skeletrack 0.1.0__tar.gz

skeletrack-0.1.0/.gitignore

@@ -0,0 +1,42 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Claude Code
+.claude/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Testing / Coverage
+.pytest_cache/
+htmlcov/
+.coverage
+
+# YOLO model weights
+*.pt
+
+# Directories not yet ready
+docs/
+examples/
+
+# Output data
+*.npy
+
+# Others
+CLAUDE.md

skeletrack-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Steven Tse
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

skeletrack-0.1.0/PKG-INFO

@@ -0,0 +1,151 @@
+Metadata-Version: 2.4
+Name: skeletrack
+Version: 0.1.0
+Summary: Extract multi-person skeleton trajectories from videos with one line of code.
+Project-URL: Homepage, https://github.com/StevenUST/skeletrack
+Author: Steven Tse
+License-Expression: MIT
+License-File: LICENSE
+Keywords: keypoints,pose,skeleton,tracking,video
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: numpy>=1.20
+Requires-Dist: opencv-python>=4.5
+Requires-Dist: tqdm>=4.0
+Provides-Extra: age
+Requires-Dist: torch>=1.9; extra == 'age'
+Requires-Dist: torchvision>=0.10; extra == 'age'
+Provides-Extra: dev
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: full
+Requires-Dist: mediapipe>=0.10; extra == 'full'
+Requires-Dist: pandas>=1.3; extra == 'full'
+Requires-Dist: torch>=1.9; extra == 'full'
+Requires-Dist: torchvision>=0.10; extra == 'full'
+Requires-Dist: ultralytics>=8.0; extra == 'full'
+Provides-Extra: mediapipe
+Requires-Dist: mediapipe>=0.10; extra == 'mediapipe'
+Provides-Extra: pandas
+Requires-Dist: pandas>=1.3; extra == 'pandas'
+Provides-Extra: yolo
+Requires-Dist: ultralytics>=8.0; extra == 'yolo'
+Description-Content-Type: text/markdown
+
+# Skeletrack
+
+Extract multi-person skeleton trajectories from videos with one line of code.
+
+```python
+import skeletrack
+
+tracks = skeletrack.extract("video.mp4")
+tracks.filter(min_duration=2.0).save("output.npy")
+```
+
+## How It Works
+
+```
+Video → Person Detection (YOLO) → Multi-Person Tracking (ByteTrack) → Pose Estimation (MediaPipe) → Skeleton Trajectories
+```
+
+Pose estimation runs **after** tracking and filtering, so compute is only spent on valid tracks.
+
+## Installation
+
+```bash
+pip install skeletrack
+pip install skeletrack[yolo,mediapipe]  # with detection + pose backends
+```
+
+### Optional dependencies
+
+| Extra | Packages | Purpose |
+|-------|----------|---------|
+| `yolo` | ultralytics | Person detection + tracking |
+| `mediapipe` | mediapipe | Pose estimation |
+| `pandas` | pandas | DataFrame export |
+| `full` | all of the above | Everything |
+
+## Quick Start
+
+```python
+import skeletrack
+
+# Extract with default settings
+tracks = skeletrack.extract("video.mp4")
+
+# Extract with options
+tracks = skeletrack.extract(
+    "video.mp4",
+    device="cuda",
+    detector="yolo:yolov8s.pt",
+    frame_skip=3,
+    min_duration=1.0,
+)
+
+# Filter + save
+tracks.filter(min_duration=2.0, min_frames=10).save("output.npy")
+
+# Access skeleton data
+for track in tracks:
+    pose = track.skeleton_array("pose")  # shape: (T, 33, 4)
+    print(f"Track {track.track_id}: {track.duration:.1f}s, {track.num_frames} frames")
+
+# Load saved tracks
+tracks = skeletrack.load("output.npy")
+
+# Export to DataFrame
+df = tracks.to_dataframe()
+```
+
+## Reusable Pipeline
+
+For processing multiple videos, create a `Pipeline` to avoid reloading models:
+
+```python
+from skeletrack import Pipeline
+
+pipeline = Pipeline(device="cuda")
+for video in video_list:
+    tracks = pipeline.run(video)
+    tracks.save(f"{video}.npy")
+pipeline.close()
+```
+
+## Output Format
+
+Each track contains:
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `track_id` | int | Unique person ID |
+| `frames` | list[int] | Frame indices |
+| `timestamps` | list[float] | Timestamps (seconds) |
+| `bboxes` | list[BBox] | Bounding boxes (x, y, w, h) |
+| `skeletons` | list[Skeleton] | Keypoints per frame |
+
+Skeleton keypoint groups (via MediaPipe Holistic):
+
+| Group | Keypoints | Dimensions |
+|-------|-----------|------------|
+| `pose` | 33 | x, y, z, visibility |
+| `left_hand` | 21 | x, y, z, visibility |
+| `right_hand` | 21 | x, y, z, visibility |
+| `face` | 468 | x, y, z |
+
+## Requirements
+
+- Python >= 3.9
+- numpy >= 1.20
+- opencv-python >= 4.5
+
+## License
+
+MIT

skeletrack-0.1.0/README.md

@@ -0,0 +1,111 @@
+# Skeletrack
+
+Extract multi-person skeleton trajectories from videos with one line of code.
+
+```python
+import skeletrack
+
+tracks = skeletrack.extract("video.mp4")
+tracks.filter(min_duration=2.0).save("output.npy")
+```
+
+## How It Works
+
+```
+Video → Person Detection (YOLO) → Multi-Person Tracking (ByteTrack) → Pose Estimation (MediaPipe) → Skeleton Trajectories
+```
+
+Pose estimation runs **after** tracking and filtering, so compute is only spent on valid tracks.
+
+## Installation
+
+```bash
+pip install skeletrack
+pip install skeletrack[yolo,mediapipe]  # with detection + pose backends
+```
+
+### Optional dependencies
+
+| Extra | Packages | Purpose |
+|-------|----------|---------|
+| `yolo` | ultralytics | Person detection + tracking |
+| `mediapipe` | mediapipe | Pose estimation |
+| `pandas` | pandas | DataFrame export |
+| `full` | all of the above | Everything |
+
+## Quick Start
+
+```python
+import skeletrack
+
+# Extract with default settings
+tracks = skeletrack.extract("video.mp4")
+
+# Extract with options
+tracks = skeletrack.extract(
+    "video.mp4",
+    device="cuda",
+    detector="yolo:yolov8s.pt",
+    frame_skip=3,
+    min_duration=1.0,
+)
+
+# Filter + save
+tracks.filter(min_duration=2.0, min_frames=10).save("output.npy")
+
+# Access skeleton data
+for track in tracks:
+    pose = track.skeleton_array("pose")  # shape: (T, 33, 4)
+    print(f"Track {track.track_id}: {track.duration:.1f}s, {track.num_frames} frames")
+
+# Load saved tracks
+tracks = skeletrack.load("output.npy")
+
+# Export to DataFrame
+df = tracks.to_dataframe()
+```
+
+## Reusable Pipeline
+
+For processing multiple videos, create a `Pipeline` to avoid reloading models:
+
+```python
+from skeletrack import Pipeline
+
+pipeline = Pipeline(device="cuda")
+for video in video_list:
+    tracks = pipeline.run(video)
+    tracks.save(f"{video}.npy")
+pipeline.close()
+```
+
+## Output Format
+
+Each track contains:
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `track_id` | int | Unique person ID |
+| `frames` | list[int] | Frame indices |
+| `timestamps` | list[float] | Timestamps (seconds) |
+| `bboxes` | list[BBox] | Bounding boxes (x, y, w, h) |
+| `skeletons` | list[Skeleton] | Keypoints per frame |
+
+Skeleton keypoint groups (via MediaPipe Holistic):
+
+| Group | Keypoints | Dimensions |
+|-------|-----------|------------|
+| `pose` | 33 | x, y, z, visibility |
+| `left_hand` | 21 | x, y, z, visibility |
+| `right_hand` | 21 | x, y, z, visibility |
+| `face` | 468 | x, y, z |
+
+## Requirements
+
+- Python >= 3.9
+- numpy >= 1.20
+- opencv-python >= 4.5
+
+## License
+
+MIT

skeletrack-0.1.0/pyproject.toml

@@ -0,0 +1,42 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "skeletrack"
+version = "0.1.0"
+description = "Extract multi-person skeleton trajectories from videos with one line of code."
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.9"
+authors = [{ name = "Steven Tse" }]
+keywords = ["skeleton", "pose", "tracking", "video", "keypoints"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "numpy>=1.20",
+    "opencv-python>=4.5",
+    "tqdm>=4.0",
+]
+
+[project.optional-dependencies]
+yolo = ["ultralytics>=8.0"]
+mediapipe = ["mediapipe>=0.10"]
+age = ["torch>=1.9", "torchvision>=0.10"]
+pandas = ["pandas>=1.3"]
+full = ["skeletrack[yolo,mediapipe,age,pandas]"]
+dev = ["pytest", "pytest-cov", "ruff", "mypy"]
+
+[project.urls]
+Homepage = "https://github.com/StevenUST/skeletrack"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/skeletrack"]
+
+[tool.ruff]
+line-length = 100

skeletrack-0.1.0/src/skeletrack/__init__.py

@@ -0,0 +1,87 @@
+"""Skeletrack: Extract multi-person skeleton trajectories from videos.
+
+Usage::
+
+    import skeletrack
+
+    # One-liner
+    tracks = skeletrack.extract("video.mp4")
+
+    # With options
+    tracks = skeletrack.extract("video.mp4", device="cuda", detector="yolo:yolov8s.pt")
+
+    # Filter + save
+    tracks.filter(min_duration=2.0).save("output.npy")
+
+    # Load previously saved tracks
+    tracks = skeletrack.load("output.npy")
+"""
+
+from ._version import __version__
+from .core.config import PipelineConfig
+from .core.pipeline import Pipeline
+from .data.collection import TrackCollection
+from .data.skeleton import Skeleton
+from .data.track import Track
+
+__all__ = [
+    "__version__",
+    "extract",
+    "load",
+    "Pipeline",
+    "PipelineConfig",
+    "Track",
+    "TrackCollection",
+    "Skeleton",
+]
+
+
+def extract(
+    source: str,
+    *,
+    device: str = "cpu",
+    detector: str = "yolo",
+    pose: str = "mediapipe",
+    frame_skip: int = 3,
+    min_duration: float = 1.0,
+    show_progress: bool = True,
+    **kwargs,
+) -> TrackCollection:
+    """Extract skeleton trajectories from a video file.
+
+    This is the main entry point for quick usage. For processing multiple
+    videos, create a :class:`Pipeline` instance for better performance.
+
+    Args:
+        source:        Path to a video file.
+        device:        "cpu" or "cuda".
+        detector:      Detector backend name (default: "yolo").
+                       Use "yolo:yolov8s.pt" to specify a model variant.
+        pose:          Pose backend name (default: "mediapipe").
+        frame_skip:    Process every N-th frame for detection (default: 3).
+        min_duration:  Discard tracks shorter than this (seconds, default: 1.0).
+        show_progress: Show a tqdm progress bar (default: True).
+        **kwargs:      Additional arguments passed to PipelineConfig.
+
+    Returns:
+        TrackCollection with all detected person skeleton trajectories.
+    """
+    config = PipelineConfig(
+        device=device,
+        frame_skip=frame_skip,
+        min_track_duration=min_duration,
+        **kwargs,
+    )
+    pipeline = Pipeline(config=config, detector=detector, pose=pose)
+    try:
+        return pipeline.run(source, show_progress=show_progress)
+    finally:
+        pipeline.close()
+
+
+def load(path: str) -> TrackCollection:
+    """Load previously saved tracks from file.
+
+    Supports .npy files (both skeletrack and legacy VideoScreener format).
+    """
+    return TrackCollection.load(path)

skeletrack-0.1.0/src/skeletrack/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

skeletrack-0.1.0/src/skeletrack/core/__init__.py

@@ -0,0 +1,5 @@
+from .config import PipelineConfig
+from .pipeline import Pipeline
+from .video import VideoReader
+
+__all__ = ["Pipeline", "PipelineConfig", "VideoReader"]

skeletrack-0.1.0/src/skeletrack/core/config.py

@@ -0,0 +1,38 @@
+"""Pipeline configuration."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class PipelineConfig:
+    """Configuration for the extraction pipeline.
+
+    All parameters have sensible defaults so you can start with just::
+
+        config = PipelineConfig()
+
+    Attributes:
+        device:                   "cpu" or "cuda".
+        frame_skip:               Process every N-th frame for detection/tracking.
+        scene_change_threshold:   Histogram correlation below this triggers a scene reset.
+        max_missing_frames:       Track dies after this many consecutive frames without detection.
+        min_confidence:           Minimum detection confidence.
+        min_track_duration:       Discard tracks shorter than this (seconds).
+        motion_threshold:         Minimum bbox center std-dev to keep a track (pixels).
+                                  Set to None to disable motion filtering.
+        pose_min_detection_conf:  MediaPipe min_detection_confidence.
+        pose_min_tracking_conf:   MediaPipe min_tracking_confidence.
+    """
+
+    device: str = "cpu"
+    frame_skip: int = 3
+    scene_change_threshold: float = 0.4
+    max_missing_frames: int = 60
+    min_confidence: float = 0.4
+    min_track_duration: float = 1.0
+    motion_threshold: float | None = 5.0
+    pose_min_detection_conf: float = 0.5
+    pose_min_tracking_conf: float = 0.5

skeletrack-0.1.0/src/skeletrack/core/pipeline.py

@@ -0,0 +1,178 @@
+"""Main extraction pipeline."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from tqdm import tqdm
+
+from ..data.bbox import BBox
+from ..data.collection import TrackCollection
+from ..data.track import Track
+from ..detection.base import DetectorBackend
+from ..detection.registry import get_detector
+from ..filters.scene import scene_changed
+from ..pose.base import PoseBackend
+from ..pose.registry import get_pose_backend
+from .config import PipelineConfig
+from .video import VideoReader
+
+
+class Pipeline:
+    """Reusable extraction pipeline.
+
+    Usage::
+
+        pipeline = Pipeline(device="cuda")
+        tracks = pipeline.run("video.mp4")
+    """
+
+    def __init__(
+        self,
+        config: PipelineConfig | None = None,
+        *,
+        detector: str | DetectorBackend = "yolo",
+        pose: str | PoseBackend = "mediapipe",
+        device: str | None = None,
+        **kwargs: Any,
+    ):
+        self.config = config or PipelineConfig()
+        if device is not None:
+            self.config.device = device
+
+        # Resolve backends (lazy — only instantiated when needed)
+        self._detector_spec = detector
+        self._pose_spec = pose
+        self._extra_kwargs = kwargs
+        self._detector: DetectorBackend | None = None
+        self._pose: PoseBackend | None = None
+
+    def _get_detector(self) -> DetectorBackend:
+        if self._detector is None:
+            self._detector = get_detector(
+                self._detector_spec,
+                device=self.config.device,
+                min_confidence=self.config.min_confidence,
+                **{k: v for k, v in self._extra_kwargs.items()
+                   if k in ("model_name",)},
+            )
+        return self._detector
+
+    def _get_pose(self) -> PoseBackend:
+        if self._pose is None:
+            self._pose = get_pose_backend(
+                self._pose_spec,
+                min_detection_confidence=self.config.pose_min_detection_conf,
+                min_tracking_confidence=self.config.pose_min_tracking_conf,
+            )
+        return self._pose
+
+    def run(self, source: str, *, show_progress: bool = True) -> TrackCollection:
+        """Extract skeleton trajectories from a video.
+
+        Args:
+            source:        Path to video file.
+            show_progress: Show a tqdm progress bar.
+
+        Returns:
+            TrackCollection with all detected person tracks and their skeletons.
+        """
+        cfg = self.config
+        detector = self._get_detector()
+
+        # -- Phase 1: Detection + Tracking ------------------------------------
+        tracks_by_id: dict[int, Track] = {}
+        yolo_to_track: dict[int, int] = {}  # yolo_track_id -> track_id
+        active_yolo: dict[int, int] = {}    # yolo_track_id -> last_frame_idx
+        next_track_id = 1
+        prev_frame: np.ndarray | None = None
+
+        with VideoReader(source, frame_skip=cfg.frame_skip) as reader:
+            video_meta = reader.metadata()
+            fps = reader.fps
+
+            frames_iter = reader
+            if show_progress:
+                frames_iter = tqdm(
+                    reader,
+                    total=reader.total_frames // cfg.frame_skip,
+                    desc="Tracking",
+                    unit="frame",
+                )
+
+            for info in frames_iter:
+                # Scene change detection
+                if prev_frame is not None:
+                    if scene_changed(prev_frame, info.frame, cfg.scene_change_threshold):
+                        detector.reset()
+                        yolo_to_track.clear()
+                        active_yolo.clear()
+
+                prev_frame = info.frame.copy()
+
+                # Detect + track
+                detections = detector.detect(info.frame, info.frame_idx)
+                current_yolo_ids = set()
+
+                for det in detections:
+                    yolo_id = det.track_id
+                    if yolo_id is None:
+                        continue
+
+                    current_yolo_ids.add(yolo_id)
+
+                    if yolo_id not in yolo_to_track:
+                        tid = next_track_id
+                        next_track_id += 1
+                        yolo_to_track[yolo_id] = tid
+                        tracks_by_id[tid] = Track(track_id=tid)
+
+                    tid = yolo_to_track[yolo_id]
+                    tracks_by_id[tid].add_frame(info.frame_idx, info.timestamp, det.bbox)
+                    active_yolo[yolo_id] = info.frame_idx
+
+                # Expire stale tracks
+                stale = [
+                    yid for yid, last_f in active_yolo.items()
+                    if info.frame_idx - last_f > cfg.max_missing_frames
+                    and yid not in current_yolo_ids
+                ]
+                for yid in stale:
+                    active_yolo.pop(yid, None)
+                    yolo_to_track.pop(yid, None)
+
+        all_tracks = list(tracks_by_id.values())
+
+        # -- Phase 2: Filter --------------------------------------------------
+        if cfg.min_track_duration > 0:
+            all_tracks = [t for t in all_tracks if t.duration >= cfg.min_track_duration]
+
+        if cfg.motion_threshold is not None:
+            kept = []
+            for t in all_tracks:
+                if t.num_frames < 3:
+                    continue
+                centers = np.array([b.center for b in t.bboxes], dtype=np.float32)
+                std = float(np.std(centers, axis=0).mean())
+                if std >= cfg.motion_threshold:
+                    kept.append(t)
+            all_tracks = kept
+
+        # -- Phase 3: Pose Estimation -----------------------------------------
+        pose_backend = self._get_pose()
+
+        with VideoReader(source, frame_skip=1) as reader:
+            for track in tqdm(all_tracks, desc="Pose estimation", disable=not show_progress):
+                for i, (frame_idx, bbox) in enumerate(zip(track.frames, track.bboxes)):
+                    for info in reader.read_range(frame_idx, frame_idx):
+                        skeleton = pose_backend.estimate(info.frame, bbox)
+                        track.skeletons[i] = skeleton
+
+        return TrackCollection(tracks=all_tracks, video_metadata=video_meta)
+
+    def close(self) -> None:
+        if self._detector is not None:
+            self._detector.close()
+        if self._pose is not None:
+            self._pose.close()