levitan_cv/figures_and_colors/main.py

import matplotlib.pyplot as plt
import numpy as np
from skimage.measure import label, regionprops
from skimage.color import rgb2hsv
from collections import defaultdict


def get_shapes(regions):
    shapes = defaultdict(lambda: 0)
    for _, region in enumerate(regions):
        eccent = region.eccentricity
        if eccent == 0:
            if region.image.size == region.area:
                key = "rect"
            else:
                key = "ball"
        else:
            key = "rect"
        shapes[key] += 1
    return shapes


file = "./balls_and_rects.png"
img = plt.imread(file)
binary = np.mean(img, 2)
binary[binary > 0] = 1
labeled = label(binary)
print(f"Total f={labeled.max()}")
hsv_image = rgb2hsv(img)

c = 0
prev = 0
for index, i in enumerate(np.linspace(0, 1, 10)):
    tmp_img = np.copy(hsv_image)
    tmp_img[tmp_img[:, :, 0] < prev] = 0
    tmp_img[tmp_img[:, :, 0] > i] = 0
    tmp_img[tmp_img[:, :, 0] > 0] = 1
    tmp_img = np.mean(tmp_img, 2)
    labeled = label(tmp_img)
    if labeled.max() > 0:
        c += 1
        print(f"Color {c}:")
        regs = regionprops(labeled)
        shapes = get_shapes(regs)
        for cur_key in shapes:
            print(f"\t{cur_key}s: {shapes[cur_key]}")

    prev = i