opensportslib 0.0.1.dev2__py3-none-any.whl

opensportslib/core/utils/checkpoint.py

@@ -0,0 +1,238 @@
+# opensportslib/core/utils/checkpoint.py
+
+import torch
+import os
+
+def localization_remap(key):
+    if key.startswith("_features"):
+        return "backbone." + key
+    elif key.startswith("_pred_fine"):
+        return "head." + key
+    return key
+
+
+def save_checkpoint(model, path, processor=None, tokenizer=None, optimizer=None, epoch=None):
+    """
+    Save model checkpoint to `path`. Uses HF save_pretrained if available, 
+    otherwise falls back to saving a PyTorch checkpoint.
+
+    Args:
+        model (torch.nn.Module or HF PreTrainedModel): model to save
+        path (str): path to save torch checkpoint (file path, e.g., /.../checkpoint.pt)
+        optimizer (torch.optim.Optimizer, optional): optimizer to save
+        epoch (int, optional): current epoch number
+        processor (optional): HF processor / feature extractor to save with model
+        tokenizer (optional): HF tokenizer to save with model
+    """
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    hf_saved = False
+
+    # 1) Try to save HuggingFace model if available
+    if hasattr(model, "save_pretrained"):
+        try:
+            model.save_pretrained(path)
+            hf_saved = True
+            print(f"[Checkpoint] HuggingFace model saved to {path}")
+        except Exception as e:
+            print(f"[Checkpoint] Warning: could not save HF model: {e}")
+            hf_saved = False
+
+    # 2) Save processor / tokenizer if provided (only if HF save succeeded)
+    if hf_saved:
+        if processor is not None:
+            try:
+                processor.save_pretrained(path)
+                print(f"[Checkpoint] Processor saved to {path}")
+            except Exception as e:
+                print(f"[Checkpoint] Warning: could not save processor: {e}")
+
+        if tokenizer is not None:
+            try:
+                tokenizer.save_pretrained(path)
+                print(f"[Checkpoint] Tokenizer saved to {path}")
+            except Exception as e:
+                print(f"[Checkpoint] Warning: could not save tokenizer: {e}")
+
+    # 3) Fallback: Save a PyTorch checkpoint if HF save is unavailable or failed
+    if not hf_saved:
+        checkpoint = {}
+        if hasattr(model, "state_dict"):
+            checkpoint["model_state_dict"] = model.state_dict()
+        else:
+            print("[Checkpoint] Warning: model has no state_dict(), skipping model_state_dict.")
+
+        if optimizer is not None:
+            checkpoint["optimizer_state_dict"] = optimizer.state_dict()
+        if epoch is not None:
+            checkpoint["epoch"] = epoch
+
+        torch.save(checkpoint, path)
+        print(f"[Checkpoint] Torch checkpoint saved at: {path}")
+
+
+def load_checkpoint(
+    model,
+    path,
+    optimizer=None,
+    scheduler=None,
+    device=None,
+    key_remap_fn=None,
+    hf_filename="model.pth.tar",   # required if loading from HF repo
+    hf_token=None,      # optional (for private repos / non-interactive envs)
+):
+    """
+    Load checkpoint from:
+      - local .pt/.pth/.tar
+      - HuggingFace repo (repo_id)
+
+    Auth behavior:
+      - If logged in -> no token needed
+      - If not logged in -> interactive prompt
+      - If non-interactive -> hf_token required
+    """
+
+    import os
+    import sys
+    import torch
+    from opensportslib.core.utils.config import expand 
+
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    ckpt_path = None
+    hf_error = None
+    is_local = os.path.exists(path)
+    # --------------------------------------------------
+    # Try Hugging Face FIRST
+    # --------------------------------------------------
+    try:
+        from huggingface_hub import hf_hub_download, whoami, login, list_repo_files
+
+        # Ensure auth if needed
+        if hf_token is None:
+            try:
+                whoami()
+            except Exception:
+                if sys.stdin.isatty():
+                    login()
+
+        if not is_local:
+            print(f"[HF] Inspecting repo: {path}")
+
+            files = list_repo_files(path, token=hf_token)
+
+            # find checkpoint file automatically
+            candidates = [
+                f for f in files
+                if f.endswith((".pt", ".pth", ".pth.tar", ".bin"))
+            ]
+
+            if not candidates:
+                raise FileNotFoundError(
+                    f"No checkpoint file found in HF repo {path}. "
+                    f"Files: {files}"
+                )
+
+            # pick first candidate
+            hf_filename = candidates[0]
+            print(f"[HF] Using checkpoint file: {hf_filename}")
+
+            ckpt_path = hf_hub_download(
+                repo_id=path,
+                filename=hf_filename,
+                token=hf_token,
+            )
+
+            print(f"[HF] Loaded from cache: {ckpt_path}")
+
+    except Exception as e:
+        hf_error = e
+
+    # --------------------------------------------------
+    # 2️⃣ Fallback to local
+    # --------------------------------------------------
+    path = expand(path)
+    if ckpt_path is None:
+        if not is_local:
+            raise FileNotFoundError(
+                f"Checkpoint not found on HuggingFace OR locally: {path}"
+            ) from hf_error
+
+        ckpt_path = path
+        print(f"[Local] Using local checkpoint: {ckpt_path}")
+    # --------------------------------------------------
+    # Load checkpoint
+    # --------------------------------------------------
+    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
+
+    # ---------------- MODEL STATE ----------------
+    if isinstance(checkpoint, dict):
+        if "model_state_dict" in checkpoint:
+            state_dict = checkpoint["model_state_dict"]
+        elif "state_dict" in checkpoint:
+            state_dict = checkpoint["state_dict"]
+        else:
+            state_dict = {
+                k: v for k, v in checkpoint.items()
+                if isinstance(v, torch.Tensor)
+            }
+    else:
+        raise ValueError("Checkpoint format not recognized")
+
+    # Clean + remap keys
+    model_keys = list(model.state_dict().keys())
+    ckpt_keys  = list(state_dict.keys())
+
+    ckpt_has_module  = ckpt_keys[0].startswith("module.")
+    model_has_module = model_keys[0].startswith("module.")
+
+    # Case 1: checkpoint has module., model doesn't
+    if ckpt_has_module and not model_has_module:
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    # Case 2: checkpoint doesn't have module., model does
+    elif not ckpt_has_module and model_has_module:
+        state_dict = {f"module.{k}": v for k, v in state_dict.items()}
+
+    # Optional custom remap
+    if key_remap_fn:
+        state_dict = {key_remap_fn(k): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    print("\n--- MISSING KEYS ---")
+    for k in missing[:20]:
+        print(k)
+
+    print("\n--- UNEXPECTED KEYS ---")
+    for k in unexpected[:20]:
+        print(k)
+    model.to(device)
+
+    # ---------------- EPOCH ----------------
+    epoch = checkpoint.get("epoch") if isinstance(checkpoint, dict) else None
+
+    # ---------------- OPTIMIZER ----------------
+    if optimizer and isinstance(checkpoint, dict):
+        opt_state = checkpoint.get("optimizer") or checkpoint.get("optimizer_state_dict")
+        if opt_state:
+            optimizer.load_state_dict(opt_state)
+
+    # ---------------- SCHEDULER ----------------
+    if scheduler and isinstance(checkpoint, dict):
+        sch_state = checkpoint.get("scheduler") or checkpoint.get("scheduler_state_dict")
+        if sch_state:
+            scheduler.load_state_dict(sch_state)
+
+    print(f"[Checkpoint] Loaded from {ckpt_path} | epoch: {epoch}")
+    print(f"Missing keys: {len(missing)}")
+    print(f"Unexpected keys: {len(unexpected)}")
+
+    return model, optimizer, scheduler, epoch
+
+
+
+def load_huggingface_checkpoint(config, path, device):
+    from opensportslib.models.base.video_mae import load_video_mae_checkpoint
+    return load_video_mae_checkpoint(config, device=device, ckpt_path=path)
+

opensportslib/core/utils/config.py

@@ -0,0 +1,199 @@
+
+import os
+import re
+import json
+import gzip
+import yaml
+
+def dict_to_namespace(d, skip_keys=("classes",)):
+    """
+    Recursively convert dict to namespace for easy access,
+    but keep certain keys (like 'classes') as raw dict/list.
+    """
+    from types import SimpleNamespace
+
+    if isinstance(d, dict):
+        out = {}
+        for k, v in d.items():
+            if k in skip_keys:
+                out[k] = v  # leave as-is
+            else:
+                out[k] = dict_to_namespace(v, skip_keys)
+        return SimpleNamespace(**out)
+    elif isinstance(d, list):
+        return [dict_to_namespace(v, skip_keys) for v in d]
+    else:
+        return d
+
+def namespace_to_dict(ns):
+    return {k: vars(v) if hasattr(v, "__dict__") else v for k, v in vars(ns).items()}
+
+def namespace_to_omegaconf(ns):
+    """
+    Recursively convert SimpleNamespace (or dict/list) back to OmegaConf
+    """
+    from omegaconf import OmegaConf
+    from types import SimpleNamespace
+
+    def to_dict(obj):
+        if isinstance(obj, SimpleNamespace):
+            return {k: to_dict(v) for k, v in vars(obj).items()}
+        elif isinstance(obj, dict):
+            return {k: to_dict(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [to_dict(v) for v in obj]
+        else:
+            return obj
+
+    return OmegaConf.create(to_dict(ns))
+    
+def load_config(config_path):
+    """
+    Loading configurations
+    """
+    print(config_path)
+    if config_path.endswith(".yaml") or config_path.endswith(".yml"):
+        with open(config_path, "r") as f:
+            cfg_dict = yaml.safe_load(f)
+    elif config_path.endswith(".json"):
+        with open(config_path, "r") as f:
+            cfg_dict = json.load(f)
+    else:
+        raise ValueError("Unsupported config format. Use YAML or JSON.")
+    return dict_to_namespace(cfg_dict)
+
+
+
+def load_config_omega(path):
+    
+    from omegaconf import OmegaConf
+    cfg = OmegaConf.load(path)
+    # OmegaConf.resolve(cfg)
+    # cfg = OmegaConf.to_container(cfg, resolve=True)
+    return dict_to_namespace(cfg)
+
+def resolve_config_omega(cfg):
+    from omegaconf import OmegaConf, DictConfig
+    #cfg = namespace_to_omegaconf(cfg)
+    #cfg = namespace_to_dict(cfg)
+    #print(type(cfg))
+    #cfg = OmegaConf.create(cfg)
+    if not isinstance(cfg, DictConfig):
+        return cfg 
+    OmegaConf.resolve(cfg)
+    cfg = dict_to_namespace(OmegaConf.to_container(cfg, resolve=True))
+    return cfg
+
+
+def expand(path):
+    return os.path.abspath(os.path.expanduser(path))
+
+
+def load_json(fpath):
+    with open(fpath) as fp:
+        return json.load(fp)
+
+def load_gz_json(fpath):
+    with gzip.open(fpath, "rt", encoding="ascii") as fp:
+        return json.load(fp)
+
+
+def store_json(fpath, obj, pretty=False):
+    kwargs = {}
+    if pretty:
+        kwargs["indent"] = 4
+        kwargs["sort_keys"] = False
+    with open(fpath, "w") as fp:
+        json.dump(obj, fp, **kwargs)
+
+
+def store_gz_json(fpath, obj):
+    with gzip.open(fpath, "wt", encoding="ascii") as fp:
+        json.dump(obj, fp)
+
+
+def load_text(fpath):
+    """Load text from a given file.
+
+    Args:
+        fpath (string): The path of the file.
+
+    Returns:
+        lines (List): List in which element is a line of the file.
+
+    """
+    lines = []
+    with open(fpath, "r") as fp:
+        for l in fp:
+            l = l.strip()
+            if l:
+                lines.append(l)
+    return lines
+
+def load_classes(input):
+    """Load classes from either list or txt file.
+
+    Args:
+        input (string): Path of the file that contains one class per line or list of classes.
+
+    Returns:
+        Dictionnary with classes associated to indexes.
+    """
+    from omegaconf import ListConfig
+    if isinstance(input, (list, ListConfig)):
+        return {x: i + 1 for i, x in enumerate(input)}
+    return {x: i + 1 for i, x in enumerate(load_text(input))}
+
+def clear_files(dir_name, re_str, exclude=[]):
+    for file_name in os.listdir(dir_name):
+        if re.match(re_str, file_name):
+            if file_name not in exclude:
+                file_path = os.path.join(dir_name, file_name)
+                os.remove(file_path)
+
+
+def _print_info_helper(src_file, labels):
+    """Print informations about videos contained in a json file.
+
+    Args:
+        src_file (string): The source file.
+        labels (list(dict)): List containing a dict fro each video.
+    """
+    num_frames = sum([x["num_frames"] for x in labels])
+    num_events = sum([len(x["events"]) for x in labels])
+    print(
+        "{} : {} videos, {} frames, {:0.5f}% non-bg".format(
+            src_file, len(labels), num_frames, num_events / num_frames * 100
+        )
+    )
+
+def select_device(config):
+    import torch
+    mode = config.device.lower()
+
+    if mode == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    elif mode == "cuda":
+        assert torch.cuda.is_available(), "CUDA requested but not available"
+        gpu_id = getattr(config, "gpu_id", 0)
+        torch.cuda.set_device(gpu_id)
+        device = torch.device(f"cuda:{gpu_id}")
+
+    elif mode == "cpu":
+        device = torch.device("cpu")
+
+    else:
+        raise ValueError(f"Unknown device mode: {mode}")
+
+    print(f"Using device: {device}")
+    if device.type == "cuda" or device.type == "auto":
+        print(f"GPU: {torch.cuda.get_device_name(device)}")
+
+    return device
+
+def is_local_path(p):
+    return p and (
+        os.path.exists(p) or
+        p.endswith((".pt", ".pth", ".tar"))
+    )

opensportslib/core/utils/data.py

@@ -0,0 +1,85 @@
+from collections import defaultdict
+from torch.utils.data import Subset
+import torch
+import numpy as np
+
+def balanced_subset(dataset, samples_per_class=5):
+    class_indices = defaultdict(list)
+
+    for idx in range(len(dataset)):
+        label = dataset.samples[idx]["label"]
+        class_indices[label].append(idx)
+
+    print(class_indices.keys())
+    selected_indices = []
+    for label, indices in class_indices.items():
+        selected_indices.extend(indices[:samples_per_class])
+
+    print(selected_indices)
+    return Subset(dataset, selected_indices)
+
+
+def batch_tensor(tensor, dim=1, squeeze=False):
+    """
+    A function to reshape PyTorch tensor `tensor` along some dimension `dim` to the batch dimension 0 such that the tensor can be processed in parallel. 
+    If `sqeeze`=True, the dimension `dim` will be removed completely, otherwise it will be of size=1. Check `unbatch_tensor()` for the reverese function.
+    """
+    batch_size, dim_size = tensor.shape[0], tensor.shape[dim]
+    returned_size = list(tensor.shape)
+    returned_size[0] = batch_size * dim_size
+    returned_size[dim] = 1
+    if squeeze:
+        return tensor.transpose(0, dim).reshape(returned_size).squeeze_(dim)
+    else:
+        return tensor.transpose(0, dim).reshape(returned_size)
+
+
+def unbatch_tensor(tensor, batch_size, dim=1, unsqueeze=False):
+    """
+    A function to chunk pytorch tensor `tensor` along the batch dimension 0 and concatenate the chuncks on dimension `dim` to recover from `batch_tensor()` function.
+    If `unsqueee`=True, it will add a dimension `dim` before the unbatching. 
+    """
+    fake_batch_size = tensor.shape[0]
+    nb_chunks = int(fake_batch_size / batch_size)
+    if unsqueeze:
+        return torch.cat(torch.chunk(tensor.unsqueeze_(dim), nb_chunks, dim=0), dim=dim).contiguous()
+    else:
+        return torch.cat(torch.chunk(tensor, nb_chunks, dim=0), dim=dim).contiguous()
+
+
+def tracking_collate_fn(batch):
+    """
+    Custom collate function for tracking data.
+    Uses PyG Batch.from_data_list for efficient C++ batching.
+    """
+    from torch_geometric.data import Batch
+    
+    batch_size = len(batch)
+    seq_len = batch[0]['seq_len']
+    
+    # flatten all graphs from all samples
+    all_graphs = []
+    for sample_idx, item in enumerate(batch):
+        for time_idx, graph in enumerate(item['graphs']):
+            all_graphs.append(graph)
+    
+    # PyG handles node offsets for edge_index automatically
+    batched_graphs = Batch.from_data_list(all_graphs)
+    
+    return {
+        'x': batched_graphs.x,
+        'edge_index': batched_graphs.edge_index,
+        'batch': batched_graphs.batch,
+        'batch_size': batch_size,
+        'seq_len': seq_len,
+        'labels': torch.tensor([item['label'] for item in batch], dtype=torch.long),
+        'id': [item['id'] for item in batch],
+    }
+
+def mixup_data(x, y, alpha=0.2):
+    """blend pairs of samples and their labels for mixup augmentation."""
+    lam = np.random.beta(alpha, alpha) if alpha > 0 else 1.0
+    index = torch.randperm(x.size(0)).to(x.device)
+    mixed_x = lam * x + (1 - lam) * x[index]
+    return mixed_x, y, y[index], lam
+    

opensportslib/core/utils/ddp.py

@@ -0,0 +1,77 @@
+import torch.distributed as dist
+def ddp_setup(rank, world_size):
+    import os
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"  # any free port
+    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+
+def ddp_cleanup():
+    dist.destroy_process_group()
+
+
+import torch
+from torch.utils.data import Sampler
+import math
+
+class DistributedWeightedSampler(Sampler):
+    def __init__(
+        self,
+        weights,
+        num_replicas=None,
+        rank=None,
+        replacement=True,
+        num_samples=None,
+        seed=0
+    ):
+        if num_replicas is None:
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            rank = dist.get_rank()
+
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.replacement = replacement
+        self.seed = seed
+
+        self.weights = torch.as_tensor(weights, dtype=torch.double)
+        self.dataset_size = len(self.weights)
+
+        # Split dataset across ranks
+        self.num_samples_per_rank = math.ceil(self.dataset_size / self.num_replicas)
+        self.total_size = self.num_samples_per_rank * self.num_replicas
+
+        # Pad if needed
+        if self.total_size > self.dataset_size:
+            padding = self.total_size - self.dataset_size
+            self.weights = torch.cat([self.weights, self.weights[:padding]])
+
+        # Indices for this rank
+        self.rank_indices = list(range(self.rank, self.total_size, self.num_replicas))
+        self.rank_weights = self.weights[self.rank_indices]
+
+        # Number of samples to draw
+        if num_samples is None:
+            self.num_samples = len(self.rank_indices)
+        else:
+            self.num_samples = num_samples // self.num_replicas
+
+        self.epoch = 0
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def __iter__(self):
+        g = torch.Generator()
+        g.manual_seed(self.seed + self.epoch)
+
+        sampled = torch.multinomial(
+            self.rank_weights,
+            self.num_samples,
+            self.replacement,
+            generator=g
+        )
+
+        return (self.rank_indices[i] % self.dataset_size for i in sampled)
+
+    def __len__(self):
+        return self.num_samples

opensportslib/core/utils/default_args.py

@@ -0,0 +1,110 @@
+def get_default_args_data_train_e2e_dali(cfg):
+    return {
+        "classes": cfg.DATA.classes,
+        "train": True,
+        "acc_grad_iter": cfg.TRAIN.acc_grad_iter,
+        "num_epochs": cfg.TRAIN.num_epochs,
+        "repartitions": cfg.TRAIN.repartitions,
+    }
+
+
+def get_default_args_data_valid_e2e_dali(cfg):
+    return {
+        "classes": cfg.DATA.classes,
+        "train": False,
+        "acc_grad_iter": cfg.TRAIN.acc_grad_iter,
+        "num_epochs": cfg.TRAIN.num_epochs,
+        "repartitions": cfg.TRAIN.repartitions,
+    }
+
+
+def get_default_args_data_train_e2e_opencv(cfg):
+    return {"classes": cfg.DATA.classes, "train": True}
+
+
+def get_default_args_data_valid_e2e_opencv(cfg):
+    return {"classes": cfg.DATA.classes, "train": False}
+
+
+def get_default_args_data_train():
+    return {"train": True}
+
+
+def get_default_args_data_valid():
+    return {"train": False}
+
+
+def get_default_args_data_valid_data_frames_e2e_dali(cfg):
+    return {"classes": cfg.DATA.classes, "repartitions": cfg.TRAIN.repartitions}
+
+
+def get_default_args_data_valid_data_frames_e2e_opencv(cfg):
+    return {"classes": cfg.DATA.classes}
+
+
+def get_default_args_dataset(split, cfg):
+    if split == "train":
+        if cfg.MODEL.runner.type == "runner_e2e":
+            if getattr(cfg, "dali", False):
+                return get_default_args_data_train_e2e_dali(cfg)
+            else:
+                return get_default_args_data_train_e2e_opencv(cfg)
+        else:
+            return get_default_args_data_train()
+
+    elif split == "valid":
+        if cfg.MODEL.runner.type == "runner_e2e":
+            if getattr(cfg, "dali", False):
+                return get_default_args_data_valid_e2e_dali(cfg)
+            else:
+                return get_default_args_data_valid_e2e_opencv(cfg)
+        else:
+            return get_default_args_data_valid()
+
+    elif split == "valid_data_frames" or split == "test" or split == "challenge":
+        if cfg.MODEL.runner.type == "runner_e2e":
+            if getattr(cfg, "dali", False):
+                return get_default_args_data_valid_data_frames_e2e_dali(cfg)
+            else:
+                return get_default_args_data_valid_data_frames_e2e_opencv(cfg)
+        else:
+            return
+    else:
+        return None
+
+
+def get_default_args_model(cfg):
+    if cfg.MODEL.type == "E2E":
+        return {"classes": cfg.DATA.classes}
+    else:
+        return None
+
+
+def get_default_args_trainer(cfg, len_train_loader):
+    if cfg.TRAIN.type == "trainer_e2e":
+        return {
+            "len_train_loader": len_train_loader,
+            "work_dir": cfg.SYSTEM.work_dir,
+            "dali": cfg.dali,
+            "repartitions": cfg.TRAIN.repartitions if cfg.dali else None,
+            "cfg_test": cfg.DATA.test,
+            #"cfg_challenge": cfg.DATA.challenge,
+            "cfg_valid_data_frames": cfg.DATA.valid_data_frames,
+        }
+    else:
+        return {"work_dir": cfg.SYSTEM.work_dir}
+
+
+def get_default_args_train(model, train_loader, valid_loader, classes, trainer_type):
+    if trainer_type == "trainer_CALF" or trainer_type == "trainer_pooling":
+        return {
+            "model": model,
+            "train_dataloaders": train_loader,
+            "val_dataloaders": valid_loader,
+        }
+    elif trainer_type == "trainer_e2e":
+        return {
+            "train_loader": train_loader,
+            "valid_loader": valid_loader,
+            "classes": classes,
+        }