import tensorflow as tf
import cv2
import numpy as np
import gradio as gr

# Load the pre-trained MobileNet SSD model
model = tf.saved_model.load("http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_fpnlite_320x320/saved_model")

# Define the label map for the MobileNet SSD model
category_index = {
    1: {'id': 1, 'name': 'person'},
    2: {'id': 2, 'name': 'bicycle'},
    3: {'id': 3, 'name': 'car'},
    # Add more label mappings as needed
}

# Function to detect objects in the image
def detect_objects(image):
    # Preprocess the image
    input_tensor = tf.convert_to_tensor(image)
    input_tensor = input_tensor[tf.newaxis,...]

    # Run the model and get detections
    detections = model(input_tensor)

    # Process detections and draw bounding boxes
    num_detections = int(detections.pop('num_detections'))
    detections = {key: value[0, :num_detections].numpy() for key, value in detections.items()}
    detection_classes = detections['detection_classes'].astype(np.int64)
    detection_boxes = detections['detection_boxes']
    detection_scores = detections['detection_scores']

    # Draw boxes on the image
    for i in range(num_detections):
        if detection_scores[i] > 0.5:  # Only consider confident detections
            class_name = category_index.get(detection_classes[i], {'name': 'N/A'})['name']
            box = detection_boxes[i]
            height, width, _ = image.shape
            ymin, xmin, ymax, xmax = box
            (startX, startY, endX, endY) = (int(xmin * width), int(ymin * height), int(xmax * width), int(ymax * height))

            # Draw bounding box and label
            cv2.rectangle(image, (startX, startY), (endX, endY), (0, 255, 0), 2)
            label = f"{class_name}: {detection_scores[i]:.2f}"
            cv2.putText(image, label, (startX, startY - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return image

# Function to handle the image from file upload or path
def gradio_interface(image):
    if isinstance(image, str):  # Check if it's a path string
        image = cv2.imread(image)
    else:
        # Convert PIL image (Gradio) to OpenCV format (numpy array)
        image = np.array(image)

    # Convert to RGB format
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Detect objects in the image
    detected_image = detect_objects(image_rgb)

    # Convert back to BGR for display in Gradio (OpenCV uses BGR)
    detected_image_bgr = cv2.cvtColor(detected_image, cv2.COLOR_RGB2BGR)
    
    return detected_image_bgr

# Create Gradio app with image input (supports path or upload)
iface = gr.Interface(fn=gradio_interface, 
                     inputs=gr.inputs.Image(type="filepath"),  # Use "filepath" to allow local path or upload
                     outputs="image", 
                     title="Object Detection with Bounding Boxes", 
                     description="Upload an image or provide a file path to detect objects.")

# Launch the Gradio app (for local or Hugging Face)
iface.launch()