SkinBurns/SkinBurns_Segmentation.py

import os
import numpy as np
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
import cv2
from skimage.color import rgb2gray

def LocalVar(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 largestRect(mask):
    h, w = mask.shape
    heights = [0] * w
    Arealargest = 0
    chosenCors = None  

    for row_idx, row in enumerate(mask):
        heights = [
            heights[col] + 1 if row[col] == 1 else 0
            for col in range(w)
        ]

        stack = []
        col = 0
        while col <= w:
            currentH = heights[col] if col < w else 0

            if not stack or currentH >= heights[stack[-1]]:
                stack.append(col)
                col += 1
            else:
                top = stack.pop()
                rectangleW = col if not stack else col - stack[-1] - 1
                area = heights[top] * rectangleW
                if area > Arealargest:
                    Arealargest = area
                    bottomY = row_idx
                    topY = row_idx - heights[top] + 1
                    leftX = col - rectangleW
                    rightX = col - 1
                    chosenCors = (topY, bottomY, leftX, rightX)
    return chosenCors

def segment_burn(patient_image_path,
                 reference_image_path,
                 threshold_texture=0.35,
                 threshold_color=10):
    
    image = io.imread(patient_image_path)
    if image.shape[-1] == 4:
        image = image[..., :3]
    image = img_as_float(image)
    height, width, _ = image.shape

   
    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

    
    image_lab_filtered = np.zeros_like(image_lab)
    for i in range(3):
        image_lab_filtered[:, :, i] = gaussian_filter(image_lab[:, :, i], sigma=1.5)


    tensor_image = tl.tensor(image_lab_filtered)
    core, factors = tucker(tensor_image, rank=[30, 30, 3])
    reconstructed = tl.tucker_to_tensor((core, factors))
    reconstructed_rgb = np.clip(color.lab2rgb(reconstructed), 0, 1)

 
    gray = rgb2gray(image)
    local_var = LocalVar(gray, window_size=9)
    vmin, vmax = np.percentile(local_var, [5, 95])
    local_var = np.clip((local_var - vmin) / (vmax - vmin), 0, 1)

  
    lab_feat = reconstructed.reshape(-1, 3)
    texture_feat = local_var.flatten()[:, None]
    features = np.hstack((lab_feat, texture_feat))

    
    n_clusters = 5
    kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
    labels = kmeans.fit_predict(features).reshape(height, width)

    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)

    burn_clusters = np.argsort(burn_scores)[-2:]
   
    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])

    
    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)

   
    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 = LocalVar(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)

        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


    segmentation_map = np.zeros((height, width, 3))
    segmentation_map[labels==skin_cluster] = [0,1,0]
    segmentation_map[final_burn==1]   = [1,0,0]

   
    coords = largestRect(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]

    
    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 burn_crop_clean, burn_crop_debug