# app.py

import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np

# Load YOLOv8 model (use 'yolov8n.pt' or your fine-tuned model path)
model = YOLO('yolov8n.pt')

# Price mapping
price_map = {
    'cell phone': '$499',
    'watch': '$199',
    'apple': '$2'
}

# Class aliases to map YOLO labels to product tags
class_aliases = {
    'cell phone': ['cell phone', 'mobile phone'],
    'watch': ['handbag', 'accessory', 'person'],  # adjust based on model behavior
    'apple': ['orange', 'fruit', 'apple']
}

# Optional: color map for visual feedback
color_map = {
    'cell phone': (255, 0, 0),
    'watch': (0, 255, 255),
    'apple': (0, 128, 0)
}

# Detection function
def detect_and_label(image):
    rgb_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = model(rgb_img)[0]

    for box in results.boxes:
        cls_id = int(box.cls[0])
        label = model.names[cls_id]
        conf = float(box.conf[0])
        x1, y1, x2, y2 = map(int, box.xyxy[0])

        for tag, aliases in class_aliases.items():
            if label in aliases:
                color = color_map.get(tag, (0, 255, 0))
                cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
                text = f"{tag}: {price_map[tag]} ({conf:.2f})"
                cv2.putText(image, text, (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

    return image

# Gradio UI
demo = gr.Interface(
    fn=detect_and_label,
    inputs=gr.Image(type="numpy", label="Upload Product Image"),
    outputs=gr.Image(type="numpy", label="Detected with Price Tags"),
    title="🛍️ Product Detector with Price Overlay",
    description="Upload an image containing cell phones, watches, or apples. The app detects them and overlays price tags."
)

demo.launch()