Create app.py

app.py

@@ -0,0 +1,84 @@
+import gradio as gr
+import torch
+from PIL import Image, ImageDraw
+
+from transformers import AutoImageProcessor
+from transformers import AutoModelForObjectDetection
+
+from PIL import Image
+
+model_save_path = "mrdbourke/detr_finetuned_trashify_box_detector"
+
+image_processor = AutoImageProcessor.from_pretrained(model_save_path)
+model = AutoModelForObjectDetection.from_pretrained(model_save_path)
+
+id2label = model.config.id2label
+color_dict = {
+    "not_trash": "red",
+    "bin": "green",
+    "trash": "blue",
+    "hand": "purple"
+}
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+
+def predict_on_image(image, conf_threshold=0.25):
+    with torch.no_grad():
+        inputs = image_processor(images=[image], return_tensors="pt")
+        outputs = model(**inputs.to(device))
+
+        target_sizes = torch.tensor([[image.size[1], image.size[0]]]) # height, width 
+
+        results = image_processor.post_process_object_detection(outputs,
+                                                                threshold=conf_threshold,
+                                                                target_sizes=target_sizes)[0]
+    # Return all items in results to CPU
+    for key, value in results.items():
+        try:
+            results[key] = value.item().cpu() # can't get scalar as .item() so add try/except block
+        except:
+            results[key] = value.cpu()
+
+    # 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 coordinates
+        x, y, x2, y2 = tuple(box.tolist())
+
+        # Get label_name
+        label_name = id2label[label.item()]
+        targ_color = color_dict[label_name]
+
+        # Draw the rectangle
+        draw.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.text(xy=(x, y),
+                  text=text_string_to_show,
+                  fill="white")
+    
+    # Remove the draw each time
+    del draw
+    
+    return image
+
+demo = gr.Interface(
+    fn=predict_on_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Target Image"),
+        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold")
+    ],
+    outputs=gr.Image(type="pil"),
+    title="🚮 Trashify Object Detection Demo",
+    description="Upload an image to detect whether there's a bin, a hand or trash in it."
+)
+
+if __name__ == "__main__":
+    demo.launch()