opensportslib 0.0.1.dev7__tar.gz → 0.0.1.dev9__tar.gz

{opensportslib-0.0.1.dev7/opensportslib.egg-info → opensportslib-0.0.1.dev9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.0.1.dev7
+Version: 0.0.1.dev9
 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.dev7 → opensportslib-0.0.1.dev9}/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.dev9/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.dev7 → opensportslib-0.0.1.dev9}/opensportslib/core/utils/load_annotations.py

@@ -386,6 +386,20 @@ def construct_labels(path, extract_fps):
         num_frames, fps, wanted_sample_fps if wanted_sample_fps < fps else fps
     )
 
+    return [
+        {
+            "video": path,
+            "path": path,
+            "num_frames": num_frames_after,
+            "num_frames_base": num_frames,
+            "num_events": 0,
+            "events": [],
+            "fps": sample_fps,
+            "width": 398,
+            "height": 224,
+        }
+    ], get_stride(fps, wanted_sample_fps if wanted_sample_fps < fps else fps)
+
 
 # def get_repartition_gpu():
 #     """Returns the distribution of gpus that will be used by pipelines for dali."""