SkinBurns Segmentation

.gitignore

@@ -3,4 +3,7 @@ venv/*
 uploads/*
 runs/*
 screenshots/*
-*.pyc
+*.pyc
+*.mp4
+*.png
+*.jpg

SkinBurns_Segmentation.py

@@ -1,68 +1,172 @@
+# ── Imports & Helper Functions ──
+import os
 import numpy as np
-from skimage import io, img_as_float, color, transform
-from scipy.ndimage import gaussian_filter
+import matplotlib.pyplot as plt
+from skimage import io, color, img_as_float
+from scipy.ndimage import gaussian_filter, uniform_filter, binary_closing, label
+from sklearn.cluster import KMeans
 import tensorly as tl
 from tensorly.decomposition import tucker
-from sklearn.cluster import KMeans
 import cv2
-import matplotlib.pyplot as plt
+from skimage.color import rgb2gray
+
+def compute_local_variance(img, window_size=9):
+    img_sq = img ** 2
+    mean = uniform_filter(img, size=window_size)
+    mean_sq = uniform_filter(img_sq, size=window_size)
+    return mean_sq - mean ** 2
 
-def segment_skin_burns(image_path, output_path='segmented_image.png'):
-    # Load and preprocess the image
-    image = io.imread(image_path)
+def largest_rectangle_in_mask(mask):
+    h, w = mask.shape
+    heights = [0] * w
+    max_area = 0
+    best_coords = None  # (y_start, y_end, x_start, x_end)
+
+    for y in range(h):
+        for x in range(w):
+            heights[x] = heights[x] + 1 if mask[y, x] == 1 else 0
+
+        stack = []
+        x = 0
+        while x <= w:
+            curr_h = heights[x] if x < w else 0
+            if not stack or curr_h >= heights[stack[-1]]:
+                stack.append(x)
+                x += 1
+            else:
+                top = stack.pop()
+                width_rect = x if not stack else x - stack[-1] - 1
+                area = heights[top] * width_rect
+                if area > max_area:
+                    max_area = area
+                    y_end = y
+                    y_start = y - heights[top] + 1
+                    x_end = x
+                    x_start = x - width_rect
+                    best_coords = (y_start, y_end, x_start, x_end)
+    return best_coords
+# ── Main Segmentation Function ──
+def segment_burn(patient_image_path,
+                 reference_image_path,
+                 threshold_texture=0.35,
+                 threshold_color=10):
+    """
+    Takes:
+      - patient_image_path: path to the RGB image
+      - reference_image_path: path to the healthy-reference RGB image
+      - threshold_texture, threshold_color: tuning parameters
+
+    Returns:
+      - burn_crop_clean: the cropped burn region (as an ndarray)
+      - burn_crop_debug: original image with the crop rectangle overlaid
+    """
+    # --- Load patient image ---
+    image = io.imread(patient_image_path)
+    if image.shape[-1] == 4:
+        image = image[..., :3]
     image = img_as_float(image)
-    image_lab = color.rgb2lab(image)  # LAB color space
+    height, width, _ = image.shape
 
-    # Adaptive Illumination Correction (CLAHE)
+    # --- Preprocessing & feature extraction ---
+    # 1. Convert to Lab + CLAHE on L
+    image_lab = color.rgb2lab(image)
     clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
     image_lab[:, :, 0] = clahe.apply((image_lab[:, :, 0] * 255).astype(np.uint8)) / 255.0
 
-    # Gaussian filtering
-    sigma_value = 1.5
+    # 2. Gaussian smoothing
     image_lab_filtered = np.zeros_like(image_lab)
     for i in range(3):
-        image_lab_filtered[:, :, i] = gaussian_filter(image_lab[:, :, i], sigma=sigma_value)
+        image_lab_filtered[:, :, i] = gaussian_filter(image_lab[:, :, i], sigma=1.5)
 
-    # Third-order tensor
+    # 3. Tucker decomposition + reconstruction
     tensor_image = tl.tensor(image_lab_filtered)
+    core, factors = tucker(tensor_image, rank=[30, 30, 3])
+    reconstructed = tl.tucker_to_tensor((core, factors))
 
-    # Tucker3 tensor decomposition
-    tucker_rank = [30, 30, 3]
-    core, factors = tucker(tensor_image, rank=tucker_rank)
+    # 4. Texture map (local variance)
+    gray = rgb2gray(image)
+    local_var = compute_local_variance(gray, window_size=9)
+    vmin, vmax = np.percentile(local_var, [5, 95])
+    local_var = np.clip((local_var - vmin) / (vmax - vmin), 0, 1)
 
-    # Resize core tensor
-    core_resized = transform.resize(core, (image.shape[0], image.shape[1], 3), anti_aliasing=True)
+    # 5. Build feature vectors
+    lab_feat = reconstructed.reshape(-1, 3)
+    texture_feat = local_var.flatten()[:, None]
+    features = np.hstack((lab_feat, texture_feat))
 
-    # Convert to RGB
-    core_resized_rgb = np.clip(color.lab2rgb(core_resized), 0, 1)
+    # --- KMeans clustering & burn scoring ---
+    n_clusters = 5
+    kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
+    labels = kmeans.fit_predict(features).reshape(height, width)
 
-    # Reconstruct image
-    reconstructed_image = tl.tucker_to_tensor((core, factors))
-    
-    # Features for clustering
-    height, width, _ = image.shape
-    feature_vectors = reconstructed_image.reshape(height * width, -1)
+    # Score clusters for burn likelihood
+    cluster_means = [features[labels.flatten()==i].mean(axis=0) for i in range(n_clusters)]
+    burn_scores = []
+    for L, A, B, T in cluster_means:
+        sat = np.sqrt(A**2 + B**2)
+        red = A
+        dark = (1 - L) + (1 - sat)
+        burn_scores.append(0.4*red + 0.2*dark + 0.4*T)
 
-    # K-means clustering
-    n_clusters = 3  # Burn wound, healthy skin, background
-    kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
-    kmeans.fit(feature_vectors)
-    labels = kmeans.labels_.reshape(height, width)
+    burn_clusters = np.argsort(burn_scores)[-2:]
+    # skin cluster = cluster closest to healthy skin color
+    skin_cluster = np.argmin([np.linalg.norm(np.array([L,A,B]) - np.array([0.7,0,0]))
+                              for L,A,B,T in cluster_means])
+
+    # Initial burn mask + morphology
+    burn_mask = np.isin(labels, burn_clusters).astype(np.uint8)
+    closed = binary_closing(burn_mask, structure=np.ones((10,10)))
+    labeled, _ = label(closed)
+    sizes = np.bincount(labeled.ravel()); sizes[0]=0
+    largest = (labeled == sizes.argmax()).astype(np.uint8)
+
+    # --- Reference-based refinement ---
+    if os.path.exists(reference_image_path):
+        ref = io.imread(reference_image_path)
+        if ref.shape[-1]==4:
+            ref = ref[..., :3]
+        ref = img_as_float(ref)
+        ref_lab = color.rgb2lab(ref)
+        ref_lab[:, :, 0] = clahe.apply((ref_lab[:, :, 0]*255).astype(np.uint8))/255.0
+
+        ref_lab_f = np.zeros_like(ref_lab)
+        for i in range(3):
+            ref_lab_f[:, :, i] = gaussian_filter(ref_lab[:, :, i], sigma=1.5)
+
+        ref_gray = rgb2gray(ref)
+        ref_var = compute_local_variance(ref_gray, window_size=9)
+        rvmin, rvmax = np.percentile(ref_var, [5,95])
+        ref_var = np.clip((ref_var - rvmin)/(rvmax - rvmin),0,1)
+
+        ref_feat = np.hstack((ref_lab_f.reshape(-1,3), ref_var.flatten()[:,None]))
+        healthy_vec = ref_feat.mean(axis=0)
 
-    # Identify burn cluster
-    cluster_means = [np.mean(feature_vectors[labels.flatten() == i], axis=0) for i in range(n_clusters)]
-    burn_cluster = np.argmax([mean[1] for mean in cluster_means])  # LAB 'a' channel: higher values → red tones
+        final_burn = np.zeros_like(largest)
+        for y in range(height):
+            for x in range(width):
+                if largest[y,x]:
+                    pix_vec = np.append(reconstructed[y,x,:], local_var[y,x])
+                    dist = np.linalg.norm(pix_vec - healthy_vec)
+                    if not (dist < threshold_color and local_var[y,x] < threshold_texture):
+                        final_burn[y,x] = 1
+    else:
+        final_burn = largest
 
-    # Color-mapped segmentation image
+    # (Optional) segmentation map if you need it later
     segmentation_map = np.zeros((height, width, 3))
-    colors = {burn_cluster: [1, 0, 0],  # Red for burn
-              (burn_cluster + 1) % 3: [0, 1, 0],  # Green for healthy skin
-              (burn_cluster + 2) % 3: [0, 0, 1]}  # Blue for background
+    segmentation_map[labels==skin_cluster] = [0,1,0]
+    segmentation_map[final_burn==1]   = [1,0,0]
 
-    for cluster, colour in colors.items():
-        segmentation_map[labels == cluster] = colour
+    # --- Crop the largest solid rectangle within the burn mask ---
+    coords = largest_rectangle_in_mask(final_burn)
+    if coords is None:
+        raise ValueError("No solid rectangular burn area found.")
+    y1,y2,x1,x2 = coords
+    burn_crop_clean = image[y1:y2+1, x1:x2+1]
 
-    # Save the segmented image
-    plt.imsave(output_path, segmentation_map)
+    # --- Build debug overlay image ---
+    debug_img = (image * 255).astype(np.uint8).copy()
+    cv2.rectangle(debug_img, (x1,y1), (x2,y2), (0,255,255), thickness=3)
+    burn_crop_debug = debug_img
 
-    return output_path
+    return burn_crop_clean, burn_crop_debug

app.py

@@ -24,6 +24,8 @@ import cv2
 import time
 from CPR.CPRAnalyzer import CPRAnalyzer
 import tempfile
+import matplotlib.pyplot as plt
+
 
 
 
@@ -145,6 +147,52 @@ async def predict_burn(file: UploadFile = File(...)):
         return JSONResponse(content={"error": str(e)}, status_code=500)
 
 
+@app.post("/segment_burn")
+async def segment_burn_endpoint(reference: UploadFile = File(...), patient: UploadFile = File(...)):
+    try:
+        # Save the reference image temporarily
+        reference_path = f"temp_ref_{reference.filename}"
+        reference_bytes = await reference.read()
+        with open(reference_path, "wb") as ref_file:
+            ref_file.write(reference_bytes)
+
+        # Save the patient image temporarily
+        patient_path = f"temp_patient_{patient.filename}"
+        patient_bytes = await patient.read()
+        with open(patient_path, "wb") as pat_file:
+            pat_file.write(patient_bytes)
+
+        # Call the segmentation logic
+        burn_crop_clean, burn_crop_debug = segment_burn(patient_path, reference_path)
+
+        # Save the cropped outputs
+        burn_crop_clean_path = f"temp_burn_crop_clean_{uuid.uuid4()}.png"
+        burn_crop_debug_path = f"temp_burn_crop_debug_{uuid.uuid4()}.png"
+  
+
+        plt.imsave(burn_crop_clean_path, burn_crop_clean)
+        plt.imsave(burn_crop_debug_path, burn_crop_debug)
+
+        # Upload to Cloudinary
+        crop_clean_upload = cloudinary.uploader.upload(burn_crop_clean_path, public_id=f"ref_{reference.filename}")
+        crop_debug_upload = cloudinary.uploader.upload(burn_crop_debug_path, public_id=f"pat_{patient.filename}")
+        crop_clean_url = crop_clean_upload["secure_url"]
+        crop_debug_url = crop_debug_upload["secure_url"]
+
+        # Clean up temp files
+        
+        os.remove(burn_crop_clean_path)
+        os.remove(burn_crop_debug_path)
+        
+
+        return {
+            "crop_clean_url": crop_clean_url,
+            "crop_debug_url": crop_debug_url
+        }
+
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+
 
 # ✅ Optimize and transform image URL
 @app.get("/cloudinary/transform")