Uploading Trashify box detection model v3 app.py with NMS post processing

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/trashify_example_2.jpeg filter=lfs diff=lfs merge=lfs -text

.gradio/cached_examples/21/log.csv

@@ -0,0 +1 @@
+Image Output (no filtering),Text Output (no filtering),Image Output (with max score per class box filtering),Text Output (with max score per class box filtering),timestamp

README.md

@@ -1,13 +1,25 @@
 ---
-title: Trashify V3
-emoji: ⚡
-colorFrom: red
-colorTo: pink
+title: Trashify Demo V3 🚮
+emoji: 🗑️
+colorFrom: purple
+colorTo: blue
 sdk: gradio
-sdk_version: 4.41.0
+sdk_version: 4.40.0
 app_file: app.py
 pinned: false
 license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🚮 Trashify Object Detector Demo V3
+
+Object detection demo to detect `trash`, `bin`, `hand`, `trash_arm`, `not_trash`, `not_bin`, `not_hand`. 
+
+Used as example for encouraging people to cleanup their local area.
+
+If `trash`, `hand`, `bin` all detected = +1 point.
+
+* V1 = model trained *without* data augmentation
+* V2 = model trained *with* data augmentation
+* V3 = model trained *with* data augmentation & NMS ([Non Maximum Suppression](https://paperswithcode.com/method/non-maximum-suppression)) post processing step
+
+TK - finish the README.md + update with links to materials

app.py

@@ -1,29 +1,105 @@
 import gradio as gr
 import torch
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 
 from transformers import AutoImageProcessor
 from transformers import AutoModelForObjectDetection
 
-from PIL import Image
-
-model_save_path = "mrdbourke/detr_finetuned_trashify_box_detector_synthetic_and_real_data"
+# Note: Can load from Hugging Face or can load from local.
+# You will have to replace {mrdbourke} for your own username if the model is on your Hugging Face account.
+model_save_path = "mrdbourke/detr_finetuned_trashify_box_detector_with_data_aug" 
 
+# Load the model and preprocessor
 image_processor = AutoImageProcessor.from_pretrained(model_save_path)
 model = AutoModelForObjectDetection.from_pretrained(model_save_path)
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+
+# Get the id2label dictionary from the model
 id2label = model.config.id2label
-color_dict = {
-    "not_trash": "red",
+
+# Set up a colour dictionary for plotting boxes with different colours
+color_dict = {   
     "bin": "green",
     "trash": "blue",
-    "hand": "purple"
+    "hand": "purple",
+    "trash_arm": "yellow",
+    "not_trash": "red",
+    "not_bin": "red",
+    "not_hand": "red",
 }
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
+# Create helper functions for seeing if items from one list are in another 
+def any_in_list(list_a, list_b):
+    "Returns True if any item from list_a is in list_b, otherwise False."
+    return any(item in list_b for item in list_a)
+
+def all_in_list(list_a, list_b):
+    "Returns True if all items from list_a are in list_b, otherwise False."
+    return all(item in list_b for item in list_a)
+
+def filter_highest_scoring_box_per_class(boxes, labels, scores):
+    """
+    Perform NMS (Non-max Supression) to only keep the top scoring box per class.
+
+    Args:
+        boxes: tensor of shape (N, 4)
+        labels: tensor of shape (N,)
+        scores: tensor of shape (N,)
+    Returns:
+        boxes: tensor of shape (N, 4) filtered for max scoring item per class
+        labels: tensor of shape (N,) filtered for max scoring item per class
+        scores: tensor of shape (N,) filtered for max scoring item per class
+    """
+    # Start with a blank keep mask (e.g. all False and then update the boxes to keep with True)
+    keep_mask = torch.zeros(len(boxes), dtype=torch.bool)
+
+    # For each unique class
+    for class_id in labels.unique():
+        # Get the indicies for the target class
+        class_mask = labels == class_id
+
+        # If any of the labels match the current class_id
+        if class_mask.any():
+            # Find the index of highest scoring box for this specific class
+            class_scores = scores[class_mask]
+            highest_score_idx = class_scores.argmax()
+
+            # Convert back to the original index
+            original_idx = torch.where(class_mask)[0][highest_score_idx]
+
+            # Update the index in the keep mask to keep the highest scoring box 
+            keep_mask[original_idx] = True
+        
+    return boxes[keep_mask], labels[keep_mask], scores[keep_mask]
+
+def create_return_string(list_of_predicted_labels, target_items=["trash", "bin", "hand"]):
+     # Setup blank string to print out
+    return_string = ""
+
+    # If no items detected or trash, bin, hand not in list, return notification 
+    if (len(list_of_predicted_labels) == 0) or not (any_in_list(list_a=target_items, list_b=list_of_predicted_labels)):
+        return_string = f"No trash, bin or hand detected at confidence threshold {conf_threshold}. Try another image or lowering the confidence threshold."
+        return return_string
+
+    # If there are some missing, print the ones which are missing
+    elif not all_in_list(list_a=target_items, list_b=list_of_predicted_labels):
+        missing_items = []
+        for item in target_items:
+            if item not in list_of_predicted_labels:
+                missing_items.append(item)
+        return_string = f"Detected the following items: {list_of_predicted_labels} (total: {len(list_of_predicted_labels)}). But missing the following in order to get +1: {missing_items}. If this is an error, try another image or altering the confidence threshold. Otherwise, the model may need to be updated with better data."
+        
+    # If all 3 trash, bin, hand occur = + 1
+    if all_in_list(list_a=target_items, list_b=list_of_predicted_labels):
+        return_string = f"+1! Found the following items: {list_of_predicted_labels} (total: {len(list_of_predicted_labels)}), thank you for cleaning up the area!"
+
+    print(return_string)
+
+    return return_string
 
-def predict_on_image(image, conf_threshold=0.25):
+def predict_on_image(image, conf_threshold):
     with torch.no_grad():
         inputs = image_processor(images=[image], return_tensors="pt")
         outputs = model(**inputs.to(device))
@@ -43,13 +119,37 @@ def predict_on_image(image, conf_threshold=0.25):
     # Can return results as plotted on a PIL image (then display the image)
     draw = ImageDraw.Draw(image)
 
-    for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
+    # Create a copy of the image to draw on it for NMS
+    image_nms = image.copy()
+    draw_nms = ImageDraw.Draw(image_nms)
+
+    # Get a font from ImageFont
+    font = ImageFont.load_default(size=20)
+
+    # Get class names as text for print out
+    class_name_text_labels = []
+
+    # TK - update this for NMS
+    class_name_text_labels_nms = []
+
+    # Get original boxes, scores, labels
+    original_boxes = results["boxes"]
+    original_labels = results["labels"]
+    original_scores = results["scores"]
+
+    # Filter boxes and only keep 1x of each label with highest score
+    filtered_boxes, filtered_labels, filtered_scores = filter_highest_scoring_box_per_class(boxes=original_boxes,
+                                                                                            labels=original_labels,
+                                                                                            scores=original_scores)
+    # TODO: turn this into a function so it's cleaner?
+    for box, label, score in zip(original_boxes, original_labels, original_scores):
         # Create coordinates
         x, y, x2, y2 = tuple(box.tolist())
 
         # Get label_name
         label_name = id2label[label.item()]
         targ_color = color_dict[label_name]
+        class_name_text_labels.append(label_name)
 
         # Draw the rectangle
         draw.rectangle(xy=(x, y, x2, y2), 
@@ -62,23 +162,70 @@ def predict_on_image(image, conf_threshold=0.25):
         # Draw the text on the image
         draw.text(xy=(x, y),
                   text=text_string_to_show,
-                  fill="white")
+                  fill="white",
+                  font=font)
+    
+    # TODO: turn this into a function so it's cleaner?
+    for box, label, score in zip(filtered_boxes, filtered_labels, filtered_scores):
+        # Create coordinates
+        x, y, x2, y2 = tuple(box.tolist())
+
+        # Get label_name
+        label_name = id2label[label.item()]
+        targ_color = color_dict[label_name]
+        class_name_text_labels_nms.append(label_name)
+
+        # Draw the rectangle
+        draw_nms.rectangle(xy=(x, y, x2, y2), 
+                       outline=targ_color,
+                       width=3)
+        
+        # Create a text string to display
+        text_string_to_show = f"{label_name} ({round(score.item(), 3)})"
+
+        # Draw the text on the image
+        draw_nms.text(xy=(x, y),
+                  text=text_string_to_show,
+                  fill="white",
+                  font=font)
+    
     
     # Remove the draw each time
     del draw
+    del draw_nms
+
+    # Create the return string
+    return_string = create_return_string(list_of_predicted_labels=class_name_text_labels)
+    return_string_nms = create_return_string(list_of_predicted_labels=class_name_text_labels_nms)
     
-    return image
+    return image, return_string, image_nms, return_string_nms
 
+# Create the interface
 demo = gr.Interface(
     fn=predict_on_image,
     inputs=[
-        gr.Image(type="pil", label="Upload Target Image"),
+        gr.Image(type="pil", label="Target Image"),
         gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold")
     ],
-    outputs=gr.Image(type="pil"),
-    title="🚮 Trashify Object Detection Demo (real and synthetic data)",
-    description="Upload an image to detect whether there's a bin, a hand or trash in it. Trained on a mixture of real and synthetic data."
+    outputs=[
+        gr.Image(type="pil", label="Image Output (no filtering)"),
+        gr.Text(label="Text Output (no filtering)"),
+        gr.Image(type="pil", label="Image Output (with max score per class box filtering)"),
+        gr.Text(label="Text Output (with max score per class box filtering)")
+        
+    ],
+    title="🚮 Trashify Object Detection Demo V3",
+    description="""Help clean up your local area! Upload an image and get +1 if there is all of the following items detected: trash, bin, hand.
+    Model in V3 has been trained with data augmentation and has an additional post-processing step to filter classes for only the highest scoring box of each class. (tk - add link to model). 
+    """,
+    # Examples come in the form of a list of lists, where each inner list contains elements to prefill the `inputs` parameter with
+    examples=[
+        ["examples/trashify_example_1.jpeg", 0.25],
+        ["examples/trashify_example_2.jpeg", 0.25],
+        ["examples/trashify_example_3.jpeg", 0.25]
+    ],
+    cache_examples=True
 )
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the demo
+demo.launch()