singlebehaviorlab 2.0.0__py3-none-any.whl

sam2/utils/transforms.py

@@ -0,0 +1,118 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms import Normalize, Resize, ToTensor
+
+
+class SAM2Transforms(nn.Module):
+    def __init__(
+        self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0
+    ):
+        """
+        Transforms for SAM2.
+        """
+        super().__init__()
+        self.resolution = resolution
+        self.mask_threshold = mask_threshold
+        self.max_hole_area = max_hole_area
+        self.max_sprinkle_area = max_sprinkle_area
+        self.mean = [0.485, 0.456, 0.406]
+        self.std = [0.229, 0.224, 0.225]
+        self.to_tensor = ToTensor()
+        self.transforms = torch.jit.script(
+            nn.Sequential(
+                Resize((self.resolution, self.resolution)),
+                Normalize(self.mean, self.std),
+            )
+        )
+
+    def __call__(self, x):
+        x = self.to_tensor(x)
+        return self.transforms(x)
+
+    def forward_batch(self, img_list):
+        img_batch = [self.transforms(self.to_tensor(img)) for img in img_list]
+        img_batch = torch.stack(img_batch, dim=0)
+        return img_batch
+
+    def transform_coords(
+        self, coords: torch.Tensor, normalize=False, orig_hw=None
+    ) -> torch.Tensor:
+        """
+        Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates,
+        If the coords are in absolute image coordinates, normalize should be set to True and original image size is required.
+
+        Returns
+            Un-normalized coordinates in the range of [0, 1] which is expected by the SAM2 model.
+        """
+        if normalize:
+            assert orig_hw is not None
+            h, w = orig_hw
+            coords = coords.clone()
+            coords[..., 0] = coords[..., 0] / w
+            coords[..., 1] = coords[..., 1] / h
+
+        coords = coords * self.resolution  # unnormalize coords
+        return coords
+
+    def transform_boxes(
+        self, boxes: torch.Tensor, normalize=False, orig_hw=None
+    ) -> torch.Tensor:
+        """
+        Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates,
+        if the coords are in absolute image coordinates, normalize should be set to True and original image size is required.
+        """
+        boxes = self.transform_coords(boxes.reshape(-1, 2, 2), normalize, orig_hw)
+        return boxes
+
+    def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor:
+        """
+        Perform PostProcessing on output masks.
+        """
+        from sam2.utils.misc import get_connected_components
+
+        masks = masks.float()
+        input_masks = masks
+        mask_flat = masks.flatten(0, 1).unsqueeze(1)  # flatten as 1-channel image
+        try:
+            if self.max_hole_area > 0:
+                # Holes are those connected components in background with area <= self.fill_hole_area
+                # (background regions are those with mask scores <= self.mask_threshold)
+                labels, areas = get_connected_components(
+                    mask_flat <= self.mask_threshold
+                )
+                is_hole = (labels > 0) & (areas <= self.max_hole_area)
+                is_hole = is_hole.reshape_as(masks)
+                # We fill holes with a small positive mask score (10.0) to change them to foreground.
+                masks = torch.where(is_hole, self.mask_threshold + 10.0, masks)
+
+            if self.max_sprinkle_area > 0:
+                labels, areas = get_connected_components(
+                    mask_flat > self.mask_threshold
+                )
+                is_hole = (labels > 0) & (areas <= self.max_sprinkle_area)
+                is_hole = is_hole.reshape_as(masks)
+                # We fill holes with negative mask score (-10.0) to change them to background.
+                masks = torch.where(is_hole, self.mask_threshold - 10.0, masks)
+        except Exception as e:
+            # Skip the post-processing step if the CUDA kernel fails
+            warnings.warn(
+                f"{e}\n\nSkipping the post-processing step due to the error above. You can "
+                "still use SAM 2 and it's OK to ignore the error above, although some post-processing "
+                "functionality may be limited (which doesn't affect the results in most cases; see "
+                "https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).",
+                category=UserWarning,
+                stacklevel=2,
+            )
+            masks = input_masks
+
+        masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False)
+        return masks

singlebehaviorlab/__init__.py

@@ -0,0 +1,4 @@
+"""SingleBehaviorLab — behavioral video annotation and training tool."""
+
+__version__ = "1.3.2"
+__author__ = "Almir Aljovic"

singlebehaviorlab/__main__.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+"""
+Main entry point for SingleBehaviorLab.
+Runs when invoked as:
+  python -m singlebehaviorlab
+  singlebehaviorlab          (pip entry point)
+"""
+
+import logging
+import sys
+import os
+
+os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
+os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")
+os.environ.setdefault("GRPC_VERBOSITY", "ERROR")
+os.environ.setdefault("GLOG_minloglevel", "2")
+
+# Let JAX grow GPU memory on demand and leave headroom for PyTorch.
+os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.45"
+# Fall back to driver JIT compilation when ptxas/nvlink is unavailable.
+os.environ["XLA_FLAGS"] = "--xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found"
+
+import yaml
+from singlebehaviorlab._paths import get_default_config_path, get_experiments_dir
+from singlebehaviorlab.gui.main_window import MainWindow
+
+
+def load_config(config_path: str = None) -> dict:
+    """Load configuration from YAML file."""
+    if config_path is None:
+        config_path = str(get_default_config_path())
+
+    if os.path.exists(config_path):
+        with open(config_path, 'r') as f:
+            config = yaml.safe_load(f) or {}
+    else:
+        config = {}
+
+    # For path keys that are missing or blank, resolve relative to the
+    # experiments directory (pip install) or the package parent (source install).
+    from singlebehaviorlab._paths import get_package_dir
+    base_dir = str(get_package_dir().parent)
+
+    defaults = {
+        "data_dir":        os.path.join(base_dir, "data"),
+        "raw_videos_dir":  os.path.join(base_dir, "data", "raw_videos"),
+        "clips_dir":       os.path.join(base_dir, "data", "clips"),
+        "annotations_dir": os.path.join(base_dir, "data", "annotations"),
+        "models_dir":      os.path.join(base_dir, "models", "behavior_heads"),
+        "backbone_dir":    os.path.join(base_dir, "models", "videoprism_backbone"),
+        "annotation_file": os.path.join(base_dir, "data", "annotations", "annotations.json"),
+    }
+
+    for key, value in defaults.items():
+        if not config.get(key):
+            config[key] = value
+        elif not os.path.isabs(config[key]):
+            config[key] = os.path.join(base_dir, config[key])
+
+    experiments_dir = str(get_experiments_dir())
+    if not config.get("experiments_dir"):
+        config["experiments_dir"] = experiments_dir
+    os.makedirs(config["experiments_dir"], exist_ok=True)
+
+    config["config_path"] = config_path
+    config.setdefault("experiment_name", None)
+    config.setdefault("experiment_path", None)
+
+    return config
+
+
+def main():
+    """Application entry point."""
+    logging.basicConfig(
+        level=logging.WARNING,
+        format="%(levelname)s [%(name)s] %(message)s",
+    )
+    from PyQt6.QtWidgets import (
+        QApplication, QDialog, QVBoxLayout, QPushButton, QLabel,
+    )
+
+    app = QApplication(sys.argv)
+    app.setApplicationName("SingleBehaviorLab")
+
+    config = load_config()
+
+    startup_dialog = QDialog()
+    startup_dialog.setWindowTitle("Welcome - Experiment Management")
+    startup_dialog.setMinimumSize(400, 200)
+    startup_dialog.setModal(True)
+
+    layout = QVBoxLayout()
+
+    welcome_label = QLabel(
+        "<h2>Welcome to SingleBehaviorLab</h2>"
+        "<p>Please choose an option to get started:</p>"
+    )
+    welcome_label.setWordWrap(True)
+    layout.addWidget(welcome_label)
+
+    create_btn = QPushButton("Create New Experiment")
+    create_btn.setMinimumHeight(40)
+    create_btn.setStyleSheet("font-size: 12px; font-weight: bold;")
+    create_btn.clicked.connect(startup_dialog.accept)
+    layout.addWidget(create_btn)
+
+    load_btn = QPushButton("Load Existing Experiment")
+    load_btn.setMinimumHeight(40)
+    load_btn.setStyleSheet("font-size: 12px; font-weight: bold;")
+    load_btn.clicked.connect(startup_dialog.reject)
+    layout.addWidget(load_btn)
+
+    startup_dialog.setLayout(layout)
+
+    result = startup_dialog.exec()
+
+    window = MainWindow(config)
+    window.show()
+
+    if result == QDialog.DialogCode.Accepted:
+        window._create_experiment()
+    elif result == QDialog.DialogCode.Rejected:
+        window._load_experiment()
+
+    sys.exit(app.exec())
+
+
+if __name__ == "__main__":
+    main()

singlebehaviorlab/_paths.py

@@ -0,0 +1,100 @@
+"""
+Central path resolver for SingleBehaviorLab.
+
+Handles two install modes transparently:
+
+  Source / zip distribution
+    SingleBehaviorLab/
+      singlebehaviorlab/   ← this file lives here
+      sam2_backend/
+      sam2_checkpoints/
+      experiments/
+
+  pip install (site-packages)
+    site-packages/singlebehaviorlab/  ← this file lives here
+    ~/SingleBehaviorLab/              ← user data lives here
+"""
+
+from pathlib import Path
+
+# Directory where this file (and the package) lives.
+_PKG_DIR = Path(__file__).parent
+
+# One level up: the SingleBehaviorLab root when running from source/zip,
+# or site-packages when pip-installed.
+_PKG_PARENT = _PKG_DIR.parent
+
+# Standard user data directory — always writable.
+USER_DATA_DIR = Path.home() / "SingleBehaviorLab"
+
+
+def _first_existing(*candidates: Path) -> Path:
+    """Return the first candidate path that exists, otherwise the first candidate."""
+    for p in candidates:
+        if p.exists():
+            return p
+    return candidates[0]
+
+
+# Public helpers.
+
+def get_package_dir() -> Path:
+    """The installed singlebehaviorlab package directory."""
+    return _PKG_DIR
+
+
+def get_sam2_backend_dir() -> Path:
+    """Return the directory containing the sam2 Python package."""
+    try:
+        import sam2  # type: ignore
+        sam2_pkg = Path(sam2.__file__).resolve().parent
+        if (sam2_pkg / "configs").exists():
+            return sam2_pkg.parent
+    except Exception:
+        pass
+    return _first_existing(
+        _PKG_PARENT / "sam2_backend",
+        _PKG_DIR / "sam2_backend",
+        USER_DATA_DIR / "sam2_backend",
+    )
+
+
+def get_sam2_checkpoints_dir() -> Path:
+    """Return the directory where SAM2 checkpoint files live, creating it if absent."""
+    source_dir = _PKG_PARENT / "sam2_checkpoints"
+    if source_dir.exists():
+        return source_dir
+    user_dir = USER_DATA_DIR / "sam2_checkpoints"
+    user_dir.mkdir(parents=True, exist_ok=True)
+    return user_dir
+
+
+def get_training_profiles_path() -> Path:
+    """Locate training_profiles.json."""
+    return _first_existing(
+        _PKG_DIR / "data" / "training_profiles.json",          # package data
+        USER_DATA_DIR / "training_profiles.json",              # user copy
+    )
+
+
+def get_default_config_path() -> Path:
+    """Locate the default (template) config.yaml."""
+    return _first_existing(
+        _PKG_PARENT / "config" / "config.yaml",                # source / zip install
+        _PKG_DIR / "data" / "config" / "config.yaml",          # package data
+    )
+
+
+def get_experiments_dir() -> Path:
+    """
+    Default experiments root directory.
+
+    - Source / zip install: SingleBehaviorLab/experiments/
+    - pip install: ~/SingleBehaviorLab/experiments/
+    """
+    local = _PKG_PARENT / "experiments"
+    if local.exists():
+        return local
+    user_exp = USER_DATA_DIR / "experiments"
+    user_exp.mkdir(parents=True, exist_ok=True)
+    return user_exp

singlebehaviorlab/backend/__init__.py

@@ -0,0 +1,2 @@
+# Backend modules
+

singlebehaviorlab/backend/augmentations.py

@@ -0,0 +1,320 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as F
+from typing import Optional
+
+
+class ClipAugment(nn.Module):
+    """Applies identical augmentation to every frame; no cropping to avoid cutting off tracked subjects."""
+    
+    def __init__(
+        self,
+        use_horizontal_flip: bool = True,
+        use_vertical_flip: bool = False,
+        use_color_jitter: bool = True,
+        use_gaussian_blur: bool = False,
+        use_random_noise: bool = False,
+        use_small_rotation: bool = False,
+        use_speed_perturb: bool = False,
+        use_random_shapes: bool = False,
+        use_grayscale: bool = False,
+        use_lighting_robustness: bool = False,
+        color_jitter_brightness: float = 0.2,
+        color_jitter_contrast: float = 0.2,
+        color_jitter_saturation: float = 0.2,
+        color_jitter_hue: float = 0.1,
+        gaussian_blur_sigma: tuple[float, float] = (0.1, 0.5),
+        noise_std: float = 0.01,
+        rotation_degrees: float = 2.0,
+        speed_range: tuple[float, float] = (0.7, 1.3),
+        random_shapes_max: int = 3,
+        random_shapes_max_size: float = 0.15,
+        grayscale_prob: float = 0.5,
+        gamma_range: tuple[float, float] = (0.75, 1.35),
+        channel_gain_range: tuple[float, float] = (0.85, 1.15),
+    ):
+        super().__init__()
+        self.use_horizontal_flip = use_horizontal_flip
+        self.use_vertical_flip = use_vertical_flip
+        self.use_color_jitter = use_color_jitter
+        self.use_gaussian_blur = use_gaussian_blur
+        self.use_random_noise = use_random_noise
+        self.use_small_rotation = use_small_rotation
+        self.use_speed_perturb = use_speed_perturb
+        self.use_random_shapes = use_random_shapes
+        self.use_grayscale = use_grayscale
+        self.use_lighting_robustness = use_lighting_robustness
+        
+        self.color_jitter_brightness = color_jitter_brightness
+        self.color_jitter_contrast = color_jitter_contrast
+        self.color_jitter_saturation = color_jitter_saturation
+        self.color_jitter_hue = color_jitter_hue
+        self.gaussian_blur_sigma = gaussian_blur_sigma
+        self.noise_std = noise_std
+        self.rotation_degrees = rotation_degrees
+        self.speed_range = speed_range
+        self.random_shapes_max = max(1, random_shapes_max)
+        self.random_shapes_max_size = random_shapes_max_size
+        self.grayscale_prob = grayscale_prob
+        self.gamma_range = gamma_range
+        self.channel_gain_range = channel_gain_range
+    
+    def _sample_params(self) -> dict:
+        """Sample one set of augmentation params shared by all frames."""
+        hflip = self.use_horizontal_flip and (torch.rand(1).item() < 0.5)
+        vflip = self.use_vertical_flip and (torch.rand(1).item() < 0.5)
+
+        if self.use_color_jitter:
+            brightness_factor = torch.empty(1).uniform_(
+                1 - self.color_jitter_brightness,
+                1 + self.color_jitter_brightness
+            ).item() if self.color_jitter_brightness > 0 else 1.0
+            
+            contrast_factor = torch.empty(1).uniform_(
+                1 - self.color_jitter_contrast,
+                1 + self.color_jitter_contrast
+            ).item() if self.color_jitter_contrast > 0 else 1.0
+            
+            saturation_factor = torch.empty(1).uniform_(
+                1 - self.color_jitter_saturation,
+                1 + self.color_jitter_saturation
+            ).item() if self.color_jitter_saturation > 0 else 1.0
+            
+            hue_factor = torch.empty(1).uniform_(
+                -self.color_jitter_hue,
+                self.color_jitter_hue
+            ).item() if self.color_jitter_hue > 0 else 0.0
+        else:
+            brightness_factor = contrast_factor = saturation_factor = hue_factor = None
+        
+        if self.use_gaussian_blur:
+            blur_sigma = torch.empty(1).uniform_(
+                self.gaussian_blur_sigma[0],
+                self.gaussian_blur_sigma[1]
+            ).item()
+        else:
+            blur_sigma = None
+        
+        if self.use_small_rotation:
+            rotation_angle = torch.empty(1).uniform_(
+                -self.rotation_degrees,
+                self.rotation_degrees
+            ).item()
+        else:
+            rotation_angle = None
+        
+        if self.use_random_noise:
+            noise_std = self.noise_std
+        else:
+            noise_std = None
+
+        gamma_factor = None
+        channel_gains = None
+        if self.use_lighting_robustness:
+            gamma_factor = torch.empty(1).uniform_(
+                self.gamma_range[0],
+                self.gamma_range[1]
+            ).item()
+            channel_gains = torch.empty(3).uniform_(
+                self.channel_gain_range[0],
+                self.channel_gain_range[1]
+            ).tolist()
+
+        # Speed perturbation: sample a speed factor once per clip
+        speed_factor = None
+        if self.use_speed_perturb:
+            speed_factor = torch.empty(1).uniform_(
+                self.speed_range[0], self.speed_range[1]
+            ).item()
+
+        # Random shapes: sample positions, sizes, colors, types once per clip
+        shapes = None
+        if self.use_random_shapes:
+            n_shapes = torch.randint(1, self.random_shapes_max + 1, (1,)).item()
+            shapes = []
+            for _ in range(n_shapes):
+                shape_type = torch.randint(0, 3, (1,)).item()  # 0=rect, 1=ellipse, 2=triangle
+                cx = torch.rand(1).item()
+                cy = torch.rand(1).item()
+                sw = torch.empty(1).uniform_(0.03, self.random_shapes_max_size).item()
+                sh = torch.empty(1).uniform_(0.03, self.random_shapes_max_size).item()
+                color = torch.rand(3).tolist()
+                shapes.append({
+                    "type": shape_type, "cx": cx, "cy": cy,
+                    "sw": sw, "sh": sh, "color": color,
+                })
+
+        do_grayscale = False
+        if self.use_grayscale:
+            do_grayscale = torch.rand(1).item() < self.grayscale_prob
+
+        return {
+            "hflip": bool(hflip),
+            "vflip": bool(vflip),
+            "brightness_factor": brightness_factor,
+            "contrast_factor": contrast_factor,
+            "saturation_factor": saturation_factor,
+            "hue_factor": hue_factor,
+            "blur_sigma": blur_sigma,
+            "rotation_angle": rotation_angle,
+            "noise_std": noise_std,
+            "gamma_factor": gamma_factor,
+            "channel_gains": channel_gains,
+            "speed_factor": speed_factor,
+            "shapes": shapes,
+            "grayscale": do_grayscale,
+        }
+
+    @staticmethod
+    def _resample_temporal(clip: torch.Tensor, speed_factor: float) -> torch.Tensor:
+        """Resample clip frames to simulate speed change. Output has same T."""
+        T = clip.shape[0]
+        if T <= 1 or abs(speed_factor - 1.0) < 0.01:
+            return clip
+        # At speed_factor > 1 we want to cover more source frames (speed up),
+        # so we sample from a wider window (indices can exceed T-1 → clamp).
+        # At speed_factor < 1 we cover fewer source frames (slow down),
+        # so indices cluster in the center.
+        src_indices = torch.linspace(0, (T - 1) * speed_factor, T)
+        src_indices = src_indices.clamp(0, T - 1)
+        idx = src_indices.round().long()
+        return clip[idx]
+
+    @staticmethod
+    def _draw_shapes_on_frame(frame: torch.Tensor, shapes: list, H: int, W: int) -> torch.Tensor:
+        """Draw pre-sampled shapes onto a single frame [C, H, W]."""
+        frame = frame.clone()
+        for s in shapes:
+            cx_px = int(s["cx"] * W)
+            cy_px = int(s["cy"] * H)
+            half_w = max(1, int(s["sw"] * W / 2))
+            half_h = max(1, int(s["sh"] * H / 2))
+            color = s["color"]
+
+            x1 = max(0, cx_px - half_w)
+            x2 = min(W, cx_px + half_w)
+            y1 = max(0, cy_px - half_h)
+            y2 = min(H, cy_px + half_h)
+            if x2 <= x1 or y2 <= y1:
+                continue
+
+            stype = s["type"]
+            if stype == 0:
+                for c_i in range(min(3, frame.shape[0])):
+                    frame[c_i, y1:y2, x1:x2] = color[c_i]
+            elif stype == 1:
+                yy = torch.arange(y1, y2, device=frame.device).float()
+                xx = torch.arange(x1, x2, device=frame.device).float()
+                gy, gx = torch.meshgrid(yy, xx, indexing="ij")
+                ey = (gy - cy_px) / max(half_h, 1)
+                ex = (gx - cx_px) / max(half_w, 1)
+                mask = (ex ** 2 + ey ** 2) <= 1.0
+                for c_i in range(min(3, frame.shape[0])):
+                    region = frame[c_i, y1:y2, x1:x2]
+                    region[mask] = color[c_i]
+            elif stype == 2:
+                # Triangle pointing up: apex at top center, base at bottom
+                yy = torch.arange(y1, y2, device=frame.device).float()
+                xx = torch.arange(x1, x2, device=frame.device).float()
+                gy, gx = torch.meshgrid(yy, xx, indexing="ij")
+                ny = (gy - y1) / max(y2 - y1 - 1, 1)
+                nx = (gx - x1) / max(x2 - x1 - 1, 1)
+                mask = (nx >= 0.5 - 0.5 * ny) & (nx <= 0.5 + 0.5 * ny)
+                for c_i in range(min(3, frame.shape[0])):
+                    region = frame[c_i, y1:y2, x1:x2]
+                    region[mask] = color[c_i]
+        return frame
+
+    def _apply_with_params(self, clip: torch.Tensor, params: dict) -> torch.Tensor:
+        T, C, H, W = clip.shape
+
+        # Speed perturbation (resamples frames, applied before per-frame ops)
+        speed_factor = params.get("speed_factor", None)
+        if speed_factor is not None:
+            clip = self._resample_temporal(clip, speed_factor)
+
+        hflip = bool(params.get("hflip", params.get("flip", False)))
+        vflip = bool(params.get("vflip", False))
+        brightness_factor = params.get("brightness_factor", None)
+        contrast_factor = params.get("contrast_factor", None)
+        saturation_factor = params.get("saturation_factor", None)
+        hue_factor = params.get("hue_factor", None)
+        blur_sigma = params.get("blur_sigma", None)
+        rotation_angle = params.get("rotation_angle", None)
+        noise_std = params.get("noise_std", None)
+        gamma_factor = params.get("gamma_factor", None)
+        channel_gains = params.get("channel_gains", None)
+        shapes = params.get("shapes", None)
+        do_grayscale = params.get("grayscale", False)
+
+        augmented_frames = []
+        for t in range(T):
+            frame = clip[t]
+
+            if hflip:
+                frame = F.hflip(frame)
+            if vflip:
+                frame = F.vflip(frame)
+
+            if do_grayscale:
+                frame = F.rgb_to_grayscale(frame, num_output_channels=C)
+
+            if brightness_factor is not None:
+                frame = F.adjust_brightness(frame, brightness_factor)
+                frame = F.adjust_contrast(frame, contrast_factor)
+                frame = F.adjust_saturation(frame, saturation_factor)
+                frame = F.adjust_hue(frame, hue_factor)
+
+            # Gamma correction + per-channel gain shifts for lighting robustness
+            if gamma_factor is not None:
+                frame = torch.clamp(frame, 0.0, 1.0)
+                frame = torch.pow(frame + 1e-6, gamma_factor)
+                if channel_gains is not None and C >= 3:
+                    gains = torch.tensor(channel_gains[:3], device=frame.device, dtype=frame.dtype).view(3, 1, 1)
+                    frame[:3] = frame[:3] * gains
+                frame = torch.clamp(frame, 0.0, 1.0)
+            
+            if rotation_angle is not None and abs(rotation_angle) > 0.01:
+                frame = F.rotate(frame, rotation_angle, interpolation=F.InterpolationMode.BILINEAR, fill=0.0)
+
+            if blur_sigma is not None and blur_sigma > 0.01:
+                kernel_size = int(2 * int(4 * blur_sigma + 0.5) + 1)
+                if kernel_size >= 3:
+                    frame = F.gaussian_blur(frame, kernel_size=[kernel_size, kernel_size], sigma=[blur_sigma, blur_sigma])
+            
+            # Random shapes (same position/color for every frame in the clip)
+            if shapes:
+                frame = self._draw_shapes_on_frame(frame, shapes, H, W)
+
+            if noise_std is not None and noise_std > 0:
+                noise = torch.randn_like(frame) * noise_std
+                frame = torch.clamp(frame + noise, 0.0, 1.0)
+            
+            augmented_frames.append(frame)
+
+        return torch.stack(augmented_frames)
+
+    def augment_with_params(self, clip: torch.Tensor):
+        """Return augmented clip and parameter dict for synchronizing spatial label transforms (e.g., bboxes, masks)."""
+        params = self._sample_params()
+        return self._apply_with_params(clip, params), params
+
+    def forward(self, clip: torch.Tensor) -> torch.Tensor:
+        """clip: [T, C, H, W] float tensor in [0, 1]."""
+        augmented_clip, _ = self.augment_with_params(clip)
+        return augmented_clip
+    
+    def __repr__(self):
+        parts = [
+            f"hflip={self.use_horizontal_flip}",
+            f"vflip={self.use_vertical_flip}",
+            f"color_jitter={self.use_color_jitter}",
+            f"gaussian_blur={self.use_gaussian_blur}",
+            f"random_noise={self.use_random_noise}",
+            f"small_rotation={self.use_small_rotation}",
+            f"speed_perturb={self.use_speed_perturb}",
+            f"random_shapes={self.use_random_shapes}",
+            f"grayscale={self.use_grayscale}",
+            f"lighting_robustness={self.use_lighting_robustness}",
+        ]
+        return f"ClipAugment({', '.join(parts)})"