Cviceni 7 UZO

segmentation.py

import cv2 as cv
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt

def filter_out_green_color(image: np.array) -> np.array:
    R = np.float32(image[:, :, 0])
    G = np.float32(image[:, :, 1])
    B = np.float32(image[:, :, 2])

    g = 255 - ((G * 255) / (R + G + B + 1e-6))
    threshold = 150

    seg = g.copy()
    seg = np.where((seg < threshold), 0, 1).astype(np.uint8)
    return seg


def get_neighbours(mask:np.array, row:int, col:int):
    neightbours = []
    if (row-1 >= 0):
        if col-1 >=0:
            neightbours.append(mask[row-1, col-1])
        neightbours.append(mask[row-1, col])
        if col+1 < mask.shape[1]:
            neightbours.append(mask[row-1, col+1])
    if (col-1>=0):
        neightbours.append(mask[row, col-1])
    return sorted(neightbours)

def test_get_neighbours():
    test_array = np.array([[0, 2, 3], [4, 5, 6], [7, 8, 9]])
    result = get_neighbours(test_array, 1, 1) 
    assert result == [1, 2, 3, 4]

def coloring_segments(image:np.array):
    counter = 1
    mask = np.zeros_like(image)
    substitions = {}
    for row in range(image.shape[0]):
        for col in range(image.shape[1]):
            if image[row, col] > 0:
                neighbours = get_neighbours(mask, row, col) 
                non_zero_neighbours = [x for x in neighbours if x > 0]
                lowest = min(non_zero_neighbours, default=0)
                if lowest == 0:
                    mask[row, col] = counter
                    counter += 1
                else:
                    mask[row, col] = lowest
                    for x in non_zero_neighbours:
                        if x != lowest:
                            substitions[x] = lowest
    def resolve_substitutions(substitions):
        for key in list(substitions.keys()):
            while substitions[key] in substitions:
                substitions[key] = substitions[substitions[key]]
        return substitions

    substitions=resolve_substitutions(substitions)

    for row in range(image.shape[0]):
        for col in range(image.shape[1]):
            if mask[row, col] in substitions:
                mask[row, col] = substitions[mask[row, col]]
    return mask

def display(image:np.array):
    plt.imshow(image)
    plt.show()
    cv.waitKey(0)
    cv.destroyAllWindows()

def plot_mask_different_color_for_every_label(mask: np.array):
    unique_labels = np.unique(mask)
    color_map = plt.cm.get_cmap('hsv', len(unique_labels))

    colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
    for i, label in enumerate(unique_labels):
        if label > 0:  # nechci pozadi
            colored_mask[mask == label] = (np.array(color_map(i)[:3]) * 255).astype(np.uint8)

    plt.imshow(colored_mask)
    plt.title("Mask with Different Colors for Each Label")
    plt.axis('off')
    plt.show()

def test_coloring():
    image = cv.imread(Path("data/cv07_barveni.bmp").as_posix())
    binary_image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
    mask = coloring_segments(binary_image)
    plot_mask_different_color_for_every_label(mask)

def evaluate_mask(mask: np.array):
    """Vyhodnoceni minci"""
    unique_labels, counts = np.unique(mask, return_counts=True)
    
    labels = {}
    for label, count in zip(unique_labels, counts):
        if label == 0:  # nechci pozadi
            continue
        rows, cols = np.where(mask == label)
        centre_of_gravity = (int(np.floor(np.mean(rows))), int(np.floor(np.mean(cols))))
        labels[label] = {
            "count": count,
            "center": centre_of_gravity,
            "type": None  
        }
    
    koruny = 0
    peti_koruny = 0
    for label, data in labels.items():
        if data["count"] > 4000:  
            peti_koruny += 1
            data["type"] = 5
        else:  
            koruny += 1
            data["type"] = 1
    
    for label, data in labels.items():
        center = data["center"]
        typ = data["type"]
        print(f"Na souradnici {center} se nachazi {typ} Kč")
    
    print("Celkova hodnota: " + str(koruny * 1 + peti_koruny * 5) + " Kč")
    return labels

    

def draw_bounding_boxes(image: np.array, labels: dict):
    output_image = image.copy()

    for label, data in labels.items():
        # Get the center and bounding box coordinates
        rows, cols = np.where(mask == label)
        top_left = (np.min(cols), np.min(rows))
        bottom_right = (np.max(cols), np.max(rows))
        center = (int(data["center"][1]), int(data["center"][0]))  # Flip (row, col) to (x, y)

        # Draw the bounding box
        cv.rectangle(output_image, top_left, bottom_right, (255, 0, 0), 2)

        # Add label type and center as text
        text = f"{data['type']} Kc"
        cv.putText(output_image, text, center, cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

    # Display the image with bounding boxes
    plt.imshow(cv.cvtColor(output_image, cv.COLOR_BGR2RGB))
    plt.title("Bounding Boxes Around Labels")
    plt.axis('off')
    plt.show()

def calculate_centre_of_gravity(mask:np.array):
    unique_labels = np.unique(mask)
    for label in unique_labels:
        if label == 0:
            continue
        rows, cols = np.where(mask == label)
        centre_of_gravity = (np.mean(rows), np.mean(cols))
        print(f"Centre of gravity for label {label}: {centre_of_gravity}")


if __name__ == "__main__":
    
    image = cv.imread(Path("data/cv07_segmentace.bmp").as_posix())#, cv.IMREAD_COLOR_RGB)
    image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
    # plot_rgb_histogram(image)
    # # image = image.astype(np.float32)
    filtered = filter_out_green_color(image)
    display(filtered)
    # filtered_gray = cv.cvtColor(filtered, cv.COLOR_BGR2GRAY)

    mask = coloring_segments(filtered)
    plot_mask_different_color_for_every_label(mask)
    labels =evaluate_mask(mask)
    draw_bounding_boxes(image, labels)