from pathlib import Path

import gradio as gr
from ultralytics import YOLO
from PIL import Image


# Load YOLOv8n model
MODEL = YOLO('weights/best.pt')
IMAGES_PATH = Path("images/")

INF_PARAMETERS = {
    "imgsz": 640,  # image size
    "conf": 0.8,   # confidence threshold
    "max_det": 1   # maximum number of detections
}

EXAMPLES = [path for path in IMAGES_PATH.iterdir()]


# Function to detect objects and crop the image
def detect_and_crop(image: Image.Image) -> Image.Image:
    # Perform object detection
    results = MODEL.predict(image,**INF_PARAMETERS)
    result = results[0]
    for box in result.boxes.xyxy.cpu().numpy():
        cropped_image = image.crop(box=box)
    return cropped_image


# Gradio UI
title = "Crop-Detection"
description = """## 🍋‍🟩 Automatically crop product pictures! 🍋‍🟩
When contributors use the mobile app, they are asked to take pictures of the product, then to crop it. 
To assist users during the process, we create a crop-detection model desin to detect the product edges.

We fine-tuned Yolov8n on images extracted from the Open Food Facts database.
Check the [model repo page](https://huggingface.co/openfoodfacts/crop-detection) for more information.
"""

# Gradio Interface
demo = gr.Interface(
    fn=detect_and_crop,               
    inputs=gr.Image(type="pil", width=300),    
    outputs=gr.Image(type="pil", width=300),
    title=title,
    description=description,
    allow_flagging="never",
    examples=EXAMPLES
)


# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()