Update predict.py

predict.py

@@ -1,106 +1,204 @@
-from fastapi import FastAPI, File, UploadFile, Response
-from fastapi.responses import FileResponse
+from fastapi import FastAPI, File, UploadFile, HTTPException, Response
 from tensorflow.keras.preprocessing.image import img_to_array
 import tensorflow as tf
 import cv2
 import numpy as np
 import os
-from scipy import fftpack
-from scipy import ndimage
 from ultralytics import YOLO
 from PIL import Image
 import io
 
-app = FastAPI()
-uploads_dir = 'uploads'
+# --- Configuration ---
+# IMPORTANT: Calibrate this value for accurate measurements.
+# To calibrate: Take a photo of a ruler next to a wound. In an image editor,
+# measure how many pixels correspond to 1 cm.
+# Example: If 1 centimeter is 25 pixels long in your image, set PIXELS_PER_CM = 25.0.
+PIXELS_PER_CM = 25.0 
 
-if not os.path.exists(uploads_dir):
-    os.makedirs(uploads_dir)
+# --- 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="2.0.0"
+)
 
-# Load the saved models
-segmentation_model_path = 'segmentation_model.h5'
-segmentation_model = tf.keras.models.load_model(segmentation_model_path)
-yolo_model_path = 'best.pt'
-yolo_model = YOLO(yolo_model_path)
+UPLOADS_DIR = 'uploads'
+if not os.path.exists(UPLOADS_DIR):
+    os.makedirs(UPLOADS_DIR)
+
+# --- Model Loading ---
+def load_models():
+    """Loads the machine learning models from disk."""
+    try:
+        # Ensure you have the correct paths to your model files
+        segmentation_model = tf.keras.models.load_model('segmentation_model.h5')
+        yolo_model = YOLO('best.pt')
+        print("Models loaded successfully.")
+        return segmentation_model, yolo_model
+    except Exception as e:
+        print(f"FATAL: Could not load models. Error: {e}")
+        return None, None
+
+segmentation_model, yolo_model = load_models()
+
+# --- Computer Vision and Analysis Functions ---
+
+def detect_wound_with_yolo(img: np.ndarray):
+    """Detects wound using YOLO model and returns the combined bounding box."""
+    if yolo_model is None: return None
+    results = yolo_model(img)
+    boxes = results[0].boxes.xyxy.tolist()
+    if not boxes:
+        return None
+    
+    combined_xmin = min(box[0] for box in boxes)
+    combined_ymin = min(box[1] for box in boxes)
+    combined_xmax = max(box[2] for box in boxes)
+    combined_ymax = max(box[3] for box in boxes)
+    
+    return int(combined_xmin), int(combined_ymin), int(combined_xmax), int(combined_ymax)
+
+def detect_wound_with_cv(img: np.ndarray):
+    """Fallback wound detection using traditional CV (HSV color-space analysis)."""
+    hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+    lower_range = np.array([0, 40, 40])
+    upper_range = np.array([25, 255, 255])
+    mask = cv2.inRange(hsv_img, lower_range, upper_range)
+    
+    kernel = np.ones((5, 5), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
+
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours: return None
+        
+    largest_contour = max(contours, key=cv2.contourArea)
+    if cv2.contourArea(largest_contour) < 150: return None
+
+    return cv2.boundingRect(largest_contour) # Returns (x, y, w, h)
+
+def calculate_wound_dimensions(mask: np.ndarray):
+    """Calculates area, length, and breadth from a binary segmentation mask."""
+    binary_mask = (mask.squeeze() > 0.5).astype(np.uint8) * 255
+    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return 0, 0, 0
+
+    main_contour = max(contours, key=cv2.contourArea)
+    
+    area_pixels = cv2.contourArea(main_contour)
+    area_cm2 = area_pixels / (PIXELS_PER_CM ** 2)
+
+    if len(main_contour) < 5:
+        x, y, w, h = cv2.boundingRect(main_contour)
+        length_px, breadth_px = max(w, h), min(w, h)
+    else:
+        rect = cv2.minAreaRect(main_contour)
+        (w, h) = rect[1]
+        length_px, breadth_px = max(w, h), min(w, h)
+
+    length_cm = length_px / PIXELS_PER_CM
+    breadth_cm = breadth_px / PIXELS_PER_CM
+    
+    return round(area_cm2, 3), round(length_cm, 3), round(breadth_cm, 3)
+
+def calculate_depth_estimate(image: np.ndarray, mask: np.ndarray):
+    """Estimates wound depth based on color intensity within the wound area."""
+    binary_mask = (mask.squeeze() > 0.5).astype(np.uint8)
+    if np.sum(binary_mask) == 0: return 0.0
+
+    wound_region = cv2.bitwise_and(image, image, mask=binary_mask)
+    lab_wound = cv2.cvtColor(wound_region, cv2.COLOR_BGR2LAB)
+    a_channel = lab_wound[:, :, 1]
+    
+    redness_values = a_channel[binary_mask == 1]
+    if redness_values.size == 0: return 0.0
+
+    avg_redness = np.mean(redness_values)
+    depth_cm = np.interp(avg_redness, [128, 180], [0.1, 2.0])
+    
+    return round(max(0, depth_cm), 3)
+
+def calculate_moisture_level(image: np.ndarray, mask: np.ndarray):
+    """Calculates a moisture score based on texture analysis within the wound area."""
+    binary_mask = (mask.squeeze() > 0.5).astype(np.uint8)
+    if np.sum(binary_mask) == 0: return 0.0
 
-def calculate_moisture_and_texture(image):
     gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    fft_image = fftpack.fft2(gray_image)
-    fft_shifted = fftpack.fftshift(fft_image)
-    magnitude_spectrum = 20 * np.log(np.abs(fft_shifted))
-    height, width = magnitude_spectrum.shape
-    center_x, center_y = width // 2, height // 2
-    radius = min(center_x, center_y) // 2
-    moisture_region = magnitude_spectrum[center_y - radius:center_y + radius, center_x - radius:center_x + radius]
-    moisture_level = np.mean(moisture_region)
-    return moisture_level
-
-def calculate_wound_dimensions(mask):
-    labeled_mask, num_labels = ndimage.label(mask > 0.5) 
-    label_count = np.bincount(labeled_mask.ravel())
-    wound_label = np.argmax(label_count[1:]) + 1  
-    wound_region = labeled_mask == wound_label   
-    rows = np.any(wound_region, axis=1)
-    cols = np.any(wound_region, axis=0)
-    rmin, rmax = np.where(rows)[0][[0, -1]]
-    cmin, cmax = np.where(cols)[0][[0, -1]]
-    length_pixels = rmax - rmin
-    breadth_pixels = cmax - cmin
-    pixel_to_cm_ratio = 0.1
-    length_cm = length_pixels * pixel_to_cm_ratio
-    breadth_cm = breadth_pixels * pixel_to_cm_ratio
-    depth_cm = np.mean(mask[wound_region]) * pixel_to_cm_ratio
-    length_cm = round(length_cm, 3)
-    breadth_cm = round(breadth_cm, 3)
-    depth_cm = round(depth_cm, 3)
-    area_cm2 = length_cm * breadth_cm
-    return length_cm, breadth_cm, depth_cm, area_cm2
+    wound_pixels = gray_image[binary_mask == 1]
+    if wound_pixels.size < 2: return 0.0
+
+    texture_metric = np.std(wound_pixels)
+    moisture_score = np.interp(texture_metric, [0, 60], [100, 0])
+    
+    return round(np.clip(moisture_score, 0, 100), 3)
+
+# --- API Endpoint ---
 
 @app.post("/analyze_wound")
 async def analyze_wounds(file: UploadFile = File(...)):
-    if file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
-        contents = await file.read()
-        nparr = np.fromstring(contents, np.uint8)
-        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-        results = yolo_model(img) 
-        combined_xmin = float('inf')
-        combined_ymin = float('inf')
-        combined_xmax = float('-inf')
-        combined_ymax = float('-inf')
-        for detection in results[0].boxes.xyxy.tolist():
-            xmin, ymin, xmax, ymax = detection
-            combined_xmin = min(combined_xmin, xmin)
-            combined_ymin = min(combined_ymin, ymin)
-            combined_xmax = max(combined_xmax, xmax)
-            combined_ymax = max(combined_ymax, ymax)
-        combined_xmin = int(combined_xmin)
-        combined_ymin = int(combined_ymin)
-        combined_xmax = int(combined_xmax)
-        combined_ymax = int(combined_ymax)
-        combined_img = img[combined_ymin:combined_ymax, combined_xmin:combined_xmax]
-        combined_img_resized = cv2.resize(combined_img, (224, 224))
-        img_array = img_to_array(combined_img_resized) / 255.0
-        img_array = np.expand_dims(img_array, axis=0)
-        output = segmentation_model.predict(img_array)
-        predicted_mask = output[0]
-        mask_overlay = (predicted_mask.squeeze() * 255).astype(np.uint8)
-        mask_overlay_colored = np.zeros((mask_overlay.shape[0], mask_overlay.shape[1], 3), dtype=np.uint8)
-        mask_overlay_colored[mask_overlay > 200] = [255, 0, 0]  # Red
-        mask_overlay_colored[(mask_overlay > 100) & (mask_overlay <= 200)] = [0, 255, 0]  # Green
-        mask_overlay_colored[mask_overlay <= 100] = [0, 0, 255]  # Blue
-        mask_overlay_colored = cv2.resize(mask_overlay_colored, (224, 224))
-        blended_image = cv2.addWeighted(combined_img_resized.astype(np.uint8), 0.6, mask_overlay_colored, 0.4, 0)
-        segmented_image = Image.fromarray(cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB))
-        img_byte_arr = io.BytesIO()
-        segmented_image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0)
-        length_cm, breadth_cm, depth_cm, area_cm2 = calculate_wound_dimensions(predicted_mask)
-        moisture = calculate_moisture_and_texture(combined_img)
-        response = Response(img_byte_arr.getvalue(), media_type='image/png')
-        response.headers['X-Length-Cm'] = str(length_cm)
-        response.headers['X-Breadth-Cm'] = str(breadth_cm)
-        response.headers['X-Depth-Cm'] = str(depth_cm)
-        response.headers['X-Area-Cm2'] = str(area_cm2)
-        response.headers['X-Moisture'] = str(moisture)
-        return response
-    return {'error': 'Invalid file format'}
+    if yolo_model is None or segmentation_model is None:
+        raise HTTPException(status_code=503, detail="Models are not loaded. The service is unavailable.")
+
+    if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
+        # Returning a dictionary for error is consistent with the original code's error path
+        return {'error': 'Invalid file format'}
+
+    contents = await file.read()
+    nparr = np.frombuffer(contents, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    if img is None:
+        return {'error': 'Could not decode image file.'}
+
+    # 1. Detect Wound Bounding Box (YOLO with CV fallback)
+    bbox = detect_wound_with_yolo(img)
+    if bbox is None:
+        cv_bbox = detect_wound_with_cv(img)
+        if cv_bbox is None:
+            raise HTTPException(status_code=404, detail="Wound not detected.")
+        x, y, w, h = cv_bbox
+        bbox = (x, y, x + w, y + h)
+
+    xmin, ymin, xmax, ymax = bbox
+    cropped_img = img[ymin:ymax, xmin:xmax]
+    if cropped_img.size == 0:
+        raise HTTPException(status_code=400, detail="Wound detection resulted in an empty crop.")
+
+    # 2. Perform Segmentation
+    resized_for_model = cv2.resize(cropped_img, (224, 224))
+    img_array = img_to_array(resized_for_model) / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
+    predicted_mask = segmentation_model.predict(img_array)[0]
+
+    # 3. Calculate Dimensions and Properties
+    area_cm2, length_cm, breadth_cm = calculate_wound_dimensions(predicted_mask)
+    depth_cm = calculate_depth_estimate(resized_for_model, predicted_mask)
+    moisture = calculate_moisture_level(resized_for_model, predicted_mask)
+
+    # 4. Create Visualization (as in original code)
+    mask_overlay = (predicted_mask.squeeze() * 255).astype(np.uint8)
+    # Using a colormap for better visualization of the segmentation mask
+    mask_overlay_colored = cv2.applyColorMap(mask_overlay, cv2.COLORMAP_JET)
+    blended_image = cv2.addWeighted(resized_for_model, 0.6, mask_overlay_colored, 0.4, 0)
+    
+    segmented_image_pil = Image.fromarray(cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB))
+    img_byte_arr = io.BytesIO()
+    segmented_image_pil.save(img_byte_arr, format='PNG')
+    img_byte_arr.seek(0)
+
+    # 5. Create Response with headers (as in original code)
+    response = Response(img_byte_arr.getvalue(), media_type='image/png')
+    response.headers['X-Area-Cm2'] = str(area_cm2)
+    response.headers['X-Length-Cm'] = str(length_cm)
+    response.headers['X-Breadth-Cm'] = str(breadth_cm)
+    response.headers['X-Depth-Cm'] = str(depth_cm)
+    response.headers['X-Moisture'] = str(moisture)
+    
+    return response
+
+@app.get("/", include_in_schema=False)
+async def root():
+    return {"message": "Welcome to the Wound Analysis API. Use the /analyze_wound endpoint to process an image."}
+
+# To run this app, save it as `main.py` and execute:
+# uvicorn main:app --reload