544 lines
20 KiB
Python
544 lines
20 KiB
Python
"""
|
||
modules/planarian_tracker.py
|
||
|
||
Détection et suivi multi-individus de planaires dans un tube.
|
||
Supporte de 1 à MAX_PLANARIANS planaires par tube.
|
||
|
||
Stratégie :
|
||
- Soustraction de fond MOG2 (léger sur Raspberry Pi 4)
|
||
- Détection de tous les contours valides (surface >= min_area_px)
|
||
- Association frame-à-frame par distance euclidienne minimale
|
||
via algorithme hongrois (scipy.optimize.linear_sum_assignment)
|
||
- Un état inter-frame indépendant par individu (PlanarianState)
|
||
- Retourne une liste de résultats, un par individu suivi
|
||
|
||
Created on 25 avr. 2026
|
||
@author: denis
|
||
"""
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
from scipy.optimize import linear_sum_assignment
|
||
import logging
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Nombre maximum de planaires suivis simultanément par tube
|
||
MAX_PLANARIANS = 10
|
||
|
||
# Distance maximale en pixels entre deux positions consécutives
|
||
# pour qu'une association soit acceptée (évite les sauts aberrants)
|
||
MAX_ASSOC_DIST_PX = 80
|
||
|
||
# Couleurs d'annotation par individu (BGR)
|
||
# Cycle automatique si plus de planaires que de couleurs
|
||
INDIVIDUAL_COLORS = [
|
||
(255, 255, 0), # cyan
|
||
( 0, 165, 255), # orange
|
||
(255, 0, 255), # magenta
|
||
( 0, 255, 255), # jaune
|
||
(128, 0, 255), # violet
|
||
( 0, 255, 128), # vert clair
|
||
(255, 128, 0), # bleu clair
|
||
( 0, 128, 255), # orange foncé
|
||
(128, 255, 0), # vert-jaune
|
||
(255, 0, 128), # rose
|
||
]
|
||
|
||
# Couleur du contour principal (individu le plus grand)
|
||
COLOR_LARGEST = (255, 255, 0) # cyan
|
||
COLOR_OTHER = ( 0, 255, 0) # vert
|
||
COLOR_CENTER = ( 0, 0, 255) # rouge
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# État inter-frame d'un individu
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class PlanarianState:
|
||
"""
|
||
Mémorise la position et le timestamp de la dernière détection
|
||
pour un planaire individuel.
|
||
|
||
Un PlanarianState par slot (index 0 à max_planarians-1).
|
||
Quand un slot n'est pas associé à un contour sur plusieurs frames
|
||
consécutives, il est marqué comme perdu (lost).
|
||
"""
|
||
|
||
# Nombre de frames sans détection avant de considérer l'individu perdu
|
||
MAX_LOST_FRAMES = 5
|
||
|
||
def __init__(self, idx: int):
|
||
"""
|
||
Args:
|
||
idx : index de l'individu (0-based)
|
||
"""
|
||
self.idx: int = idx
|
||
self.cx: int | None = None
|
||
self.cy: int | None = None
|
||
self.ts: float | None = None
|
||
self.lost: int = 0 # compteur de frames sans détection
|
||
self.active: bool = False # vrai si l'individu a été détecté au moins une fois
|
||
|
||
def update(self, cx: int, cy: int, ts: float):
|
||
"""
|
||
Met à jour la position suite à une association réussie.
|
||
|
||
Args:
|
||
cx, cy : position du centre de masse en pixels
|
||
ts : timestamp de la frame
|
||
"""
|
||
self.cx = cx
|
||
self.cy = cy
|
||
self.ts = ts
|
||
self.lost = 0
|
||
self.active = True
|
||
|
||
def mark_lost(self):
|
||
"""Incrémente le compteur de perte — appelé quand aucun contour n'est associé."""
|
||
self.lost += 1
|
||
|
||
@property
|
||
def is_lost(self) -> bool:
|
||
"""Retourne True si l'individu est considéré perdu (trop de frames sans détection)."""
|
||
return self.lost >= self.MAX_LOST_FRAMES
|
||
|
||
def compute_speed(self, cx: int, cy: int, ts: float, tube_axis: str) -> tuple:
|
||
"""
|
||
Calcule la vitesse instantanée depuis la position précédente.
|
||
|
||
Args:
|
||
cx, cy : position courante en pixels
|
||
ts : timestamp courant
|
||
tube_axis : "vertical" ou "horizontal"
|
||
|
||
Returns:
|
||
tuple (speed_px_s, axial_speed) ou (0.0, 0.0) si état vide
|
||
"""
|
||
if self.cx is None or self.cy is None or self.ts is None:
|
||
return 0.0, 0.0
|
||
|
||
dt = ts - self.ts
|
||
if dt <= 0:
|
||
return 0.0, 0.0
|
||
|
||
dx = cx - self.cx
|
||
dy = cy - self.cy
|
||
speed_px_s = float(np.sqrt(dx**2 + dy**2) / dt)
|
||
axial_speed = float((dy / dt) if tube_axis == "vertical" else (dx / dt))
|
||
|
||
return speed_px_s, axial_speed
|
||
|
||
def reset(self):
|
||
"""Réinitialise l'état de cet individu."""
|
||
self.cx = None
|
||
self.cy = None
|
||
self.ts = None
|
||
self.lost = 0
|
||
self.active = False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tracker multi-individus
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class PlanarianTracker:
|
||
"""
|
||
Détection et suivi multi-individus de planaires dans un tube.
|
||
|
||
Instancié une fois par caméra active, réutilisé frame à frame.
|
||
Utilise la soustraction de fond MOG2 — léger sur Raspberry Pi 4.
|
||
Association frame-à-frame par algorithme hongrois (distance euclidienne).
|
||
|
||
Usage :
|
||
tracker = PlanarianTracker(tube_axis="vertical", max_planarians=3)
|
||
while capturing:
|
||
frame_out, results = tracker.process(frame, ts)
|
||
# results : liste de dicts, un par individu détecté
|
||
for r in results:
|
||
metrics.update(r, planarian_id=r["planarian_id"])
|
||
"""
|
||
|
||
# Nombre de frames d'initialisation MOG2 ignorées (fond non appris)
|
||
WARMUP_FRAMES = 10
|
||
|
||
def __init__(
|
||
self,
|
||
tube_axis: str = "vertical",
|
||
min_area_px: int = 20,
|
||
max_area_ratio: float = 0.10,
|
||
max_planarians: int = 1,
|
||
merge_kernel_size: int = 15,
|
||
min_contour_dist_px:int = 40,
|
||
draw_contours: bool = True,
|
||
):
|
||
"""
|
||
Args:
|
||
tube_axis : axe principal du tube — "vertical" (cy) ou "horizontal" (cx)
|
||
min_area_px : surface minimale d'un contour pour être considéré valide (px²)
|
||
max_area_ratio : surface maximale d'un contour en fraction de la frame (défaut 10%)
|
||
filtre les faux positifs du fond non encore appris par MOG2
|
||
max_planarians : nombre maximum de planaires à suivre simultanément (1-10)
|
||
merge_kernel_size : taille du kernel elliptique de fusion des fragments (px).
|
||
Régler ≈ largeur du planaire en pixels. Défaut : 15.
|
||
min_contour_dist_px : distance min entre deux contours pour les considérer
|
||
comme individus distincts. Défaut : 40px.
|
||
"""
|
||
self.tube_axis = tube_axis
|
||
self.min_area_px = min_area_px
|
||
self.max_area_ratio = max_area_ratio
|
||
self.max_planarians = max(1, min(max_planarians, MAX_PLANARIANS))
|
||
self.draw_contours = draw_contours
|
||
|
||
# Un état inter-frame par slot individu
|
||
self._states = [PlanarianState(i) for i in range(self.max_planarians)]
|
||
|
||
# Taille du kernel de fusion morphologique (pixels) —
|
||
# doit être proche de la largeur du planaire en pixels.
|
||
# Trop petit : fragments non fusionnés → IDs multiples.
|
||
# Trop grand : deux planaires proches fusionnés en un seul.
|
||
self.merge_kernel_size = merge_kernel_size
|
||
|
||
# Distance minimale en pixels entre deux contours distincts.
|
||
# En-dessous : le plus petit est considéré comme fragment du plus grand.
|
||
self.min_contour_dist_px = min_contour_dist_px
|
||
|
||
# Soustracteur de fond adaptatif MOG2
|
||
self._bg_sub = self._make_bg_sub()
|
||
|
||
# Compteur de frames d'initialisation — MOG2 retourne du bruit
|
||
# pendant les premières WARMUP_FRAMES frames
|
||
self._warmup_count = 0
|
||
|
||
@staticmethod
|
||
def _make_bg_sub():
|
||
"""Crée et retourne un soustracteur de fond MOG2."""
|
||
return cv2.createBackgroundSubtractorMOG2(
|
||
history = 50,
|
||
varThreshold = 25,
|
||
detectShadows= False,
|
||
)
|
||
|
||
def reset(self):
|
||
"""
|
||
Réinitialise l'état inter-frame complet.
|
||
À appeler lors du changement de puits.
|
||
"""
|
||
for s in self._states:
|
||
s.reset()
|
||
self._bg_sub = self._make_bg_sub()
|
||
self._warmup_count = 0
|
||
# Les paramètres morphologiques (merge_kernel_size, min_contour_dist_px)
|
||
# sont conservés — ils ne dépendent pas du puits
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Interface principale
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def process(self, frame: np.ndarray, ts: float) -> tuple:
|
||
"""
|
||
Analyse une frame, associe les contours aux individus connus,
|
||
dessine les annotations et retourne les métriques.
|
||
|
||
Args:
|
||
frame : image BGR (numpy array)
|
||
ts : timestamp de la frame (float, secondes epoch)
|
||
|
||
Returns:
|
||
tuple (frame_annotée, results)
|
||
|
||
frame_annotée : copie BGR avec contours, croix et textes
|
||
results : liste de dicts — un dict par planaire actif détecté.
|
||
Chaque dict contient :
|
||
planarian_id int index de l'individu (0-based)
|
||
detected bool True si détecté cette frame
|
||
cx, cy int centre de masse en pixels
|
||
area_px int surface du contour (px²)
|
||
speed_px_s float vitesse totale (px/s)
|
||
axial_speed float vitesse axiale (px/s)
|
||
axial_pos float position axiale normalisée (0-1)
|
||
timestamp float ts de la frame
|
||
"""
|
||
frame_out = frame.copy()
|
||
h, w = frame.shape[:2]
|
||
|
||
# --- Extraction du premier plan ---
|
||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||
fg_mask = self._bg_sub.apply(gray)
|
||
|
||
# --- Morphologie : fusion des fragments du corps ---
|
||
# Un planaire ondulant est souvent segmenté en plusieurs contours
|
||
# (tête, milieu, queue). La dilatation fusionne les fragments proches
|
||
# avant la détection des contours.
|
||
# Kernel 3×3 : supprime le bruit fin (OPEN)
|
||
# Kernel merge_kernel : fusionne les fragments du corps (CLOSE + DILATE)
|
||
noise_kernel = np.ones((3, 3), np.uint8)
|
||
merge_kernel = cv2.getStructuringElement(
|
||
cv2.MORPH_ELLIPSE, (self.merge_kernel_size, self.merge_kernel_size)
|
||
)
|
||
fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_OPEN, noise_kernel)
|
||
fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, merge_kernel)
|
||
fg_mask = cv2.dilate(fg_mask, merge_kernel, iterations=1)
|
||
|
||
# Warmup MOG2 : les premières WARMUP_FRAMES frames retournent du bruit
|
||
# (fond non encore appris) — on les alimente mais on ne détecte rien
|
||
self._warmup_count += 1
|
||
if self._warmup_count <= self.WARMUP_FRAMES:
|
||
return frame_out, []
|
||
|
||
contours, _ = cv2.findContours(
|
||
fg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
|
||
)
|
||
|
||
# Surface maximale admissible : fraction de la frame
|
||
# Filtre les faux positifs du fond (contours couvrant toute l'image)
|
||
max_area_px = h * w * self.max_area_ratio
|
||
|
||
# Filtrage : surface min ET surface max, triés par surface décroissante
|
||
valid = sorted(
|
||
[c for c in contours
|
||
if self.min_area_px <= cv2.contourArea(c) <= max_area_px],
|
||
key=cv2.contourArea,
|
||
reverse=True,
|
||
)
|
||
|
||
# --- Suppression des fragments résiduels trop proches ---
|
||
# Si plusieurs contours valides ont leur centre à moins de
|
||
# min_contour_dist_px les uns des autres, seul le plus grand est conservé.
|
||
# Évite les cas où la fusion morphologique est incomplète.
|
||
filtered = []
|
||
for c in valid:
|
||
M = cv2.moments(c)
|
||
if M["m00"] == 0:
|
||
continue
|
||
cx_c = int(M["m10"] / M["m00"])
|
||
cy_c = int(M["m01"] / M["m00"])
|
||
too_close = False
|
||
for kept in filtered:
|
||
Mk = cv2.moments(kept)
|
||
if Mk["m00"] == 0:
|
||
continue
|
||
cx_k = int(Mk["m10"] / Mk["m00"])
|
||
cy_k = int(Mk["m01"] / Mk["m00"])
|
||
dist = np.sqrt((cx_c - cx_k)**2 + (cy_c - cy_k)**2)
|
||
if dist < self.min_contour_dist_px:
|
||
too_close = True
|
||
break
|
||
if not too_close:
|
||
filtered.append(c)
|
||
|
||
# Limiter au nombre maximum de planaires attendus
|
||
valid = filtered[:self.max_planarians]
|
||
|
||
# --- Calcul des centres de masse des contours détectés ---
|
||
detections = [] # liste de (cx, cy, area, contour)
|
||
for c in valid:
|
||
M = cv2.moments(c)
|
||
if M["m00"] == 0:
|
||
continue
|
||
cx = int(M["m10"] / M["m00"])
|
||
cy = int(M["m01"] / M["m00"])
|
||
area = cv2.contourArea(c)
|
||
detections.append((cx, cy, int(area), c))
|
||
|
||
# --- Association hongroise détections → slots individus ---
|
||
assignments = self._hungarian_assign(detections)
|
||
|
||
# --- Mise à jour des états et construction des résultats ---
|
||
results = []
|
||
|
||
for slot_idx, det_idx in assignments.items():
|
||
state = self._states[slot_idx]
|
||
|
||
if det_idx is None:
|
||
# Aucune détection associée à ce slot
|
||
state.mark_lost()
|
||
if state.active and not state.is_lost:
|
||
# L'individu était suivi : on retourne un résultat "perdu"
|
||
results.append(self._lost_result(slot_idx, ts))
|
||
continue
|
||
|
||
cx, cy, area, contour = detections[det_idx]
|
||
|
||
# Calcul de la vitesse depuis la position précédente
|
||
speed_px_s, axial_speed = state.compute_speed(cx, cy, ts, self.tube_axis)
|
||
|
||
axial_pos = (cy / h) if self.tube_axis == "vertical" else (cx / w)
|
||
|
||
# Mise à jour de l'état
|
||
state.update(cx, cy, ts)
|
||
|
||
|
||
if self.draw_contours:
|
||
|
||
# Annotation visuelle
|
||
color = INDIVIDUAL_COLORS[slot_idx % len(INDIVIDUAL_COLORS)]
|
||
cv2.drawContours(frame_out, [contour], -1, color, 2)
|
||
self._draw_center(frame_out, cx, cy, slot_idx, speed_px_s, axial_pos, color)
|
||
|
||
results.append({
|
||
"planarian_id": slot_idx,
|
||
"detected": True,
|
||
"cx": cx,
|
||
"cy": cy,
|
||
"area_px": area,
|
||
"speed_px_s": round(speed_px_s, 3),
|
||
"axial_speed": round(axial_speed, 3),
|
||
"axial_pos": round(axial_pos, 4),
|
||
"timestamp": ts,
|
||
})
|
||
|
||
# Marquer les slots non présents dans les assignments comme perdus
|
||
assigned_slots = set(assignments.keys())
|
||
for state in self._states:
|
||
if state.idx not in assigned_slots:
|
||
state.mark_lost()
|
||
return frame_out, results
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Association hongroise
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def _hungarian_assign(self, detections: list) -> dict:
|
||
"""
|
||
Associe les détections courantes aux slots individus connus
|
||
via l'algorithme hongrois (coût = distance euclidienne).
|
||
|
||
Contrainte : une association n'est acceptée que si la distance
|
||
est inférieure à MAX_ASSOC_DIST_PX (évite les sauts aberrants).
|
||
|
||
Args:
|
||
detections : liste de (cx, cy, area, contour)
|
||
|
||
Returns:
|
||
dict {slot_idx: det_idx | None}
|
||
det_idx = None si aucune détection assignée à ce slot
|
||
"""
|
||
n_slots = self.max_planarians
|
||
n_dets = len(detections)
|
||
|
||
if n_dets == 0:
|
||
# Aucune détection : tous les slots sont "perdus"
|
||
return {i: None for i in range(n_slots)}
|
||
|
||
# Slots actifs (déjà vus au moins une fois et non perdus)
|
||
active_slots = [s for s in self._states if s.active and not s.is_lost]
|
||
|
||
if not active_slots:
|
||
# Première frame ou tous perdus : attribution séquentielle simple
|
||
assignment = {}
|
||
for i in range(n_slots):
|
||
assignment[i] = i if i < n_dets else None
|
||
return assignment
|
||
|
||
# --- Construction de la matrice de coût (distance euclidienne) ---
|
||
cost = np.full((len(active_slots), n_dets), fill_value=1e6)
|
||
|
||
for si, state in enumerate(active_slots):
|
||
for di, (cx, cy, _, _) in enumerate(detections):
|
||
dist = np.sqrt((cx - state.cx)**2 + (cy - state.cy)**2)
|
||
cost[si, di] = dist
|
||
|
||
# --- Algorithme hongrois ---
|
||
row_ind, col_ind = linear_sum_assignment(cost)
|
||
|
||
# Construire le dict d'association
|
||
assignment: dict[int, int | None] = {i: None for i in range(n_slots)}
|
||
|
||
assigned_dets = set()
|
||
for ri, ci in zip(row_ind, col_ind):
|
||
if cost[ri, ci] <= MAX_ASSOC_DIST_PX:
|
||
slot_idx = active_slots[ri].idx
|
||
assignment[slot_idx] = ci
|
||
assigned_dets.add(ci)
|
||
|
||
# --- Nouvelles détections non assignées → slots inactifs libres ---
|
||
free_slots = [s for s in self._states if not s.active or s.is_lost]
|
||
new_dets = [di for di in range(n_dets) if di not in assigned_dets]
|
||
|
||
for state, det_idx in zip(free_slots, new_dets):
|
||
assignment[state.idx] = det_idx
|
||
|
||
return assignment
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Dessin des annotations
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def _draw_center(
|
||
self,
|
||
frame: np.ndarray,
|
||
cx: int,
|
||
cy: int,
|
||
idx: int,
|
||
speed_px_s: float,
|
||
axial_pos: float,
|
||
color: tuple,
|
||
):
|
||
"""
|
||
Dessine la croix, le cercle et le label de vitesse/position
|
||
pour un individu.
|
||
|
||
Args:
|
||
frame : image à annoter (en place)
|
||
cx, cy : centre de masse en pixels
|
||
idx : index de l'individu
|
||
speed_px_s : vitesse en px/s
|
||
axial_pos : position axiale normalisée
|
||
color : couleur BGR de l'individu
|
||
"""
|
||
cross = 8
|
||
cv2.line(frame, (cx - cross, cy), (cx + cross, cy), COLOR_CENTER, 1)
|
||
cv2.line(frame, (cx, cy - cross), (cx, cy + cross), COLOR_CENTER, 1)
|
||
cv2.circle(frame, (cx, cy), 12, color, 1)
|
||
|
||
# Badge numéro individu
|
||
cv2.circle(frame, (cx + 14, cy - 14), 8, color, -1)
|
||
cv2.putText(
|
||
frame, str(idx),
|
||
(cx + 10, cy - 10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.35, (0, 0, 0), 1, cv2.LINE_AA,
|
||
)
|
||
|
||
# Texte vitesse + position axiale
|
||
label = (
|
||
f"#{idx} v={speed_px_s:.1f}px/s ax={axial_pos:.2f}"
|
||
if speed_px_s > 0
|
||
else f"#{idx} ax={axial_pos:.2f}"
|
||
)
|
||
cv2.putText(
|
||
frame, label,
|
||
(max(cx - 60, 0), max(cy - 22, 12)),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.35, (255, 255, 255), 1, cv2.LINE_AA,
|
||
)
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Résultats vides / perdus
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def _lost_result(self, planarian_id: int, ts: float) -> dict:
|
||
"""
|
||
Retourne un résultat pour un individu temporairement non détecté.
|
||
|
||
Args:
|
||
planarian_id : index de l'individu
|
||
ts : timestamp de la frame courante
|
||
|
||
Returns:
|
||
dict avec detected=False et les dernières coordonnées connues
|
||
"""
|
||
state = self._states[planarian_id]
|
||
return {
|
||
"planarian_id": planarian_id,
|
||
"detected": False,
|
||
"cx": state.cx or 0,
|
||
"cy": state.cy or 0,
|
||
"area_px": 0,
|
||
"speed_px_s": 0.0,
|
||
"axial_speed": 0.0,
|
||
"axial_pos": 0.0,
|
||
"timestamp": ts,
|
||
}
|