import streamlit as st
from ultralytics import YOLO
from PIL import Image
import numpy as np
import io

# --- Page Config ---
st.set_page_config(page_title="Palm Oil Ripeness AI", layout="wide")
st.title("🌴 Palm Oil FFB Ripeness Detector")
st.markdown("### R&D Proof of Concept: Automated Maturity Grading")

# --- Load Model (Cached for performance) ---
@st.cache_resource
def load_model():
    return YOLO("best.pt")

model = load_model()

# --- Sidebar Controls ---
st.sidebar.header("Settings")
conf_threshold = st.sidebar.slider("Confidence Threshold", 0.1, 1.0, 0.5)

# --- Image Upload ---
uploaded_file = st.file_uploader("Drag and drop a Palm Oil FFB image here...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Convert uploaded file to PIL Image
    image = Image.open(uploaded_file)
    
    # Layout: Original vs Predicted
    col1, col2 = st.columns(2)
    
    with col1:
        st.image(image, caption="Uploaded Image", use_container_width=True)
        
    with col2:
        with st.spinner('Analyzing ripeness...'):
            # Run Inference
            results = model.predict(source=image, conf=conf_threshold)
            
            # The .plot() method returns a BGR numpy array with boxes drawn
            annotated_img = results[0].plot()
            
            # Convert BGR (OpenCV format) to RGB (Streamlit/PIL format)
            annotated_img_rgb = annotated_img[:, :, ::-1]
            
            st.image(annotated_img_rgb, caption="AI Analysis Result", use_container_width=True)

    # --- Metrics Section ---
    st.divider()
    st.subheader("Analysis Summary")
    
    detections = results[0].boxes
    if len(detections) > 0:
        for box in detections:
            label = model.names[int(box.cls)]
            conf = float(box.conf)
            st.success(f"**Detected:** {label} | **Confidence:** {conf:.2%}")
    else:
        st.warning("No fruit bunches detected. Try adjusting the confidence slider.")