Update predict.py

predict.py

@@ -1,194 +1,148 @@
-# main.py
-
 from fastapi import FastAPI, File, UploadFile, HTTPException, Response
 import cv2
 import numpy as np
 import io
 from typing import Union
 
-# --- Model Loading ---
+# --- Load Models ---
 def load_models():
-    """Loads TensorFlow and YOLO models using your specified filenames."""
-    segmentation_model, yolo_detector = None, None
-    
-    try:
-        from ultralytics import YOLO
-        yolo_detector = YOLO("best.pt")
-        print("YOLOv8 detection model 'best.pt' loaded successfully.")
-    except (ImportError, IOError, Exception) as e:
-        print(f"Warning: YOLOv8 model not loaded. Using contour-based region detection. Error: {e}")
-
+    segmentation_model, yolo_model = None, None
     try:
         import tensorflow as tf
         segmentation_model = tf.keras.models.load_model("segmentation_model.h5")
-        print("TensorFlow segmentation model 'segmentation_model.h5' loaded successfully.")
-    except (ImportError, IOError, Exception) as e:
-        print(f"Warning: TensorFlow segmentation model not loaded. Using OpenCV fallback. Error: {e}")
-        
-    return segmentation_model, yolo_detector
+        print("✅ Segmentation model loaded.")
+    except Exception as e:
+        print(f"⚠️ Failed to load segmentation model: {e}")
+    try:
+        from ultralytics import YOLO
+        yolo_model = YOLO("best.pt")
+        print("✅ YOLO model loaded.")
+    except Exception as e:
+        print(f"⚠️ Failed to load YOLO model: {e}")
+    return segmentation_model, yolo_model
 
 segmentation_model, yolo_model = load_models()
 
-
-# --- Configuration ---
 PIXELS_PER_CM = 50.0
+app = FastAPI(title="Wound Analyzer", version="10.0")
 
-# --- App Initialization ---
-app = FastAPI(
-    title="Wound Analysis API",
-    description="A comprehensive API to analyze wound images using deep learning and computer vision techniques.",
-    version="9.1.0" # Version with fix for model prediction output format
-)
-
-
-# --- Helper Functions ---
-
+# --- Preprocessing ---
 def preprocess_image(image: np.ndarray) -> np.ndarray:
-    """Applies the full preprocessing pipeline: Denoise -> CLAHE -> Gamma Correction."""
     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)
+    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)
     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)
+    return np.clip((result / 255.0) ** gamma * 255, 0, 255).astype(np.uint8)
 
+# --- Segmentation ---
 def segment_wound(image: np.ndarray) -> np.ndarray:
-    """Segments the wound using the TF model if available, otherwise falls back to color clustering."""
-    if segmentation_model:
-        try:
-            orig_h, orig_w = image.shape[:2]
-            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)
-            
-            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:
-                return mask
-        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)
+    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
+                prediction = prediction[0]
+            prediction = prediction[0]  # remove batch dim
+
+            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)
     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)
-        refined_mask = np.zeros_like(mask)
-        cv2.drawContours(refined_mask, [largest_contour], -1, 255, cv2.FILLED)
-        return refined_mask
+    wound_idx = np.argmax(centers_lab[:, 1])
+    mask = (labels.reshape(image.shape[:2]) == wound_idx).astype(np.uint8) * 255
     return mask
 
-def calculate_all_metrics(mask: np.ndarray, image: np.ndarray) -> dict:
-    """Computes all specified wound metrics from the mask and original image."""
-    wound_pixels = cv2.countNonZero(mask)
-    if wound_pixels == 0:
-        return {"area_cm2": 0.0, "length_cm": 0.0, "breadth_cm": 0.0, "depth_cm": 0.0, "moisture": 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
-    mask_bool = mask.astype(bool)
-    lab_img = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
-    mean_a = np.mean(lab_img[:, :, 1][mask_bool])
-    depth_score = mean_a - 128.0
-    gray_img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    texture_std = np.std(gray_img[mask_bool])
-    moisture_score = max(0.0, 100.0 * (1.0 - texture_std / 127.0))
+# --- Metrics ---
+def calculate_metrics(mask: np.ndarray, image: np.ndarray):
+    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)
     
-    return {
-        "area_cm2": f"{area_cm2:.2f}", "length_cm": f"{length_cm:.2f}", "breadth_cm": f"{breadth_cm:.2f}",
-        "depth_cm": f"{depth_score:.1f}", "moisture": f"{moisture_score:.0f}"
-    }
+    area_cm2 = area_px / (PIXELS_PER_CM ** 2)
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    rect = cv2.minAreaRect(max(contours, key=cv2.contourArea))
+    (w, h) = rect[1]
+    length_cm, breadth_cm = max(w, h) / PIXELS_PER_CM, min(w, h) / PIXELS_PER_CM
 
-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."""
+    mask_bool = mask.astype(bool)
+    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
+    gray = cv2.cvtColor(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)
+    )
+
+# --- Overlay ---
+def draw_overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
     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, 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)]
+    heatmap = np.zeros_like(image)
 
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    cv2.drawContours(final_image, contours, -1, (255, 255, 255), 1)
+    heatmap[dist >= 0.66] = (0, 0, 255)        # Red - 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
+
+    blended = cv2.addWeighted(image, 0.7, heatmap, 0.3, 0)
+    annotated = image.copy()
+    annotated[mask.astype(bool)] = blended[mask.astype(bool)]
 
-    return final_image
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    cv2.drawContours(annotated, contours, -1, (255, 255, 255), 2)  # White outline
+    return annotated
 
-# --- Main API Endpoint ---
+# --- API Endpoint ---
 @app.post("/analyze_wound")
-async def analyze_wound(file: UploadFile = File(...)):
-    contents = await file.read()
-    image_array = np.frombuffer(contents, np.uint8)
-    original_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
-    if original_image is None:
-        raise HTTPException(status_code=400, detail="Invalid or corrupt image file.")
-
-    processed_image = preprocess_image(original_image)
-    
-    roi_image = processed_image
-    original_roi = original_image
+async def analyze(file: UploadFile = File(...)):
+    image = cv2.imdecode(np.frombuffer(await file.read(), np.uint8), cv2.IMREAD_COLOR)
+    if image is None:
+        raise HTTPException(status_code=400, detail="Invalid image file.")
+
+    image = preprocess_image(image)
+    crop = image.copy()
+
     if yolo_model:
         try:
-            results = yolo_model.predict(processed_image, verbose=False)
-            if results and results[0].boxes and len(results[0].boxes) > 0:
-                best_box = sorted(results[0].boxes, key=lambda b: b.conf[0], reverse=True)[0]
-                coords = best_box.xyxy[0].cpu().numpy()
-                x1, y1, x2, y2 = map(int, coords)
-                roi_image = processed_image[y1:y2, x1:x2]
-                original_roi = original_image[y1:y2, x1:x2]
+            results = yolo_model.predict(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]
         except Exception as e:
-            print(f"YOLO prediction failed, analyzing full image. Error: {e}")
+            print(f"⚠️ YOLO detection failed: {e}")
 
-    wound_mask = segment_wound(roi_image)
-    if cv2.countNonZero(wound_mask) == 0:
-        _, png_data = cv2.imencode(".png", original_image)
-        headers = {
-            'X-Length-Cm': '0.0', 'X-Breadth-Cm': '0.0', 'X-Depth-Cm': '0.0',
-            'X-Area-Cm2': '0.0', 'X-Moisture': '0.0'
-        }
-        return Response(content=png_data.tobytes(), media_type="image/png", headers=headers)
+    mask = segment_wound(crop)
+    metrics = calculate_metrics(mask, crop)
+    annotated = draw_overlay(crop, mask)
 
-    metrics = calculate_all_metrics(wound_mask, original_roi)
-    annotated_image = create_visual_overlay(original_roi, wound_mask)
-
-    success, png_data = cv2.imencode(".png", annotated_image)
+    success, out = cv2.imencode(".png", annotated)
     if not success:
-        raise HTTPException(status_code=500, detail="Failed to encode output image.")
+        raise HTTPException(status_code=500, detail="Failed to encode image.")
 
     headers = {
-        'X-Length-Cm': metrics['length_cm'],
-        'X-Breadth-Cm': metrics['breadth_cm'],
-        'X-Depth-Cm': metrics['depth_cm'],
-        'X-Area-Cm2': metrics['area_cm2'],
-        'X-Moisture': metrics['moisture']
+        "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=png_data.tobytes(), media_type="image/png", headers=headers)
+    return Response(content=out.tobytes(), media_type="image/png", headers=headers)