166 lines
5.6 KiB
Python
166 lines
5.6 KiB
Python
# modules/planarian_tracker.py
|
|
'''
|
|
Created on 16 avr. 2026
|
|
|
|
@author: denis
|
|
'''
|
|
|
|
import cv2
|
|
import logging
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PlanarianTracker:
|
|
"""
|
|
Détection et suivi d'une planaire 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.
|
|
"""
|
|
|
|
def __init__(self, tube_axis: str = "vertical", min_area_px: int = 20):
|
|
# Axe du tube : "vertical" (cy) ou "horizontal" (cx)
|
|
self.tube_axis = tube_axis
|
|
self.min_area_px = min_area_px
|
|
|
|
# Etat inter-frame
|
|
self._prev_cx = None
|
|
self._prev_cy = None
|
|
self._prev_ts = None
|
|
|
|
# Soustracteur de fond adaptatif MOG2
|
|
self._bg_sub = cv2.createBackgroundSubtractorMOG2(
|
|
history = 50,
|
|
varThreshold = 25,
|
|
detectShadows= False,
|
|
)
|
|
|
|
def reset(self):
|
|
"""
|
|
Réinitialise l'état inter-frame — appeler lors du changement de puits.
|
|
"""
|
|
self._prev_cx = None
|
|
self._prev_cy = None
|
|
self._prev_ts = None
|
|
# Réinitialise le fond appris
|
|
self._bg_sub = cv2.createBackgroundSubtractorMOG2(
|
|
history = 50,
|
|
varThreshold = 25,
|
|
detectShadows= False,
|
|
)
|
|
|
|
def process(self, frame: np.ndarray, ts: float) -> tuple[np.ndarray, dict]:
|
|
"""
|
|
Analyse une frame et dessine les contours détectés directement sur l'image.
|
|
Retourne (frame_annotée, métriques).
|
|
|
|
Contours fins Vert (0,255,0) Tous les contours valides détectés
|
|
Contour épais Cyan (255,255,0) Planaire principale (plus grand contour)
|
|
Croix + cercle Rouge (0,0,255) Centre de masse exact
|
|
Texte Blanc Vitesse px/s + position axiale normalisée
|
|
"""
|
|
result = self._empty_result(ts)
|
|
frame_out = frame.copy() # copie pour ne pas modifier l'original
|
|
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
fg_mask = self._bg_sub.apply(gray)
|
|
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_OPEN, kernel)
|
|
fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, kernel)
|
|
|
|
contours, _ = cv2.findContours(
|
|
fg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
|
|
)
|
|
|
|
if not contours:
|
|
self._update_prev(None, None, ts)
|
|
return frame_out, result
|
|
|
|
# Filtre les contours significatifs
|
|
valid_contours = [c for c in contours if cv2.contourArea(c) >= self.min_area_px]
|
|
|
|
if not valid_contours:
|
|
self._update_prev(None, None, ts)
|
|
return frame_out, result
|
|
|
|
# Dessine tous les contours valides en vert fin
|
|
cv2.drawContours(frame_out, valid_contours, -1, (0, 255, 0), 1)
|
|
|
|
# Plus grand contour = planaire principale
|
|
largest = max(valid_contours, key=cv2.contourArea)
|
|
area = cv2.contourArea(largest)
|
|
|
|
# Contour principal en cyan plus épais
|
|
cv2.drawContours(frame_out, [largest], -1, (255, 255, 0), 2)
|
|
|
|
M = cv2.moments(largest)
|
|
if M["m00"] == 0:
|
|
return frame_out, result
|
|
|
|
cx = int(M["m10"] / M["m00"])
|
|
cy = int(M["m01"] / M["m00"])
|
|
h, w = frame.shape[:2]
|
|
|
|
axial_pos = (cy / h) if self.tube_axis == "vertical" else (cx / w)
|
|
speed_px_s = None
|
|
axial_speed = None
|
|
|
|
if self._prev_cx is not None and self._prev_ts is not None:
|
|
dt = ts - self._prev_ts
|
|
if dt > 0:
|
|
dx = cx - self._prev_cx
|
|
dy = cy - self._prev_cy
|
|
speed_px_s = float(np.sqrt(dx**2 + dy**2) / dt)
|
|
axial_speed = float((dy / dt) if self.tube_axis == "vertical" else (dx / dt))
|
|
|
|
# Croix sur le centre de masse
|
|
cross_size = 8
|
|
cv2.line(frame_out, (cx - cross_size, cy), (cx + cross_size, cy), (0, 0, 255), 1)
|
|
cv2.line(frame_out, (cx, cy - cross_size), (cx, cy + cross_size), (0, 0, 255), 1)
|
|
|
|
# Cercle centré sur la planaire
|
|
cv2.circle(frame_out, (cx, cy), 12, (0, 0, 255), 1)
|
|
|
|
# Texte vitesse + position axiale
|
|
label = f"v={speed_px_s:.1f}px/s ax={axial_pos:.2f}" if speed_px_s is not None else f"ax={axial_pos:.2f}"
|
|
cv2.putText(
|
|
frame_out, label,
|
|
(max(cx - 60, 0), max(cy - 18, 12)),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv2.LINE_AA,
|
|
)
|
|
|
|
result.update({
|
|
"detected" : True,
|
|
"cx" : cx,
|
|
"cy" : cy,
|
|
"area_px" : int(area),
|
|
"speed_px_s" : round(speed_px_s, 3) if speed_px_s is not None else 0.0,
|
|
"axial_speed" : round(axial_speed, 3) if axial_speed is not None else 0.0,
|
|
"axial_pos" : round(axial_pos, 4),
|
|
})
|
|
|
|
self._update_prev(cx, cy, ts)
|
|
return frame_out, result
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
def _empty_result(self, ts: float) -> dict:
|
|
return {
|
|
"timestamp" : ts,
|
|
"detected" : False,
|
|
"cx" : 0,
|
|
"cy" : 0,
|
|
"area_px" : 0,
|
|
"speed_px_s" : 0.0,
|
|
"axial_speed": 0.0,
|
|
"axial_pos" : 0.0,
|
|
}
|
|
|
|
def _update_prev(self, cx, cy, ts):
|
|
self._prev_cx = cx
|
|
self._prev_cy = cy
|
|
self._prev_ts = ts
|
|
|
|
|