diff --git a/crabs/detector/models.py b/crabs/detector/models.py index a6bc8a4b..4b1f93c0 100644 --- a/crabs/detector/models.py +++ b/crabs/detector/models.py @@ -243,3 +243,37 @@ def configure_optimizers(self) -> dict[str, torch.optim.Optimizer]: weight_decay=self.config["wdecay"], ) return {"optimizer": optimizer} + + # def predict_step( + # self, images: torch.Tensor, dataloader_idx=0 + # ) -> torch.Tensor: + # """Perform inference on a batch of images.""" + + # # ensure batch is on device + # # batch = batch.to(self.device) + # # torch.Size([2160, 4096, 3]) + + # # apply inference transform + # # test_val_transforms = CrabsDataModule._get_test_val_transform(None) + # # inference_transforms = transforms.Compose( + # # [ + # # transforms.ToImage(), + # # transforms.ToDtype(torch.float32, scale=True), + # # ] + # # ) + # # batch_transformed = inference_transforms(batch) # B, C, H, W + # # breakpoint() + + # # # run prediction + # # predictions = self.model(images) + # transform = transforms.Compose( + # [ + # transforms.ToImage(), + # transforms.ToDtype(torch.float32, scale=True), + # ] + # ) + # image_tensors = transform(images).to(self.device) + # # image_tensors = image_tensors.unsqueeze(0) + # with torch.no_grad(): + # predictions = self.model(image_tensors) + # return predictions diff --git a/crabs/tracker/track_video.py b/crabs/tracker/track_video.py index 7971ef3e..b3d4efcf 100644 --- a/crabs/tracker/track_video.py +++ b/crabs/tracker/track_video.py @@ -7,12 +7,18 @@ from pathlib import Path import cv2 +import lightning import numpy as np import torch import torchvision.transforms.v2 as transforms import yaml # type: ignore from crabs.detector.models import FasterRCNN +from crabs.detector.utils.detection import ( + log_mlflow_metadata_as_info, + set_mlflow_run_name, + setup_mlflow_logger, +) from crabs.tracker.evaluate_tracker import TrackerEvaluate from crabs.tracker.sort import Sort from crabs.tracker.utils.io import ( @@ -61,6 +67,41 @@ def __init__(self, args: argparse.Namespace) -> None: iou_threshold=self.config["iou_threshold"], ) + # MLflow experiment name and run name + self.experiment_name = args.experiment_name + self.run_name = set_mlflow_run_name() + self.mlflow_folder = args.mlflow_folder + + # Log MLflow information to screen + log_mlflow_metadata_as_info(self) + + def setup_trainer(self): + """Set up trainer object with logging for testing.""" + # Setup logger + mlf_logger = setup_mlflow_logger( + experiment_name=self.experiment_name, + run_name=self.run_name, + mlflow_folder=self.mlflow_folder, + cli_args=self.args, + ) + + # Add trained model section to MLflow hyperparameters + mlf_logger.log_hyperparams( + { + "trained_model/experiment_name": self.trained_model_expt_name, + "trained_model/run_name": self.trained_model_run_name, + "trained_model/ckpt_file": Path(self.trained_model_path).name, + } + ) + + # Add other unlogged information from init? + + # Return trainer linked to logger + return lightning.Trainer( + accelerator=self.accelerator, # lightning accelerators + logger=mlf_logger, + ) + def setup(self): """Load tracking config, trained model and input video path.""" with open(self.config_file) as f: @@ -229,6 +270,12 @@ def run_tracking(self): if self.args.save_frames: close_csv_file(self.csv_file) + # if this is a slurm job: add slurm logs as artifacts + # slurm_job_id = os.environ.get("SLURM_JOB_ID") + # slurm_job_name = os.environ.get("SLURM_JOB_NAME") + # if slurm_job_id and (slurm_job_name != "bash"): + # slurm_logs_as_artifacts(trainer.logger, slurm_job_id) + def main(args) -> None: """Run detection+tracking inference on video. @@ -319,6 +366,25 @@ def tracking_parse_args(args): "Valid inputs are: cpu or gpu. Default: gpu." ), ) + parser.add_argument( + "--experiment_name", + type=str, + default="Inference", + help=( + "Name of the experiment in MLflow, under which the current run " + "will be logged. " + "By default: Inference." + ), + ) + parser.add_argument( + "--mlflow_folder", + type=str, + default="./ml-runs", + help=( + "Path to MLflow directory where to log the evaluation data. " + "Default: 'ml-runs' directory under the current working directory." + ), + ) parser.add_argument( "--max_frames_to_read", type=int, diff --git a/tests/test_unit/test_track_video.py b/tests/test_unit/test_track_video.py index 7d7ffa89..31f5d92c 100644 --- a/tests/test_unit/test_track_video.py +++ b/tests/test_unit/test_track_video.py @@ -20,6 +20,8 @@ def mock_args(): accelerator="gpu", annotations_file=None, save_video=None, + experiment_name=None, + mlflow_folder=Path(temp_dir) / "ml-runs", )