Update app.py

app.py

@@ -11,25 +11,27 @@ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 def predict_image(image: Image.Image):
     try:
         # Step 1: Use DETR for object detection
+        # Ensure padding is enabled when processing single images
         inputs = processor(images=image, return_tensors="pt", padding=True)
-        
+
+        # Step 2: Run inference
         with torch.no_grad():
             outputs = model(**inputs)
-        
-        # Post-process the outputs for object detection
-        target_sizes = torch.tensor([image.size[::-1]])  # Convert to (height, width)
+
+        # Step 3: Post-process the outputs for object detection
+        target_sizes = torch.tensor([image.size[::-1]])  # Convert image size to (height, width)
         results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.5)[0]
         
-        # Draw bounding boxes on the image
+        # Step 4: Draw bounding boxes on the image
         draw = ImageDraw.Draw(image)
         for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
             box = [round(i, 2) for i in box.tolist()]
             draw.rectangle(box, outline="red", width=2)
             draw.text((box[0], box[1]), f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}", fill="red")
         
-        # Return the image with bounding boxes and description text
+        # Step 5: Return the processed image and description
         return image, "Object detection complete."
-    
+
     except Exception as e:
         return None, f"Error: {str(e)}"