Update predict.py

predict.py

@@ -63,34 +63,29 @@ def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
     return None
 
 def segment_wound_with_model(image: np.ndarray) -> Union[np.ndarray, None]:
-    """Segments the wound using the U-Net model."""
     if not segmentation_model:
         return None
     try:
-        input_shape = segmentation_model.input.shape[1:3]
+        input_shape = segmentation_model.input_shape[1:3]
         img_resized = cv2.resize(image, (input_shape[1], input_shape[0]))
         img_norm = np.expand_dims(img_resized.astype(np.float32) / 255.0, axis=0)
-        
+
         prediction = segmentation_model.predict(img_norm, verbose=0)
-        
-        # --- FIX APPLIED HERE ---
-        # Handle cases where the model returns a list of outputs
-        if isinstance(prediction, list):
-            pred_mask = prediction[0]
-        else:
-            pred_mask = prediction
-        
-        # The output of predict() is batched, so get the first item.
-        pred_mask = pred_mask[0]
-        # --- END OF FIX ---
 
+        # FIX: Handle nested list output or Tensor
+        while isinstance(prediction, list):
+            prediction = prediction[0]
+        if isinstance(prediction, tf.Tensor):
+            prediction = prediction.numpy()
+
+        pred_mask = prediction[0]
         pred_mask_resized = cv2.resize(pred_mask, (image.shape[1], image.shape[0]))
-        
         return (pred_mask_resized.squeeze() >= 0.5).astype(np.uint8) * 255
     except Exception as e:
         print(f"Segmentation model prediction failed: {e}")
     return None
 
+
 def segment_wound_with_fallback(image: np.ndarray) -> np.ndarray:
     pixels = image.reshape((-1, 3)).astype(np.float32)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)