117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
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()
|