yiking/process.py

from operator import itemgetter

from PIL import Image, ImageTk

import math
from dataclasses import dataclass

import numpy as np
import cv2
from pathlib import Path

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


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


def process_frame(params):
    """
    Simulates OpenCV processing using parameters from the GUI.

    Args:
        params (dict): A dictionary of variable values passed from the GUI.

    Returns:
        ImageTk.PhotoImage: A Tkinter-compatible image.
        str: A result text description.
    """
    # Simulate processing: for now, return a dummy image and text.
    # width, height = 400, 300
    # image = Image.new("RGB", (width, height),
    #                   color=(params["color1_R"] * 4, params["color1_V"] * 4, params["color1_B"] * 4))
    # image_tk = ImageTk.PhotoImage(image)
    #
    # result_text = f"Processed image with params: {params}"
    # return image_tk, result_text

    (minDist, param1, param2, minRadius, maxRadius,
     color1_R, color1_V, color1_B, color2_R, color2_V, color2_B) = itemgetter(
        'minDist', 'param1', 'param2', 'minRadius', 'maxRadius',
        'color1_R', 'color1_V', 'color1_B', 'color2_R', 'color2_V', 'color2_B'
    )(params)

    # 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)

    circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist, param1=param1, param2=param2,
                               minRadius=minRadius,
                               maxRadius=maxRadius)

    min_rayon = 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]), rayon=int(i[2]))

            # 3. Détection de la box la plus petite : cochonnet
            if boule.rayon < min_rayon:
                min_rayon = boule.rayon
                if cochonnet is not None:
                    boules.append(cochonnet)
                cochonnet = boule
            else:
                boules.append(boule)

    # 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])
    boules_couleurs = []
    for boule in boules:
        half_diametre = int(boule.rayon / 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)

    # 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 - -
    return_text = ""
    img_final = raw_image.copy()
    for i in boules_proches:
        boule = boules[i]
        return_text += f"{boules_couleurs[i]}\n"
        cv2.circle(img_final, (boule.x, boule.y), boule.rayon, (0, 255, 0), 2)

    return img_final, return_text