Create app.py

app.py

@@ -0,0 +1,94 @@
+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)
+    
+
+
+demo = gr.Interface(
+    
+    fn=run_detection,
+    inputs=gr.Image(type="filepath", label="Upload Image"),
+    outputs=gr.Image(type="pil",label="Output Image"),
+
+    title="FasterRCNN_ResNet50_FPN model detects 91 classes. For demo purpose, it's been adjusted to only detect car and traffic light"
+)
+
+demo.launch()