Add causal person fusion and MuVAP stage-3 baselines

2026-09-10 06:54:22 +08:00
parent 0534c9d3d0
commit d426b4740a
225 changed files with 27716 additions and 119 deletions
@@ -0,0 +1,67 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import time
+
+import wandb
+from rekognition_online_action_detection.criterions import build_criterion
+from rekognition_online_action_detection.datasets import build_data_loader
+from rekognition_online_action_detection.engines import do_train
+from rekognition_online_action_detection.models import build_model
+from rekognition_online_action_detection.optimizers import (build_optimizer,
+                                                            build_scheduler)
+from rekognition_online_action_detection.utils.checkpointer import \
+    setup_checkpointer
+from rekognition_online_action_detection.utils.env import setup_environment
+from rekognition_online_action_detection.utils.logger import setup_logger
+from rekognition_online_action_detection.utils.parser import load_cfg
+from utils import parse_config
+
+
+def main(cfg):
+    parse_config(cfg)
+    # Setup configurations
+    device = setup_environment(cfg)
+    checkpointer = setup_checkpointer(cfg, phase="train")
+    logger = setup_logger(cfg, phase="train")
+
+    print(cfg)
+    print(f"cfg['SOLVER']['NUM_EPOCHS']: {cfg['SOLVER']['NUM_EPOCHS']}")
+    print(f"cfg['MODEL']['CHECKPOINT']: {cfg['MODEL']['CHECKPOINT']}")
+
+    # Build data loaders
+    data_loaders = {phase: build_data_loader(cfg, phase) for phase in cfg.SOLVER.PHASES}
+
+    # Build model
+    model = build_model(cfg, device)
+
+    # Build criterion
+    criterion = build_criterion(cfg, device)
+
+    # Build optimizer
+    optimizer = build_optimizer(cfg, model)
+
+    # Load pretrained model and optimizer
+    checkpointer.load(model, optimizer)
+
+    # Build scheduler
+    scheduler = build_scheduler(cfg, optimizer, len(data_loaders["train"]))
+
+    do_train(
+        cfg,
+        data_loaders,
+        model,
+        criterion,
+        optimizer,
+        scheduler,
+        device,
+        checkpointer,
+        logger,
+    )
+
+    wandb.finish() if cfg.USE_WANDB else None
+
+
+if __name__ == "__main__":
+    main(load_cfg())