anyantudre commited on
Commit
d336d1e
·
1 Parent(s): 2ee155f
app.py CHANGED
@@ -18,13 +18,19 @@ AVAILABLE_MODELS = [
18
  "ucf",
19
  ]
20
 
21
- # load the model from HF Model Registry
22
  def load_model(model_name, config_path, weights_path):
23
  with open(config_path, 'r') as f:
24
  config = yaml.safe_load(f)
25
 
26
  config['model_name'] = model_name
27
 
 
 
 
 
 
 
 
28
  model_class = DETECTOR[model_name]
29
  model = model_class(config).to(device)
30
 
@@ -33,6 +39,7 @@ def load_model(model_name, config_path, weights_path):
33
  model.eval()
34
  return model
35
 
 
36
  # preprocess a single video
37
  def preprocess_video(video_path, output_dir, frame_num=32):
38
  os.makedirs(output_dir, exist_ok=True)
 
18
  "ucf",
19
  ]
20
 
 
21
  def load_model(model_name, config_path, weights_path):
22
  with open(config_path, 'r') as f:
23
  config = yaml.safe_load(f)
24
 
25
  config['model_name'] = model_name
26
 
27
+ # download the pretrained model from Hugging Face
28
+ if 'pretrained' in config:
29
+ pretrained_filename = f"{model_name}_best.pth"
30
+ repo_id = "ArissBandoss/deepfake-video-classifier"
31
+ pretrained_path = hf_hub_download(repo_id=repo_id, filename=pretrained_filename)
32
+ config['pretrained'] = pretrained_path
33
+
34
  model_class = DETECTOR[model_name]
35
  model = model_class(config).to(device)
36
 
 
39
  model.eval()
40
  return model
41
 
42
+
43
  # preprocess a single video
44
  def preprocess_video(video_path, output_dir, frame_num=32):
45
  os.makedirs(output_dir, exist_ok=True)
training/config/detector/ucf.yaml CHANGED
@@ -2,7 +2,7 @@
2
  log_dir: /data/home/zhiyuanyan/DeepfakeBench/debug_logs/ucf
3
 
4
  # model setting
5
- pretrained: /teamspace/studios/this_studio/DeepfakeBench/training/pretrained/xception-b5690688.pth # path to a pre-trained model, if using one
6
  # pretrained: '/home/zhiyuanyan/.cache/torch/hub/checkpoints/resnet34-b627a593.pth' # path to a pre-trained model, if using one
7
  model_name: ucf # model name
8
  backbone_name: xception # backbone name
 
2
  log_dir: /data/home/zhiyuanyan/DeepfakeBench/debug_logs/ucf
3
 
4
  # model setting
5
+ pretrained: xception_best.pth # path to a pre-trained model, if using one
6
  # pretrained: '/home/zhiyuanyan/.cache/torch/hub/checkpoints/resnet34-b627a593.pth' # path to a pre-trained model, if using one
7
  model_name: ucf # model name
8
  backbone_name: xception # backbone name
training/config/detector/xception.yaml CHANGED
@@ -2,7 +2,7 @@
2
  log_dir: /teamspace/studios/this_studio/DeepfakeBench/logs/testing_bench
3
 
4
  # model setting
5
- pretrained: /teamspace/studios/this_studio/DeepfakeBench/training/pretrained/xception-b5690688.pth # path to a pre-trained model, if using one
6
  model_name: xception # model name
7
  backbone_name: xception # backbone name
8
 
 
2
  log_dir: /teamspace/studios/this_studio/DeepfakeBench/logs/testing_bench
3
 
4
  # model setting
5
+ pretrained: xception_best.pth # path to a pre-trained model, if using one
6
  model_name: xception # model name
7
  backbone_name: xception # backbone name
8