Fix metadata.min.json

Author: ControlNet

evaluate.py

@@ -1,4 +1,5 @@
 import argparse
+import os
 
 import toml
 import torch
@@ -8,7 +9,7 @@
 from metrics import AP, AR
 from model import Batfd, BatfdPlus
 from post_process import post_process
-from utils import read_json
+from utils import generate_metadata_min, read_json
 
 parser = argparse.ArgumentParser(description="BATFD evaluation")
 parser.add_argument("--config", type=str)
@@ -125,6 +126,10 @@ def evaluate_lavdf(config, args):
 
 if __name__ == '__main__':
     args = parser.parse_args()
+    
+    if os.path.exists(os.path.join(args.data_root, "metadata.min.json")):
+        generate_metadata_min(args.data_root)
+
     config = toml.load(args.config)
     torch.backends.cudnn.benchmark = True
     if config["dataset"] == "lavdf":

train.py

@@ -1,12 +1,13 @@
 import argparse
+import os
 
 import toml
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
 
 from dataset.lavdf import LavdfDataModule
 from model import Batfd, BatfdPlus
-from utils import LrLogger, EarlyStoppingLR
+from utils import LrLogger, EarlyStoppingLR, generate_metadata_min
 
 parser = argparse.ArgumentParser(description="BATFD training")
 parser.add_argument("--config", type=str)
@@ -24,6 +25,9 @@
     args = parser.parse_args()
     config = toml.load(args.config)
 
+    if os.path.exists(os.path.join(args.data_root, "metadata.min.json")):
+        generate_metadata_min(args.data_root)
+
     learning_rate = config["optimizer"]["learning_rate"]
     gpus = args.gpus
     total_batch_size = args.batch_size * gpus

utils.py

@@ -1,4 +1,6 @@
+from importlib import metadata
 import json
+import os
 import re
 from abc import ABC
 from typing import List, Tuple, Optional
@@ -219,3 +221,12 @@ def _run_early_stop_checking(self, trainer: Trainer) -> None:
         elif self.mode == "any":
             if any(lr <= self.lr_threshold for lr in all_lr):
                 trainer.should_stop = True
+
+
+def generate_metadata_min(data_root: str):
+    metadata_full = read_json(os.path.join(data_root, "metadata.json"))
+    metadata_min = []
+    for meta in metadata_full:
+        del meta["timestamps"]
+    with open(os.path.join(data_root, "metadata.min.json"), "w") as f:
+        json.dump(metadata_min, f)