Update predict.py

predict.py

@@ -4,25 +4,24 @@ import numpy as np
 from ultralytics import YOLO
 import tensorflow as tf
 import io
-from typing import Union # <--- IMPORT THIS
+from typing import Union
 
 # --- Configuration ---
-PIXELS_PER_CM = 50.0  # Calibrate this value: pixels per centimeter
+PIXELS_PER_CM = 50.0
 
 # --- App Initialization ---
 app = FastAPI(
     title="Wound Analysis API",
     description="An API to analyze wound images and return an annotated image with data in headers.",
-    version="3.3.0" # Version updated for Python 3.9 compatibility
+    version="3.4.0" # Version updated for prediction output fix
 )
 
 # --- Model Loading ---
 def load_models():
-    """Loads ML models safely, allowing the app to run even if a model is missing."""
     segmentation_model, yolo_model = None, None
     try:
-        segmentation_model = tf.keras.models.load_model("segmentation_model.h5")
-        print("Segmentation model 'segmentation_model.h5' loaded successfully.")
+        segmentation_model = tf.keras.models.load_model("segmentation model.h5")
+        print("Segmentation model 'segmentation model.h5' loaded successfully.")
     except Exception as e:
         print(f"Warning: Could not load segmentation model. Using fallback. Error: {e}")
 
@@ -39,26 +38,20 @@ segmentation_model, yolo_model = load_models()
 # --- Helper Functions ---
 
 def preprocess_image(image: np.ndarray) -> np.ndarray:
-    """Applies denoising, CLAHE, and gamma correction to the image."""
     img_denoised = cv2.medianBlur(image, 3)
-    
     lab = cv2.cvtColor(img_denoised, cv2.COLOR_BGR2LAB)
     l_channel, a_channel, b_channel = cv2.split(lab)
     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
     l_clahe = clahe.apply(l_channel)
     lab_clahe = cv2.merge([l_clahe, a_channel, b_channel])
     img_clahe = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2BGR)
-    
     gamma = 1.2
     img_float = img_clahe.astype(np.float32) / 255.0
     img_gamma = np.power(img_float, gamma)
     return (img_gamma * 255).astype(np.uint8)
 
-# --- FIX APPLIED HERE ---
 def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
-    """Detects the primary wound region using YOLO and returns its bounding box."""
-    if not yolo_model:
-        return None
+    if not yolo_model: return None
     try:
         results = yolo_model.predict(image, verbose=False)
         if results and results[0].boxes:
@@ -69,7 +62,6 @@ def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
         print(f"YOLO prediction failed: {e}")
     return None
 
-# --- FIX APPLIED HERE ---
 def segment_wound_with_model(image: np.ndarray) -> Union[np.ndarray, None]:
     """Segments the wound using the U-Net model."""
     if not segmentation_model:
@@ -79,7 +71,19 @@ def segment_wound_with_model(image: np.ndarray) -> Union[np.ndarray, None]:
         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)
         
-        pred_mask = segmentation_model.predict(img_norm, verbose=0)[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 ---
+
         pred_mask_resized = cv2.resize(pred_mask, (image.shape[1], image.shape[0]))
         
         return (pred_mask_resized.squeeze() >= 0.5).astype(np.uint8) * 255
@@ -88,16 +92,12 @@ def segment_wound_with_model(image: np.ndarray) -> Union[np.ndarray, None]:
     return None
 
 def segment_wound_with_fallback(image: np.ndarray) -> np.ndarray:
-    """Segments the wound using K-means color clustering as a fallback."""
     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)
-    
     centers_lab = cv2.cvtColor(centers.reshape(1, -1, 3).astype(np.uint8), cv2.COLOR_BGR2LAB)[0]
     wound_cluster_idx = np.argmax(centers_lab[:, 1])
-    
     mask = (labels.reshape(image.shape[:2]) == wound_cluster_idx).astype(np.uint8) * 255
-    
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if contours:
         largest_contour = max(contours, key=cv2.contourArea)
@@ -107,19 +107,13 @@ def segment_wound_with_fallback(image: np.ndarray) -> np.ndarray:
     return mask
 
 def calculate_metrics(mask: np.ndarray, image: np.ndarray) -> dict:
-    """Calculates area, dimensions, depth, and moisture from the final mask and image."""
     wound_pixels = cv2.countNonZero(mask)
     if wound_pixels == 0:
-        return {
-            "area_cm2": 0.0, "length_cm": 0.0, "breadth_cm": 0.0,
-            "depth_score": 0.0, "moisture_score": 0.0
-        }
+        return {"area_cm2": 0.0, "length_cm": 0.0, "breadth_cm": 0.0, "depth_score": 0.0, "moisture_score": 0.0}
 
     area_cm2 = wound_pixels / (PIXELS_PER_CM ** 2)
-    
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     largest_contour = max(contours, key=cv2.contourArea)
-    
     (_, (width, height), _) = cv2.minAreaRect(largest_contour)
     length_cm = max(width, height) / PIXELS_PER_CM
     breadth_cm = min(width, height) / PIXELS_PER_CM
@@ -130,28 +124,19 @@ def calculate_metrics(mask: np.ndarray, image: np.ndarray) -> dict:
     
     mean_a = np.mean(lab_img[:, :, 1][mask_bool])
     depth_score = mean_a - 128.0
-    
     texture_std = np.std(gray_img[mask_bool])
     moisture_score = max(0.0, 100.0 * (1.0 - texture_std / 127.0))
 
-    return {
-        "area_cm2": area_cm2, "length_cm": length_cm, "breadth_cm": breadth_cm,
-        "depth_score": depth_score, "moisture_score": moisture_score
-    }
+    return {"area_cm2": area_cm2, "length_cm": length_cm, "breadth_cm": breadth_cm, "depth_score": depth_score, "moisture_score": moisture_score}
 
 def create_visual_overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
-    """Generates a colored overlay on the image based on the wound mask."""
-    if cv2.countNonZero(mask) == 0:
-        return image
-
+    if cv2.countNonZero(mask) == 0: return image
     dist = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
     cv2.normalize(dist, dist, 0, 1.0, cv2.NORM_MINMAX)
-    
     overlay = np.zeros_like(image)
     overlay[dist >= 0.66] = (0, 0, 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)
     annotated_img = image.copy()
     annotated_img[mask.astype(bool)] = blended[mask.astype(bool)]
@@ -167,7 +152,6 @@ async def analyze_wound(file: UploadFile = File(...)):
         raise HTTPException(status_code=400, detail="Invalid or corrupt image file")
 
     processed_image = preprocess_image(original_image)
-
     bbox = detect_wound_region_yolo(processed_image)
     if bbox:
         xmin, ymin, xmax, ymax = bbox
@@ -180,7 +164,6 @@ async def analyze_wound(file: UploadFile = File(...)):
         mask = segment_wound_with_fallback(cropped_image)
 
     metrics = calculate_metrics(mask, cropped_image)
-
     full_mask = np.zeros(original_image.shape[:2], dtype=np.uint8)
     if bbox:
         full_mask[ymin:ymax, xmin:xmax] = mask
@@ -188,7 +171,6 @@ async def analyze_wound(file: UploadFile = File(...)):
         full_mask = mask
     
     annotated_image = create_visual_overlay(original_image, full_mask)
-
     success, png_data = cv2.imencode(".png", annotated_image)
     if not success:
         raise HTTPException(status_code=500, detail="Failed to encode output image")
@@ -200,4 +182,4 @@ async def analyze_wound(file: UploadFile = File(...)):
         "Wound-Depth-Score": f"{metrics['depth_score']:.1f}",
         "Wound-Moisture-Score": f"{metrics['moisture_score']:.0f}"
     }
-    return Response(content=png_data.tobytes(), media_type="image/png", headers=headers)
+    return Response(content=png_data.tobytes(), media_type="image/png", headers=headers)