main.py

import numpy as np
import cv2
from skimage import color
import time
import os
import numpy as np

def ncu_line_match_filter_direc_org(image, deg):
    """ This version is closer to the one in paper code-wise not performance or quality"""

    m, n = image.shape
    image = image.astype(float)
    j_out = np.zeros((m, n))
    rad = np.deg2rad(deg)
    
    sigmas = [1, np.sqrt(2), 2, 2 * np.sqrt(2)]
    
    for sigma in sigmas:
        mf0, s = g0_filter(sigma, [np.cos(rad), np.sin(rad)])
        a0 = np.sum(mf0 * mf0)
        
        for i in range(m):
            for j in range(n):
                flag = False
                
                i_min = i - s 
                i_max = i + s

                x_start = 0 
                x_end = 2 * s + 1
                
                if i_min < 0:
                    x_start = -i_min
                    i_min = 0
                    flag = True
                
                if i_max >= m:
                    overflow = i_max - (m - 1)
                    x_end = x_end - overflow
                    i_max = m - 1
                    flag = True

                j_min = j - s
                j_max = j + s
                
                y_start = 0
                y_end = 2 * s + 1
                
                if j_min < 0:
                    y_start = -j_min
                    j_min = 0
                    flag = True
                
                if j_max >= n:
                    overflow = j_max - (n - 1)
                    y_end = y_end - overflow
                    j_max = n - 1
                    flag = True
                
                i0_patch = image[i_min : i_max + 1, j_min : j_max + 1]
                
                i0_patch = i0_patch - np.max(i0_patch)
                i0_patch = -i0_patch
                
                if flag:
                    mf = mf0[x_start : x_end, y_start : y_end]
                    a = np.sum(mf * mf)
                else:
                    a = a0
                    mf = mf0
                
                b = np.sum(i0_patch * i0_patch)
                m0 = np.sum(i0_patch * mf)
                
                if a != 0 and b != 0:
                    t = (m0**2) / (a * b)
                else:
                    t = 0 # Fallback if energy is strictly zero to avoid error
                    print("energy was set to zero")

                if t > j_out[i, j]:
                    j_out[i, j] = t

    return j_out

def ncu_line_match_filter_direc(image, deg):
    m, n = image.shape
    image = image.astype(float)
    j_out = np.zeros((m, n))
    rad = np.deg2rad(deg)
    v = [np.cos(rad), np.sin(rad)]

    for sigma in [1, np.sqrt(2), 2, 2 * np.sqrt(2)]:
        mf0, s = g0_filter(sigma, v)
        a0 = np.sum(mf0 * mf0) # filter energy

        # Manually slide window to match local normalization
        for i in range(s, m - s):
            for j in range(s, n - s):
                # Local neighborhood normalization
                i0 = image[i-s:i+s+1, j-s:j+s+1]
                i0_norm = -(i0 - np.max(i0)) # Invert and shift

                b = np.sum(i0_norm * i0_norm)
                m0 = np.sum(i0_norm * mf0)

                if a0 > 0 and b > 0:
                    t = (m0**2) / (a0 * b) # Eq (2) from paper
                    if t > j_out[i, j]:
                        j_out[i, j] = t
    return j_out

def g0_filter(sigma, v):
    d = int(np.ceil(sigma * 3))
    x, y = np.meshgrid(np.arange(-d, d + 1), np.arange(-d, d + 1), indexing='ij')
    t = x * v[0] + y * v[1]
    f0 = np.exp(-(x**2 + y**2 - t**2) / (2 * sigma**2))
    return f0, d

def stdDilateDarkest(I):
    high = 0.75*225
    low = 0.65*225
    n1 = 40

    dx = [1, 1, 0, -1, -1, -1, 0, 1]
    dy = [0, 1, 1, 1, 0, -1, -1, -1]

    rows, cols = I.shape[0], I.shape[1]
    
    if I.ndim == 3:
        img_ch1 = I[:, :, 0]
    else:
        img_ch1 = I

    M = np.zeros((rows, cols), dtype=int)

    J = np.zeros((rows, cols, 3), dtype=np.uint8)

    for r in range(rows):
        for c in range(cols):
            if img_ch1[r, c] >= high:
                M[r, c] = 1
            else:
                J[r, c, :] = 255  # White background for non-seeds

    idx = 1
    
    for r in range(rows):
        for c in range(cols):
            
            if M[r, c] == 1:
                idx = idx + 1
                
                region_pixels = []
                
                M[r, c] = idx
                region_pixels.append((r, c))
                
                ptr = 0
                while ptr < len(region_pixels):
                    curr_r, curr_c = region_pixels[ptr]
                    ptr += 1
                    
                    for a in range(8):
                        nr = curr_r + dx[a]
                        nc = curr_c + dy[a]
                        
                        if 0 <= nr < rows and 0 <= nc < cols:
                            if M[nr, nc] == 1:
                                M[nr, nc] = idx
                                region_pixels.append((nr, nc))

                count = len(region_pixels)
                
                if count <= 2:
                    idx = idx - 1
                    for (pr, pc) in region_pixels:
                        M[pr, pc] = 0
                        J[pr, pc, :] = 255
                else:
                    nh = [] 
                    
                    for k in range(len(region_pixels)):
                        curr_r, curr_c = region_pixels[k]
                        for a in range(8):
                            nr = curr_r + dx[a]
                            nc = curr_c + dy[a]
                            
                            if 0 <= nr < rows and 0 <= nc < cols:
                                if M[nr, nc] <= 1 and img_ch1[nr, nc] >= low:
                                    M[nr, nc] = idx 
                                    nh.append([float(img_ch1[nr, nc]), nr, nc])

                    if len(nh) >= 1:
                        dlen = 0
                        
                        while dlen <= n1 and len(nh) > 0:
                            best_idx = -1
                            max_val = -1.0
                            
                            for k in range(len(nh)):
                                if nh[k][0] > max_val:
                                    max_val = nh[k][0]
                                    best_idx = k
                            
                            if best_idx == -1 or max_val < low:
                                break
                                
                            dlen += 1
                            
                            val, best_r, best_c = nh.pop(best_idx)
                            
                            region_pixels.append((best_r, best_c))
                            J[best_r, best_c, :] = 150 
                            
                            for a in range(8):
                                nr = best_r + dx[a]
                                nc = best_c + dy[a]
                                if 0 <= nr < rows and 0 <= nc < cols:
                                    if M[nr, nc] <= 1 and img_ch1[nr, nc] >= low:
                                        M[nr, nc] = idx
                                        nh.append([float(img_ch1[nr, nc]), nr, nc])

                    for item in nh:
                         M[item[1], item[2]] = 0

                    final_count = len(region_pixels)
                    
                    should_remove = False
                    
                    if final_count < 12:
                        should_remove = True
                    else:
                        intensities = []
                        for (pr, pc) in region_pixels:
                            intensities.append(float(img_ch1[pr, pc]))
                        
                        intensities.sort(reverse=True)
                        
                        val_10th = intensities[9]
                        val_median = np.median(intensities)
                        
                        if val_10th < (0.750 * 255):
                            if val_median < (0.725 * 255):
                                should_remove = True
                    
                    if should_remove:
                        for (pr, pc) in region_pixels:
                            M[pr, pc] = 0
                            J[pr, pc, :] = 255 # Back to white
    
    # Convert labeled map M to binary mask (0 and 1)
    binary_mask = (M > 0).astype(int)
    
    return binary_mask

def stdDilateColorDist3(I_rgb, M_binary, high_thresh=0.5, s=25):
    rows, cols = M_binary.shape
    lab = color.rgb2lab(I_rgb)
    
    M = M_binary.copy().astype(int)
    
    dx = [1, 1, 0, -1, -1, -1, 0, 1]
    dy = [0, 1, 1, 1, 0, -1, -1, -1]
    
    idx = 1
    
    for r in range(rows):
        for c in range(cols):
            if M[r, c] == 1:
                idx += 1

                region_pixels = []
                M[r, c] = idx
                region_pixels.append((r, c))
                
                ptr = 0
                while ptr < len(region_pixels):
                    curr_r, curr_c = region_pixels[ptr]
                    ptr += 1
                    
                    for k in range(8):
                        nr, nc = curr_r + dx[k], curr_c + dy[k]
                        if 0 <= nr < rows and 0 <= nc < cols:
                            if M[nr, nc] == 1:
                                M[nr, nc] = idx
                                region_pixels.append((nr, nc))
                
                count = len(region_pixels)
                
                if count <= 2:
                    idx -= 1
                    for (pr, pc) in region_pixels:
                        M[pr, pc] = 0
                else:
                    nh = []
                    
                    for k in range(count):
                        curr_r, curr_c = region_pixels[k]
                        for a in range(8):
                            nr, nc = curr_r + dx[a], curr_c + dy[a]
                            if 0 <= nr < rows and 0 <= nc < cols:
                                if M[nr, nc] == 0:
                                    M[nr, nc] = idx 
                                    
                                    x0, x1 = max(0, nr-s), min(rows, nr+s+1)
                                    y0, y1 = max(0, nc-s), min(cols, nc+s+1)
                                    
                                    mask_roi = M[x0:x1, y0:y1]
                                    lab_roi = lab[x0:x1, y0:y1]
                                    target_px = lab[nr, nc]
                                    
                                    dist = find_lda_dist(mask_roi, lab_roi, target_px)
                                    nh.append([dist, nr, nc])

                    if len(nh) > 0:
                        dlen = 0
                        n0 = count / 2 
                        
                        while dlen <= n0:
                            best_idx = -1
                            min_dist = 99999.0
                            
                            for k in range(len(nh)):
                                d = nh[k][0]
                                if d < 255.0 and d < min_dist:
                                    min_dist = d
                                    best_idx = k
                            
                            if best_idx == -1 or min_dist > high_thresh:
                                break
                                
                            dlen += 1
                            dist_val, br, bc = nh[best_idx]
                            nh[best_idx][0] = 255.0 
                            region_pixels.append((br, bc))
                            
                            for a in range(8):
                                nr, nc = br + dx[a], bc + dy[a]
                                if 0 <= nr < rows and 0 <= nc < cols:
                                    if M[nr, nc] == 0:
                                        M[nr, nc] = idx
                                        
                                        x0, x1 = max(0, nr-s), min(rows, nr+s+1)
                                        y0, y1 = max(0, nc-s), min(cols, nc+s+1)
                                        
                                        mask_roi = M[x0:x1, y0:y1]
                                        lab_roi = lab[x0:x1, y0:y1]
                                        target_px = lab[nr, nc]
                                        
                                        new_dist = find_lda_dist(mask_roi, lab_roi, target_px)
                                        nh.append([new_dist, nr, nc])

                    for item in nh:
                        if item[0] < 255.0:
                            M[item[1], item[2]] = 0

                    final_count = len(region_pixels)
                    if final_count < 6:
                        for (pr, pc) in region_pixels:
                            M[pr, pc] = 0

    return (M > 0)

def find_lda_dist(mask_roi, lab_roi, target_color):
    sub_mask = mask_roi[::2, ::2]
    sub_lab  = lab_roi[::2, ::2]
    
    hair_pixels = sub_lab[sub_mask > 0]
    skin_pixels = sub_lab[sub_mask == 0]
    
    if len(hair_pixels) == 0 or len(skin_pixels) == 0:
        return 255.0

    m1 = np.mean(hair_pixels, axis=0)
    m0 = np.mean(skin_pixels, axis=0)

    diff1 = hair_pixels - m1
    s1 = np.dot(diff1.T, diff1)
    
    diff0 = skin_pixels - m0
    s0 = np.dot(diff0.T, diff0)

    try:
        w = np.linalg.solve(s1 + s0, m1 - m0)
    except np.linalg.LinAlgError:
        return 255.0 # Fallback for singular matrix
        
    t0 = np.dot(w, m0)
    t1 = np.dot(w, m1)
    t  = np.dot(w, target_color)
    
    if t0 == t1:
        return 0.0
    
    return (t - t1) / (t0 - t1)

def hair_remov_med(image, mask, s=15):
    result = image.copy()
    m, n = mask.shape
    for i in range(m):
        for j in range(n):
            if not mask[i, j]:
                x0, x1 = max(0, i-s), min(m, i+s+1)
                y0, y1 = max(0, j-s), min(n, j+s+1)

                roi_mask = mask[x0:x1, y0:y1]
                roi_img = image[x0:x1, y0:y1]

                valid_pixels = roi_img[roi_mask]
                if len(valid_pixels) > 0:
                    result[i, j] = np.median(valid_pixels, axis=0)
    return result

def process_image(image):
    I_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    M = None
    print("Processing filter responses...")
    M = np.zeros_like(gray, dtype=float)
    angles = np.linspace(-90, 90, 12, endpoint=False)

    total_time = 0

    for deg in angles:
        start_time = time.time()
        ji = ncu_line_match_filter_direc(gray, deg)
        end_time = time.time()
        angle_time = end_time - start_time
        total_time += angle_time

        M = np.maximum(M, ji)

        print(f"Angle {deg:.1f}°: {angle_time:.4f}s (total: {total_time:.4f}s)")
        print(f"\nTotal filter processing time: {total_time:.4f} seconds")

    print("Applying stdDilateDarkest...")
    start_time_dilate = time.time()
    if np.max(M) > 0:
        M_scaled = (M / np.max(M)) * 255.0
    else:
        M_scaled = M

    binary_mask = stdDilateDarkest(M_scaled)
    end_time_dilate = time.time()
    print(f"stdDilateDarkest processing time: {end_time_dilate - start_time_dilate:.4f} seconds")
    print("Applying stdDilateColorDist3...")
    start_time_lda = time.time()
    K_binary = stdDilateColorDist3(I_rgb, binary_mask, 0.5, 25)
    end_time_lda = time.time()
    print(f"stdDilateColorDist3 processing time: {end_time_lda - start_time_lda:.4f} seconds")

    result = hair_remov_med(I_rgb, ~K_binary, 15)
    hair_mask = (K_binary * 255).astype(np.uint8)

    return cv2.cvtColor(result.astype(np.uint8), cv2.COLOR_RGB2BGR), hair_mask

if __name__ == "__main__":
    input_image_path = "2006-124-2.png"
    image = cv2.imread(input_image_path)

    if image is None:
        print(f"Error: Could not load image from {input_image_path}")
        exit(1)

    start_time_overall = time.time()
    result_image, hair_mask = process_image(image)#, use_cache=True
    end_time_overall = time.time()

    cv2.imwrite("result.png", result_image)
    cv2.imwrite("hair_mask.png", hair_mask)

    print(f"\nOverall processing time: {end_time_overall - start_time_overall:.4f} seconds")
    print(f"Processing complete. Results saved to result.png and hair_mask.png")


#if you want to experment with std* methods
# def process_image(image, use_cache=True, cache_file="cached_filter_response.npy"):
#     I_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
#     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
#     cache_file_path = cache_file
#     M = None

#     if use_cache and os.path.exists(cache_file_path):
#         print("Loading cached filter responses...")
#         start_time = time.time()
#         M = np.load(cache_file_path)
#         end_time = time.time()
#         print(f"Loaded cache in {end_time - start_time:.4f} seconds")
#     else:
#         print("Processing filter responses...")
#         M = np.zeros_like(gray, dtype=float)
#         angles = np.linspace(-90, 90, 12, endpoint=False)

#         total_time = 0

#         for deg in angles:
#             start_time = time.time()
#             ji = ncu_line_match_filter_direc(gray, deg)
#             end_time = time.time()
#             angle_time = end_time - start_time
#             total_time += angle_time

#             M = np.maximum(M, ji)

#             print(f"Angle {deg:.1f}°: {angle_time:.4f}s (total: {total_time:.4f}s)")

#         print(f"\nTotal filter processing time: {total_time:.4f} seconds")

#         # Save cache for future use
#         if use_cache:
#             print("Saving filter responses to cache...")
#             start_time = time.time()
#             np.save(cache_file_path, M)
#             end_time = time.time()
#             print(f"Saved cache in {end_time - start_time:.4f} seconds")

#     print("Applying stdDilateDarkest...")
#     start_time_dilate = time.time()
#     if np.max(M) > 0:
#         M_scaled = (M / np.max(M)) * 255.0
#     else:
#         M_scaled = M

#     binary_mask = stdDilateDarkest(M_scaled)
#     end_time_dilate = time.time()
#     print(f"stdDilateDarkest processing time: {end_time_dilate - start_time_dilate:.4f} seconds")

#     print("Applying stdDilateColorDist3...")
#     start_time_lda = time.time()
#     K_binary = stdDilateColorDist3(I_rgb, binary_mask, 0.5, 25)
#     end_time_lda = time.time()
#     print(f"stdDilateColorDist3_v2 processing time: {end_time_lda - start_time_lda:.4f} seconds")

#     result = hair_remov_med(I_rgb, ~K_binary, 15)
#     hair_mask = (K_binary * 255).astype(np.uint8)

#     return cv2.cvtColor(result.astype(np.uint8), cv2.COLOR_RGB2BGR), hair_mask