Update app.py

app.py

@@ -1,47 +1,39 @@
 import gradio as gr
-from transformers import DetrFeatureExtractor, DetrForObjectDetection, DetrImageProcessor  # Use DetrImageProcessor
+from transformers import DetrImageProcessor, DetrForObjectDetection
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw
 
-# Load the model and processor (suppressing warnings and using DetrImageProcessor)
+# Model loading (same as before - with error handling)
 try:
-    feature_extractor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") # Correct image processor
-    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", ignore_mismatched_sizes=True) # Suppress warnings
-except Exception as e: # Catch potential errors during model loading
-    print(f"Error loading model: {e}") # So you can see in the logs what’s happening
-    raise e
-
+    feature_extractor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", ignore_mismatched_sizes=True)
+except Exception as e:  #  Error handling during model loading
+    print(f"Error loading model: {e}")  # Log the error so you can see in HF logs
+    raise e  # Re-raise for Space to report it
 
 def predict(image):
-
     inputs = feature_extractor(images=image, return_tensors="pt")
     outputs = model(**inputs)
 
-
     target_sizes = torch.tensor([image.size[::-1]])
-    results = feature_extractor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.7)[0] # Simplify this line
+    results = feature_extractor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.7)[0]
 
-    potholes = []
+    # Draw bounding boxes on the image
+    draw = ImageDraw.Draw(image) # Create a drawing object
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        box = [round(i, 2) for i in box.tolist()]
-        potholes.append({
-                "box": box,
-                "score": round(score.item(), 3),
-                "label": model.config.id2label[label.item()]
-            })
-
-    return potholes
-
-
+        box = [round(i) for i in box.tolist()]  # Convert to integers for drawing
+        draw.rectangle(box, outline="red", width=2) # Outline
+        draw.text((box[0], box[1]), model.config.id2label[label.item()], fill="red") # Add a label
 
+    return image  # Return the image with the bounding boxes drawn
 
+# Gradio Interface (updated output type)
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs=gr.JSON(label="Detected Potholes"), #  Correct for Gradio 3.0+
+    outputs=gr.Image(type="pil", label="Detected Potholes (Image)"), # Updated
     title="Pothole Detection POC",
     description="Upload an image to detect potholes."
 )
 
-
 iface.launch()