adapt to test set

Author: ControlNet

examples/batfd/batfd/inference.py

@@ -1,12 +1,13 @@
 import os.path
 from typing import Any, List, Optional
+import torch
 from torch import Tensor
 import pandas as pd
 from pathlib import Path
 from lightning.pytorch import LightningModule, Trainer, Callback
+from torch.utils.data import DataLoader
 
 from avdeepfake1m.loader import Metadata
-from torch.utils.data import DataLoader
 
 
 def nullable_index(obj, index):
@@ -38,9 +39,12 @@ def on_predict_batch_end(
             batch_size = fusion_bm_map.shape[0]
 
             for i in range(batch_size):
-                temporal_size = batch[3][i]
+                temporal_size = torch.tensor(100) # the first value of `Batfd.get_meta_attr`
                 video_name = self.metadata[batch_idx * batch_size + i].file
                 n_frames = self.metadata[batch_idx * batch_size + i].video_frames
+                # if n_frames is not available, it should be in test set, and we can get it from the batch
+                if n_frames == -1:
+                    n_frames = batch[-1][i].cpu().numpy().item()
 
                 assert isinstance(video_name, str)
                 self.gen_df_for_batfd(fusion_bm_map[i], temporal_size, n_frames, os.path.join(
@@ -52,9 +56,12 @@ def on_predict_batch_end(
             batch_size = fusion_bm_map.shape[0]
 
             for i in range(batch_size):
-                temporal_size = batch[3][i]
+                temporal_size = torch.tensor(100) # the first value of `BatfdPlus.get_meta_attr`
                 video_name = self.metadata[batch_idx * batch_size + i].file
                 n_frames = self.metadata[batch_idx * batch_size + i].video_frames
+                # if n_frames is not available, it should be in test set, and we can get it from the batch
+                if n_frames == -1:
+                    n_frames = batch[-1][i].cpu().numpy().item()
                 assert isinstance(video_name, str)
 
                 self.gen_df_for_batfd_plus(fusion_bm_map[i], nullable_index(fusion_start, i),
@@ -112,11 +119,9 @@ def inference_model(model_name: str, model: LightningModule, dataloader: DataLoa
     metadata: List[Metadata],
     max_duration: int, model_type: str,
     gpus: int = 1,
-    temp_dir: str = "output/",
-    subset: str = "test"
+    temp_dir: str = "output/"
 ) -> List[Metadata]:
     Path(os.path.join(temp_dir, model_name)).mkdir(parents=True, exist_ok=True)
-    assert subset in ["test", "val"]
 
     model.eval()
 

examples/batfd/infer.py

@@ -20,7 +20,7 @@ def main():
                         help="Root directory of the dataset.")
     parser.add_argument("--num_workers", type=int, default=8,
                         help="Number of workers for data loading.")
-    parser.add_argument("--subset", type=str, choices=["val", "test"], 
+    parser.add_argument("--subset", type=str, choices=["val", "test", "testA", "testB"], 
                         default="test", help="Dataset subset.")
     parser.add_argument("--gpus", type=int, default=1,
                         help="Number of GPUs. Set to 0 for CPU.")
@@ -62,23 +62,41 @@ def main():
         num_workers=args.num_workers,
         get_meta_attr=model.get_meta_attr,
         return_file_name=True,
-        is_plusplus=is_plusplus
+        is_plusplus=is_plusplus,
+        test_subset=args.subset if args.subset in ("test", "testA", "testB") else None
     )
     dm.setup()
 
     Path(output_file).parent.mkdir(parents=True, exist_ok=True)
     Path(temp_dir).mkdir(parents=True, exist_ok=True)
 
-    if args.subset == "test":
+    if args.subset in ("test", "testA", "testB"):
         dataloader = dm.test_dataloader()
-        metadata_path = os.path.join(dm.root, "test_metadata.json")
+        metadata_path = os.path.join(dm.root, f"{args.subset}_metadata.json")
     elif args.subset == "val":
         dataloader = dm.val_dataloader()
         metadata_path = os.path.join(dm.root, "val_metadata.json")
     else:
         raise ValueError("Invalid subset")
 
-    metadata = [Metadata(**each, fps=25) for each in read_json(metadata_path)]
+    if os.path.exists(metadata_path):
+        metadata = [Metadata(**each, fps=25) for each in read_json(metadata_path)]
+    else:
+        metadata = [
+            Metadata(file=file_name, 
+                     original=None,
+                     split=args.subset,
+                     fake_segments=[],
+                     fps=25,
+                     visual_fake_segments=[],
+                     audio_fake_segments=[],
+                     audio_model="",
+                     modify_type="",
+                     # handle by the predictor in `inference_model`
+                     video_frames=-1,
+                     audio_frames=-1)
+            for file_name in dataloader.dataset.file_list
+        ]
 
     inference_model(
         model_name=config["name"], 
@@ -88,8 +106,7 @@ def main():
         max_duration=config["max_duration"], 
         model_type=config["model_type"], 
         gpus=args.gpus, 
-        temp_dir=temp_dir, 
-        subset=args.subset
+        temp_dir=temp_dir
     )
 
     post_process(

python/avdeepfake1m/loader.py

@@ -105,6 +105,7 @@ def __getitem__(self, index: int) -> List[Union[Tensor, str, int]]:
         file = self.file_list[index]
 
         video, audio, _ = read_video(os.path.join(self.root, self.subset, file))
+        n_frames = video.shape[0]
         video = F.interpolate(video.float().permute(1, 0, 2, 3)[None], size=(self.temporal_size, 96, 96))[0]
         audio = F.interpolate(audio.float().permute(1, 0)[None], size=self.audio_temporal_size, mode="linear")[0].permute(1, 0)
         video = self.video_transform(video)
@@ -113,7 +114,7 @@ def __getitem__(self, index: int) -> List[Union[Tensor, str, int]]:
 
         outputs = [video, audio]
 
-        if self.subset != "test":
+        if self.subset not in ("test", "testA", "testB"):
             if self.is_plusplus:
                 subset_folder = self.subset
             else:
@@ -131,6 +132,9 @@ def __getitem__(self, index: int) -> List[Union[Tensor, str, int]]:
             if self.return_file_name:
                 outputs.append(meta.file)
 
+        else:
+            outputs = outputs + [n_frames]
+
         return outputs
 
     def get_label(self, file: str, meta: Metadata) -> tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
@@ -244,6 +248,7 @@ def __init__(self, root: str = "data", temporal_size: int = 100,
         get_meta_attr: Callable[[Metadata, Tensor, Tensor, Tensor], List[Any]] = _default_get_meta_attr,
         return_file_name: bool = False,
         is_plusplus: bool = False,
+        test_subset: Optional[str] = None
     ):
         super().__init__()
         self.root = root
@@ -260,11 +265,15 @@ def __init__(self, root: str = "data", temporal_size: int = 100,
         self.return_file_name = return_file_name
         self.is_plusplus = is_plusplus
         self.Dataset = AVDeepfake1m
+        if test_subset is None:
+            self.test_subset = "test" if not self.is_plusplus else "testA"
+        else:
+            self.test_subset = test_subset
 
     def setup(self, stage: Optional[str] = None) -> None:
         train_file_list = [meta["file"] for meta in read_json(os.path.join(self.root, "train_metadata.json"))]
         val_file_list = [meta["file"] for meta in read_json(os.path.join(self.root, "val_metadata.json"))]
-        with open(os.path.join(self.root, "test_files.txt"), "r") as f:
+        with open(os.path.join(self.root, f"{self.test_subset}_files.txt"), "r") as f:
             test_file_list = list(filter(lambda x: x != "", f.read().split("\n")))
 
         if self.take_val is not None:
@@ -285,7 +294,7 @@ def setup(self, stage: Optional[str] = None) -> None:
             return_file_name=self.return_file_name,
             is_plusplus=self.is_plusplus
         )
-        self.test_dataset = self.Dataset("test", self.root, self.temporal_size, self.max_duration, self.fps,
+        self.test_dataset = self.Dataset(self.test_subset, self.root, self.temporal_size, self.max_duration, self.fps,
             file_list=test_file_list, get_meta_attr=self.get_meta_attr,
             require_match_scores=self.require_match_scores,
             return_file_name=self.return_file_name,