impl xception infer for avdf1m++

Author: ControlNet

examples/xception/infer.py

@@ -12,7 +12,7 @@
 parser.add_argument("--checkpoint", type=str)
 parser.add_argument("--model", type=str)
 parser.add_argument("--batch_size", type=int, default=128)
-parser.add_argument("--subset", type=str, choices=["train", "val", "test"])
+parser.add_argument("--subset", type=str, choices=["train", "val", "test", "testA", "testB"])
 parser.add_argument("--gpus", type=int, default=1)
 parser.add_argument("--take_num", type=int, default=None)
 
@@ -27,14 +27,14 @@
         raise ValueError(f"Unknown model: {args.model}")
 
     model.to(device)
-    model.train()
-    test_dataset = AVDeepfake1mPlusPlusVideo(args.subset, args.data_root, take_num=args.take_num)
+    model.train()  # not sure why but eval mode will generate nonsense output
+    test_dataset = AVDeepfake1mPlusPlusVideo(args.subset, args.data_root, take_num=args.take_num, pred_mode=True)
 
     save_path = f"output/{args.model}_{args.subset}.txt"
     Path(save_path).parent.mkdir(parents=True, exist_ok=True)
     with open(save_path, "w") as f:
         with torch.inference_mode():
-            for i, (video, _, label) in enumerate(tqdm(test_dataset)):
+            for i, (video, _, _) in enumerate(tqdm(test_dataset)):
                 # batch video as frames use batch_size
                 preds_video = []
                 for j in range(0, len(video), args.batch_size):
@@ -47,3 +47,4 @@
 
                 file_name = test_dataset.metadata[i].file
                 f.write(f"{file_name};{pred}\n")
+                f.flush()

python/avdeepfake1m/loader.py

@@ -562,13 +562,34 @@ def __init__(self, subset: str, data_root: str = "data",
         image_size: int = 96,
         take_num: Optional[int] = None,
         metadata: Optional[List[Metadata]] = None,
+        pred_mode: bool = False
     ):
         self.subset = subset
         self.data_root = data_root
         self.image_size = image_size
+        self.pred_mode = pred_mode
+
         if metadata is None:
-            metadata_json = read_json(os.path.join(self.data_root, f"{subset}_metadata.json"))
-            self.metadata = [Metadata(**meta, fps=25) for meta in metadata_json]
+            if self.pred_mode:
+                with open(os.path.join(self.data_root, f"{subset}_files.txt"), "r") as f:
+                    files = [line.strip() for line in f.readlines() if line.strip() != ""]
+                self.metadata = [ # dummy metadata for prediction
+                    Metadata(file=file_name, 
+                            original=None,
+                            split=subset,
+                            fake_segments=[],
+                            fps=25,
+                            visual_fake_segments=[],
+                            audio_fake_segments=[],
+                            audio_model="",
+                            modify_type="",
+                            video_frames=-1,
+                            audio_frames=-1)
+                    for file_name in files
+                ]
+            else:
+                metadata_json = read_json(os.path.join(self.data_root, f"{subset}_metadata.json"))
+                self.metadata = [Metadata(**meta, fps=25) for meta in metadata_json]
         else:
             self.metadata = metadata
 
@@ -584,5 +605,5 @@ def __getitem__(self, index):
         video, audio, _ = read_video(os.path.join(self.data_root, self.subset, meta.file))
         if self.image_size != 224:
             video = resize_video(video, (self.image_size, self.image_size))
-        label = len(meta.fake_periods) > 0
-        return video, audio, label
+        label = len(meta.fake_periods) > 0 if not self.pred_mode else False
+        return video, audio, label