目标跟踪：SORT, DeepSORT, ByteTrack

# 目标跟踪基本数据结构
from dataclasses import dataclass
from typing import List, Optional
import numpy as np

@dataclass
class BoundingBox:
    """边界框：[x_center, y_center, width, height]"""
    x: float
    y: float
    w: float
    h: float

    @property
    def area(self) -> float:
        return self.w * self.h

    def iou(self, other: 'BoundingBox') -> float:
        """计算两个框的 IoU（交并比）"""
        x1 = max(self.x - self.w/2, other.x - other.w/2)
        y1 = max(self.y - self.h/2, other.y - other.h/2)
        x2 = min(self.x + self.w/2, other.x + other.w/2)
        y2 = min(self.y + self.h/2, other.y + other.h/2)
        inter = max(0, x2 - x1) * max(0, y2 - y1)
        union = self.area + other.area - inter
        return inter / union if union > 0 else 0.0

@dataclass
class Track:
    """跟踪轨迹"""
    track_id: int
    state: np.ndarray  # 卡尔曼滤波状态 [x, y, s, r, vx, vy, vs, vr]
    hits: int          # 连续匹配帧数
    age: int           # 自创建以来的帧数
    time_since_update: int  # 距上次更新的帧数

# MOT 跟踪主循环伪代码
def mot_tracking(video_path, detector, tracker):
    """Tracking-by-Detecting 主流程"""
    cap = cv2.VideoCapture(video_path)
    tracks = []  # 活跃的跟踪轨迹

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Step 1: 检测当前帧所有目标
        detections = detector.detect(frame)  # 返回 [BBox, conf] 列表

        # Step 2: 预测已有轨迹的当前位置
        tracker.predict()

        # Step 3: 将检测结果与已有轨迹关联
        tracks = tracker.update(tracks, detections)

        # Step 4: 可视化
        for track in tracks:
            if track.is_confirmed():
                draw_box(frame, track.bbox, track.track_id)

    return tracks

import numpy as np
from scipy.optimize import linear_sum_assignment

class KalmanFilter:
    """8 维匀速卡尔曼滤波器"""
    def __init__(self):
        # 状态维度: [x, y, s, r, vx, vy, vs, vr]
        ndim = 8
        dt = 1.0  # 帧间时间间隔

        # 状态转移矩阵 F（匀速模型）
        self.F = np.eye(ndim)
        for i in range(4):
            self.F[i, i + 4] = dt  # 位置 += 速度 * dt

        # 观测矩阵 H（只观测位置，不观测速度）
        self.H = np.eye(4, ndim)

        # 过程噪声协方差 Q
        self.Q = np.eye(ndim) * 0.01

        # 观测噪声协方差 R
        self.R = np.eye(4) * 10.0

    def init_state(self, bbox):
        """用初始检测框初始化状态 [x, y, s, r]"""
        x, y, w, h = bbox
        state = np.array([x, y, w * h, w / h, 0, 0, 0, 0])
        return state.reshape(8, 1)

    def predict(self, state, cov):
        """预测步：x_k = F * x_{k-1}"""
        new_state = self.F @ state
        new_cov = self.F @ cov @ self.F.T + self.Q
        return new_state, new_cov

    def update(self, state, cov, measurement):
        """更新步：融合观测 z_k"""
        # 卡尔曼增益: K = P*H^T * (H*P*H^T + R)^{-1}
        y = measurement - self.H @ state       # 新息 (innovation)
        S = self.H @ cov @ self.H.T + self.R   # 新息协方差
        K = cov @ self.H.T @ np.linalg.inv(S)  # 卡尔曼增益

        new_state = state + K @ y
        I_k = np.eye(8)
        new_cov = (I_k - K @ self.H) @ cov
        return new_state, new_cov

# 使用示例
kf = KalmanFilter()
state = kf.init_state([300, 200, 80, 120])  # [x, y, w, h]
cov = np.eye(8) * 100
state, cov = kf.predict(state, cov)
print(f"预测位置: x={state[0,0]:.0f}, y={state[1,0]:.0f}")

def associate_detections_to_trackers(detections, trackers, iou_threshold=0.3):
    """
    SORT 核心关联：IoU + 匈牙利算法
    detections: N×4 检测框数组
    trackers:   M×4 预测轨迹框数组
    返回: (matched, unmatched_dets, unmatched_trks)
    """
    if len(trackers) == 0:
        return np.empty((0, 2), dtype=int), np.arange(len(detections)), []

    # 计算 IoU 代价矩阵 (N×M)
    iou_matrix = np.zeros((len(detections), len(trackers)))
    for d, det in enumerate(detections):
        for t, trk in enumerate(trackers):
            iou_matrix[d, t] = iou(det, trk)

    # 匈牙利算法求解最小代价匹配
    # 将 IoU 转为代价: cost = 1 - IoU
    matched_indices = linear_sum_assignment(1 - iou_matrix)
    matched_indices = np.array(matched_indices).T

    # 过滤低于 IoU 阈值的匹配
    unmatched_detections = []
    for d in range(len(detections)):
        if d not in matched_indices[:, 0]:
            unmatched_detections.append(d)

    unmatched_trackers = []
    for t in range(len(trackers)):
        if t not in matched_indices[:, 1]:
            unmatched_trackers.append(t)

    # 移除低 IoU 匹配
    matches = []
    for m in matched_indices:
        if iou_matrix[m[0], m[1]] < iou_threshold:
            unmatched_detections.append(m[0])
            unmatched_trackers.append(m[1])
        else:
            matches.append(m.reshape(1, 2))

    return matches, unmatched_detections, unmatched_trackers

class SortTracker:
    """SORT 跟踪器主类"""
    def __init__(self, max_age=1, min_hits=3, iou_threshold=0.3):
        self.max_age = max_age          # 轨迹最大丢失帧数
        self.min_hits = min_hits        # 确认轨迹所需最小匹配次数
        self.iou_threshold = iou_threshold
        self.kf = KalmanFilter()
        self.trackers = []              # 活跃的卡尔曼滤波器列表
        self.frame_count = 0

    def update(self, dets):
        """
        dets: N×5 检测框数组 [x1, y1, x2, y2, confidence]
        返回: M×6 跟踪结果 [x1, y1, x2, y2, track_id, confidence]
        """
        self.frame_count += 1

        # Step 1: 用卡尔曼滤波预测所有现有轨迹
        trks = np.zeros((len(self.trackers), 5))
        to_del = []
        for t, trk in enumerate(trks):
            pos = self.trackers[t].predict()
            trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
            if np.any(np.isnan(pos)):
                to_del.append(t)
        trks = np.ma.compress_rows(np.ma.masked_invalid(trks))

        # Step 2: IoU 关联
        matched, unmatched_dets, unmatched_trks = \
            associate_detections_to_trackers(dets, trks, self.iou_threshold)

        # Step 3: 更新匹配的轨迹
        for m in matched:
            self.trackers[m[1]].update(dets[m[0], :])

        # Step 4: 创建新轨迹
        for i in unmatched_dets:
            trk = self.kf.init_bbox(dets[i, :4])
            self.trackers.append(trk)

        return self._get_results()

import torch
import torch.nn as nn
import torchvision.models as models

class ReIDFeatureExtractor(nn.Module):
    """DeepSORT 外观特征提取器"""
    def __init__(self, feature_dim=128):
        super().__init__()
        # 使用 ResNet-50 作为 backbone
        backbone = models.resnet50(weights=None)
        # 移除最后的分类层
        self.backbone = nn.Sequential(*list(backbone.children())[:-1])
        # 添加投影头输出 128 维特征
        self.projection = nn.Linear(2048, feature_dim)
        self.l2_norm = nn.LayerNorm(feature_dim)

    def forward(self, x):
        """输入: 裁剪的目标图像 [N, 3, 256, 128]"""
        features = self.backbone(x)          # [N, 2048, 1, 1]
        features = features.view(features.size(0), -1)  # [N, 2048]
        features = self.projection(features)  # [N, 128]
        features = self.l2_norm(features)
        return features

# 提取检测框的外观特征
def extract_features(extractor, frame, bboxes, device='cuda'):
    """从帧中裁剪目标区域并提取 Re-ID 特征"""
    crops = []
    for x1, y1, x2, y2 in bboxes:
        crop = frame[int(y1):int(y2), int(x1):int(x2)]
        crop = cv2.resize(crop, (128, 256))
        crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
        crops.append(crop)
    crops = np.stack(crops).transpose(0, 3, 1, 2) / 255.0
    tensor = torch.from_numpy(crops).float().to(device)
    with torch.no_grad():
        features = extractor(tensor)
    return features.cpu().numpy()  # [N, 128]

def gated_sigmoid_distance(cost_matrix, detections, trackers, threshold=0.2):
    """级联匹配 + 外观距离 + 马氏距离门控"""
    from scipy.spatial.distance import cosine

    # 计算外观余弦距离矩阵
    feature_matrix = np.zeros((len(detections), len(trackers)))
    for i, det in enumerate(detections):
        for j, trk in enumerate(trackers):
            # 取轨迹特征库中最近的特征
            nearest_feat = min(
                trk.features[-100:],  # 最多用最近 100 个
                key=lambda f: cosine(det.feature, f)
            )
            feature_matrix[i, j] = cosine(det.feature, nearest_feat)

    # 计算马氏距离（运动一致性）
    mahalanobis_dist = compute_mahalanobis(detections, trackers)

    # 门控: 马氏距离超过阈值 → 外观距离设为无穷大（不允许匹配）
    gated_cost = np.where(
        mahalanobis_dist <= 9.4877,  # 卡方分布 4 自由度的 95% 分位
        feature_matrix,
        float('inf')
    )
    return gated_cost

# 级联匹配: 优先匹配丢失时间短的轨迹
def cascade_matching(detections, trackers, max_age=30):
    matches, unmatched_dets, unmatched_trks = [], [], list(range(len(trackers)))
    for age in range(max_age):
        # 只考虑丢失 age 帧的轨迹
        candidates = [t for t in unmatched_trks if trackers[t].time_since_update == age]
        if not candidates:
            continue
        # 在候选轨迹上做匈牙利匹配
        cost = compute_cost(detections, [trackers[t] for t in candidates])
        matched = hungarian_matching(cost)
        # 更新匹配结果...
    return matches, unmatched_dets, unmatched_trks

class ByteTrack:
    """ByteTrack 核心：两轮关联策略"""
    def __init__(self, track_thresh=0.5, high_thresh=0.6,
                 match_thresh=0.8, new_track_thresh=0.7):
        self.track_thresh = track_thresh    # 检测阈值
        self.high_thresh = high_thresh      # 高置信度阈值
        self.match_thresh = match_thresh    # IoU 匹配阈值
        self.new_track_thresh = new_track_thresh  # 新轨迹阈值
        self.tracker = BYTETracker(max_age=30)

    def update(self, detections):
        """
        detections: [x1, y1, x2, y2, score] 数组
        """
        # 分类检测框
        remain_inds = detections[:, 4] > self.high_thresh
        inds_low = (detections[:, 4] > self.track_thresh) & \
                   (detections[:, 4] <= self.high_thresh)
        inds_high = detections[:, 4] > self.high_thresh

        dets_high = detections[inds_high]
        dets_low = detections[inds_low]

        # 第一轮: 高置信度检测框关联
        online_targets = self.tracker.update(
            dets_high, match_thresh=self.match_thresh
        )

        # 第二轮: 低置信度检测框与未匹配轨迹关联
        activated_targets = [t for t in online_targets if t.is_activated]
        unconfirmed_targets = [t for t in online_targets if not t.is_activated]
        refind_targets = self.tracker.refind(
            unconfirmed_targets, dets_low
        )

        # 合并结果
        return activated_targets + refind_targets

    def forward(self, frame):
        """完整的前向传播"""
        raw_dets = self.detector(frame)
        tracks = self.update(raw_dets)
        return tracks

# 两轮 IoU 匹配的代价矩阵构建
def byte_track_match(tracks, dets, match_thresh=0.8):
    """
    ByteTrack 风格 IoU 匹配
    返回: matches (list of (track_idx, det_idx))
    """
    if len(tracks) == 0 or len(dets) == 0:
        return [], list(range(len(tracks))), list(range(len(dets)))

    # 计算 IoU 矩阵
    iou_mat = np.zeros((len(tracks), len(dets)))
    for i, trk in enumerate(tracks):
        for j, det in enumerate(dets):
            iou_mat[i, j] = compute_iou(trk.predicted_bbox, det.bbox)

    # 匈牙利匹配（代价 = 1 - IoU）
    row_ind, col_ind = linear_sum_assignment(1 - iou_mat)

    matches = []
    unmatched_tracks = list(range(len(tracks)))
    unmatched_dets = list(range(len(dets)))

    for r, c in zip(row_ind, col_ind):
        if iou_mat[r, c] >= match_thresh:
            matches.append((r, c))
            unmatched_tracks.remove(r)
            unmatched_dets.remove(c)

    return matches, unmatched_tracks, unmatched_dets

# 对比: 传统方法 vs ByteTrack
print("传统方法: 丢弃 score < 0.5 的检测 → 丢失遮挡目标")
print("ByteTrack:  保留所有 score > 0.1 的检测 → 两轮匹配回收真实目标")

def compute_mota(tp, fp, fn, id_switches, n_gt):
    """
    MOTA = 1 - (FN + FP + IDSW) / GT
    """
    if n_gt == 0:
        return 1.0
    mota = 1 - (fn + fp + id_switches) / n_gt
    return max(0, mota)

def compute_motp(tp, total_distance):
    """
    MOTP = Σd(i,t) / Σc_t
    平均定位精度：匹配框的平均 IoU 距离
    """
    if tp == 0:
        return 0.0
    return total_distance / tp

# 示例：某跟踪器在 1000 帧上的统计
tp = 8500       # 正确匹配
fp = 800        # 误检
fn = 1500       # 漏检
idsw = 200      # ID 切换
n_gt = 10000    # 总 GT 数

mota = compute_mota(tp, fp, fn, idsw, n_gt)
print(f"MOTA: {mota*100:.1f}%")  # 1 - (1500+800+200)/10000 = 75.0%

def compute_idf1(id_tp, id_fp, id_fn):
    """
    IDF1 = 2 * ID_TP / (2 * ID_TP + ID_FP + ID_FN)
    衡量 ID 一致性：同一个 ID 被正确分配了多少次
    """
    denom = 2 * id_tp + id_fp + id_fn
    if denom == 0:
        return 0.0
    return 2 * id_tp / denom

def compute_hota(det_accuracy, assoc_accuracy, alpha=0.5):
    """
    HOTA = (DetA^alpha * AssA^(1-alpha))^0.5
    统一检测和关联两个维度的评估
    """
    if det_accuracy <= 0 or assoc_accuracy <= 0:
        return 0.0
    return (det_accuracy  alpha * assoc_accuracy  (1 - alpha)) ** 0.5

# 对比三个指标对同一跟踪器的评估
results = {
    "Tracker_A": {"MOTA": 0.75, "IDF1": 0.72, "HOTA": 0.62},
    "Tracker_B": {"MOTA": 0.78, "IDF1": 0.65, "HOTA": 0.59},
    "Tracker_C": {"MOTA": 0.70, "IDF1": 0.80, "HOTA": 0.63},
}
print(f"{'Tracker':<12} {'MOTA':>8} {'IDF1':>8} {'HOTA':>8}")
for name, metrics in results.items():
    print(f"{name:<12} {metrics['MOTA']*100:>7.1f}% {metrics['IDF1']*100:>7.1f}% {metrics['HOTA']*100:>7.1f}%")
# Tracker_B MOTA 最高但 IDF1 最低 → 检测准但 ID 切换多
# Tracker_C IDF1 最高但 MOTA 最低 → ID 保持好但有漏检

import cv2
import numpy as np
from deep_sort_realtime.deepsort_tracker import DeepSort

class VideoTracker:
    """YOLO + DeepSORT 完整视频跟踪系统"""
    def __init__(self, video_path, output_path, conf_thresh=0.45):
        self.cap = cv2.VideoCapture(video_path)
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # YOLO 检测器（使用 Ultralytics）
        from ultralytics import YOLO
        self.model = YOLO("yolov8n.pt")
        self.conf_thresh = conf_thresh

        # DeepSORT 跟踪器
        self.tracker = DeepSort(
            max_age=30,
            n_init=3,
            nms_max_overlap=1.0,
            max_cosine_distance=0.2,
            nn_budget=100,
            override_track_class=None
        )

        # 输出视频写入
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        self.out = cv2.VideoWriter(output_path, fourcc, self.fps,
                                    (self.width, self.height))

        # 颜色表
        np.random.seed(42)
        self.colors = np.random.randint(0, 255, size=(100, 3))

    def process_frame(self, frame):
        """处理单帧"""
        # YOLO 检测
        results = self.model(frame, verbose=False)
        detections = []
        for r in results:
            for box in r.boxes:
                if box.conf[0] >= self.conf_thresh:
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    conf = float(box.conf[0])
                    cls = int(box.cls[0])
                    # DeepSORT 格式: ([x1,y1,w,h], conf, class)
                    w, h = x2 - x1, y2 - y1
                    detections.append(([x1, y1, w, h], conf, cls))

        # DeepSORT 更新
        tracks = self.tracker.update_tracks(detections, frame=frame)

        # 渲染
        for track in tracks:
            if not track.is_confirmed():
                continue
            track_id = track.track_id
            ltrb = track.to_ltrb()
            x1, y1, x2, y2 = map(int, ltrb)
            color = self.colors[track_id % 100].tolist()
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            cv2.putText(frame, f"ID:{track_id}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        return frame

    def run(self):
        """处理整个视频"""
        frame_count = 0
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break

            frame = self.process_frame(frame)
            self.out.write(frame)
            frame_count += 1

            if frame_count % 100 == 0:
                print(f"已处理 {frame_count} 帧")

        self.cap.release()
        self.out.release()
        print(f"完成! 共处理 {frame_count} 帧, "
              f"FPS: {frame_count / (frame_count / self.fps):.1f}")

# 使用示例
if __name__ == "__main__":
    tracker = VideoTracker(
        video_path="input.mp4",
        output_path="output.mp4",
        conf_thresh=0.45
    )
    tracker.run()

# 性能优化建议:
# 1. YOLO 使用 n/s/m/l/x 不同尺寸模型平衡速度和精度
# 2. 高分辨率视频可以先 resize 到 640x640 再检测
# 3. 使用 TensorRT 推理可以加速 3-5 倍
# 4. 跟踪器可以跑在 CPU 上，节省 GPU 给检测器