tracking.py

from collections import deque

import cv2
import numpy as np

def create_kalman():
    kf = cv2.KalmanFilter(4, 2)  # 4 стани, 2 виміри

    # Стан: [x, y, dx, dy]
    kf.transitionMatrix = np.array([
        [1, 0, 1, 0],  # x = x + dx
        [0, 1, 0, 1],  # y = y + dy
        [0, 0, 1, 0],  # dx = dx
        [0, 0, 0, 1],  # dy = dy
    ], dtype=np.float32)

    # Вимірюємо тільки x, y (центр bbox)
    kf.measurementMatrix = np.array([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
    ], dtype=np.float32)

    # Шум процесу і вимірювання (можна тюнити)
    kf.processNoiseCov = np.eye(4, dtype=np.float32) * 1e-2
    kf.measurementNoiseCov = np.eye(2, dtype=np.float32) * 5e-2

    return kf

def bbox_to_measurement(bbox):
    x, y, w, h = [int(v) for v in bbox]
    cx = x + w // 2
    cy = y + h // 2
    return np.array([cx, cy], dtype=np.float32).reshape(2, 1)

def prediction_to_bbox(pred, w, h):
    cx, cy = int(pred[0]), int(pred[1])
    return (cx - w // 2, cy - h // 2, w, h)

def track(video):
    tracker = cv2.TrackerCSRT_create()
    kf = create_kalman()

    ok, frame=video.read()
    bbox = cv2.selectROI(frame)
    ok = tracker.init(frame, bbox)
    bbox_history = deque(maxlen=20)

    # Ініціалізуємо стан Кальмана першим bbox
    kf.statePost = np.array(
        [bbox[0] + bbox[2] // 2,  # cx
         bbox[1] + bbox[3] // 2,  # cy
         0, 0],                 # dx, dy
        dtype=np.float32
    ).reshape(-1, 1)

    while True:
       ok,frame=video.read()
       if not ok:
           break

       # 1. Kalman передбачає
       predicted = kf.predict().flatten()

       pred_bbox = prediction_to_bbox(predicted, bbox[2], bbox[3])

       ok, bbox = tracker.update(frame)

       if ok:
           measurement = bbox_to_measurement(bbox)
           kf.correct(measurement)
           bbox_history.append(bbox)
           x, y, w, h = [int(v) for v in bbox]
           px, py, pw, ph = [int(v) for v in pred_bbox]

           cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2, 1)
           cv2.putText(frame, f'CSRT - {bbox}', (x, y - 10),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

           cv2.rectangle(frame, (px, py), (px + pw, py + ph), (0, 165, 255), 2, 1)
           cv2.putText(frame, f'Kalman - {pred_bbox}', (px, py - 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 165, 255), 2)

       else:
           # намагаємось перезапустити трекер, беремо одну із попередніх позицій
           tracker = cv2.TrackerCSRT_create()
           px, py, pw, ph = [int(v) for v in bbox_history[5]]
           tracker.init(frame, (px, py, pw, ph))

           # пробуємо взяте значення яке дає Калман
           # cv2.rectangle(frame, (px, py), (px + pw, py + ph), (0, 165, 255), 2, 1)
           # cv2.putText(frame, f'Kalman - {pred_bbox}', (px, py - 30),
           #             cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 165, 255), 2)

       cv2.imshow('Tracking',frame)
       if cv2.waitKey(1) & 0XFF==27:
            break
    cv2.destroyAllWindows()


# Я спробував декілька алгоритмів (розташовую їх у порядку якості трекінгу):
# MeanShift, CamShift, KCF та в кінці зупинився на CSRT.
# KCF майже ідеально спрацював, але на відео race2, де одночасно є багато машин,
# іноді трекер переструбвав на іншу машину. В CSRT це пофіксилось.
# Тут залишається одна проблема, яку не вийшло вирішити:
# коли машина заїзжає під міст, або під дорожні інформаційні щитки що вісять над дорогоую,
# то трекінг зникає. Це гарно видно на відео race1.
# Я намагався спочатку вирішити це задопомоги збереження попередніх позицій bbox, але не вийшло.
# Далі знайшов інфу про Фільтр Калмана. Додав його, і виглядає що він схоже зможе вирішити подібну проблему,
# якщо його доналаштувати.
# Плюс у мене є підозра, що проблема саме із моїм відео, бо трекінг зниках в момент, коли машина раптово
# йде на обгон і одразу потрапляє під дорожній щит, а так як Калман запамʼятовує останній вектор руху,
# то можливо саме через це у мене і не вийшло на даному відео вирішити цю проблему.
# Для захвату обличчя працює теж непогано, якщо обличчя не закривати тільки нічим.

if __name__ == '__main__':

    # Машину треба виділити рамкою та нажати Enter. Esc для закінчення роботи.

    video = cv2.VideoCapture('race1.mp4')
    # video = cv2.VideoCapture('race2.mp4')
    track(video)