# Install dependencies manually if not found
# import os; os.system("pip install gradio pillow transformers")
#As of April 2025, version 4.39.3 of transformers is stable with DETR and avoids the init_empty_weights bug.
import os
os.system("pip install --upgrade torch torchvision torchaudio transformers==4.39.3 gradio pillow")



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

# for local in cache file
# model_path = "../models/models--facebook--detr-resnet-50/snapshots/1d5f47bd3bdd2c4bbfa585418ffe6da5028b4c0b"

pipe = pipeline("object-detection", model="facebook/detr-resnet-50")

object_detector = pipeline("object-detection", model="facebook/detr-resnet-50")

def draw_detections_on_image(pil_image, detections):
    draw = ImageDraw.Draw(pil_image)

    try:
        font = ImageFont.truetype("arial.ttf", 14)
    except:
        font = ImageFont.load_default()

    for det in detections:
        score = det.get('score', 0)
        label = det.get('label', 'Object')
        box = det['box']
        xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']

        draw.rectangle([xmin, ymin, xmax, ymax], outline='red', width=2)
        text = f"{label} {score:.2f}"
        bbox = font.getbbox(text)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        text_background = [xmin, ymin - text_height, xmin + text_width, ymin]
        draw.rectangle(text_background, fill='red')
        draw.text((xmin, ymin - text_height), text, fill='white', font=font)

    return pil_image

def detect_object(image):
    try:
        row_image = image
        output = object_detector(row_image)
        processed_image = draw_detections_on_image(row_image, output)
        return processed_image
    except Exception as e:
        # Handle network issues or other exceptions
        print(f"Error during object detection: {e}")
        return "An error occurred during object detection. Please check your network connection and try again."

demo = gr.Interface(
    fn=detect_object,
    inputs=gr.Image(label="Select to upload image", type="pil"),
    outputs=gr.Image(label="Processed Image", type="pil"),
    title="@cygon: Object detector",
    description="THIS SIMPLE APP CAN BE USED TO SAMPLE AND DETECT VARIOUS OBJECTS USING IMAGES.",
)

if __name__ == "__main__":
    demo.launch(share=True)  # Enables public sharing