anyantudre
commited on
Commit
·
d336d1e
1
Parent(s):
2ee155f
fixing CI
Browse files- app.py +8 -1
- training/config/detector/ucf.yaml +1 -1
- training/config/detector/xception.yaml +1 -1
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:
|
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:
|
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 |
|