import os
import gradio as gr
from transformers import pipeline, DetrForObjectDetection, DetrConfig, DetrImageProcessor
import numpy as np
import cv2
from PIL import Image

def draw_detections(image, detections):
    # Convert PIL image to a numpy array
    np_image = np.array(image)
    
    # Convert RGB to BGR for OpenCV
    np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
    
    for detection in detections:
        # Each detection includes ['score', 'label', 'box']
        score = detection['score']
        label = detection['label']
        box = detection['box']
        x_min, y_min, x_max, y_max = map(int, box)
        cv2.rectangle(np_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        cv2.putText(np_image, f'{label} {score:.2f}', (x_min, max(y_min - 10, 0)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    
    # Convert BGR to RGB for displaying
    final_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
    # Convert the numpy array to PIL Image
    final_pil_image = Image.fromarray(final_image)
    return final_pil_image

# Initialize objects from transformers
config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=config)
image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
od_pipe = pipeline(task='object-detection', model=model, image_processor=image_processor)

def get_pipeline_prediction(pil_image):
    # Run the object detection pipeline
    pipeline_output = od_pipe(pil_image)
    
    # Draw the detection results on the image
    processed_image = draw_detections(pil_image, pipeline_output)
    
    # Provide both the image and the JSON detection results
    return processed_image, pipeline_output

demo = gr.Interface(
    fn=get_pipeline_prediction,
    inputs=gr.Image(label="Input image", type="pil"),
    outputs=[
        gr.Image(label="Annotated Image"),
        gr.JSON(label="Detected Objects")
    ]
)

demo.launch()