opensportslib 0.0.1.dev8__tar.gz → 0.0.1.dev10__tar.gz

{opensportslib-0.0.1.dev8/opensportslib.egg-info → opensportslib-0.0.1.dev10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.0.1.dev8
+Version: 0.0.1.dev10
 Summary: OpenSportsLib is the professional library, designed for advanced video understanding in sports. It provides state-of-the-art tools for action recognition, spotting, retrieval, and captioning, making it ideal for researchers, analysts, and developers working with sports video data.
 Author: Jeet Vora
 Requires-Python: >=3.12
@@ -25,6 +25,7 @@ Provides-Extra: localization
 Requires-Dist: nvidia-dali-cuda120; extra == "localization"
 Requires-Dist: cupy-cuda12x; extra == "localization"
 Requires-Dist: tabulate; extra == "localization"
+Requires-Dist: pytorch-lightning; extra == "localization"
 Provides-Extra: py-geometric
 Requires-Dist: torch-geometric; extra == "py-geometric"
 Requires-Dist: torch-scatter; extra == "py-geometric"

{opensportslib-0.0.1.dev8 → opensportslib-0.0.1.dev10}/opensportslib/apis/localization.py

@@ -112,7 +112,7 @@ class LocalizationAPI:
             gpu=self.config.SYSTEM.GPU,
             default_args=data_obj_train.default_args,
         )
-        train_loader = data_obj_train.building_dataloader(dataset_Train, cfg=data_obj_train.cfg.dataloader, gpu=self.config.SYSTEM.GPU, dali=True)
+        train_loader = data_obj_train.building_dataloader(dataset_Train, cfg=data_obj_train.cfg.dataloader, gpu=self.config.SYSTEM.GPU, dali=self.config.dali)
         print(len(train_loader))
         # Valid
         data_obj_valid = build_dataset(self.config,split="valid")
@@ -121,7 +121,7 @@ class LocalizationAPI:
             gpu= self.config.SYSTEM.GPU,
             default_args=data_obj_valid.default_args,
         )
-        valid_loader = data_obj_valid.building_dataloader(dataset_Valid, cfg=data_obj_valid.cfg.dataloader, gpu=self.config.SYSTEM.GPU, dali=True)
+        valid_loader = data_obj_valid.building_dataloader(dataset_Valid, cfg=data_obj_valid.cfg.dataloader, gpu=self.config.SYSTEM.GPU, dali=self.config.dali)
         print(len(valid_loader))
 
         # Trainer
@@ -200,7 +200,7 @@ class LocalizationAPI:
                 gpu=self.config.SYSTEM.GPU,
                 default_args=data_obj_test.default_args,
             )
-            test_loader = data_obj_test.building_dataloader(dataset_Test, cfg=data_obj_test.cfg.dataloader, gpu=self.config.SYSTEM.GPU, dali=True)
+            test_loader = data_obj_test.building_dataloader(dataset_Test, cfg=data_obj_test.cfg.dataloader, gpu=self.config.SYSTEM.GPU, dali=self.config.dali)
             print(len(test_loader))
 
             # # Inference

{opensportslib-0.0.1.dev8 → opensportslib-0.0.1.dev10}/opensportslib/core/trainer/localization_trainer.py

@@ -135,6 +135,9 @@ def build_trainer(cfg, model=None, default_args=None, resume_from=None):
             trainer.best_criterion_valid = checkpoint.get('best_criterion_valid', 
                 0 if cfg.TRAIN.criterion_valid == "map" else float("inf"))
             logging.info(f"Restored best epoch: {trainer.best_epoch}")
+    
+    else:
+        trainer = Trainer_pl(cfg, default_args["work_dir"])
             
 
     return trainer
@@ -147,6 +150,37 @@ class Trainer(ABC):
     def train(self):
         pass
 
+class Trainer_pl(Trainer):
+    """Trainer class used for models that rely on lightning modules.
+
+    Args:
+        cfg (dict): Dict config. It should contain the key 'max_epochs' and the key 'GPU'.
+    """
+
+    def __init__(self, cfg, work_dir):
+        from opensportslib.core.utils.lightning import CustomProgressBar, MyCallback
+        import pytorch_lightning as pl
+
+        self.work_dir = work_dir
+        call = MyCallback()
+        self.trainer = pl.Trainer(
+            max_epochs=cfg.max_epochs,
+            devices=[cfg.GPU],
+            callbacks=[call, CustomProgressBar(refresh_rate=1)],
+            num_sanity_val_steps=0,
+        )
+
+    def train(self, **kwargs):
+        self.trainer.fit(**kwargs)
+
+        best_model = kwargs["model"].best_state
+
+        logging.info("Done training")
+        logging.info("Best epoch: {}".format(best_model.get("epoch")))
+        torch.save(best_model, os.path.join(self.work_dir, "model.pth.tar"))
+
+        logging.info("Model saved")
+        logging.info(os.path.join(self.work_dir, "model.pth.tar"))
 
 
 class Trainer_e2e(Trainer):

opensportslib-0.0.1.dev10/opensportslib/core/utils/lightning.py

@@ -0,0 +1,52 @@
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
+import logging
+
+
+class CustomProgressBar(TQDMProgressBar):
+    """Override the custom progress bar used by pytorch lightning to change some attributes."""
+
+    def get_metrics(self, trainer, pl_module):
+        """Override the method to don't show the version number in the progress bar."""
+        items = super().get_metrics(trainer, pl_module)
+        items.pop("v_num", None)
+        return items
+
+
+class MyCallback(pl.Callback):
+    """Override the Callback class of pl to change the behaviour on validation epoch end."""
+
+    def __init__(self):
+        super().__init__()
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        loss_validation = pl_module.losses.avg
+        state = {
+            "epoch": trainer.current_epoch + 1,
+            "state_dict": pl_module.model.state_dict(),
+            "best_loss": pl_module.best_loss,
+            "optimizer": pl_module.optimizer.state_dict(),
+        }
+
+        # remember best prec@1 and save checkpoint
+        is_better = loss_validation < pl_module.best_loss
+        pl_module.best_loss = min(loss_validation, pl_module.best_loss)
+
+        # Save the best model based on loss only if the evaluation frequency too long
+        if is_better:
+            pl_module.best_state = state
+            # torch.save(state, best_model_path)
+
+        # Reduce LR on Plateau after patience reached
+        prevLR = pl_module.optimizer.param_groups[0]["lr"]
+        pl_module.scheduler.step(loss_validation)
+        currLR = pl_module.optimizer.param_groups[0]["lr"]
+
+        if currLR is not prevLR and pl_module.scheduler.num_bad_epochs == 0:
+            logging.info("\nPlateau Reached!")
+        if (
+            prevLR < 2 * pl_module.scheduler.eps
+            and pl_module.scheduler.num_bad_epochs >= pl_module.scheduler.patience
+        ):
+            logging.info("\nPlateau Reached and no more reduction -> Exiting Loop")
+            trainer.should_stop = True

{opensportslib-0.0.1.dev8 → opensportslib-0.0.1.dev10}/opensportslib/core/utils/video_processing.py

@@ -1,6 +1,10 @@
 import torch
 import numpy as np
 import math
+import random
+import torch.nn as nn
+import torchvision.transforms as T
+import torchvision.transforms.functional as F
 
 try:
     import decord
@@ -301,12 +305,6 @@ def get_remaining(data_len, batch_size):
     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 RandomHorizontalFlipFLow(nn.Module):
 
     def __init__(self, p=0.5):
@@ -719,6 +717,131 @@ def oneHotToShifts(onehot, params):
         Shifts[:, i] = shifts
 
     return Shifts
+
+def timestamps2long(output_spotting, video_size, chunk_size, receptive_field):
+    """Method to transform the timestamps to vectors"""
+    start = 0
+    last = False
+    receptive_field = receptive_field // 2
+
+    timestamps_long = (
+        torch.zeros(
+            [video_size, output_spotting.size()[-1] - 2],
+            dtype=torch.float,
+            device=output_spotting.device,
+        )
+        - 1
+    )
+
+    for batch in np.arange(output_spotting.size()[0]):
+
+        tmp_timestamps = (
+            torch.zeros(
+                [chunk_size, output_spotting.size()[-1] - 2],
+                dtype=torch.float,
+                device=output_spotting.device,
+            )
+            - 1
+        )
+
+        for i in np.arange(output_spotting.size()[1]):
+            tmp_timestamps[
+                torch.floor(output_spotting[batch, i, 1] * (chunk_size - 1)).type(
+                    torch.int
+                ),
+                torch.argmax(output_spotting[batch, i, 2:]).type(torch.int),
+            ] = output_spotting[batch, i, 0]
+
+        # ------------------------------------------
+        # Store the result of the chunk in the video
+        # ------------------------------------------
+
+        # For the first chunk
+        if start == 0:
+            timestamps_long[0 : chunk_size - receptive_field] = tmp_timestamps[
+                0 : chunk_size - receptive_field
+            ]
+
+        # For the last chunk
+        elif last:
+            timestamps_long[start + receptive_field : start + chunk_size] = (
+                tmp_timestamps[receptive_field:]
+            )
+            break
+
+        # For every other chunk
+        else:
+            timestamps_long[
+                start + receptive_field : start + chunk_size - receptive_field
+            ] = tmp_timestamps[receptive_field : chunk_size - receptive_field]
+
+        # ---------------
+        # Loop Management
+        # ---------------
+
+        # Update the index
+        start += chunk_size - 2 * receptive_field
+        # Check if we are at the last index of the game
+        if start + chunk_size >= video_size:
+            start = video_size - chunk_size
+            last = True
+    return timestamps_long
+
+
+def batch2long(output_segmentation, video_size, chunk_size, receptive_field):
+    """Method to transform the batches to vectors."""
+    start = 0
+    last = False
+    receptive_field = receptive_field // 2
+
+    segmentation_long = torch.zeros(
+        [video_size, output_segmentation.size()[-1]],
+        dtype=torch.float,
+        device=output_segmentation.device,
+    )
+
+    for batch in np.arange(output_segmentation.size()[0]):
+
+        tmp_segmentation = torch.nn.functional.one_hot(
+            torch.argmax(output_segmentation[batch], dim=-1),
+            num_classes=output_segmentation.size()[-1],
+        )
+
+        # ------------------------------------------
+        # Store the result of the chunk in the video
+        # ------------------------------------------
+
+        # For the first chunk
+        if start == 0:
+            segmentation_long[0 : chunk_size - receptive_field] = tmp_segmentation[
+                0 : chunk_size - receptive_field
+            ]
+
+        # For the last chunk
+        elif last:
+            segmentation_long[start + receptive_field : start + chunk_size] = (
+                tmp_segmentation[receptive_field:]
+            )
+            break
+
+        # For every other chunk
+        else:
+            segmentation_long[
+                start + receptive_field : start + chunk_size - receptive_field
+            ] = tmp_segmentation[receptive_field : chunk_size - receptive_field]
+
+        # ---------------
+        # Loop Management
+        # ---------------
+
+        # Update the index
+        start += chunk_size - 2 * receptive_field
+        # Check if we are at the last index of the game
+        if start + chunk_size >= video_size:
+            start = video_size - chunk_size
+            last = True
+    return segmentation_long
+
 # import torch
 # import numpy as np
 # import decord

{opensportslib-0.0.1.dev8 → opensportslib-0.0.1.dev10}/opensportslib/datasets/localization_dataset.py

@@ -16,7 +16,7 @@ import logging
 import tqdm
 from opensportslib.core.utils.default_args import get_default_args_dataset
 from opensportslib.core.utils.load_annotations import get_repartition_gpu
-from opensportslib.opensportslib.core.utils.video_processing import feats2clip, getChunks_anchors, getTimestampTargets, oneHotToShifts
+from opensportslib.core.utils.video_processing import feats2clip, getChunks_anchors, getTimestampTargets, oneHotToShifts
 from SoccerNet.Downloader import getListGames
 from SoccerNet.Downloader import SoccerNetDownloader
 from SoccerNet.Evaluation.utils import (
@@ -246,7 +246,7 @@ class LocalizationDataset(Dataset):
             num_workers=cfg.num_workers if gpu >= 0 else 0,
             pin_memory=cfg.pin_memory if gpu >= 0 else False,
             prefetch_factor=(
-                cfg.prefetch_factor if "prefetch_factor" in cfg.keys() else None
+                getattr(cfg, "prefetch_factor", None)
             ),
             worker_init_fn=worker_init_fn
         )
@@ -457,7 +457,7 @@ class ActionSpotDataset(Dataset):
         from opensportslib.core.utils.video_processing import _get_deferred_rgb_transform, _get_img_transforms
 
         self._src_file = label_file
-        self._labels = annotationstoe2eformat(
+        self._labels, self.task_name = annotationstoe2eformat(
             label_file, video_dir, input_fps, extract_fps, False
         )
         # self._labels = load_json(label_file)
@@ -491,17 +491,17 @@ class ActionSpotDataset(Dataset):
 
         self._mixup = mixup
 
+        self.IMAGENET_MEAN = IMAGENET_MEAN
+        self.IMAGENET_STD = IMAGENET_STD
+        self.TARGET_HEIGHT = TARGET_HEIGHT
+        self.TARGET_WIDTH = TARGET_WIDTH
         # Try to do defer the latter half of the transforms to the GPU
         self._gpu_transform = None
         if not is_eval and same_transform:
             if modality == "rgb":
                 print("=> Deferring some RGB transforms to the GPU!")
-                self._gpu_transform = _get_deferred_rgb_transform()
+                self._gpu_transform = _get_deferred_rgb_transform(self.IMAGENET_MEAN, self.IMAGENET_STD)
 
-        self.IMAGENET_MEAN = IMAGENET_MEAN
-        self.IMAGENET_STD = IMAGENET_STD
-        self.TARGET_HEIGHT = TARGET_HEIGHT
-        self.TARGET_WIDTH = TARGET_WIDTH
 
         crop_transform, img_transform = _get_img_transforms(
             self.IMAGENET_MEAN,
@@ -771,7 +771,7 @@ class ActionSpotVideoDataset(Dataset, DatasetVideoSharedMethods):
     
         self._src_file = label_file
         if label_file.endswith(".json"):
-            self._labels = annotationstoe2eformat(
+            self._labels, self.task_name = annotationstoe2eformat(
                 label_file, video_dir, input_fps, extract_fps, False
             )
             # self._labels = load_json(label_file)