# app.py for Hugging Face Spaces

import gradio as gr
from ultralyticsplus import YOLO
from PIL import Image
import numpy as np

# Load the model
model = YOLO('foduucom/product-detection-in-shelf-yolov8')

# Set model parameters
model.overrides['conf'] = 0.25  # NMS confidence threshold
model.overrides['iou'] = 0.45  # NMS IoU threshold
model.overrides['agnostic_nms'] = False  # NMS class-agnostic
model.overrides['max_det'] = 1000  # maximum number of detections per image

def detect_skus(image):
    # Convert Gradio image to numpy array if needed, but YOLO accepts PIL
    results = model(image)
    
    # Extract unique SKU names (assuming classes are SKU names)
    sku_set = set()
    for result in results:
        for box in result.boxes:
            class_id = int(box.cls)
            sku_name = result.names[class_id]
            sku_set.add(sku_name)
    
    sku_list = list(sku_set)
    return "\n".join(sku_list) if sku_list else "No SKUs detected."

# Create Gradio interface
iface = gr.Interface(
    fn=detect_skus,
    inputs=gr.Image(type="pil", label="Upload Shelf Image"),
    outputs=gr.Textbox(label="Detected SKUs"),
    title="Shelf SKU Detector",
    description="Upload an image of a cigarette shelf to detect and list the SKUs."
)

if __name__ == "__main__":
    iface.launch()