#!/usr/bin/env python
# coding: utf-8

# In[ ]:


from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
from huggingface_hub import snapshot_download
import os
import numpy as np

# Best Yolov8 trained model
best_yolo_model = "best.pt"

def load_model(repo_id):
    download_dir = snapshot_download(repo_id)
    print(download_dir)
    path  = os.path.join(download_dir, "best_int8_openvino_model")
    print(path)
    detection_model = YOLO(path, task='detect')
    return detection_model


def predict(pilimg):
 
    detection_model = YOLO(best_yolo_model, task='detect')
    
    result = detection_model.predict(pilimg, conf=0.4, iou=0.5)
    #result = detection_model.predict(pilimg)
        
    obb = result[0].obb
    confs = obb.conf
    #print("obb : ", obb)
    #print( confidence : ", confs)
              
    img_bgr = result[0].plot()
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
   
    # Show warning if no object is detected
    if len(result[0].obb.cls) == 0:
        gr.Warning("No object detected!")   
         
    return out_pilimg


REPO_ID = "ITI107-2024S2/8035531F"
#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 Detection With Confidence Level"),
                title=title,
             )

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


# In[ ]: