Update predict.py

predict.py

@@ -1,132 +1,221 @@
-from fastapi import FastAPI, File, UploadFile, Response, HTTPException
+from fastapi import FastAPI, File, UploadFile, HTTPException, Response
 import cv2
 import numpy as np
 from ultralytics import YOLO
 import tensorflow as tf
-import os
+import io
 from typing import Union
 
-app = FastAPI()
-
+# --- Configuration ---
 PIXELS_PER_CM = 50.0
 
-# --- 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}")
+# --- App Initialization ---
+app = FastAPI(
+    title="Wound Analysis API",
+    description="An API to analyze wound images, zoom to the wound, and return an annotated image with data in headers.",
+    version="3.5.0" # Version updated for polygon and zoom features
+)
 
-try:
-    yolo_model = YOLO("best.pt")
-except Exception as e:
-    print(f"YOLO model not loaded: {e}")
+# --- Model Loading ---
+def load_models():
+    """Loads the segmentation and YOLO models, handling potential errors."""
+    segmentation_model, yolo_model = None, None
+    try:
+        # Load your trained segmentation model
+        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}")
+
+    try:
+        # Load your trained YOLO model for wound detection
+        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
+
+segmentation_model, yolo_model = load_models()
+
+# --- Helper Functions ---
 
-# --- Helpers ---
 def preprocess_image(image: np.ndarray) -> np.ndarray:
-    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
-    l, a, b = cv2.split(lab)
+    """Applies a series of preprocessing steps to enhance the image for analysis."""
+    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))
-    cl = clahe.apply(l)
-    limg = cv2.merge((cl, a, b))
-    return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
+    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)
 
-def detect_with_yolo(image: np.ndarray) -> Union[tuple, None]:
+def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
+    """Detects the wound bounding box using the YOLO model."""
+    if not yolo_model: return None
     try:
         results = yolo_model.predict(image, verbose=False)
-        if results and results[0].boxes:
+        if results and results[0].boxes and len(results[0].boxes) > 0:
+            # Get the box with the highest confidence
             best_box = sorted(results[0].boxes, key=lambda b: b.conf[0], reverse=True)[0]
             coords = best_box.xyxy[0].cpu().numpy()
             return tuple(map(int, coords))
     except Exception as e:
-        print(f"YOLO error: {e}")
+        print(f"YOLO prediction failed: {e}")
+    return None
+
+def segment_wound_with_model(image: np.ndarray) -> Union[np.ndarray, None]:
+    """Segments the wound from the image using the primary segmentation model."""
+    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)
+
+        # Handle nested list or Tensor output from some model versions
+        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}")
     return None
 
-def fallback_segmentation(image: np.ndarray) -> np.ndarray:
-    Z = image.reshape((-1, 3)).astype(np.float32)
+def segment_wound_with_fallback(image: np.ndarray) -> np.ndarray:
+    """A fallback segmentation method using k-means clustering if the primary model fails."""
+    pixels = image.reshape((-1, 3)).astype(np.float32)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
-    _, 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:
-    area_px = cv2.countNonZero(mask)
-    area_cm2 = area_px / (PIXELS_PER_CM ** 2)
+    _, 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]) # 'a' channel in LAB is good for redness
+    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
+
+def calculate_metrics(mask: np.ndarray) -> dict:
+    """Calculates dimensional and analytical metrics from the wound mask."""
+    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)
     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
-    }
+        return {"area_cm2": area_cm2, "length_cm": 0.0, "breadth_cm": 0.0, "depth_cm": 0.0, "moisture": 0.0}
+
+    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
+
+    # Placeholder for depth and moisture calculation.
+    # These would typically require more advanced sensors or algorithms.
+    depth_cm = 0.1 # Placeholder value
+    moisture = 75.0 # Placeholder value
+
+    return {"area_cm2": area_cm2, "length_cm": length_cm, "breadth_cm": breadth_cm, "depth_cm": depth_cm, "moisture": moisture}
+
+def create_visual_overlay_and_zoom(image: np.ndarray, mask: np.ndarray, bbox: tuple = None) -> np.ndarray:
+    """Draws a polygon around the wound, applies a color overlay, and zooms to the region."""
+    annotated_img = image.copy()
+    
+    # Find contours to draw the polygon
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if contours:
+        # Draw a distinct polygon outline around the wound
+        cv2.drawContours(annotated_img, contours, -1, (0, 255, 255), 2) # Yellow, 2px thick
+
+    # Zoom to the wound area if a bounding box is available
+    if bbox:
+        xmin, ymin, xmax, ymax = bbox
+        # Add a 10% padding around the bounding box for better context
+        padding_x = int((xmax - xmin) * 0.10)
+        padding_y = int((ymax - ymin) * 0.10)
+        
+        # Ensure coordinates are within image bounds
+        zoom_xmin = max(0, xmin - padding_x)
+        zoom_ymin = max(0, ymin - padding_y)
+        zoom_xmax = min(image.shape[1], xmax + padding_x)
+        zoom_ymax = min(image.shape[0], ymax + padding_y)
+        
+        return annotated_img[zoom_ymin:zoom_ymax, zoom_xmin:zoom_xmax]
+    
+    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
 
-# --- API ---
+# --- Main API Endpoint ---
 @app.post("/analyze_wound")
 async def analyze_wound(file: UploadFile = File(...)):
+    """
+    Receives an image, analyzes the wound, and returns a zoomed,
+    annotated image with metrics in the response headers.
+    """
     contents = await file.read()
-    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)
+    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)
+    
+    # Use YOLO to find the general wound region
+    bbox = detect_wound_region_yolo(processed_image)
+    
     if bbox:
-        full_mask[bbox[1]:bbox[3], bbox[0]:bbox[2]] = mask
+        xmin, ymin, xmax, ymax = bbox
+        # Crop to the detected region for more accurate segmentation
+        cropped_for_segmentation = processed_image[ymin:ymax, xmin:xmax]
     else:
-        full_mask = mask
+        # If no wound is detected, analyze the whole image as a fallback
+        cropped_for_segmentation = processed_image
+
+    # Segment the wound within the cropped region
+    mask = segment_wound_with_model(cropped_for_segmentation)
+    if mask is None:
+        mask = segment_wound_with_fallback(cropped_for_segmentation)
+
+    # Calculate metrics based on the precise mask
+    metrics = calculate_metrics(mask)
 
-    final_image = annotate(image, full_mask, metrics['contour'])
-    _, buf = cv2.imencode(".png", final_image)
+    # Create a full-sized mask to pass to the visualization function
+    full_mask = np.zeros(original_image.shape[:2], dtype=np.uint8)
+    if bbox:
+        full_mask[ymin:ymax, xmin:xmax] = mask
+    else:
+        full_mask = mask
+    
+    # Generate the final visual output: draw polygon and zoom
+    final_image = create_visual_overlay_and_zoom(original_image, full_mask, bbox)
+    
+    success, png_data = cv2.imencode(".png", final_image)
+    if not success:
+        raise HTTPException(status_code=500, detail="Failed to encode output 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
+    # Set the custom headers as requested
+    headers = {
+        'X-Length-Cm': f"{metrics['length_cm']:.2f}",
+        'X-Breadth-Cm': f"{metrics['breadth_cm']:.2f}",
+        'X-Depth-Cm': f"{metrics['depth_cm']:.2f}",
+        'X-Area-Cm2': f"{metrics['area_cm2']:.2f}",
+        'X-Moisture': f"{metrics['moisture']:.1f}"
+    }
+    
+    return Response(content=png_data.tobytes(), media_type="image/png", headers=headers)