Update predict.py

predict.py

@@ -1,57 +1,37 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException, Response
+from fastapi import FastAPI, File, UploadFile, Response, HTTPException
 import cv2
 import numpy as np
 from ultralytics import YOLO
 import tensorflow as tf
-import io
+import os
 from typing import Union
 
-# --- Configuration ---
-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.4.0" # Version updated for prediction output fix
-)
-
-# --- Model Loading ---
-def load_models():
-    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.")
-    except Exception as e:
-        print(f"Warning: Could not load segmentation model. Using fallback. Error: {e}")
+app = FastAPI()
 
-    try:
-        yolo_model = YOLO("best.pt")
-        print("YOLO model 'best.pt' loaded successfully.")
-    except Exception as e:
-        print(f"Warning: Could not load YOLO model. Using fallback. Error: {e}")
-        
-    return segmentation_model, yolo_model
+PIXELS_PER_CM = 50.0
 
-segmentation_model, yolo_model = load_models()
+# --- Model loading ---
+segmentation_model, yolo_model = None, None
+try:
+    segmentation_model = tf.keras.models.load_model("segmentation_model.h5")
+except Exception as e:
+    print(f"Segmentation model not loaded: {e}")
 
-# --- Helper Functions ---
+try:
+    yolo_model = YOLO("best.pt")
+except Exception as e:
+    print(f"YOLO model not loaded: {e}")
 
+# --- Helpers ---
 def preprocess_image(image: np.ndarray) -> np.ndarray:
-    img_denoised = cv2.medianBlur(image, 3)
-    lab = cv2.cvtColor(img_denoised, cv2.COLOR_BGR2LAB)
-    l_channel, a_channel, b_channel = cv2.split(lab)
+    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
+    l, a, b = 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)
+    cl = clahe.apply(l)
+    limg = cv2.merge((cl, a, b))
+    return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
 
-def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
-    if not yolo_model: return None
+def detect_with_yolo(image: np.ndarray) -> Union[tuple, None]:
     try:
         results = yolo_model.predict(image, verbose=False)
         if results and results[0].boxes:
@@ -59,122 +39,94 @@ def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
             coords = best_box.xyxy[0].cpu().numpy()
             return tuple(map(int, coords))
     except Exception as e:
-        print(f"YOLO prediction failed: {e}")
-    return None
-
-def segment_wound_with_model(image: np.ndarray) -> Union[np.ndarray, None]:
-    if not segmentation_model:
-        return None
-    try:
-        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: 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}")
+        print(f"YOLO error: {e}")
     return None
 
-
-def segment_wound_with_fallback(image: np.ndarray) -> np.ndarray:
-    pixels = image.reshape((-1, 3)).astype(np.float32)
+def fallback_segmentation(image: np.ndarray) -> np.ndarray:
+    Z = 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)
-        refined_mask = np.zeros_like(mask)
-        cv2.drawContours(refined_mask, [largest_contour], -1, 255, cv2.FILLED)
-        return refined_mask
-    return mask
+    _, label, center = cv2.kmeans(Z, 2, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
+    label = label.reshape(image.shape[:2])
+    unique_vals = np.unique(label)
+    if len(unique_vals) > 1:
+        wound_label = np.argmax([np.sum(label == val) for val in unique_vals])
+    else:
+        wound_label = unique_vals[0]
+    return (label == wound_label).astype(np.uint8) * 255
+
+def segment(image: np.ndarray) -> np.ndarray:
+    if segmentation_model is not None:
+        try:
+            input_shape = segmentation_model.input.shape[1:3]
+            resized = cv2.resize(image, (input_shape[1], input_shape[0]))
+            norm = np.expand_dims(resized / 255.0, axis=0)
+            prediction = segmentation_model.predict(norm)
+            if isinstance(prediction, list):
+                prediction = prediction[0]
+            mask = (prediction[0].squeeze() >= 0.5).astype(np.uint8) * 255
+            return cv2.resize(mask, (image.shape[1], image.shape[0]))
+        except Exception as e:
+            print(f"Segmentation model failed: {e}")
+    return fallback_segmentation(image)
 
 def calculate_metrics(mask: np.ndarray, image: np.ndarray) -> dict:
-    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}
-
-    area_cm2 = wound_pixels / (PIXELS_PER_CM ** 2)
+    area_px = cv2.countNonZero(mask)
+    area_cm2 = area_px / (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)
-    gray_img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    
-    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}
+    if not contours:
+        return {"length": 0, "breadth": 0, "area": 0, "depth": 0, "moisture": 0}
+    c = max(contours, key=cv2.contourArea)
+    rect = cv2.minAreaRect(c)
+    length, breadth = max(rect[1]) / PIXELS_PER_CM, min(rect[1]) / PIXELS_PER_CM
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    texture_std = np.std(gray[mask.astype(bool)])
+    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
+    mean_a = np.mean(lab[:, :, 1][mask.astype(bool)])
+    depth = mean_a - 128
+    moisture = max(0, 100 * (1.0 - texture_std / 127.0))
+    return {
+        "area": area_cm2,
+        "length": length,
+        "breadth": breadth,
+        "depth": depth,
+        "moisture": moisture,
+        "contour": c
+    }
 
-def create_visual_overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
-    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)]
-    return annotated_img
+def annotate(image: np.ndarray, mask: np.ndarray, contour) -> np.ndarray:
+    poly_image = image.copy()
+    if contour is not None:
+        cv2.drawContours(poly_image, [contour], -1, (0, 255, 0), 2)
+    return poly_image
 
-# --- Main API Endpoint ---
+# --- API ---
 @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)
-    bbox = detect_wound_region_yolo(processed_image)
+    arr = np.frombuffer(contents, np.uint8)
+    image = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    if image is None:
+        raise HTTPException(status_code=400, detail="Invalid image")
+
+    image = preprocess_image(image)
+    bbox = detect_with_yolo(image)
+    cropped = image[bbox[1]:bbox[3], bbox[0]:bbox[2]] if bbox else image
+    mask = segment(cropped)
+
+    metrics = calculate_metrics(mask, cropped)
+    full_mask = np.zeros(image.shape[:2], dtype=np.uint8)
     if bbox:
-        xmin, ymin, xmax, ymax = bbox
-        cropped_image = processed_image[ymin:ymax, xmin:xmax]
-    else:
-        cropped_image = processed_image
-
-    mask = segment_wound_with_model(cropped_image)
-    if mask is None:
-        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
+        full_mask[bbox[1]:bbox[3], bbox[0]:bbox[2]] = mask
     else:
         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")
 
-    headers = {
-        "Wound-Area-cm2": f"{metrics['area_cm2']:.2f}",
-        "Wound-Length-cm": f"{metrics['length_cm']:.2f}",
-        "Wound-Breadth-cm": f"{metrics['breadth_cm']:.2f}",
-        "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)
+    final_image = annotate(image, full_mask, metrics['contour'])
+    _, buf = cv2.imencode(".png", final_image)
+
+    response = Response(content=buf.tobytes(), media_type="image/png")
+    response.headers['X-Length-Cm'] = str(metrics['length'])
+    response.headers['X-Breadth-Cm'] = str(metrics['breadth'])
+    response.headers['X-Depth-Cm'] = str(metrics['depth'])
+    response.headers['X-Area-Cm2'] = str(metrics['area'])
+    response.headers['X-Moisture'] = str(metrics['moisture'])
+    return response