Upload potholedetector.py

potholedetector.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+"""Untitled1.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1ZlMhMCHSQcigQYuNzsGtJAda3vE_K5R5
+"""
+
+!pip install transformers datasets gradio
+
+from transformers import DetrFeatureExtractor, DetrForObjectDetection
+
+feature_extractor = DetrFeatureExtractor.from_pretrained("facebook/detr-resnet-50")
+model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+
+import torch
+from PIL import Image
+import requests
+
+def predict(image_url):
+    image = Image.open(requests.get(image_url, stream=True).raw)
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # convert outputs (bounding boxes and class logits) to COCO API output
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = feature_extractor.post_process(outputs, target_sizes=target_sizes)[0]
+
+    potholes = []
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [round(i, 2) for i in box.tolist()]
+        if score > 0.7: # Adjust threshold as needed
+            #  Here,  we'll simply identify everything.  Ideally, you'd filter by the "pothole" class.
+            potholes.append({
+                "box": box,
+                "score": round(score.item(), 3),
+                "label": model.config.id2label[label.item()]
+            })
+
+    return potholes
+
+import gradio as gr
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(label="Image URL"),  # Input is a text box for the image URL
+    outputs=gr.JSON(label="Detected Potholes"), # Output the results as JSON
+    title="Pothole Detection POC",
+    description="Enter the URL of an image to detect potholes."
+)
+
+iface.launch(debug=True)  # Launch the interface within the Colab notebook (for testing)