yiking/main.py

import math
from dataclasses import dataclass

import numpy as np
import cv2
import tkinter
import os
import sys
from pathlib import Path

dir_path = Path(".").absolute()
TYPE_1 = "_________"
TYPE_2 = "___   ___"


@dataclass
class Object:
    x: int
    y: int
    diametre: int


# 1. Acquisition de l'image
src = dir_path.joinpath('tests/images/balls-full-small.jpg')
raw_image = cv2.imread(str(src))

# 2. Boxing des objets via opencv
gray = cv2.cvtColor(raw_image, cv2.COLOR_BGR2GRAY)
blurred = cv2.medianBlur(gray, 25)
minDist = 100
param1 = 30  # 500
param2 = 25  # 200 #smaller value-> more false circles
minRadius = 5
maxRadius = 1000  # 10
circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist, param1=param1, param2=param2, minRadius=minRadius,
                           maxRadius=maxRadius)

min_diameter = 9999
cochonnet = None
boules = []
if circles is not None:
    circles = np.uint16(np.around(circles))
    for i in circles[0, :]:
        boule = Object(x=int(i[0]), y=int(i[1]), diametre=int(i[2]))
        # cv2.circle(img, (boule.x, boule.y), boule.diametre, (0, 255, 0), 2)

        # 3. Détection de la box la plus petite : cochonnet
        if boule.diametre < min_diameter:
            min_diameter = boule.diametre
            if cochonnet != None:
                boules.append(cochonnet)
            cochonnet = boule
        else:
            boules.append(boule)

img_check_shapes = raw_image.copy()
for boule in boules:
    cv2.circle(img_check_shapes, (boule.x, boule.y), boule.diametre, (0, 255, 0), 2)
cv2.circle(img_check_shapes, (cochonnet.x, cochonnet.y), cochonnet.diametre, (255, 255, 0), -1)
# cv2.imshow('img_check_shapes', img_check_shapes)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

# 4. Regroupement en liste de boules 1 ou 2 selon la couleur principale de chaque box restante

hsv = cv2.cvtColor(raw_image, cv2.COLOR_BGR2HSV)
(h, s, v) = cv2.split(hsv)
s = s * 2
s = np.clip(s, 0, 255)
imghsv = cv2.merge([h, s, v])
# cv2.imshow('imghsv', imghsv)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

boules_couleurs = []
for boule in boules:
    half_diametre = int(boule.diametre / 2)
    crop = imghsv[
           boule.y - half_diametre:boule.y + half_diametre,
           boule.x - half_diametre:boule.x + half_diametre,
           ].copy()
    pixels = np.float32(crop.reshape(-1, 3))
    n_colors = 2
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, .1)
    _, labels, palette = cv2.kmeans(pixels, n_colors, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    _, counts = np.unique(labels, return_counts=True)
    (b, g, r) = palette[np.argmax(counts)] / 16

    # A modulariser
    boules_couleurs.append(TYPE_1 if b > 4 else TYPE_2)

    # cv2.imshow('crop', crop)
    # cv2.waitKey(0)

# 5. Calcul des distances entre chaque boule et le cochonnet selon le centre des boxs
boules_distance = {}
for i, boule in enumerate(boules):
    dist = int(math.sqrt(math.pow(cochonnet.x - boule.x, 2) + math.pow(cochonnet.y - boule.y, 2)))
    boules_distance[i] = dist
boules_distance = dict(sorted(boules_distance.items(), key=lambda item: item[1]))

# 6. Liste ordonnée des 6 distances les plus faibles

boules_proches = [x for x in list(boules_distance)[0:6]]

# 7. Sortie des 6 couleurs en  --- ou - -
img_final = raw_image.copy()
for i in boules_proches:
    boule = boules[i]
    print(boules_couleurs[i])
    cv2.circle(img_final, (boule.x, boule.y), boule.diametre, (0, 255, 0), 2)

# Show result for testing:
cv2.imshow('img_final', img_final)
cv2.waitKey(0)
cv2.destroyAllWindows()