Share
## https://sploitus.com/exploit?id=A3729249-BCF1-5E47-BD8F-8B2605EB9053
# SwimTrack PoC

This directory contains a proof of concept for swimmer detection and single-target tracking.

## What it does

- Uses an off-the-shelf `RT-DETRv2-S` checkpoint for person detection.
- Uses TensorRT for GPU inference (FP16 by default).
- Tracks detections across frames (ByteTrack when available, with a built-in fallback tracker).
- Selects one swimmer track and renders bounding box, center point, and path history.
- Processes all configured input videos and writes annotated outputs to `output/`.

## Prerequisites

1. **Linux x86_64 with NVIDIA GPU**
2. **Python 3.12**
3. **CUDA 13 runtime/toolkit compatibility**
4. **TensorRT 10.x** (Python package install is supported in this repo)
5. **ffmpeg with `libx264`** (required when `output.codec: x264`)

Install ffmpeg on Ubuntu/Debian:

```bash
sudo apt-get install ffmpeg
```

## Installation

Run from the `swimtrack/` directory.

Using `uv` (recommended):

```bash
# Install uv if needed
curl -LsSf https://astral.sh/uv/install.sh | sh

# Create/update environment from pyproject.toml + uv.lock
uv sync
```

Using `pip` (manual, less reproducible):

```bash
# Install CUDA 13 PyTorch wheels first
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu130

# Install project dependencies
pip install -e .
```

## Run

Full pipeline (prepare artifacts + process videos):

```bash
uv run python run_poc.py --config example_config.yaml
```

Prepare artifacts only (checkpoint, ONNX, TensorRT engine):

```bash
uv run python run_poc.py --config example_config.yaml --prepare-only
```

## Input and output defaults

- Default input pattern is `input/*.mp4`.
- Default output directory is `output/`.
- Output names are `_tracked.mp4`.

## YAML configuration reference

All tunable parameters live in `example_config.yaml`.

### `runtime`

- `runtime.device` (`str`): CUDA device string, e.g. `cuda:0`.
- `runtime.expected_video_size` (`[int, int]`): required frame size as `[width, height]`.
- `runtime.expected_fps` (`float`): expected frame rate.
- `runtime.fps_tolerance` (`float`): allowed absolute FPS deviation.

### `artifacts`

- `artifacts.checkpoint_url` (`str`): download URL for RT-DETRv2 checkpoint.
- `artifacts.checkpoint_path` (`str`): local checkpoint path.
- `artifacts.rtdetr_config_path` (`str`): RT-DETRv2 model config used for ONNX export.
- `artifacts.onnx_path` (`str`): local ONNX model path.
- `artifacts.engine_path` (`str`): local TensorRT engine path.
- `artifacts.rebuild_onnx` (`bool`): force ONNX re-export even if file exists.
- `artifacts.rebuild_engine` (`bool`): force TensorRT engine rebuild even if file exists.
- `artifacts.use_fp16` (`bool`): build TensorRT engine with FP16 enabled.
- `artifacts.trt_workspace_gb` (`float`): TensorRT workspace memory limit in GB.

### `input`

- `input.videos` (`list[str]`): list of video paths and/or glob patterns.

### `output`

- `output.directory` (`str`): output directory.
- `output.suffix` (`str`): suffix appended to output filename stem.
- `output.codec` (`str`): final output codec mode. `x264` triggers ffmpeg transcode.
- `output.temp_codec` (`str`): OpenCV codec used for temporary video before transcode.
- `output.require_ffmpeg` (`bool`): if `true`, fail when ffmpeg is unavailable for x264.
- `output.keep_temp_video` (`bool`): keep `.tmp.mp4` intermediate file after transcode.

### `detector`

- `detector.input_size` (`[int, int]`): inference input size as `[width, height]`.
- `detector.score_threshold` (`float`): minimum detection score.
- `detector.person_label` (`int`): class id for person in detector output.
- `detector.max_detections` (`int`): max detections retained per frame.
- `detector.min_box_area` (`float`): minimum bounding-box area filter.
- `detector.roi.enabled` (`bool`): enable ROI center filtering.
- `detector.roi.normalized_xyxy` (`[float, float, float, float]`): ROI in normalized `[x1, y1, x2, y2]`.

### `tracking`

- `tracking.track_thresh` (`float`): minimum confidence for tracker updates.
- `tracking.track_buffer` (`int`): number of frames to keep unmatched tracks.
- `tracking.match_thresh` (`float`): match threshold for data association.
- `tracking.mot20` (`bool`): ByteTrack MOT20 mode flag.

### `single_target`

- `single_target.enabled` (`bool`): if `true`, keep one active swimmer track.
- `single_target.max_missing_frames` (`int`): frames allowed before dropping active track.
- `single_target.center_distance_weight` (`float`): weight for center distance continuity.
- `single_target.score_weight` (`float`): weight favoring higher-confidence tracks.
- `single_target.area_weight` (`float`): weight favoring larger-area tracks.

### `visualization`

- `visualization.draw_roi` (`bool`): draw detector ROI rectangle overlay when ROI filtering is enabled.
- `visualization.draw_center_point` (`bool`): draw center marker.
- `visualization.draw_path` (`bool`): draw path polyline.
- `visualization.path_history` (`int`): max center points kept per track.
- `visualization.roi_color_bgr` (`[int, int, int]`): ROI rectangle color in BGR.
- `visualization.box_color_bgr` (`[int, int, int]`): bounding-box color in BGR.
- `visualization.path_color_bgr` (`[int, int, int]`): path color in BGR.
- `visualization.center_color_bgr` (`[int, int, int]`): center color in BGR.
- `visualization.roi_line_thickness` (`int`): ROI rectangle line thickness.
- `visualization.line_thickness` (`int`): box/path line thickness.
- `visualization.center_radius` (`int`): center circle radius.
- `visualization.font_scale` (`float`): label text scale.

## Notes

- Relative paths in YAML are resolved from this repository directory (`swimtrack/`).
- TensorRT engine files are platform-specific; moving between machines/drivers may require rebuild.
- If an existing engine cannot be deserialized, the script automatically rebuilds it from ONNX.
- `RT-DETRv2/` (sibling directory) is required when checkpoint/ONNX export is needed.
- `ByteTrack/` (sibling directory) is used when available; if missing, the fallback tracker is used.