from pprint import pprint

import torch
import torchvision.models as tv_models
import torchvision.transforms as transforms
import torchvision.utils as vutils
from IPython.display import Image as DisplayImage
from PIL import Image
from torchvision.io import decode_image

import matplotlib.pyplot as plt
import numpy as np

from torchvision.transforms.functional import to_pil_image
from IPython.display import display
import gradio as gr

bb_model_weights = tv_models.detection.FasterRCNN_ResNet50_FPN_Weights.DEFAULT
model = tv_models.detection.fasterrcnn_resnet50_fpn(weights=bb_model_weights).eval()


def detect_and_draw_bboxes(model, image_path, object_indices, labels, bbox_colors, threshold, bbox_width=5):

    # Open and transform the image, and prepare the result tensor
    pil_image = Image.open(image_path).convert("RGB")
    transform_to_tensor = transforms.Compose([transforms.ToTensor()])
    tensor_image_batch = transform_to_tensor(pil_image).unsqueeze(0)
    result_image_tensor = (tensor_image_batch.squeeze(0) * 255).byte()

    # Perform inference to get predictions for all possible objects
    with torch.no_grad():
        prediction = model(tensor_image_batch)[0]

    # Initialize lists to collect all boxes, labels, and colors that meet the criteria
    all_boxes_to_draw = []
    all_labels_to_draw = []
    all_colors_to_draw = []

    # Loop through each target class to find its boxes
    for index, label, color in zip(object_indices, labels, bbox_colors):
        # Filter predictions for the current class index and confidence threshold
        class_mask = (prediction['labels'] == index) & (prediction['scores'] > threshold)
        
        # Get the boxes for the current class
        boxes_for_this_class = prediction['boxes'][class_mask]

        if boxes_for_this_class.nelement() > 0:
            # Add the found boxes to our master list
            all_boxes_to_draw.extend(boxes_for_this_class.tolist())
            # Create and add corresponding labels and colors
            all_labels_to_draw.extend([label] * len(boxes_for_this_class))
            all_colors_to_draw.extend([color] * len(boxes_for_this_class))

    # After checking all classes, draw all collected boxes at once if any were found
    if all_boxes_to_draw:
        result_image_tensor = vutils.draw_bounding_boxes(
            result_image_tensor,
            torch.tensor(all_boxes_to_draw),
            labels=all_labels_to_draw,
            colors=all_colors_to_draw,
            width=bbox_width
        )
    else:
        # If the list of boxes to draw is empty, print this information.
        print(f"No objects from the list {labels} were found with a confidence score above {threshold}.\n")
    
    return result_image_tensor


def run_detection(image_path):
    result_image_tensor = detect_and_draw_bboxes(
        model=model,
        image_path=image_path,
        object_indices=[3, 10],
        labels=['car', 'traffic light'],
        bbox_colors=['red', 'blue'],
        threshold=0.7,
    )

    
    return to_pil_image(result_image_tensor)
    

with gr.Blocks() as demo:

    iface = gr.Interface(
        fn=run_detection,
        inputs=gr.Image(type="filepath", label="Upload Image"),
        outputs=gr.Image(type="pil", label="Output Image"),
        title="<br>This app uses:<br>PyTorch Vision <br>PyTorch Model: FasterRCNN_ResNet50_FPN<br>PS: click submit after image uploaded"
    )
    gr.Markdown("### The model is adjusted to only detect car and traffic light. Example output:")
    gr.Image(
        value="exap1.JPG",  
        show_label=False,
        width=450
    )

demo.launch()