Update predict.py

predict.py

@@ -24,50 +24,48 @@ def load_models():
 segmentation_model, yolo_model = load_models()
 
 PIXELS_PER_CM = 50.0
-app = FastAPI(title="Wound Analyzer", version="10.0")
+app = FastAPI(title="Wound Analyzer", version="10.1") # Version with depth calculation fix
 
 # --- Preprocessing ---
 def preprocess_image(image: np.ndarray) -> np.ndarray:
+    """Enhances image for better segmentation."""
     img_denoised = cv2.medianBlur(image, 3)
     lab = cv2.cvtColor(img_denoised, cv2.COLOR_BGR2LAB)
     l, a, b = cv2.split(lab)
     clahe = cv2.createCLAHE(2.0, (8, 8))
-    l = clahe.apply(l)
-    lab = cv2.merge((l, a, b))
-    result = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
+    l_clahe = clahe.apply(l)
+    lab_clahe = cv2.merge((l_clahe, a, b))
+    result = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2BGR)
     gamma = 1.2
     return np.clip((result / 255.0) ** gamma * 255, 0, 255).astype(np.uint8)
 
 # --- Segmentation ---
 def segment_wound(image: np.ndarray) -> np.ndarray:
+    """Segments wound from a preprocessed image, with a fallback."""
     try:
         if segmentation_model:
             input_size = segmentation_model.input_shape[1:3]
             resized = cv2.resize(image, (input_size[1], input_size[0]))
             norm = np.expand_dims(resized / 255.0, axis=0)
-
             prediction = segmentation_model.predict(norm, verbose=0)
-            if isinstance(prediction, list):  # <-- Fix
+            if isinstance(prediction, list):
                 prediction = prediction[0]
-            prediction = prediction[0]  # remove batch dim
-
+            prediction = prediction[0]
             mask = cv2.resize(prediction.squeeze(), (image.shape[1], image.shape[0]))
             return (mask >= 0.5).astype(np.uint8) * 255
     except Exception as e:
         print(f"⚠️ Model prediction failed: {e}")
 
-    # Fallback segmentation
     Z = image.reshape((-1, 3)).astype(np.float32)
-    _, labels, centers = cv2.kmeans(Z, 2, None,
-                                    (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0),
-                                    5, cv2.KMEANS_PP_CENTERS)
+    _, labels, centers = cv2.kmeans(Z, 2, None, (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0), 5, cv2.KMEANS_PP_CENTERS)
     centers_lab = cv2.cvtColor(centers.reshape(1, -1, 3).astype(np.uint8), cv2.COLOR_BGR2LAB)[0]
     wound_idx = np.argmax(centers_lab[:, 1])
     mask = (labels.reshape(image.shape[:2]) == wound_idx).astype(np.uint8) * 255
     return mask
 
 # --- Metrics ---
-def calculate_metrics(mask: np.ndarray, image: np.ndarray):
+def calculate_metrics(mask: np.ndarray, original_image: np.ndarray):
+    """Calculates metrics using the mask and the ORIGINAL image for color accuracy."""
     area_px = cv2.countNonZero(mask)
     if area_px == 0:
         return dict(area_cm2=0.0, length_cm=0.0, breadth_cm=0.0, depth_cm=0.0, moisture=0.0)
@@ -79,26 +77,24 @@ def calculate_metrics(mask: np.ndarray, image: np.ndarray):
     length_cm, breadth_cm = max(w, h) / PIXELS_PER_CM, min(w, h) / PIXELS_PER_CM
 
     mask_bool = mask.astype(bool)
-    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    # **FIX**: Use the original, unaltered image for color-based metrics.
+    lab = cv2.cvtColor(original_image, cv2.COLOR_BGR2LAB)
+    gray = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)
+    
     depth = np.mean(lab[:, :, 1][mask_bool]) - 128.0
     moisture = max(0.0, 100.0 * (1 - np.std(gray[mask_bool]) / 127.0))
 
-    return dict(
-        area_cm2=round(area_cm2, 2),
-        length_cm=round(length_cm, 2),
-        breadth_cm=round(breadth_cm, 2),
-        depth_cm=round(depth, 1),
-        moisture=round(moisture, 0)
-    )
+    return dict(area_cm2=round(area_cm2, 2), length_cm=round(length_cm, 2), breadth_cm=round(breadth_cm, 2), depth_cm=round(depth, 1), moisture=round(moisture, 0))
 
 # --- Overlay ---
 def draw_overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    """Draws heatmap and boundary on the image."""
+    # **CHANGE**: Use Yellow for high-intensity areas for better visibility
     dist = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
     cv2.normalize(dist, dist, 0, 1.0, cv2.NORM_MINMAX)
     heatmap = np.zeros_like(image)
 
-    heatmap[dist >= 0.66] = (0, 0, 255)        # Red - Most Affected
+    heatmap[dist >= 0.66] = (0, 255, 255)      # Yellow - Most Affected
     heatmap[(dist >= 0.33) & (dist < 0.66)] = (255, 0, 0)  # Blue - Moderate
     heatmap[(dist > 0) & (dist < 0.33)] = (0, 255, 0)      # Green - Least
 
@@ -107,42 +103,52 @@ def draw_overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
     annotated[mask.astype(bool)] = blended[mask.astype(bool)]
 
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    cv2.drawContours(annotated, contours, -1, (255, 255, 255), 2)  # White outline
+    cv2.drawContours(annotated, contours, -1, (255, 255, 255), 2)
     return annotated
 
 # --- API Endpoint ---
 @app.post("/analyze_wound")
 async def analyze(file: UploadFile = File(...)):
-    image = cv2.imdecode(np.frombuffer(await file.read(), np.uint8), cv2.IMREAD_COLOR)
-    if image is None:
+    contents = await file.read()
+    original_image = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
+    if original_image is None:
         raise HTTPException(status_code=400, detail="Invalid image file.")
 
-    image = preprocess_image(image)
-    crop = image.copy()
+    # 1. Create a preprocessed version for segmentation
+    preprocessed_image = preprocess_image(original_image)
+    
+    # 2. Define Regions of Interest (ROI) for both original and preprocessed images
+    original_roi = original_image.copy()
+    preprocessed_roi = preprocessed_image.copy()
 
     if yolo_model:
         try:
-            results = yolo_model.predict(image, verbose=False)
+            results = yolo_model.predict(preprocessed_image, verbose=False)
             if results and results[0].boxes:
                 coords = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
                 x1, y1, x2, y2 = coords
-                crop = image[y1:y2, x1:x2]
+                # Crop both the original and preprocessed images to the same ROI
+                original_roi = original_image[y1:y2, x1:x2]
+                preprocessed_roi = preprocessed_image[y1:y2, x1:x2]
         except Exception as e:
             print(f"⚠️ YOLO detection failed: {e}")
 
-    mask = segment_wound(crop)
-    metrics = calculate_metrics(mask, crop)
-    annotated = draw_overlay(crop, mask)
+    # 3. Get mask from the preprocessed ROI
+    mask = segment_wound(preprocessed_roi)
+    
+    # 4. Calculate metrics using the ORIGINAL ROI for color accuracy
+    metrics = calculate_metrics(mask, original_roi)
+    
+    # 5. Draw overlay on the ORIGINAL ROI for correct visualization
+    annotated_image = draw_overlay(original_roi, mask)
 
-    success, out = cv2.imencode(".png", annotated)
+    success, out_bytes = cv2.imencode(".png", annotated_image)
     if not success:
         raise HTTPException(status_code=500, detail="Failed to encode image.")
 
     headers = {
-        "X-Length-Cm": str(metrics["length_cm"]),
-        "X-Breadth-Cm": str(metrics["breadth_cm"]),
-        "X-Depth-Cm": str(metrics["depth_cm"]),
-        "X-Area-Cm2": str(metrics["area_cm2"]),
+        "X-Length-Cm": str(metrics["length_cm"]), "X-Breadth-Cm": str(metrics["breadth_cm"]),
+        "X-Depth-Cm": str(metrics["depth_cm"]), "X-Area-Cm2": str(metrics["area_cm2"]),
         "X-Moisture": str(metrics["moisture"]),
     }
-    return Response(content=out.tobytes(), media_type="image/png", headers=headers)
+    return Response(content=out_bytes.tobytes(), media_type="image/png", headers=headers)