# Description: This is the main file to run the Gradio interface for the object detection model.
from ultralytics import YOLO
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
import os

model_path = "best_int8_openvino_model"

# Example paths for Gradio
image_examples = [["DurianMangosteen1.jpg"], ["DurianMangosteen2.jpg"]]

# Load the model    
def load_model(repo_id):
    download_dir = snapshot_download(repo_id) # download the model from the Hugging Face Hub
    print(download_dir)
    path  = os.path.join(download_dir, "best_int8_openvino_model") # path to the model
    print(path)
    detection_model = YOLO(path, task='detect') # load the model
    return detection_model

# Predict the image
def predict(pilimg):
    source = pilimg
    # x = np.asarray(pilimg)
    # print(x.shape)
    result = detection_model.predict(source, conf=0.4, iou=0.6) # confidence threshold, intersection over union threshold
    
    #print("Result: ", result)
    if not result or len(result[0].boxes) == 0: # if no object detected
        gr.Warning("No object detected in the image!")
    else:
        img_bgr = result[0].plot() # plot the image
        out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
    
        return out_pilimg

REPO_ID = "ITI107-2024S2/8035531F"  # The repo ID of the model
detection_model = load_model(REPO_ID)

title = "Detect Durian and Mangosteen (King and Queen of Fruits) In The Image"
interface = gr.Interface(
                fn=predict,
                inputs=gr.Image(type="pil", label="Input Image"),
                outputs=gr.Image(type="pil", label="Object Detected Image"),
                title=title,
                examples=image_examples,
             )

# Launch the interface
interface.launch(share=True)