Update predict.py

predict.py

@@ -7,8 +7,6 @@ import io
 from typing import Union
 
 # --- Model Loading ---
-# This section attempts to load the specific deep learning models you are using.
-
 def load_models():
     """Loads TensorFlow and YOLO models using your specified filenames."""
     segmentation_model, yolo_detector = None, None
@@ -39,7 +37,7 @@ PIXELS_PER_CM = 50.0
 app = FastAPI(
     title="Wound Analysis API",
     description="A comprehensive API to analyze wound images using deep learning and computer vision techniques.",
-    version="9.0.0" # Version with improved visualization (Yellow Heatmap + Boundary)
+    version="9.1.0" # Version with fix for model prediction output format
 )
 
 
@@ -67,7 +65,17 @@ def segment_wound(image: np.ndarray) -> np.ndarray:
             model_input_size = segmentation_model.input.shape[1:3]
             img_resized = cv2.resize(image, (model_input_size[1], model_input_size[0]))
             img_norm = np.expand_dims(img_resized.astype(np.float32) / 255.0, axis=0)
-            pred_mask = segmentation_model.predict(img_norm, verbose=0)[0]
+            
+            prediction = segmentation_model.predict(img_norm, verbose=0)
+            
+            # --- FIX: Handle model output that is a list ---
+            # If the prediction is a list, extract the first element which is the actual numpy array.
+            if isinstance(prediction, list):
+                pred_mask = prediction[0]
+            else:
+                pred_mask = prediction
+            # --- END FIX ---
+
             pred_mask_resized = cv2.resize(pred_mask, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
             mask = (pred_mask_resized.squeeze() >= 0.5).astype(np.uint8) * 255
             if cv2.countNonZero(mask) > 0:
@@ -75,6 +83,7 @@ def segment_wound(image: np.ndarray) -> np.ndarray:
         except Exception as e:
             print(f"Model prediction failed, switching to fallback segmentation. Error: {e}")
 
+    # Fallback Method
     pixels = image.reshape((-1, 3)).astype(np.float32)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
     _, labels, centers = cv2.kmeans(pixels, 2, None, criteria, 5, cv2.KMEANS_PP_CENTERS)
@@ -115,28 +124,21 @@ def calculate_all_metrics(mask: np.ndarray, image: np.ndarray) -> dict:
     }
 
 def create_visual_overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
-    """
-    Generates a visual overlay with a Yellow/Blue/Green heatmap and a white boundary.
-    """
-    # --- 1. Create the Color Heatmap ---
+    """Generates a visual overlay with a Yellow/Blue/Green heatmap and a white boundary."""
     dist = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
     cv2.normalize(dist, dist, 0, 1.0, cv2.NORM_MINMAX)
     overlay = np.zeros_like(image)
     
-    # **CHANGE**: Use Yellow instead of Red for the most affected area for better visibility.
-    overlay[dist >= 0.66] = (0, 255, 255)      # Yellow in BGR
-    overlay[(dist >= 0.33) & (dist < 0.66)] = (255, 0, 0)  # Blue in BGR
-    overlay[(dist > 0) & (dist < 0.33)] = (0, 255, 0)      # Green in BGR
+    overlay[dist >= 0.66] = (0, 255, 255)
+    overlay[(dist >= 0.33) & (dist < 0.66)] = (255, 0, 0)
+    overlay[(dist > 0) & (dist < 0.33)] = (0, 255, 0)
     
     blended = cv2.addWeighted(image, 0.7, overlay, 0.3, 0)
     final_image = image.copy()
     final_image[mask.astype(bool)] = blended[mask.astype(bool)]
 
-    # --- 2. Draw the Boundary Contour ---
-    # Find contours from the mask to draw the boundary line.
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    # **NEW**: Draw a crisp white boundary on the final image.
-    cv2.drawContours(final_image, contours, -1, (255, 255, 255), 1) # White color, 1px thickness
+    cv2.drawContours(final_image, contours, -1, (255, 255, 255), 1)
 
     return final_image