Update predict.py

predict.py

@@ -2,6 +2,7 @@ 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
 
@@ -10,32 +11,35 @@ PIXELS_PER_CM = 50.0
 
 # --- App Initialization ---
 app = FastAPI(
-    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
+    title="High-Quality Wound Heatmap API",
+    description="Generates a high-quality wound heatmap using a DL model, preserving original image quality.",
+    version="7.0.0" # Version updated for quality preservation
 )
 
 # --- Model Loading ---
-def load_yolo_model():
-    """Loads the YOLO model for initial wound detection."""
+def load_models():
+    """Loads both the YOLO and the TensorFlow segmentation models."""
+    yolo_model, segmentation_model = None, None
     try:
         yolo_model = YOLO("best.pt")
         print("YOLO model 'best.pt' loaded successfully.")
-        return yolo_model
     except Exception as e:
-        print(f"Fatal: Could not load YOLO model. Error: {e}")
-        return None
+        print(f"Warning: Could not load YOLO model. Error: {e}")
 
-yolo_model = load_yolo_model()
+    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. Error: {e}")
+        
+    return yolo_model, segmentation_model
 
-# --- Helper Functions ---
+yolo_model, segmentation_model = load_models()
 
-def preprocess_image(image: np.ndarray) -> np.ndarray:
-    """Applies a median blur to reduce noise."""
-    return cv2.medianBlur(image, 5)
+# --- Helper Functions ---
 
 def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
-    """Detects the primary wound bounding box using the YOLO model."""
+    """Detects the primary wound bounding box using the YOLO model on the original quality image."""
     if not yolo_model: return None
     try:
         results = yolo_model.predict(image, verbose=False)
@@ -47,34 +51,42 @@ def detect_wound_region_yolo(image: np.ndarray) -> Union[tuple, None]:
         print(f"YOLO prediction failed: {e}")
     return None
 
-def create_wound_bed_mask(cropped_image: np.ndarray) -> np.ndarray:
+def segment_wound_with_model(image: np.ndarray) -> Union[np.ndarray, None]:
     """
-    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.
+    Segments the wound using the TF/Keras model.
+    It resizes a copy for the model but returns a mask matching the original image's dimensions.
     """
-    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
+    if not segmentation_model:
+        print("Segmentation model not loaded, cannot create mask.")
+        return None
+    try:
+        input_shape = segmentation_model.input_shape[1:3]
+        # A temporary, resized copy is made for the model's prediction
+        img_resized_for_model = cv2.resize(image, (input_shape[1], input_shape[0]))
+        img_norm = np.expand_dims(img_resized_for_model.astype(np.float32) / 255.0, axis=0)
+
+        prediction = segmentation_model.predict(img_norm, verbose=0)
+
+        while isinstance(prediction, list):
+            prediction = prediction[0]
+        if isinstance(prediction, tf.Tensor):
+            prediction = prediction.numpy()
+
+        pred_mask = prediction[0]
+        # The resulting mask is resized back to the original image's dimensions to ensure perfect alignment
+        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 calculate_metrics(mask: np.ndarray) -> dict:
-    """Calculates dimensional metrics from the overall wound mask."""
+    """Calculates dimensional metrics from the full-resolution 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 = area_pixels / (PIXELS_PER_CM ** 2)
-    
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if not contours:
         return {"area_cm2": area_cm2, "length_cm": 0.0, "breadth_cm": 0.0, "depth_cm": 0.0, "moisture": 0.0}
@@ -90,39 +102,25 @@ def calculate_metrics(mask: np.ndarray) -> dict:
     return {"area_cm2": area_cm2, "length_cm": length_cm, "breadth_cm": breadth_cm, "depth_cm": depth_cm, "moisture": moisture}
 
 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.
-    """
+    """Generates and overlays a three-color heatmap, preserving the underlying image quality."""
     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
+    a_channel = lab_image[:, :, 1]
     overlay = np.zeros_like(image)
     
-    # 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
+    RED_THRESHOLD = 160
+    BLUE_THRESHOLD = 145
+
+    overlay[(a_channel >= RED_THRESHOLD) & (mask == 255)] = (0, 0, 255)
+    overlay[(a_channel >= BLUE_THRESHOLD) & (a_channel < RED_THRESHOLD) & (mask == 255)] = (255, 0, 0)
+    overlay[(a_channel < BLUE_THRESHOLD) & (mask == 255)] = (0, 255, 0)
+
+    # Blend the overlay with the original image
     blended_image = cv2.addWeighted(overlay, 0.4, image, 0.6, 0)
     
-    # To make the result clean, we only apply the blended result where the mask is active
+    # Create the final image by taking the original and replacing only the masked area with the blended version
     final_image = image.copy()
     final_image[mask == 255] = blended_image[mask == 255]
 
@@ -131,33 +129,37 @@ def create_three_color_heatmap(image: np.ndarray, mask: np.ndarray) -> np.ndarra
 # --- Main API Endpoint ---
 @app.post("/analyze_wound")
 async def analyze_wound(file: UploadFile = File(...)):
-    if not yolo_model:
-        raise HTTPException(status_code=503, detail="YOLO model is not available.")
+    if not yolo_model or not segmentation_model:
+        raise HTTPException(status_code=503, detail="A required model is not available.")
 
     contents = await file.read()
+    # Read the image from buffer into a cv2 object without any lossy conversions
     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")
+        raise HTTPException(status_code=400, detail="Invalid or corrupt image file")
 
-    processed_image = preprocess_image(original_image)
-    
-    bbox = detect_wound_region_yolo(processed_image)
+    # No preprocessing is done to the original image to preserve quality.
+    # A copy of the original image is used for detection.
+    bbox = detect_wound_region_yolo(original_image.copy())
     if not bbox:
         raise HTTPException(status_code=404, detail="No wound detected in the image.")
 
     xmin, ymin, xmax, ymax = bbox
+    # Crop the region of interest from the original, high-quality image
     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)
+    # Step 1: Use the DL model on the high-quality crop to get a precise mask
+    wound_mask = segment_wound_with_model(cropped_image_roi)
+    if wound_mask is None or cv2.countNonZero(wound_mask) == 0:
+        raise HTTPException(status_code=404, detail="Segmentation model failed to identify a wound in the detected region.")
     
-    # Step 2: Calculate metrics based on this overall wound mask
+    # Step 2: Calculate metrics based on the full-resolution mask
     metrics = calculate_metrics(wound_mask)
 
-    # Step 3: Generate the three-color heatmap on the cropped image
+    # Step 3: Generate the heatmap on the high-quality cropped image
     heatmap_image = create_three_color_heatmap(cropped_image_roi, wound_mask)
     
-    # Step 4: Encode the final, annotated (and already cropped) image
+    # Step 4: Encode the final image into PNG (lossless format) to preserve quality
     success, png_data = cv2.imencode(".png", heatmap_image)
     if not success:
         raise HTTPException(status_code=500, detail="Failed to encode output image")