Update app.py

app.py

@@ -13,23 +13,26 @@ def predict_image(image: Image.Image):
         # Step 1: Preprocess the image with padding enabled
         inputs = processor(images=image, return_tensors="pt", padding=True)
         
-        # Step 2: Run inference without gradients
+        # Step 2: Ensure the model works on single-image batch by adding a batch dimension
+        inputs = {key: val.unsqueeze(0) if len(val.shape) == 3 else val for key, val in inputs.items()}
+        
+        # Step 3: Run inference without gradients
         with torch.no_grad():
             outputs = model(**inputs)
 
-        # Step 3: Post-process the outputs
+        # Step 4: 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
+        # Step 5: 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 with status
+        # Step 6: Return the processed image with status
         return image, "Object detection complete."
 
     except Exception as e: