src/util/material_property_eval.py

import cv2
import torch
import matplotlib
import tkinter as tk
import numpy as np
import matplotlib.pyplot as plt

from tkinter import filedialog, messagebox
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.widgets import RectangleSelector

crop_coords = None
pixel_size_um = None


def calculate_pixel_size(image, image_size_um=2.0):
    """Calculate the pixel size in micrometers based on the image dimensions."""
    image_width_pixels = image.shape[1]
    pixel_size_um = image_size_um / image_width_pixels
    return pixel_size_um


def process_image(data, height, width, pixel_size_um: float = 2.0) -> dict:
    """
    Process a 2D material with shape: [H, W]

    Returns
    ---
    - Dictionary of material properties
    """

    # Calculate pixel size
    pixel_size_um = calculate_pixel_size(data)

    # Normalize data
    normalized_data = cv2.normalize(data, None, 0, 255, cv2.NORM_MINMAX).astype(
        np.uint8
    )

    # Ensure the image is in float format for processing
    img = data.astype(np.float32)
    h, w = data.shape
    # Flattening: Fit a 1st-order polynomial (linear) to each row and subtract
    flattened_img = np.zeros_like(img, dtype=np.float32)

    for i in range(h):
        row = img[i, :]
        x = np.arange(w)

        # Fit a linear polynomial (1st order)
        p = np.polyfit(x, row, 1)  # Fit line y = ax + b

        # Compute background
        background = np.polyval(p, x)

        # Subtract the background from the row
        flattened_img[i, :] = row - background

        # flattened_img[i, :] -= np.min(flattened_img[i, :])  # Shift to keep variations

    # Normalize the image for display
    flattened_img = cv2.normalize(flattened_img, None, 0, 255, cv2.NORM_MINMAX)
    flattened_img = np.uint8(flattened_img)

    # Calculate the average current
    average_current = np.mean(data) / 1e-9

    # Define the coverage threshold
    threshold = 120 * 1e-12

    # Create a binary mask for coverage (1 = covered, 0 = uncovered)
    coverage_mask = (data > threshold).astype(int)

    # Calculate coverage percentage
    coverage_percentage = (np.sum(coverage_mask) / coverage_mask.size) * 100

    # Create a color map: Light blue for coverage, light grey for uncovered
    colors = ["black", "mistyrose"]
    cmap = LinearSegmentedColormap.from_list("coverage_map", colors, N=2)

    # Plot the coverage map
    x = np.linspace(0, 2, data.shape[1])  # X-axis (microns)
    y = np.linspace(0, 2, data.shape[0])  # Y-axis (microns)
    X, Y = np.meshgrid(x, y)

    # Apply Gamma Correction
    gamma = 0.7
    gamma_corrected_img = np.power(flattened_img / 255.0, gamma) * 255.0
    gamma_corrected_img = gamma_corrected_img.astype(np.uint8)

    # Apply Bilateral Filter
    bilateral_filtered = cv2.bilateralFilter(gamma_corrected_img, 9, 30, 75)

    # Apply Gaussian Blur
    blurred = cv2.GaussianBlur(bilateral_filtered, (3, 3), 3)

    # Adaptive Thresholding
    thresholded = cv2.adaptiveThreshold(
        blurred, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 7, 1
    )

    # Contour Detection
    contours, _ = cv2.findContours(
        thresholded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
    )
    output_image = cv2.cvtColor(blurred, cv2.COLOR_GRAY2BGR)

    curves_length = 0
    circular_shapes_area = 0
    extended_shapes_area = 0
    circular_shapes = []
    curved_lines = []
    extended_shapes = []

    for contour in contours:
        if cv2.contourArea(contour) > 0:
            perimeter = cv2.arcLength(contour, closed=True)
            area = cv2.contourArea(contour)
            circularity = 4 * np.pi * area / (perimeter**2) if perimeter != 0 else 0

            # Create a mask for the contour
            mask = np.zeros_like(thresholded, dtype=np.uint8)
            cv2.drawContours(mask, [contour], -1, 255, -1)
            # Check if the contour is inside a covered area (valid contour)
            if not np.any(coverage_mask[mask == 255]):
                continue  # Skip this contour if it falls in an uncovered region

            # Calculate the total number of pixels inside the contour
            total_pixels = cv2.countNonZero(mask)

            # Calculate the number of white pixels inside the contour
            white_pixels = cv2.countNonZero(cv2.bitwise_and(mask, thresholded))
            # Calculate the number of dark pixels inside the contour (where the thresholded image is 0)
            dark_pixels = total_pixels - cv2.countNonZero(
                cv2.bitwise_and(mask, thresholded)
            )

            x, y, w, h = cv2.boundingRect(contour)
            aspect_ratio = float(w) / h

            if circularity > 0.2 and area <= 1000:
                # Check for white pixel coverage only for circular shapes
                # if dark_pixels / total_pixels >= 0.8:
                # continue  # Skip this contour if 80% or more of the pixels are white

                # Process circular shapes
                circular_shapes.append(contour)
                circular_shapes_area += area
                cv2.drawContours(output_image, [contour], -1, (0, 255, 255), 1)
            elif aspect_ratio > 2 and area <= 1000:
                # Check for white pixel coverage only for extended shapes
                # if total_pixels > 0 and (white_pixels / total_pixels) >= 0.8:
                # continue  # Skip this contour if 80% or more of the pixels are white

                # Process extended shapes
                extended_shapes.append(contour)
                extended_shapes_area += area
                cv2.drawContours(output_image, [contour], -1, (0, 255, 255), 1)

            else:
                # First, calculate the area of the contour
                contour_area = cv2.contourArea(contour)

                if contour_area <= 250:  # Threshold for small areas
                    extended_shapes.append(contour)
                    extended_shapes_area += contour_area
                    cv2.drawContours(output_image, [contour], -1, (0, 255, 255), 1)
                else:
                    # Process curved lines without the white pixel check
                    curved_lines.append(contour)
                    curve_length = cv2.arcLength(contour, closed=True)
                    curves_length += curve_length
                    cv2.drawContours(output_image, [contour], -1, (0, 255, 0), 1)

    # Convert areas and lengths to micrometers
    circular_shapes_area_um2 = circular_shapes_area * (pixel_size_um**2)
    extended_shapes_area_um2 = extended_shapes_area * (pixel_size_um**2)
    curves_length_um = curves_length * pixel_size_um

    Total_Defect_Area = circular_shapes_area_um2 + extended_shapes_area_um2
    Total_Defect_Percentage = 100 * Total_Defect_Area / (height * width)

    results = {
        "average_current": average_current,
        "coverage_percentage": coverage_percentage,
        "circular_shapes_area": circular_shapes_area_um2,
        "extended_shapes_area": extended_shapes_area_um2,
        "curves_length": curves_length_um,
        "total_defect_area": Total_Defect_Area,
        "total_defect_percentage": Total_Defect_Percentage,
        "number_ext_shapes": len(extended_shapes),
    }

    return results


def calculate_roughness(height_data: np.ndarray) -> dict:

    if isinstance(height_data, torch.Tensor):
        height_data = height_data.cpu().numpy()

    # Zero-centering the data
    height_data -= np.mean(height_data)

    # RMS roughness in nm
    Sq = np.sqrt(np.mean(height_data**2)) * 1e9

    # Mean roughness in nm
    Sa = np.mean(np.abs(height_data)) * 1e9

    return {"RMS roughness (Sq)": Sq, "Mean roughness (Sa)": Sa}


def pcnt_abs_diff_surface_roughness(x1: torch.Tensor, x2: torch.Tensor) -> dict:

    x1 = x1.cpu().numpy()
    x2 = x2.cpu().numpy()

    char_1 = calculate_roughness(x1)
    char_2 = calculate_roughness(x2)

    diffs = {}
    for k in char_1:
        v1, v2 = char_1[k], char_2[k]
        err = abs(v1 - v2)
        diffs[k] = err

    return diffs


def calculate_abs_diff_between_samples(
    x1: torch.Tensor, x2: torch.Tensor, image_size_um: float
) -> dict:

    x1 = x1.cpu().numpy()
    x2 = x2.cpu().numpy()
    char_1 = process_image(x1, x1.shape[0], x1.shape[1], image_size_um)
    char_2 = process_image(x2, x2.shape[1], x2.shape[1], image_size_um)

    diffs = {}
    for k in char_1:
        v1, v2 = char_1[k], char_2[k]
        err = abs(v1 - v2)
        diffs[k] = err

    return diffs


if __name__ == "__main__":

    dummy = np.random.random((128, 128))
    res = process_image(dummy, 128, 128)

    breakpoint()