Create app.py

app.py

@@ -0,0 +1,56 @@
+import torch
+from transformers import DetrForObjectDetection, DetrImageProcessor
+from PIL import Image
+import gradio as gr
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+# Load pre-trained model and processor
+model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+
+def detect_car(image: Image.Image) -> Image.Image:
+    # Preprocess the input image
+    inputs = processor(images=image, return_tensors="pt")
+
+    # Run the model to get predictions
+    outputs = model(**inputs)
+
+    # Postprocess the outputs to get bounding boxes and labels
+    target_sizes = torch.tensor([image.size[::-1]])  # (height, width)
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]
+
+    # Plotting the image with bounding boxes for objects
+    fig, ax = plt.subplots(1, figsize=(12, 8))
+    ax.imshow(image)
+
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        if score > 0.7:  # Confidence threshold for detecting cars
+            xmin, ymin, xmax, ymax = box.detach().numpy()
+            width, height = xmax - xmin, ymax - ymin
+            rect = patches.Rectangle((xmin, ymin), width, height, linewidth=2, edgecolor='red', facecolor='none')
+            ax.add_patch(rect)
+            ax.text(xmin, ymin, f"{model.config.id2label[label.item()]}: {score:.2f}",
+                    color='white', fontsize=12, bbox=dict(facecolor='red', alpha=0.5))
+
+    # Convert the plot to an image
+    plt.axis('off')
+    plt.tight_layout()
+
+    # Save the figure to a canvas and convert to image
+    fig.canvas.draw()
+    result_img = Image.frombytes('RGB', fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
+    plt.close(fig)
+    return result_img
+
+# Gradio interface to upload images and get object detection results
+iface = gr.Interface(
+    fn=detect_car,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Car Detection with DETR",
+    description="Upload an image and the model will detect cars with bounding boxes. Only cars will be displayed."
+)
+
+if __name__ == "__main__":
+    iface.launch()