Update Space

app.py

@@ -63,12 +63,11 @@ class Model:
         outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
         
         # Process outputs
-        # Assuming outputs are in the format: [boxes, labels, scores]
-        # Adjust based on your actual model's output format
+        # Extract the first (and only) batch element
         if len(outputs) >= 3:
-            boxes = outputs[0]  # shape: [num_detections, 4]
-            labels = outputs[1].astype(int)  # shape: [num_detections]
-            scores = outputs[2]  # shape: [num_detections]
+            boxes = outputs[0][0]    # shape: [num_detections, 4]
+            labels = outputs[1][0].astype(int)  # shape: [num_detections]
+            scores = outputs[2][0]   # shape: [num_detections]
             return boxes, labels, scores
         else:
             raise ValueError("Unexpected number of outputs from the model.")
@@ -87,7 +86,9 @@ def draw_boxes(image, boxes, labels, scores, threshold=0.5):
     for box, label, score in zip(boxes, labels, scores):
         if score < threshold:
             continue
-        # Assuming box format is [xmin, ymin, xmax, ymax] normalized [0,1]
+        if len(box) != 4:
+            print(f"Invalid box format: {box}")
+            continue
         xmin, ymin, xmax, ymax = box
         width, height = image.size
         xmin = int(xmin * width)
@@ -108,10 +109,14 @@ def draw_boxes(image, boxes, labels, scores, threshold=0.5):
 
 # Prediction function for Gradio
 def predict_image(input_image):
-    boxes, labels, scores = model.predict(input_image)
-    output_image = input_image.copy()
-    output_image = draw_boxes(output_image, boxes, labels, scores, threshold=0.5)
-    return output_image
+    try:
+        boxes, labels, scores = model.predict(input_image)
+        output_image = input_image.copy()
+        output_image = draw_boxes(output_image, boxes, labels, scores, threshold=0.5)
+        return output_image
+    except Exception as e:
+        print(f"Error during prediction: {e}")
+        return input_image  # Return the original image if prediction fails
 
 # Define Gradio Interface
 def get_example_images():
@@ -136,7 +141,7 @@ iface = gr.Interface(
     examples=example_images,
     title=title,
     description=description,
-    allow_flagging="never"
+    flagging_mode="never"  # Updated from allow_flagging="never"
 )
 
 # Launch the interface