import gradio as gr
from PIL import Image

# Use a pipeline as a high-level helper
from transformers import pipeline

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

from PIL import Image, ImageDraw, ImageFont

def merge_bboxes(bboxes, threshold=50):
    """
    Merges bounding boxes that are close to each other into one.

    :param bboxes: List of bounding boxes (xmin, ymin, xmax, ymax)
    :param threshold: Maximum distance between boxes to be considered as part of the same object
    :return: List of merged bounding boxes
    """
    merged_bboxes = []
    for bbox in bboxes:
        added = False
        for i, merged_bbox in enumerate(merged_bboxes):
            if (abs(bbox[0] - merged_bbox[0]) < threshold and abs(bbox[1] - merged_bbox[1]) < threshold) or \
               (abs(bbox[2] - merged_bbox[2]) < threshold and abs(bbox[3] - merged_bbox[3]) < threshold):
                # Merge the boxes by extending the boundaries
                new_bbox = (
                    min(bbox[0], merged_bbox[0]),
                    min(bbox[1], merged_bbox[1]),
                    max(bbox[2], merged_bbox[2]),
                    max(bbox[3], merged_bbox[3])
                )
                merged_bboxes[i] = new_bbox
                added = True
                break
        if not added:
            merged_bboxes.append(bbox)
    return merged_bboxes

def draw_bounding_boxes(image, detections):
    """
    Draws bounding boxes on the given PIL image based on the object detection results.
    Merges close bounding boxes to avoid duplicate detections.

    :param image: PIL Image object
    :param detections: List of dictionaries containing object detection results
    :return: PIL Image with bounding boxes
    """
    draw = ImageDraw.Draw(image)

    try:
        font = ImageFont.truetype("arial.ttf", 40)  # Load font (adjust size as needed)
    except:
        font = ImageFont.load_default()  # Use default font if Arial is not available

    bboxes = []
    labels = []

    # Collect all bounding boxes and labels
    for obj in detections:
        box = obj["box"]
        label = f"{obj['label']} ({obj['score']:.2f})"
        xmin, ymin, xmax, ymax = box["xmin"], box["ymin"], box["xmax"], box["ymax"]

        bboxes.append((xmin, ymin, xmax, ymax))
        labels.append(label)

    # Merge close bounding boxes
    merged_bboxes = merge_bboxes(bboxes)

    # Draw bounding boxes and labels for merged boxes
    for idx, bbox in enumerate(merged_bboxes):
        xmin, ymin, xmax, ymax = bbox
        label = labels[idx]

        # Draw bounding box
        draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=5)

        # Adjust label placement to avoid overlap
        label_bbox = draw.textbbox((xmin, ymin), label, font=font)
        label_width, label_height = label_bbox[2] - label_bbox[0], label_bbox[3] - label_bbox[1]
        label_x = xmin
        label_y = ymin - label_height if ymin - label_height > 0 else ymin + 5  # Adjust above or below box

        # Draw label background to make text more visible
        draw.rectangle([label_x, label_y, label_x + label_width, label_y + label_height], fill="red")
        draw.text((label_x, label_y), label, fill="white", font=font)

    return image

# Example usage:
# image = Image.open("your_image.jpg")
# detections = [{'score': 0.92, 'label': 'dog', 'box': {'xmin': 2929, 'ymin': 1297, 'xmax': 5067, 'ymax': 3693}}, ...]
# image_with_boxes = draw_bounding_boxes(image, detections)
# image_with_boxes.show()

def detect_object(image):
    raw_imagee = image
    output=pipe(raw_imagee)
    processed_image = draw_bounding_boxes(raw_imagee, output)
    return processed_image

demo = gr.Interface(fn = detect_object,
                   inputs=[gr.Image(label='Please upload the file to detect the data', type='pil')],
                   outputs=[gr.Image(label='Detected object', type='pil')],
                   title= 'Object Detector',
                   description = 'This application will be used to detect image'
                   )

demo.launch(share=True)