opensportslib 0.0.1.dev2__py3-none-any.whl

opensportslib/core/utils/video_processing.py

@@ -0,0 +1,389 @@
+import torch
+import numpy as np
+import math
+
+try:
+    import decord
+    from decord import cpu
+    USE_DECORD = True
+except:
+    import av
+    USE_DECORD = False
+
+def read_video(video_path):
+    """Read video frames into list of HxWxC uint8 arrays for VideoMAE."""
+    if USE_DECORD:
+        vr = decord.VideoReader(video_path, ctx=cpu(0))
+        frames = vr.get_batch(range(len(vr)))  # (T, H, W, C)
+        frames = frames.asnumpy().astype(np.uint8)  # ensure uint8 for VideoMAE
+        frames_list = [frame for frame in frames]  # list of T frames
+    else:
+        container = av.open(video_path)
+        frames = []
+        for frame in container.decode(video=0):
+            frames.append(frame.to_ndarray(format="rgb24").astype(np.uint8))
+        return frames
+    return frames_list
+
+
+def resample_video_idx(num_frames, original_fps, new_fps):
+    """Return frame indices to match new fps"""
+    step = float(original_fps) / new_fps
+    if step.is_integer():
+        step = int(step)
+        return slice(None, None, step)
+    idxs = torch.arange(num_frames, dtype=torch.float32) * step
+    idxs = idxs.floor().to(torch.int64)
+    return idxs
+
+def process_frames(frames, target_num_frames, input_fps, target_fps, start_frame=0, end_frame=-1, uniform_sample=False):
+    """
+    frames: list of np arrays (H, W, C)
+    target_num_frames: int
+    Returns: list of np arrays (H, W, C) ready for processor
+    """
+
+    
+    target_fps = target_fps or input_fps
+    num_frames = len(frames)
+    end_frame = end_frame if end_frame != -1 else num_frames
+    duration = num_frames / input_fps
+
+    # unfiorm sampling throughout the video
+    if uniform_sample:
+        # Too short → resample with new fps
+        if num_frames < target_num_frames:
+            new_fps = np.ceil(target_num_frames / duration)
+            idxs = resample_video_idx(target_num_frames, input_fps, new_fps)
+            idxs = np.clip(idxs, 0, num_frames - 1)
+            frames = [frames[i] for i in idxs]
+
+        # Too long → uniform sampling
+        elif num_frames > target_num_frames:
+            idxs = np.linspace(0, num_frames - 1, target_num_frames).astype(int)
+            frames = [frames[i] for i in idxs]
+
+        # Pad if still short
+        if len(frames) < target_num_frames:
+            pad = target_num_frames - len(frames)
+            frames.extend([frames[-1]] * pad)
+    else:
+        window = frames[start_frame:end_frame]          # 24 frames
+        assert len(window) > 0, "Empty temporal window"
+        factor = input_fps / target_fps
+        idxs = [int(i * factor) for i in range(target_num_frames)]
+        idxs = [min(i, len(window) - 1) for i in idxs]
+        frames = [window[i] for i in idxs]
+        assert len(frames) == target_num_frames, f"Expected {target_num_frames} frames, got {len(frames)}"
+
+    return frames
+
+
+def get_stride(src_fps, sample_fps):
+    """Get stride to apply based on the input and output fps.
+
+    Args:
+        src_fps (int): The input fps of the video.
+        sample_fps (int): The output fps.
+    Returns:
+        stride (int): The stride to apply.
+    """
+    if sample_fps <= 0:
+        stride = 1
+    else:
+        stride = int(src_fps / sample_fps)
+    return stride
+
+
+def read_fps(fps, sample_fps):
+    """Computes the exact output fps based on input fps and wanted output fps.
+    Example: if input fps is 25 and wanted output fps is 2, the exact output fps is 2.0833333333333335.
+
+    Args:
+        fps (int): The input fps.
+        sample_fps (int): The wanted output fps.
+
+    Returns:
+        est_out_fps (float): The exact output fps.
+
+    """
+    stride = get_stride(fps, sample_fps)
+    est_out_fps = fps / stride
+    return est_out_fps
+
+def get_num_frames(num_frames, fps, sample_fps):
+    """Compute the number of frames of a video after fps changes.
+
+    Args:
+        num_frames (int): Number of frames in the base video.
+        fps (int): The input fps.
+        sample_fps (int): The output fps.
+
+    Returns:
+        (int): The number of frames with the output fps.
+    """
+    return math.ceil(num_frames / get_stride(fps, sample_fps))
+
+
+def distribute_elements(batch_size, len_devices):
+    """Return a list containing the distribution of the batch along the devices.
+
+    Args:
+        batch_size (int).
+        len_device (int).
+
+    Returns:
+        distribution (list): For example if batch size is 8 and there is 4 gpus, the distribution is [2,2,2,2], meaning that each gpu will process 2 samples.
+    """
+    quotient, remainder = divmod(batch_size, len_devices)
+    distribution = [quotient] * len_devices
+    if remainder > 0:
+        for i in range(len(distribution)):
+            distribution[i] += 1
+
+    return distribution
+
+def _get_deferred_rgb_transform(IMAGENET_MEAN, IMAGENET_STD):
+    import torchvision.transforms as T
+    import torch.nn as nn
+    img_transforms = [
+        # Jittering separately is faster (low variance)
+        T.RandomApply(
+            nn.ModuleList([T.ColorJitter(hue=0.2)]), p=0.25
+        ),
+        T.RandomApply(
+            nn.ModuleList([T.ColorJitter(saturation=(0.7, 1.2))]), p=0.25
+        ),
+        T.RandomApply(
+            nn.ModuleList([T.ColorJitter(brightness=(0.7, 1.2))]), p=0.25
+        ),
+        T.RandomApply(
+            nn.ModuleList([T.ColorJitter(contrast=(0.7, 1.2))]), p=0.25
+        ),
+        # Jittering together is slower (high variance)
+        # transforms.RandomApply(
+        #     nn.ModuleList([
+        #         transforms.ColorJitter(
+        #             brightness=(0.7, 1.2), contrast=(0.7, 1.2),
+        #             saturation=(0.7, 1.2), hue=0.2)
+        #     ]), 0.8),
+        T.RandomApply(nn.ModuleList([T.GaussianBlur(5)]), p=0.25),
+        T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+    ]
+    return torch.jit.script(nn.Sequential(*img_transforms))
+
+def get_remaining(data_len, batch_size):
+    """Return the padding that ensures that all batches have an equal number of items, which is required with the pipeline to make sur that all clips are processed.
+    Args:
+        data_len (int): The length of dataset.
+        batch_size (int).
+
+    Returns:
+        (int): The number of elements to add.
+    """
+    return (math.ceil(data_len / batch_size) * batch_size) - data_len
+
+
+import random
+import numpy as np
+import torch
+import torchvision.transforms as T
+import torchvision.transforms.functional as F
+
+
+class VideoTransform:
+    def __init__(self, config, mode="train"):
+        self.mode = mode
+        self.config = config
+
+        self.frame_height, self.frame_width = config.DATA.frame_size
+        self.augmentations = config.DATA.augmentations
+
+    def __call__(self, frames: np.ndarray):
+        """
+        frames: np.ndarray (T, H, W, C)
+        """
+
+        if self.mode != "train":
+            return frames
+
+        is_float = frames.dtype in (np.float32, np.float64)
+
+        T_, H, W, C = frames.shape
+        frames_t = torch.from_numpy(frames).permute(0, 3, 1, 2)
+
+        aug = self.augmentations
+
+        # ---------------- Random crop ----------------
+        if getattr(aug, "random_crop", False):
+            scale = getattr(aug, "scale", (0.8, 1.0))
+            ratio = getattr(aug, "ratio", (3/4, 4/3))
+
+            i, j, h, w = T.RandomResizedCrop.get_params(
+                frames_t[0], scale=scale, ratio=ratio
+            )
+
+            frames_t = torch.stack([
+                F.resized_crop(
+                    f, i, j, h, w,
+                    size=(self.frame_height, self.frame_width),
+                    interpolation=F.InterpolationMode.BILINEAR,
+                )
+                for f in frames_t
+            ])
+
+        # ---------------- Affine ----------------
+        if getattr(aug, "random_affine", False):
+            max_translate = getattr(aug, "translate", (0.1, 0.1))
+            scale_range = getattr(aug, "affine_scale", (0.9, 1.0))
+
+            tx = int(random.uniform(-max_translate[0], max_translate[0]) * W)
+            ty = int(random.uniform(-max_translate[1], max_translate[1]) * H)
+            scale_factor = random.uniform(scale_range[0], scale_range[1])
+
+            frames_t = torch.stack([
+                F.affine(
+                    f,
+                    angle=0.0,
+                    translate=[tx, ty],
+                    scale=scale_factor,
+                    shear=[0.0, 0.0],
+                    interpolation=F.InterpolationMode.BILINEAR,
+                )
+                for f in frames_t
+            ])
+
+        # ---------------- Perspective ----------------
+        if getattr(aug, "random_perspective", False):
+            distortion_scale = getattr(aug, "distortion_scale", 0.3)
+            p = getattr(aug, "perspective_prob", 0.5)
+
+            if random.random() < p:
+                startpoints, endpoints = T.RandomPerspective.get_params(
+                    width=W, height=H, distortion_scale=distortion_scale
+                )
+
+                frames_t = torch.stack([
+                    F.perspective(
+                        f,
+                        startpoints=startpoints,
+                        endpoints=endpoints,
+                        interpolation=F.InterpolationMode.BILINEAR,
+                    )
+                    for f in frames_t
+                ])
+
+        # ---------------- Rotation ----------------
+        if getattr(aug, "random_rotation", False):
+            degrees = getattr(aug, "rotation_degrees", 5)
+            angle = random.uniform(-degrees, degrees)
+
+            frames_t = torch.stack([
+                F.rotate(f, angle=angle,
+                         interpolation=F.InterpolationMode.BILINEAR)
+                for f in frames_t
+            ])
+
+        # ---------------- Color jitter ----------------
+        if getattr(aug, "color_jitter", False):
+            jitter_prob = getattr(aug, "jitter_prob", 1.0) # defaults to 1.0
+            brightness, contrast, saturation, hue = getattr(
+                aug, "jitter_params", (0.2, 0.2, 0.2, 0.05)
+            )
+
+            jitter = T.ColorJitter(brightness, contrast, saturation, hue)
+            frames_t = torch.stack([jitter(f) for f in frames_t])
+
+        # ---------------- Flip ----------------
+        if getattr(aug, "random_horizontal_flip", False):
+            if random.random() < getattr(aug, "flip_prob", 0.5):
+                frames_t = torch.flip(frames_t, dims=[3])
+
+        result = frames_t.permute(0, 2, 3, 1)
+        return result.numpy().astype(np.float32 if is_float else np.uint8)
+
+
+def build_transform(config, mode="train"):
+    return VideoTransform(config, mode)
+
+def get_transforms_model(pre_model):
+    from torchvision.models.video import R3D_18_Weights, MC3_18_Weights
+    from torchvision.models.video import R2Plus1D_18_Weights, S3D_Weights
+    from torchvision.models.video import MViT_V2_S_Weights, MViT_V1_B_Weights
+    from torchvision.models.video import mvit_v2_s, MViT_V2_S_Weights, mvit_v1_b, MViT_V1_B_Weights
+    
+    if pre_model == "r3d_18":
+        transforms_model = R3D_18_Weights.KINETICS400_V1.transforms()        
+    elif pre_model == "s3d":
+        transforms_model = S3D_Weights.KINETICS400_V1.transforms()       
+    elif pre_model == "mc3_18":
+        transforms_model = MC3_18_Weights.KINETICS400_V1.transforms()       
+    elif pre_model == "r2plus1d_18":
+        transforms_model = R2Plus1D_18_Weights.KINETICS400_V1.transforms()
+    elif pre_model == "mvit_v2_s":
+        transforms_model = MViT_V2_S_Weights.KINETICS400_V1.transforms()
+    else:
+        transforms_model = R2Plus1D_18_Weights.KINETICS400_V1.transforms()
+
+    return transforms_model
+
+
+# import torch
+# import numpy as np
+# import decord
+# from decord import cpu
+# import torchvision.transforms as T
+
+# def read_video(video_path):
+#     """Read video frames into tensor (T, C, H, W)"""
+#     vr = decord.VideoReader(video_path, ctx=cpu(0))
+#     frames = vr.get_batch(range(len(vr)))  # (T, H, W, C)
+#     frames = torch.from_numpy(frames.asnumpy()).permute(0, 3, 1, 2)  # (T, C, H, W)
+#     return frames
+
+# def resample_video_idx(num_frames, original_fps, new_fps):
+#     """Return frame indices to match new fps"""
+#     step = float(original_fps) / new_fps
+#     if step.is_integer():
+#         step = int(step)
+#         return slice(None, None, step)
+#     idxs = torch.arange(num_frames, dtype=torch.float32) * step
+#     idxs = idxs.floor().to(torch.int64)
+#     return idxs
+
+# def process_frames(video_tensor, target_num_frames, input_fps, target_fps=None):
+#     """Uniform sampling, padding, FPS adjustment"""
+#     target_fps = target_fps or input_fps
+#     num_frames = video_tensor.size(0)
+#     duration = num_frames / input_fps
+
+#     # Too short → resample
+#     if num_frames < target_num_frames:
+#         new_fps = np.ceil(target_num_frames / duration)
+#         idxs = resample_video_idx(target_num_frames, input_fps, new_fps)
+#         idxs = np.clip(idxs, 0, num_frames - 1)
+#         video_tensor = video_tensor[idxs]
+
+#     # Too long → uniform sample
+#     elif num_frames > target_num_frames:
+#         idxs = np.linspace(0, num_frames - 1, target_num_frames)
+#         video_tensor = video_tensor[idxs.astype(int)]
+
+#     # Pad if still short
+#     if video_tensor.size(0) < target_num_frames:
+#         pad = target_num_frames - video_tensor.size(0)
+#         last_frame = video_tensor[-1:].repeat(pad, 1, 1, 1)
+#         video_tensor = torch.cat([video_tensor, last_frame], dim=0)
+
+#     return video_tensor
+
+# def build_transform(frame_size=(224, 224), mean=None, std=None):
+#     """Return default frame transform"""
+#     mean = mean or [0.485, 0.456, 0.406]
+#     std = std or [0.229, 0.224, 0.225]
+#     return T.Compose([
+#         T.ConvertImageDtype(torch.float32),
+#         T.Resize(frame_size),
+#         T.Normalize(mean, std),
+#     ])
+

opensportslib/core/utils/wandb.py

@@ -0,0 +1,110 @@
+import wandb
+import matplotlib.pyplot as plt
+import numpy as np
+import logging
+
+def init_wandb(cfg, run_id, use_wandb=False):
+    """
+    Initialize Weights & Biases if enabled.
+    
+    Args:
+        cfg: config object with attributes:
+             - use_wandb (bool)
+             - project_name (str)
+             - run_name (str)
+    """
+
+    if not use_wandb:
+        logging.info("W&B disabled.")
+        return None
+
+    try:
+        import wandb
+    except ImportError:
+        logging.warning("wandb not installed. Install with `pip install wandb`.")
+        return None
+
+    if getattr(cfg.DATA, "data_modality", None):
+        run_name = f"{cfg.MODEL.backbone.type}_{cfg.DATA.data_modality}"
+    else:
+        run_name = f"{cfg.MODEL.backbone.type}"
+
+    wandb.init(
+        project=cfg.TASK,
+        name=run_name,
+        id=run_id,
+        resume="allow",
+        config=vars(cfg) if hasattr(cfg, "__dict__") else cfg,
+    )
+
+    logging.info(f"Wandb initialised")
+    return wandb
+
+def log_table_wandb(name, rows, headers):
+    """
+    Log a table to Weights & Biases.
+
+    Args:
+        name (str): Name of the table in wandb.
+        rows (list[list]): Table rows.
+        headers (list[str]): Column headers.
+    """
+    if wandb.run is None:
+        return
+
+    table = wandb.Table(columns=headers)
+
+    for row in rows:
+        table.add_data(*row)
+
+    wandb.log({name: table})
+
+def log_attention_wandb(attention, split_name):
+
+    attn = attention.detach().cpu().numpy()
+
+    fig, ax = plt.subplots(figsize=(6, 3))
+    ax.imshow(attn, aspect="auto", cmap="viridis")
+    ax.set_title(f"{split_name} Attention Map")
+    ax.set_xlabel("Views / Time")
+    ax.set_ylabel("Batch")
+
+    wandb.log({
+        f"{split_name}/attention_map": wandb.Image(fig)
+    })
+
+    plt.close(fig)
+
+
+def log_sample_videos_wandb(mvclips, preds, labels, split_name, max_samples=2, fps=5):
+
+
+    # mvclips: (B, V, C, T, H, W)
+    mvclips = mvclips.detach().cpu().numpy()
+
+    for i in range(min(len(mvclips), max_samples)):
+        views = mvclips[i]  # (V, C, T, H, W)
+
+        # Log each view separately
+        for v in range(views.shape[0]):
+            video = views[v].transpose(1, 2, 3, 0)  # (T, H, W, C)
+            video = (video * 255).astype(np.uint8) if video.max() <= 1.0 else video
+
+            wandb.log({
+                f"{split_name}/sample_{i}_view_{v}": wandb.Video(
+                    video,
+                    fps=fps,
+                    caption=f"Pred: {preds[i]}, GT: {labels[i]}"
+                )
+            })
+
+
+def log_confusion_matrix_wandb(y_true, y_pred, class_names, split_name):
+    wandb.log({
+        f"{split_name}/confusion_matrix": wandb.plot.confusion_matrix(
+            probs=None,
+            y_true=y_true,
+            preds=y_pred,
+            class_names=class_names
+        )
+    })

opensportslib/datasets/builder.py

@@ -0,0 +1,42 @@
+# opensportslib/datasets/builder.py
+# from .spotting_dataset import SpottingDataset
+
+def build_dataset(config, annotation_file=None, processor=None, split="train"):
+    """Return a dataset instance based on model type"""
+    task = config.TASK.lower()
+
+    if "classification" in task:
+        from opensportslib.datasets import classification_dataset
+        return classification_dataset.build(config, annotation_file, processor, split)
+    
+    elif "localization" in task:
+        from opensportslib.datasets.localization_dataset import LocalizationDataset
+        return LocalizationDataset(config, annotation_file, processor, split=split)
+    
+    else:
+        raise ValueError(f"No dataset found for task: {task}")
+
+
+##### --------- ####
+# def build_dataset(config, annotation_file=None, processor=None, split="train"):
+#     """Return a dataset instance based on task and modality"""
+#     task = config.TASK.lower()
+
+#     if "classification" in task:
+#         modality = config.DATA.data_modality.lower()
+        
+#         if modality == "tracking_parquet":
+#             from opensportslib.datasets.classification_dataset import TrackingDataset
+#             return TrackingDataset(config, annotation_file, split)
+#         elif modality == "video":
+#             from opensportslib.datasets.classification_dataset import VideoDataset
+#             return VideoDataset(config, annotation_file, processor, split)
+#         else:
+#             raise ValueError(f"Unknown data_modality: {modality}")
+    
+#     elif "localization" in task:
+#         from opensportslib.datasets.localization_dataset import LocalizationDataset
+#         return LocalizationDataset(config, annotation_file, processor, split=split)
+    
+#     else:
+#         raise ValueError(f"No dataset found for task: {task}")