Uploading the trash object detection system model app.py

.gitattributes

@@ -33,3 +33,5 @@ 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
+image_examples/trash_example_2.jpeg filter=lfs diff=lfs merge=lfs -text
+image_examples/trash_example_3.jpeg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,22 @@
 ---
-title: Trash Object Detection Demo
-emoji: 🏆
-colorFrom: green
-colorTo: indigo
+title: Trash Object Detection Demo 🚮
+emoji: 🗑️
+colorFrom: purple
+colorTo: blue
 sdk: gradio
-sdk_version: 5.42.0
+sdk_version: 5.34.0
 app_file: app.py
-pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🚮 Trash Object Detector
+Object detection demo to detect `trash`, `bin`, `hand`, `trash_arm`, `not_trash`, `not_bin`, `not_hand`.
+
+Used as example to encourage people to clean up their local area.
+
+If `trash`, `hand`, `bin` all dected = +1 point.
+
+## Dataset
+The model is trained on a custom dataset, hand-labelled of people picking up trash and placing it in a bin.
+
+The dataset is found in Hugging Face as [`mrdbourke/trashify_manual_labelled_images`](https://huggingface.co/datasets/mrdbourke/trashify_manual_labelled_images).

app.py

@@ -0,0 +1,153 @@
+
+# 1. Importing the required libraries and packages
+import gradio as gr
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+
+# 2. Setup preprocessing and helper functions
+# The model path to load (from Hugging Face)
+model_path = "Saint5/rt_detrv2_finetuned_trash_box_detector_v1"
+
+# Loading the model and the processor
+image_processor = AutoImageProcessor.from_pretrained(model_path)
+model = AutoModelForObjectDetection.from_pretrained(model_path)
+
+# Set the target device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+
+# Get the id2label dictionary from the model
+id2label = model.config.id2label
+
+# Setting up colour dictionary for plotting boxes with different colours
+colour_dict = {
+    "bin" : "green",
+    "trash" : "blue",
+    "hand" : "purple",
+    "trash_arm" : "yellow",
+    "not_trash" : "red",
+    "not_bin" : "red",
+    "not_bin" : "red",
+}
+
+# 3. Create function to predict on a given image with a given confidence threshold
+def predict_on_image(image, conf_threshold):
+  model.eval()
+
+  # Make a prediction on target image
+  with torch.no_grad():
+    inputs = image_processor(images=[image], return_tensors="pt")
+    model_outputs = model(**inputs.to(device))
+
+    target_sizes = torch.tensor([[image.size[1], image.size[0]]]) # -> [batch_size, height, width]
+
+    # Post process the raw outputs from the model
+    results = image_processor.post_process_object_detection(model_outputs,
+                                                            threshold=conf_threshold,
+                                                            target_sizes=target_sizes)[0]
+  # Return all items in results to CPU (for display with matplotlib)
+  for key, value in results.items():
+    try:
+      result[key] = value.item().cpu()
+    except:
+      results[key] = value.cpu()
+
+  # 4. Draw the predictions on the target image
+  draw = ImageDraw.Draw(image)
+
+  # Get a font from ImageFont
+  font = ImageFont.load_default(size=20)
+
+  # Get a class name as text for print out
+  detected_class_name_text_labels = []
+
+  # Iterate through the predictions of the model and draw them on the target image
+  for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
+    # Create coordinates
+    x, y, x2, y2 = tuple(box.tolist())
+
+    # Get label name
+    label_name = id2label[label.item()]
+    targ_colour = colour_dict[label_name]
+    detected_class_name_text_labels.append(label_name)
+
+    # Draw the rectangle
+    draw.Rectangle(xy=(x, y, x2, y2),
+                   outline=targ_colour,
+                   width=3)
+
+    # Create a text string to display
+    text_string_to_show = f"{label_name} ({round(score.item(), 3)})"
+
+    # Draw the text string on the image
+    draw.text(xy=(x, y),
+              text=text_string_to_show,
+              fill="white",
+              font=font)
+
+  # Remove the draw each time
+  del draw
+
+  # 5. Create logic for outputting information message
+  # Setup set of targets to discover
+  target_items = {"trash", "bin", "hand"}
+  detected_items = set(detected_class_name_text_labels)
+
+  # If the items are not detected, return notification
+  if not detected_items and target_items:
+    return_string = (
+        f"""No trash, bin or hand detected at confidence threshold {conf_threshold}.
+        Try another image or lowering the confidence threshold."""
+    )
+    print(return_string)
+    return image, return_string
+
+  # If there are missing items, say what the missing items are
+  missing_items = target_items - detected_items
+  if missing_items:
+    return_string = (
+        f"""Detected the following items: {sorted(detected_items and target_items)} but missing the following in order to get 1 point: {sorted(missing_items)}.
+        If this is an error, try another image or alter the confidence threshold.
+        Otherwise, the model may need to be updated with better data."""
+    )
+    print(return_string)
+    return image, return_string
+
+  # If all target items are present
+  return_string = f"+1 Point!🪙 Found the following items: {sorted(detected_items)}, thank you for cleaning up the area!"
+  print(return_string)
+  return image, return_string
+
+# 6. Setup the demo application to take in image, make a prediction with the model, return the image with drawn predicitons
+# Write a description for the gradio interface
+description = """
+An object detection system that lets the user upload a picture of them holding trash in their hand and placing it in a bin. The system will be able to detect the hand, trash and the bin. If all the three items are available,
+the user get 1 point!
+
+Model used for the system is a finetuned version of [RT-DETRv2](https://huggingface.co/docs/transformers/main/en/model_doc/rt_detr_v2#transformers.RTDetrV2Config) on the manually hand labelled [dataset](https://huggingface.co/datasets/mrdbourke/trashify_manual_labelled_images).
+"""
+
+# Create the gradio interface
+demo = gr.Interface(
+    fn = predict_on_image,
+    inputs = [
+        gr.Image(type="pil", label="Target Image"),
+        gr.Slider(minimum=0, maximum=1, value=0.3, label="Confidence Threshold")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Image Output"),
+        gr.Text(label="Text Output")
+    ],
+    title = "🗑️🚮 Trash Object Detection Model Demo",
+    description=description,
+    examples=[
+        ["image_examples/trash_example_1.jpeg", 0.3],
+        ["image_examples/trash_example_2.jpeg", 0.3],
+        ["image_examples/trash_example_3.jpeg", 0.3],
+    ],
+    cache_examples=True
+)
+
+# Launch the demo
+demo.launch()

requirements.txt

@@ -0,0 +1,4 @@
+timm
+gradio
+torch
+transformers