from PIL import Image, ImageDraw,ImageFont
import gradio as gr
import torch
from transformers import pipeline

modelname= "SenseTime/deformable-detr-with-box-refine"
modelpath = 'models/models--SenseTime--deformable-detr-with-box-refine/snapshots/2e9e461623a8fdc296e19666c46c8a4389a3a6fe'

def draw_bounding_boxes(image, detections, color=(0, 255, 0), thickness=2, font_path=None, font_size=12):
  draw = ImageDraw.Draw(image)
  for detection in detections:
    xmin, ymin, xmax, ymax = detection['box']['xmin'], detection['box']['ymin'], detection['box']['xmax'], detection['box']['ymax']
    draw.rectangle(((xmin, ymin), (xmax, ymax)), outline=color, width=thickness)
    if font_path:
      try:
        font = ImageFont.truetype(font_path, font_size)
        label_text = f"{detection['label']}: {detection['score']:.2f}"
        text_width, text_height = draw.textsize(label_text, font=font)
        draw.rectangle(((xmin, ymin), (xmin + text_width + 5, ymin + text_height + 5)), fill=(0, 0, 0, 0.5))  # Semi-transparent black background
        draw.text((xmin, ymin), label_text, fill=color, font=font)
      except (IOError, OSError):
        print(f"Warning: Could not load font '{font_path}'. Labels will not be drawn.")
  return image


def hf_pipeline(model_name=None,model_path=None):
    model = model_path if model_name == None else model_name
    print(f"=============model: {model} =============")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    image_detector = pipeline("object-detection", model=model,device=device)
    return image_detector


def detect_image_withbox(image):
    obj_detector = hf_pipeline(modelname)
    detections = obj_detector(image)
    image_with_boxes = draw_bounding_boxes(image.copy(), detections)
    print(detections)
    return image_with_boxes



demo = gr.Interface(fn=detect_image_withbox,
                    inputs=[gr.Image(label="Select Image",type="pil")],
                    outputs=[gr.Image(label="Processed Image With Boxes", type="pil")],
                    title="@SmartChoiceLearningHub HF Project 2 : Object Detector With Box",
                    description="This app detects objects in an image and draws bounding boxes around them.")
demo.launch()