shape_clustering.py

import cv2
import numpy as np
import os
import random
import math
from multiprocessing import Pool, cpu_count
from functools import partial

#--------AXIS AND CENTER FINDING---------------------------#
def rotating_calipers(points: np.ndarray) -> tuple:
    """
    Finds the length, width, and center of a cluster using rotating calipers.

    Args:
        points (np.ndarray): Points in the cluster, represented as an array of shape (N, 1, 2).

    Returns:
        tuple:
            - Length information (start point, end point, length, angle).
            - Width information (start point, end point, width).
            - Center of the cluster.
    """
    #---------CENTER CLUSTER---------#
    # calculate center of cluster based on the pixels from the clusters
    center_x = np.mean(points[:, 0, 0])
    center_y = np.mean(points[:, 0, 1])
    #--center cluster:
    center_cluster = (center_x, center_y)

    # ---------CLUSTER POINTS---------#
    # Extract cluster points into a usable format (N x 2 array)
    if points.ndim == 3:  # If the input is in N x 1 x 2 format
        cluster_points = points[:, 0, :]
    elif points.ndim == 2:  # If the input is already in N x 2 format
        cluster_points = points
    else:
        raise ValueError("Invalid shape for points. Expected (N x 1 x 2) or (N x 2).")



    #----------LENGTH AND ANGLE---------#
    #initilaize the points from the cluster:
    hull_points = cv2.convexHull(points, returnPoints=True)
    hull_points = hull_points.reshape(-1, 2)
    num_points = len(hull_points)

    if num_points < 2:
        return ((None, None, 0, 0), (None, None, 0), center_cluster)
    # initialize to find the length
    max_length = 0
    length_start_point = None
    length_end_point = None
    for i in range(num_points):
        for j in range(i + 1, num_points):
            pt1 = hull_points[i]
            pt2 = hull_points[j]
            distance = np.linalg.norm(pt1 - pt2)
            if distance > max_length:
                max_length = distance
                length_start_point = tuple(pt1)
                length_end_point = tuple(pt2)

    #--length points:
    dx_l = length_end_point[0] - length_start_point[0]
    dy_l = length_end_point[1] - length_start_point[1]
    #--angle:
    angle = abs(math.degrees(math.atan2(dy_l, dx_l)))

    #---------WIDTH---------#
    # Initialization of the axis
    perp_dx = -dy_l
    perp_dy = dx_l
    perp_length = np.sqrt(perp_dx ** 2 + perp_dy ** 2)
    perp_unit_dx = perp_dx / perp_length
    perp_unit_dy = perp_dy / perp_length

    # initialize the points for the normal to the length: the width
    min_projection = float('inf')
    max_projection = float('-inf')
    width_start_point = None
    width_end_point = None
    # Initialize variables for minimal normal: the min width
    min_width = float('inf')

    for pt in hull_points:
        projection = (pt[0] - center_x) * perp_unit_dx + \
                     (pt[1] - center_y) * perp_unit_dy

        normal_length = np.linalg.norm((pt[0] - center_x, pt[1] - center_y))

        if normal_length < min_width:
            min_width = normal_length

        if projection < min_projection:
            min_projection = projection
            width_start_point = (center_x + min_projection * perp_unit_dx,
                                 center_y + min_projection * perp_unit_dy)
        if projection > max_projection:
            max_projection = projection
            width_end_point = (center_x + max_projection * perp_unit_dx,
                               center_y + max_projection * perp_unit_dy)
    #--width distance:
    width = max_projection - min_projection
    #mean_distance_width= abs((width + min_width) / 2)
    mean_distance_width= abs((max_projection + min_width) / 2)

    dx_w = width_end_point[0] - width_start_point[0]
    dy_w = width_end_point[1] - width_start_point[1]

    # Calculate the midpoint
    midpoint_width_x = (width_start_point[0] + width_end_point[0]) / 2
    midpoint_width_y = (width_start_point[1] + width_end_point[1]) / 2
    midpoint_width = (midpoint_width_x, midpoint_width_y)

    #---------FINDING MIN DISTANCE IN LENGTH FROM CENTER---------#
    #Calculate distance between center cluster and start and end points in length:
    dist_center_cluster_extreme_1 = math.sqrt(
        (center_cluster[0] - length_end_point[0]) ** 2 + (center_cluster[1] - length_end_point[1]) ** 2)
    dist_center_cluster_extreme_2 = math.sqrt(
        (center_cluster[0] - length_start_point[0]) ** 2 + (center_cluster[1] - length_start_point[1]) ** 2)
    shorter_distance_length = min(dist_center_cluster_extreme_1, dist_center_cluster_extreme_2)
    length = math.sqrt(
        (length_start_point[0] - length_end_point[0]) ** 2 + (length_start_point[1] - length_end_point[1]) ** 2)

    #--print
    #print(f"Center cluster: {center_point}, Center point length: {center_length_point}")
    #print(f"shortest dist: {2*shorter_distance}")
    #print(f"Length: {max_length}, Width: {width}, Angle: {angle}")


    #---------RETURN---------#
    #2 * shorter_distance_length     instead of length
    #mean_distance_width
    return (
        (length_start_point, length_end_point, max_length, angle),
        (width_start_point, width_end_point, width, cluster_points),
        center_cluster
    )


def apply_mosaic_ratio(image, block_size: int, ratio: float) -> list[list[int]]:
    """
    Divides an image into blocks and determines whether each block is selected
    based on the ratio of white pixels.

    Args:
        image (np.ndarray): Input grayscale image.
        block_size (int): Size of the blocks for mosaic processing.
        ratio (float): Minimum ratio of white pixels required for a block to be selected.

    Returns:
        list[list[int]]: A grid representing the image with 1s for selected blocks
                         and 0s for unselected blocks.
    """
    if block_size < 1:
        block_size = 1

    grid: list[list[int]] = []

    (h, w) = image.shape[:2]
    for y in range(0, h, block_size):
        row = []
        for x in range(0, w, block_size):
            block = image[y:y + block_size, x:x + block_size]
            white_pixels = np.sum(block >= 255)  # get white pixel
            total_pixels = block.size
            white_ratio = white_pixels / total_pixels

            if white_ratio >= ratio:
                row.append(1)
            else:
                row.append(0)

        grid.append(row)

    return grid

def find_cluster(grid: list[list[int]], pos: tuple[int, int], value: int) -> tuple[list[list[int]], list[list[int]]]:
    """
    Finds a connected cluster of blocks with a specific value starting from a given position.

    Args:
        grid (list[list[int]]): A 2D grid representing the image.
        pos (tuple[int, int]): Starting position (x, y) in the grid.
        value (int): The value of the cluster to find.

    Returns:
        tuple[list[list[int]], list[list[int]]]:
            - Modified grid with the cluster removed.
            - Grid containing only the found cluster.
    """
    stack = [pos]
    cluster_grid = [[0 for _ in range(len(grid[0]))] for _ in range(len(grid))]

    while stack:
        x, y = stack.pop()
        if grid[x][y] == value:
            grid[x][y] = 0
            cluster_grid[x][y] = value

            for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                nx, ny = x + dx, y + dy
                if 0 <= nx < len(grid) and 0 <= ny < len(grid[0]) and grid[nx][ny] == value:
                    stack.append((nx, ny))

    return grid, cluster_grid

def grid_clustering(grid: list[list[int]], value: int = 1) -> list[list[list[int]]]:
    """
    Identifies clusters of blocks with a specific value in a grid.

    Args:
        grid (list[list[int]]): A 2D grid representing the image.
        value (int): The value of the clusters to find (default is 1).

    Returns:
        list[list[list[int]]]: A list of grids, each representing a separate cluster.
    """
    clusters = []

    for i in range(len(grid)):
        for j in range(len(grid[0])):
            if grid[i][j] == value:
                grid, cluster = find_cluster(grid, (i, j), value)
                clusters.append(cluster)

    return clusters


def image_clustering(image, clustered_grid: list[list[list[int]]], block_size: int) -> list[np.ndarray]:
    """
    Generates binary images for each cluster based on the clustered grid.

    Args:
        image (np.ndarray): Input grayscale image.
        clustered_grid (list[list[list[int]]]): List of clustered grids.
        block_size (int): Size of the blocks in the grid.

    Returns:
        list[np.ndarray]: A list of binary images, one for each cluster.
    """
    (h, w) = image.shape[:2]
    cluster_images = []

    for cluster in clustered_grid:
        cluster_image = np.zeros((h, w), dtype=np.uint8)

        for i in range(len(cluster)):
            for j in range(len(cluster[0])):
                if cluster[i][j] == 1:
                    y_start, y_end = i * block_size, min((i + 1) * block_size, h)
                    x_start, x_end = j * block_size, min((j + 1) * block_size, w)
                    cluster_image[y_start:y_end, x_start:x_end] = 255

        cluster_images.append(cluster_image)

    return cluster_images


def edge_expand(image, clusters: list[list[list[int]]], block_size: int, ratio: float, distance: int) -> list[
    list[list[int]]]:
    """
    Expands the edges of clusters based on the ratio of white pixels in neighboring blocks.

    Args:
        image (np.ndarray): Input grayscale image.
        clusters (list[list[list[int]]]): List of clustered grids.
        block_size (int): Size of the blocks in the grid.
        ratio (float): Minimum ratio of white pixels required for expansion.
        distance (int): Maximum distance for edge expansion.

    Returns:
        list[list[list[int]]]: List of expanded clusters.
    """
    (h, w) = image.shape[:2]
    expanded_clusters = []

    for cluster in clusters:
        expanded_cluster = [row[:] for row in cluster]  # Deep copy

        for i in range(len(cluster)):
            for j in range(len(cluster[0])):
                if cluster[i][j] == 1:
                    for dx in range(-distance, distance + 1):
                        for dy in range(-distance, distance + 1):
                            ni, nj = i + dx, j + dy
                            if 0 <= ni < len(cluster) and 0 <= nj < len(cluster[0]) and expanded_cluster[ni][nj] == 0:
                                y_start, y_end = ni * block_size, min((ni + 1) * block_size, h)
                                x_start, x_end = nj * block_size, min((nj + 1) * block_size, w)
                                block = image[y_start:y_end, x_start:x_end]
                                white_pixels = np.sum(block >= 255)
                                total_pixels = block.size
                                white_ratio = white_pixels / total_pixels

                                if white_ratio >= ratio:
                                    expanded_cluster[ni][nj] = 1

        expanded_clusters.append(expanded_cluster)

    return expanded_clusters


def combine_clusters_in_distance(clusters: list[list[list[int]]], dist: int) -> list[list[list[int]]]:
    """
    Combines clusters that are within a specified distance of each other.

    Args:
        clusters (list[list[list[int]]]): List of clustered grids.
        dist (int): Maximum distance between clusters to combine them.

    Returns:
        list[list[list[int]]]: List of combined clusters.
    """
    def clusters_are_close(cluster1, cluster2, dist):
        for i in range(len(cluster1)):
            for j in range(len(cluster1[0])):
                if cluster1[i][j] == 1:
                    for dx in range(-dist, dist + 1):
                        for dy in range(-dist, dist + 1):
                            ni, nj = i + dx, j + dy
                            if 0 <= ni < len(cluster2) and 0 <= nj < len(cluster2[0]) and cluster2[ni][nj] == 1:
                                return True
        return False

    combined = []
    while clusters:
        base_cluster = clusters.pop(0)
        merge_indices = []
        for idx, other_cluster in enumerate(clusters):
            if clusters_are_close(base_cluster, other_cluster, dist):
                merge_indices.append(idx)

        for idx in sorted(merge_indices, reverse=True):
            other_cluster = clusters.pop(idx)
            for i in range(len(base_cluster)):
                for j in range(len(base_cluster[0])):
                    if other_cluster[i][j] == 1:
                        base_cluster[i][j] = 1

        combined.append(base_cluster)

    return combined


def filter_biggest_clusters(clusters: list[list[list[int]]], num: int) -> list[list[list[int]]]:
    """
    Filters and returns the largest clusters by size.

    Args:
        clusters (list[list[list[int]]]): List of clustered grids.
        num (int): Number of largest clusters to retain.

    Returns:
        list[list[list[int]]]: List of the largest clusters.
    """
    #-------CLUSTER SIZE---------
    # calculate the size of each cluster
    cluster_sizes = [(cluster, sum(sum(row) for row in cluster))
                     for cluster in clusters]
    # sort the clusters by size
    sorted_clusters = sorted(cluster_sizes, key=lambda x: x[1], reverse=True)
    # Find the size of the biggest cluster
    max_size = max(size for _, size in cluster_sizes) if cluster_sizes else 1
    print("Max cluster size=", max_size)

    #-------CLUSTER WEIGHT  #Verify this method:
    # Instead of taking the 3 biggest clusters we do percentage and we take the 5 or 2 biggest clusters
    #PLEASE double check the calculation.
    # Calculate the weight for each cluster as a percentage of the biggest cluster
    weighted_clusters = [(cluster, (size / max_size) * 100) for cluster, size in cluster_sizes]

    # Sort clusters by weight in descending order
    sorted_clusters = sorted(weighted_clusters, key=lambda x: x[1], reverse=True)

    #-------RETURN
    # return the num-th biggest cluster
    return [cluster for cluster, _ in sorted_clusters[:num]]


def merge_clusters(clusters: list[list[list[int]]]) -> list[list[int]]:
    """
    Merges all clusters into a single grid.

    Args:
        clusters (list[list[list[int]]]): List of clustered grids.

    Returns:
        list[list[int]]: A single grid containing all merged clusters.
    """
    if not clusters:
        return []

    rows, cols = len(clusters[0]), len(clusters[0][0])

    combined_cluster = [[0] * cols for _ in range(rows)]

    for cluster in clusters:
        for i in range(rows):
            for j in range(cols):
                if cluster[i][j] == 1:
                    combined_cluster[i][j] = 1

    return combined_cluster

def eliminate_clusters(clusters: list[list[list[int]]], smallest_cluster: int) -> list[list[list[int]]]:
    """
    Removes clusters smaller than a specified size.

    Args:
        clusters (list[list[list[int]]]): List of clustered grids.
        smallest_cluster (int): Minimum size of clusters to retain.

    Returns:
        list[list[list[int]]]: List of clusters that meet the size requirement.
    """
    filtered_clusters = []
    for cluster in clusters:
        cluster_size = sum(sum(row) for row in cluster)
        if cluster_size >= smallest_cluster:
            filtered_clusters.append(cluster)
    return filtered_clusters


#TODO modify
def do_clustering(image, block_size: int, ratio: float, new_ratio: float, n: int, dist: int, smallest_cluster: int,percent_ac):
    """
    Performs the full clustering process and identifies shapes within clusters.

    Args:
        image (np.ndarray): Input grayscale image.
        block_size (int): Size of the blocks for mosaic processing.
        ratio (float): Minimum ratio of white pixels required for a block to be selected.
        new_ratio (float): Ratio used during edge expansion.
        n (int): Number of blocks for edge expansion.
        dist (int): Maximum distance between clusters to combine them.
        smallest_cluster (int): Minimum size of clusters to retain.
        percent_ac (int): Percentage of largest clusters to retain.

    Returns:
        tuple:
            - cluster_images (list[np.ndarray]): Binary images for each cluster.
            - cluster_dimensions (list): Dimensions of each cluster (length, width, center).
            - best_shape_masks (list): Best matching shapes for each cluster.
            - center_points (list[tuple]): Center points for each cluster.
    """
    grid = apply_mosaic_ratio(image, block_size, ratio)
    # TODO: be replaced
    clusters = grid_clustering(grid)
    expanded_clusters = edge_expand(image, clusters, block_size, new_ratio, n)
    combined_clusters_list = combine_clusters_in_distance(expanded_clusters, dist)
    # TODO: new function
    # clusters = dbscan_clustering(grid, block_size, eps=15, min_samples=5)
    eliminated_clusters = eliminate_clusters(
        combined_clusters_list, smallest_cluster)
    # we take here the percent_ac biggest clusters
    filtered_clusters = filter_biggest_clusters(eliminated_clusters, percent_ac)
    big_cluster = merge_clusters(filtered_clusters)

    # split the real image
    cluster_images = image_clustering(image, [big_cluster], block_size)

    # find the longest dist between 2 arbitrary points
    cluster_dimensions = []
    center_points = []  # List to hold center points
    for cluster_image in cluster_images:
        points = cv2.findNonZero(cluster_image)
        if points is not None and len(points) >= 3:
            length_info, width_info, center_point = rotating_calipers(points)
            cluster_dimensions.append((length_info, width_info, center_point))
            center_points.append(center_point)  # Store the center point
        else:
            cluster_dimensions.append(((None, None, 0, 0), (None, None, 0), (0, 0)))
            center_points.append((0, 0))  # If no points, default to (0, 0)

    return cluster_images, cluster_dimensions, center_points