Update app.py

app.py

@@ -30,26 +30,40 @@ id_to_names = {
 }
 
 def recognize_image(input_img, conf_threshold, iou_threshold):
-    det_res = model.predict(
-        input_img,
-        imgsz=1024,
-        conf=conf_threshold,
-        device=device,
-    )[0]
-    boxes = det_res.__dict__['boxes'].xyxy
-    classes = det_res.__dict__['boxes'].cls
-    scores = det_res.__dict__['boxes'].conf
+    try:
+        det_res = model.predict(
+            input_img,
+            imgsz=1024,
+            conf=conf_threshold,
+            device=device,
+        )[0]
+        boxes = det_res.__dict__['boxes'].xyxy
+        classes = det_res.__dict__['boxes'].cls
+        scores = det_res.__dict__['boxes'].conf
 
-    indices = torchvision.ops.nms(boxes=torch.Tensor(boxes), scores=torch.Tensor(scores),iou_threshold=iou_threshold)
-    boxes, scores, classes = boxes[indices], scores[indices], classes[indices]
-    if len(boxes.shape) == 1:
-        boxes = np.expand_dims(boxes, 0)
-        scores = np.expand_dims(scores, 0)
-        classes = np.expand_dims(classes, 0)
+        indices = torchvision.ops.nms(
+            boxes=torch.Tensor(boxes), 
+            scores=torch.Tensor(scores),
+            iou_threshold=iou_threshold
+        )
+
+        boxes, scores, classes = boxes[indices], scores[indices], classes[indices]
+
+        if len(boxes.shape) == 1:
+            boxes = np.expand_dims(boxes, 0)
+            scores = np.expand_dims(scores, 0)
+            classes = np.expand_dims(classes, 0)
+
+        output = visualize_bbox(input_img, boxes, classes, scores, id_to_names)
+        if not isinstance(output, (np.ndarray, Image.Image)):
+            raise ValueError("Output is not a valid image")
+
+        return output
+    except Exception as e:
+        print(f"[ERROR] recognize_image failed: {e}")
+        # Return blank image or raise if debugging
+        return np.zeros((512, 512, 3), dtype=np.uint8)
 
-    vis_result = visualize_bbox(input_img, boxes, classes, scores, id_to_names)
-    return vis_result
-    
 def gradio_reset():
     return gr.update(value=None), gr.update(value=None)