from clearml import Model
from ultralytics import YOLO
from PIL import Image, ImageDraw
from huggingface_hub import hf_hub_download
import gradio as gr


# fetching the model from the ClearML 
def fetch_model_from_HG():
    try:
        hg_model = hf_hub_download(repo_id="manvi23/GroupG", filename="best.pt")
        return hg_model
    except Exception as e:
        print(f"Failed to fetch model from HuggingFace: {e}")
        raise
 
model_path = fetch_model_from_HG()
model = YOLO(model_path)


def predict_and_visualize(image, confidence_threshold=0.3):
    results = model.predict(image, conf=confidence_threshold)

    boxes = results[0].boxes.xyxy.cpu().numpy()  # Bounding box coordinates
    scores = results[0].boxes.conf.cpu().numpy()  # Confidence scores

    # Draw bounding boxes on the image
    draw = ImageDraw.Draw(image)
    for box, score in zip(boxes, scores):
        x_min, y_min, x_max, y_max = box
        draw.rectangle([x_min, y_min, x_max, y_max], outline="red", width=3)
        draw.text((x_min, y_min), f"{score:.2f}", fill="red")

    return image


# Gradio interface function
def gradio_app(image):
    result_image = predict_and_visualize(image)
    return result_image


# Create Gradio app
interface = gr.Interface(
    fn=gradio_app,
    inputs=[
        gr.Image(type="pil"),  # Input image
    ],
    outputs=gr.Image(type="pil"),  # Output image
    title="Object Detection with model of Group G",
    description="Upload an image to detect the objects.",
)

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