Update predict.py

predict.py

@@ -2,7 +2,6 @@ from fastapi import FastAPI, File, UploadFile, HTTPException, Response
 import cv2
 import numpy as np
 from ultralytics import YOLO
-import tensorflow as tf
 import io
 from typing import Union
 
@@ -11,56 +10,36 @@ PIXELS_PER_CM = 50.0
 
 # --- 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
+    title="Wound Heatmap Analysis API",
+    description="An API that generates a three-color heatmap (Red, Blue, Green) on a wound to show tissue characteristics.",
+    version="5.0.0" # Version updated for three-layer color heatmap
 )
 
 # --- Model Loading ---
-def load_models():
-    """Loads the segmentation and YOLO models, handling potential errors."""
-    segmentation_model, yolo_model = None, None
+def load_yolo_model():
+    """Loads the YOLO model for initial wound detection."""
     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.")
+        return yolo_model
     except Exception as e:
-        print(f"Warning: Could not load YOLO model. Using fallback. Error: {e}")
-        
-    return segmentation_model, yolo_model
+        print(f"Fatal: Could not load YOLO model. Error: {e}")
+        return None
 
-segmentation_model, yolo_model = load_models()
+yolo_model = load_yolo_model()
 
 # --- Helper Functions ---
 
 def preprocess_image(image: np.ndarray) -> np.ndarray:
-    """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))
-    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)
+    """Applies a median blur to reduce noise."""
+    return cv2.medianBlur(image, 5)
 
 def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
-    """Detects the wound bounding box using the YOLO model."""
+    """Detects the primary 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 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))
@@ -68,53 +47,33 @@ def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
         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 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)
-    _, 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 create_wound_bed_mask(cropped_image: np.ndarray) -> np.ndarray:
+    """
+    Creates a general binary mask of the entire wound bed, which will be the area for the heatmap.
+    This uses a broader color range than specific tissue segmentation.
+    """
+    lab_image = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2LAB)
+    
+    # Broad thresholds to capture the entire wound area (slough, granulation, etc.)
+    lower_bound = np.array([0, 135, 135])
+    upper_bound = np.array([255, 200, 200])
+    
+    mask = cv2.inRange(lab_image, lower_bound, upper_bound)
+    
+    # Clean up the mask
+    kernel = np.ones((5, 5), np.uint8)
+    mask_cleaned = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=2)
+    mask_cleaned = cv2.morphologyEx(mask_cleaned, cv2.MORPH_CLOSE, kernel, iterations=3)
+    
+    return mask_cleaned
 
 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:
+    """Calculates dimensional metrics from the overall wound mask."""
+    area_pixels = cv2.countNonZero(mask)
+    if area_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)
+    area_cm2 = area_pixels / (PIXELS_PER_CM ** 2)
     
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if not contours:
@@ -125,91 +84,85 @@ def calculate_metrics(mask: np.ndarray) -> dict:
     
     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
+    depth_cm = 0.1 # Placeholder
+    moisture = 75.0 # Placeholder
 
     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()
+def create_three_color_heatmap(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    """
+    Generates and overlays a three-color (Red, Blue, Green) heatmap onto the image
+    based on the intensity of the 'a' channel (redness) in the LAB color space.
+    """
+    if cv2.countNonZero(mask) == 0:
+        return image
+
+    # Convert the region of interest to LAB color space for analysis
+    lab_image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
+    a_channel = lab_image[:, :, 1] # The 'a' channel represents the green-red axis
+
+    # Create a color overlay image, initially transparent
+    overlay = np.zeros_like(image)
     
-    # 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]
+    # Define thresholds for redness intensity. These values might need tuning.
+    # Values are based on the 'a' channel, where ~128 is neutral.
+    RED_THRESHOLD = 160  # Most intense red (e.g., fresh granulation)
+    BLUE_THRESHOLD = 145 # Medium red (e.g., developing tissue)
+    # Anything above a lower bound (e.g., 135) will be green.
+
+    # Apply colors based on thresholds, only within the wound mask
+    # Red for "most" affected
+    overlay[(a_channel >= RED_THRESHOLD) & (mask == 255)] = (0, 0, 255) # BGR for Red
+    # Blue for "less" affected
+    overlay[(a_channel >= BLUE_THRESHOLD) & (a_channel < RED_THRESHOLD) & (mask == 255)] = (255, 0, 0) # BGR for Blue
+    # Green for "least" affected (but still part of the wound bed)
+    overlay[(a_channel < BLUE_THRESHOLD) & (mask == 255)] = (0, 255, 0) # BGR for Green
+
+    # Blend the original image with the color overlay
+    # A weight of 0.4 for the overlay makes it visible but not overpowering
+    blended_image = cv2.addWeighted(overlay, 0.4, image, 0.6, 0)
     
-    return annotated_img
+    # To make the result clean, we only apply the blended result where the mask is active
+    final_image = image.copy()
+    final_image[mask == 255] = blended_image[mask == 255]
 
+    return final_image
 
 # --- 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.
-    """
+    if not yolo_model:
+        raise HTTPException(status_code=503, detail="YOLO model is not available.")
+
     contents = await file.read()
-    image_array = np.frombuffer(contents, np.uint8)
-    original_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
+    original_image = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
     if original_image is None:
-        raise HTTPException(status_code=400, detail="Invalid or corrupt image file")
+        raise HTTPException(status_code=400, detail="Invalid image file")
 
     processed_image = preprocess_image(original_image)
     
-    # Use YOLO to find the general wound region
     bbox = detect_wound_region_yolo(processed_image)
+    if not bbox:
+        raise HTTPException(status_code=404, detail="No wound detected in the image.")
+
+    xmin, ymin, xmax, ymax = bbox
+    cropped_image_roi = original_image[ymin:ymax, xmin:xmax]
+
+    # Step 1: Create a mask for the entire wound bed in the cropped image
+    wound_mask = create_wound_bed_mask(cropped_image_roi)
     
-    if bbox:
-        xmin, ymin, xmax, ymax = bbox
-        # Crop to the detected region for more accurate segmentation
-        cropped_for_segmentation = processed_image[ymin:ymax, xmin:xmax]
-    else:
-        # 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)
-
-    # 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)
+    # Step 2: Calculate metrics based on this overall wound mask
+    metrics = calculate_metrics(wound_mask)
+
+    # Step 3: Generate the three-color heatmap on the cropped image
+    heatmap_image = create_three_color_heatmap(cropped_image_roi, wound_mask)
     
-    success, png_data = cv2.imencode(".png", final_image)
+    # Step 4: Encode the final, annotated (and already cropped) image
+    success, png_data = cv2.imencode(".png", heatmap_image)
     if not success:
         raise HTTPException(status_code=500, detail="Failed to encode output image")
 
-    # Set the custom headers as requested
+    # Step 5: Set the custom headers
     headers = {
         'X-Length-Cm': f"{metrics['length_cm']:.2f}",
         'X-Breadth-Cm': f"{metrics['breadth_cm']:.2f}",