# Install required libraries if not already installed
# !pip install gradio opencv-python torch torchvision

import gradio as gr
import cv2
import torch
from torchvision import models, transforms
from PIL import Image

# Load the pre-trained Faster R-CNN model
model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

# Define the transformation for the input image
transform = transforms.Compose([
    transforms.ToTensor()
])

# Function to perform object detection
def detect_objects(input_image):
    # Convert the Gradio image to PIL Image
    image_pil = Image.fromarray(input_image.astype('uint8'), 'RGB')
    
    # Apply transformations
    image = transform(image_pil)
    image = image.unsqueeze(0)  # Add batch dimension
    
    # Get predictions
    with torch.no_grad():
        predictions = model(image)
    
    # Process predictions
    boxes = predictions[0]['boxes'].detach().numpy()
    labels = predictions[0]['labels'].detach().numpy()
    scores = predictions[0]['scores'].detach().numpy()
    
    # Convert PIL Image to OpenCV format
    image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    
    # Draw bounding boxes on the image
    for box, label, score in zip(boxes, labels, scores):
        if score < 0.5:
            continue  # Skip detections with low confidence
        x1, y1, x2, y2 = box.astype(int)
        cv2.rectangle(image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image_cv, f'{label}: {score:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    
    # Convert back to RGB for Gradio
    image_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
    return image_rgb

# Create the Gradio interface
app = gr.Interface(
    fn=detect_objects,
    inputs="image",
    outputs="image",
    title="Object Detection using Faster R-CNN",
    description="Upload an image and the model will detect objects and draw bounding boxes around them."
)

# Launch the app
app.launch()