opensportslib 0.0.1.dev2__py3-none-any.whl

opensportslib/models/base/e2e.py

@@ -0,0 +1,252 @@
+"""
+Copyright 2022 James Hong, Haotian Zhang, Matthew Fisher, Michael Gharbi,
+Kayvon Fatahalian
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+this list of conditions and the following disclaimer in the documentation and/or
+other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors
+may be used to endorse or promote products derived from this software without
+specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+import logging
+from opensportslib.core.loss.builder import build_criterion
+from opensportslib.models.backbones.builder import build_backbone
+from opensportslib.models.utils.common import step, BaseRGBModel
+from opensportslib.models.heads.builder import build_head
+from contextlib import nullcontext
+import torch
+from tqdm import tqdm
+import torch.nn as nn
+
+
+class E2EModel(BaseRGBModel):
+    """Class used to handle the model for the E2E method. This class is not the model but contains the model as a variable.
+    This class is used to initialize the model, perform an epoch and predict.
+    Args:
+        cfg (dict): Dict of config.
+        num_classes (int): Number of classes.
+        backbone (dict): Dict used to build the backbone.
+        head (dict): Dict used to build the head.
+        clip_len (int): the length of the clip.
+        device (string): The device to use.
+            Default :'cuda'.
+        multi_gpu: Whether to use one gpu or more for the model.
+            Default: False.
+    """
+
+    class Impl(nn.Module):
+        """Model for the E2E method. Model is build of a backbone and a head.
+
+        Args:
+            num_classes (int): Number of classes.
+            backbone (dict): Dict used to build the backbone.
+            head (dict): Dict used to build the head.
+            clip_len (int): the length of the clip.
+        """
+
+        def __init__(self, num_classes, backbone, head, clip_len, modality):
+            super().__init__()
+            is_rgb = modality == "rgb"
+            in_channels = {"flow": 2, "bw": 1, "rgb": 3}[modality]
+
+            backbone.clip_len = clip_len
+            backbone.is_rgb = is_rgb
+            backbone.in_channels = in_channels
+
+            self.backbone = build_backbone(backbone)
+
+            head.num_classes = num_classes
+            head.feat_dim = self.backbone._feat_dim
+
+            self.head = build_head(head)
+
+        def forward(self, x):
+            im_feat = self.backbone(x)
+            return self.head(im_feat)
+
+        def print_stats(self):
+            print("Model params:", sum(p.numel() for p in self.parameters()))
+            print(
+                "  CNN features:",
+                sum(p.numel() for p in self.backbone._features.parameters()),
+            )
+            print(
+                "  Temporal:", sum(p.numel() for p in self.head._pred_fine.parameters())
+            )
+
+    def __init__(
+        self,
+        cfg,
+        num_classes,
+        backbone,
+        head,
+        clip_len,
+        modality,
+        device="cuda",
+        multi_gpu=False,
+    ):
+
+        last_gpu_index = torch.cuda.device_count() - 1
+
+        # self.device = torch.device('cuda:{}'.format(0))
+        self.device = device
+        self._multi_gpu = multi_gpu
+        self._model = E2EModel.Impl(num_classes, backbone, head, clip_len, modality)
+        self._model.print_stats()
+        logging.info("Build criterion")
+        self.criterion = build_criterion(cfg.TRAIN.criterion)
+        logging.info(self.criterion)
+        if multi_gpu:
+            self._model = nn.DataParallel(self._model)
+            self._model.to(device)
+            # self._model = nn.DataParallel(self._model,device_ids = [1, 0, 2, 3] if torch.cuda.device_count()==4 else [1, 0])
+            # self._model.to(f'cuda:{self._model.device_ids[0]}')
+            # self.device = torch.device('cuda:{}'.format(1))
+        else:
+            self._model.to(device)
+
+        self._multi_gpu = multi_gpu
+        self._num_classes = num_classes
+
+    def epoch(
+        self,
+        loader,
+        dali,
+        optimizer=None,
+        scaler=None,
+        lr_scheduler=None,
+        acc_grad_iter=1,
+        fg_weight=5,
+    ):
+        """Performs an epoch for training/validating.
+
+        Args:
+            loader: The dataloader.
+            dali (bool): Whether dali has been used or opencv for data processing.
+            optimizer (torch.optim.Optimizer): The optimizer to update model parameters. Set to None if validation epoch.
+                Default: None.
+            scaler (torch.cuda.amp.GradScaler): The gradient scaler for mixed precision training.
+                Default: None.
+            lr_scheduler : The learning rate scheduler.
+                Default: None.
+            acc_grad_iter (int): Use gradient accumulation.
+                Default: 1.
+            fg_weight (int): Used to build tensor of weights for the criterion.
+                Default: 5.
+        Returns:
+            (float): The average loss over the batches.
+        """
+        if optimizer is None:
+            self._model.eval()
+        else:
+            optimizer.zero_grad()
+            self._model.train()
+
+        ce_kwargs = {}
+        if fg_weight != 1:
+            ce_kwargs["weight"] = torch.FloatTensor(
+                [1] + [fg_weight] * (self._num_classes - 1)
+            ).to(self.device)
+
+        epoch_loss = 0.0
+
+        times = []
+        import timeit
+
+        with torch.no_grad() if optimizer is None else nullcontext():
+            for batch_idx, batch in enumerate(tqdm(loader)):
+                if dali:
+                    frame = batch["frame"].to(self.device)
+                    label = batch["label"].to(self.device)
+                else:
+                    frame = loader.dataset.load_frame_gpu(batch, self.device)
+                    label = batch["label"].to(self.device)
+
+                label = (
+                    label.flatten()
+                    if len(label.shape) == 2
+                    else label.view(-1, label.shape[-1])
+                )
+
+                with torch.cuda.amp.autocast():
+                    pred = self._model(frame)
+
+                    pred = pred.to(self.device)
+
+                    loss = 0.0
+                    if len(pred.shape) == 3:
+                        pred = pred.unsqueeze(0)
+
+                    # label=label.to(self.device)
+
+                    for i in range(pred.shape[0]):
+                        loss += self.criterion(
+                            pred[i].reshape(-1, self._num_classes), label, **ce_kwargs
+                        )
+                        # loss += F.cross_entropy(
+                        #     pred[i].reshape(-1, self._num_classes), label,
+                        #     **ce_kwargs)
+
+                if optimizer is not None:
+                    step(
+                        optimizer,
+                        scaler,
+                        loss / acc_grad_iter,
+                        lr_scheduler=lr_scheduler,
+                        backward_only=(batch_idx + 1) % acc_grad_iter != 0,
+                    )
+
+                epoch_loss += loss.detach().item()
+
+        print(epoch_loss, len(loader), epoch_loss / len(loader))
+        return epoch_loss / len(loader)  # Avg loss
+
+    def predict(self, seq, use_amp=True):
+        """Perform prediction on the input.
+
+        Args:
+            seq (torch.tensor): The input.
+            use_amp (bool): Whether to use automatic precision.
+                Default: True.
+        Returns:
+            pred_cls (numpy.ndarray): Predicted class indices.
+            pred (numpy.ndarray): Predicted probabilities.
+        """
+        if not isinstance(seq, torch.Tensor):
+            seq = torch.FloatTensor(seq)
+        if len(seq.shape) == 4:  # (L, C, H, W)
+            seq = seq.unsqueeze(0)
+        if seq.device != self.device:
+            seq = seq.to(self.device)
+
+        self._model.eval()
+        with torch.no_grad():
+            with torch.cuda.amp.autocast() if use_amp else nullcontext():
+                pred = self._model(seq)
+            if isinstance(pred, tuple):
+                pred = pred[0]
+            if len(pred.shape) > 3:
+                pred = pred[-1]
+            pred = torch.softmax(pred, axis=2)
+            pred_cls = torch.argmax(pred, axis=2)
+            return pred_cls.cpu().numpy(), pred.cpu().numpy()

opensportslib/models/base/tracking.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn import global_mean_pool
+
+from opensportslib.models.backbones.builder import build_backbone
+from opensportslib.models.neck.builder import build_neck
+from opensportslib.models.heads.builder import build_head
+from opensportslib.datasets.utils.tracking import FEATURE_DIM
+
+
+class TrackingModel(nn.Module):
+    """
+    Tracking-based classification model.
+    Combines graph backbone, temporal neck, and classification head.
+    """
+    
+    def __init__(self, config, device):
+        super().__init__()
+        print("Building TrackingModel")
+        
+        self.device = device
+        self.num_frames = config.DATA.num_frames
+        
+        # backbone: graph encoder
+        self.backbone = build_backbone(
+            config.MODEL.backbone,
+            default_args={"input_dim": FEATURE_DIM}
+        )
+        
+        # neck: temporal aggregation
+        self.neck = build_neck(
+            config.MODEL.neck,
+            default_args={"window_size": self.num_frames}
+        )
+        
+        # head: classifier
+        self.head = build_head(
+            config.MODEL.head,
+            default_args={"input_dim": self.neck.feat_dim}
+        )
+    
+    def forward(self, batch):
+        """
+        Args:
+            batch: dict with keys:
+                - x: (B*T*N, F) all node features batched
+                - edge_index: (2, E) all edges with proper offsets
+                - batch: (B*T*N,) graph assignment per node
+                - batch_size: int
+                - seq_len: int
+        
+        Returns:
+            logits: (B, num_classes)
+        """
+        x = batch['x']
+        edge_index = batch['edge_index']
+        batch_idx = batch['batch']
+        batch_size = batch['batch_size']
+        seq_len = batch['seq_len']
+        
+        # single forward through backbone for all B*T graphs
+        graph_emb = self.backbone(x, edge_index, batch_idx)  # (B*T, H)
+        
+        # reshape to (B, T, H)
+        x = graph_emb.view(batch_size, seq_len, -1)
+        
+        # temporal aggregation
+        x = self.neck(x)  # (B, H)
+        
+        # classification
+        logits = self.head(x)  # (B, num_classes)
+        
+        return logits

opensportslib/models/base/vars.py

@@ -0,0 +1,29 @@
+# opensportslib/models/base/vars.py
+
+import __future__
+import torch
+from opensportslib.models.backbones.builder import build_backbone
+from opensportslib.models.neck.builder import build_neck
+from opensportslib.models.heads.builder import build_head
+
+class MVNetwork(torch.nn.Module):
+
+    def __init__(self, config, backbone, neck, head):
+        super().__init__()
+        print("Building MVNetwork Model")
+        self.lifting_net = torch.nn.Sequential()
+        
+        self.backbone = build_backbone(backbone)
+        self.mvaggregate = build_neck(neck, default_args=dict(
+                model=self.backbone,
+                feat_dim=self.backbone.feat_dim,
+                lifting_net=self.lifting_net
+            )
+        )
+        head.num_classes = config.DATA.num_classes
+        head.feat_dim = self.backbone.feat_dim
+        self.head = build_head(head)
+
+    def forward(self, mvimages):
+        features, attention = self.mvaggregate(mvimages)
+        return self.head(features), attention

opensportslib/models/base/video.py

@@ -0,0 +1,130 @@
+# opensportslib/models/base/video.py
+
+"""video backbone and model for frames_npy modality.
+
+this file contains two independent things:
+
+1. the existing VideoMAE HuggingFace full-model builder functions
+   (build_video_mae_backbone, load_video_mae_checkpoint).
+   these are left exactly as they were and route through MODEL.type == "huggingface".
+
+2. the new VideoBackbone + VideoModel classes for the custom frames_npy path,
+   supporting dinov3, clip, videomae, videomae2 as pure feature extractors
+   wired to the library's existing TemporalAggregation neck and
+   TrackingClassifierHead head.
+"""
+
+import torch
+import torch.nn as nn
+
+from opensportslib.models.backbones.builder import build_backbone
+from opensportslib.models.neck.builder import build_neck
+from opensportslib.models.heads.builder import build_head
+
+
+# -----------------------------------------------------------------------
+# video mae backbone for MVFoul
+# -----------------------------------------------------------------------
+
+def build_video_mae_backbone(config, device, ckpt_path=None, infer=False):
+    """
+    Build HuggingFace VideoMAE model for video classification.
+    This includes both backbone and classification head.
+    """
+    from transformers import VideoMAEForVideoClassification, VideoMAEImageProcessor
+
+    num_classes = config.MODEL.num_classes
+    pretrained_model_name = ckpt_path if ckpt_path else config.MODEL.pretrained_model
+    processor = VideoMAEImageProcessor.from_pretrained(config.MODEL.pretrained_model)
+    model = VideoMAEForVideoClassification.from_pretrained(
+        pretrained_model_name,
+        num_labels=num_classes,
+        ignore_mismatched_sizes=True,
+        trust_remote_code=True,
+        device_map=device
+    )
+
+    for param in model.parameters():
+        param.requires_grad = False
+
+    if not infer:
+        if config.MODEL.unfreeze_head:
+            for p in model.classifier.parameters():
+                p.requires_grad = True
+
+        n_unfreeze = getattr(config.MODEL, "unfreeze_last_n_layers", 0)
+        if n_unfreeze > 0:
+            for layer in model.videomae.encoder.layer[-n_unfreeze:]:
+                for p in layer.parameters():
+                    p.requires_grad = True
+
+    trainable = [name for name, p in model.named_parameters() if p.requires_grad]
+    print("Number of trainable params:", len(trainable))
+    for n in trainable:
+        print(n)
+
+    return model, processor
+
+
+def load_video_mae_checkpoint(config, device, ckpt_path, infer=True):
+    """
+    Load fine-tuned VideoMAE checkpoint from a HuggingFace-style directory.
+    """
+    return build_video_mae_backbone(config, device, ckpt_path, infer=infer)
+
+
+# -----------------------------------------------------------------------
+# new custom path: full model
+# -----------------------------------------------------------------------
+
+class VideoModel(nn.Module):
+    """Video classification model for the frames_npy modality.
+
+    follows the same backbone -> neck -> head pattern as TrackingModel.
+
+    the backbone is a VideoBackbone (pure feature extractor).
+    the neck is TemporalAggregation.
+    the head is TrackingClassifierHead.
+
+    Args:
+        config: full YAML config
+        device: torch device string
+    """
+
+    def __init__(self, config, device):
+        super().__init__()
+        print("Building VideoModel")
+
+        self.device = device
+        self.num_frames = config.DATA.num_frames
+
+        # backbone: pure feature extractor
+        self.backbone = build_backbone(config.MODEL.backbone)
+
+        # neck: temporal aggregation over the frame sequence
+        self.neck = build_neck(
+            config.MODEL.neck,
+            default_args={"window_size": self.num_frames}
+        )
+
+        # head: linear classifier
+        self.head = build_head(
+            config.MODEL.head,
+            default_args={"input_dim": self.neck.feat_dim}
+        )
+
+    def forward(self, batch):
+        """
+        Args:
+            batch: dict with key "pixel_values" of shape (B, T, H, W, C).
+
+        Returns:
+            logits: (B, num_classes)
+        """
+        x = batch["pixel_values"]           # (B, T, H, W, C)
+
+        x = self.backbone(x)                # (B, T, hidden_dim) or (B, 1, hidden_dim)
+        x = self.neck(x)                    # (B, hidden_dim)
+        logits = self.head(x)               # (B, num_classes)
+
+        return logits

opensportslib/models/base/video_mae.py

@@ -0,0 +1,60 @@
+# models/backbones/video_mae.py
+
+from transformers import VideoMAEForVideoClassification, VideoMAEImageProcessor
+import os
+
+def build_video_mae_backbone(config, device, ckpt_path=None, infer=False):
+    """
+    Build HuggingFace VideoMAE model for video classification.
+    This includes both backbone and classification head.
+    """
+    num_classes = config.MODEL.num_classes
+    pretrained_model_name = ckpt_path if ckpt_path else config.MODEL.pretrained_model
+    processor = VideoMAEImageProcessor.from_pretrained(config.MODEL.pretrained_model)
+    model = VideoMAEForVideoClassification.from_pretrained(
+        pretrained_model_name,
+        num_labels=num_classes,
+        ignore_mismatched_sizes=True,
+        trust_remote_code=True,
+        device_map=device
+    )
+
+    # freeze everything 
+    for param in model.parameters():
+        param.requires_grad = False
+    
+    if not infer:
+        # Unfreeze classification head 
+        if config.MODEL.unfreeze_head:
+            for p in model.classifier.parameters():
+                p.requires_grad = True
+
+        # -------- Unfreeze last N VideoMAE encoder layers --------
+        n_unfreeze = getattr(config.MODEL, "unfreeze_last_n_layers", 0)
+        # unfreeze last encoder layer
+        if n_unfreeze > 0:
+            for layer in model.videomae.encoder.layer[-n_unfreeze:]:
+                for p in layer.parameters():
+                    p.requires_grad = True
+
+    trainable = []
+    for name, p in model.named_parameters():
+        if p.requires_grad:
+            trainable.append(name)
+
+    print("Number of trainable params:", len(trainable))
+    for n in trainable:
+        print(n)
+    return model, processor
+
+
+def load_video_mae_checkpoint(config, device, ckpt_path, infer=True):
+    """
+    Load fine-tuned VideoMAE checkpoint from a HuggingFace-style directory.
+
+    Supports:
+      - model.safetensors
+      - pytorch_model.bin
+      - config.json
+    """
+    return build_video_mae_backbone(config, device, ckpt_path, infer=infer)

opensportslib/models/builder.py

@@ -0,0 +1,43 @@
+# opensportslib/models/builder.py
+
+def build_model(config, device):
+    """
+    Dispatch model builder based on cfg.MODEL.task
+    """
+    task = config.TASK.lower()
+    
+    if task == "classification":
+        # return model, processor
+        if config.MODEL.backbone.type == "video_mae":
+            from opensportslib.models.base.video import build_video_mae_backbone
+            return build_video_mae_backbone(config, device)
+        
+        elif config.MODEL.backbone.type in ["r3d_18", "mc3_18", "r2plus1d_18", "s3d", "mvit_v2_s"]:
+            from opensportslib.models.base.vars import MVNetwork
+            return MVNetwork(config, config.MODEL.backbone, config.MODEL.neck, config.MODEL.head), None
+        
+        elif config.MODEL.backbone.type == "graph_conv":
+            from opensportslib.models.base.tracking import TrackingModel
+            return TrackingModel(config, device), None
+
+        elif config.MODEL.backbone.type in ("dinov3", "clip", "videomae", "videomae2"):
+            from opensportslib.models.base.video import VideoModel
+            return VideoModel(config, device), None
+        
+        else:
+            raise ValueError(f"Unsupported backbone type: {config.MODEL.backbone.type}")
+    
+    if task == "localization":
+        from opensportslib.models.base.e2e import E2EModel
+        if config.MODEL.type == "E2E":
+            model = E2EModel(config, 
+                            len(config.DATA.classes)+1,
+                            config.MODEL.backbone,
+                            config.MODEL.head,
+                            clip_len=config.DATA.clip_len,
+                            modality=config.DATA.modality,
+                            device=device,
+                            multi_gpu=config.MODEL.multi_gpu)
+            return model
+    else:
+        raise ValueError(f"Unsupported model type: {config.MODEL.backbone} for task: {task}")