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

# Initialize the model
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")

def load_model(threshold):
    # Reinitialize the model with the desired detection threshold
    config = DetrConfig.from_pretrained("facebook/detr-resnet-50", threshold=threshold)
    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=config)
    image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
    return pipeline(task='object-detection', model=model, image_processor=image_processor)

od_pipe = load_model(0.5)  # Default threshold

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)
    
    # Draw detections
    for detection in detections:
        score = detection['score']
        label = detection['label']
        box = detection['box']
        x_min = box['xmin']
        y_min = box['ymin']
        x_max = box['xmax']
        y_max = box['ymax']

        # Increase font size for better visibility
        cv2.rectangle(np_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        label_text = f'{label} {score:.2f}'
        cv2.putText(np_image, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 4)
    
    # Convert BGR to RGB for displaying
    final_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
    final_pil_image = Image.fromarray(final_image)
    return final_pil_image

def get_pipeline_prediction(threshold, pil_image):
    global od_pipe
    try:
        # Check if the model threshold needs adjusting
        if od_pipe.config.threshold != threshold:
            od_pipe = load_model(threshold)
            print("Model reloaded with new threshold:", threshold)
        
        # Ensure input is a PIL image
        if not isinstance(pil_image, Image.Image):
            pil_image = Image.fromarray(np.array(pil_image).astype('uint8'), 'RGB')
        
        # Run detection and return annotated image and results
        pipeline_output = od_pipe(pil_image)
        processed_image = draw_detections(pil_image, pipeline_output)
        return processed_image, pipeline_output
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        print(error_message)
        return pil_image, {"error": error_message}

# Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            inp_image = gr.Image(label="Input image")
            slider = gr.Slider(minimum=0, maximum=1, step=0.05, label="Detection Sensitivity", value=0.5)
            gr.Markdown("Adjust the slider to change detection sensitivity.")
            btn_run = gr.Button('Run Detection')
        with gr.Column():
            with gr.Tab("Annotated Image"):
                out_image = gr.Image()
            with gr.Tab("Detection Results"):
                out_json = gr.JSON()

    btn_run.click(get_pipeline_prediction, inputs=[slider, inp_image], outputs=[out_image, out_json])

demo.launch()