Update app.py

app.py

@@ -10,26 +10,26 @@ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 # Prediction function
 def predict_image(image: Image.Image):
     try:
-        # Step 1: Use DETR for object detection
-        # Ensure padding is enabled when processing single images
+        # Step 1: Preprocess the image with padding enabled
         inputs = processor(images=image, return_tensors="pt", padding=True)
-
-        # Step 2: Run inference
+        
+        # Step 2: Run inference without gradients
         with torch.no_grad():
             outputs = model(**inputs)
 
-        # Step 3: Post-process the outputs for object detection
-        target_sizes = torch.tensor([image.size[::-1]])  # Convert image size to (height, width)
+        # Step 3: Post-process the outputs
+        # Convert image size to (height, width)
+        target_sizes = torch.tensor([image.size[::-1]])
         results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.5)[0]
-        
+
         # 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")
-        
-        # Step 5: Return the processed image and description
+
+        # Step 5: Return the processed image with status
         return image, "Object detection complete."
 
     except Exception as e: