mrneuralnet commited on
Commit
f4c3cd9
1 Parent(s): e10b61a

Modify inference script args

Browse files
Files changed (3) hide show
  1. app.py +6 -5
  2. extract_video.py +1 -1
  3. inference.py +2 -3
app.py CHANGED
@@ -18,7 +18,7 @@ from file_picker import st_file_selector
18
 
19
  import os
20
 
21
- DEBUG = False
22
 
23
  def main():
24
  st.markdown("###")
@@ -36,6 +36,7 @@ def main():
36
  filename = "{}.{}".format(random_id, uploaded_file.type.split("/")[-1])
37
  file_type = uploaded_file.type.split("/")[0]
38
  filepath = f"{base_folder}/{filename}"
 
39
 
40
  if uploaded_file.type == 'video/mp4':
41
  with open(f"temps/{filename}", mode='wb') as f:
@@ -52,6 +53,7 @@ def main():
52
  file_type = selected_file.split(".")[-1]
53
  filename = selected_file.split("/")[-1]
54
  filepath = f"{base_folder}/{selected_file}"
 
55
 
56
  if file_type == 'mp4':
57
  video_file = open(filepath, 'rb')
@@ -71,10 +73,10 @@ def main():
71
  processing_stdout = subprocess.run(shlex.split(f"""python extract_video.py --device cpu --max_frames 50 --bs 2 --frame_interval 60 --confidence_threshold 0.997 --data_path "{filepath}" """), capture_output=True)
72
  st.text(f'1. Processing {file_type} ✅')
73
  with st.spinner(f'Analyzing {file_type}...'):
74
- analyze_stdout = subprocess.run(shlex.split(f"""python inference.py --weight weights/model_params_ffpp_c23.pickle --device cpu --image_folder "{base_folder}/images/{filename}" """), capture_output=True)
75
  st.text(f'2. Analyzing {file_type} ✅')
76
 
77
- if len(os.listdir("{}/images/{}".format(base_folder, filename.split(".")[0]))) < 1:
78
  st.text("No faces could be detected! 🚨")
79
  return
80
 
@@ -87,8 +89,7 @@ def main():
87
  st.text("fake probability {:.2f}".format(fake_probability))
88
 
89
  # os.remove(f"{base_folder}/{filename}")
90
- folder_name = ".".join(filename.split(".")[:-1])
91
- shutil.rmtree(f"{base_folder}/images/{folder_name}")
92
  except Exception as e:
93
  if DEBUG:
94
  st.text(processing_stdout.stdout.decode('utf-8'))
 
18
 
19
  import os
20
 
21
+ DEBUG = True
22
 
23
  def main():
24
  st.markdown("###")
 
36
  filename = "{}.{}".format(random_id, uploaded_file.type.split("/")[-1])
37
  file_type = uploaded_file.type.split("/")[0]
38
  filepath = f"{base_folder}/{filename}"
39
+ faces_folder = f"{base_folder}/images/{random_id}"
40
 
41
  if uploaded_file.type == 'video/mp4':
42
  with open(f"temps/{filename}", mode='wb') as f:
 
53
  file_type = selected_file.split(".")[-1]
54
  filename = selected_file.split("/")[-1]
55
  filepath = f"{base_folder}/{selected_file}"
56
+ faces_folder = f"{base_folder}/images/" + selected_file.split(".")[0]
57
 
58
  if file_type == 'mp4':
59
  video_file = open(filepath, 'rb')
 
73
  processing_stdout = subprocess.run(shlex.split(f"""python extract_video.py --device cpu --max_frames 50 --bs 2 --frame_interval 60 --confidence_threshold 0.997 --data_path "{filepath}" """), capture_output=True)
74
  st.text(f'1. Processing {file_type} ✅')
75
  with st.spinner(f'Analyzing {file_type}...'):
76
+ analyze_stdout = subprocess.run(shlex.split(f"""python inference.py --weight weights/model_params_ffpp_c23.pickle --device cpu --image_folder "{faces_folder}" """), capture_output=True)
77
  st.text(f'2. Analyzing {file_type} ✅')
78
 
79
+ if len(os.listdir(faces_folder)) < 1:
80
  st.text("No faces could be detected! 🚨")
81
  return
82
 
 
89
  st.text("fake probability {:.2f}".format(fake_probability))
90
 
91
  # os.remove(f"{base_folder}/{filename}")
92
+ shutil.rmtree(faces_folder)
 
93
  except Exception as e:
94
  if DEBUG:
95
  st.text(processing_stdout.stdout.decode('utf-8'))
extract_video.py CHANGED
@@ -187,7 +187,6 @@ def extract_method_videos(data_path, interval):
187
 
188
  image_folder = video.split('.')[0]
189
  try:
190
- print(data_path)
191
  image_list = extract_frames(data_path, interval)
192
  detect(image_list, join(images_path, image_folder))
193
  except Exception as ex:
@@ -225,6 +224,7 @@ if __name__ == '__main__':
225
  print(device)
226
 
227
  # net and model
 
228
  net = RetinaFace(cfg=cfg, phase='test')
229
  net = load_model(net, pretrained_weights, load_to_cpu=False if device=='cuda' else True)
230
  net.eval()
 
187
 
188
  image_folder = video.split('.')[0]
189
  try:
 
190
  image_list = extract_frames(data_path, interval)
191
  detect(image_list, join(images_path, image_folder))
192
  except Exception as ex:
 
224
  print(device)
225
 
226
  # net and model
227
+ print('loading the model...')
228
  net = RetinaFace(cfg=cfg, phase='test')
229
  net = load_model(net, pretrained_weights, load_to_cpu=False if device=='cuda' else True)
230
  net.eval()
inference.py CHANGED
@@ -74,9 +74,8 @@ def prepare_data():
74
  images.append(preprocess(args.image))
75
  paths.append(args.image)
76
  elif args.image_folder:
77
- image_folder = '.'.join(args.image_folder.split('.')[:-1])
78
- image_paths = glob(image_folder + "/*.jpg")
79
- image_paths.extend(glob(image_folder + "/*.png"))
80
  for _ in image_paths:
81
  images.append(preprocess(_))
82
  paths.append(_)
 
74
  images.append(preprocess(args.image))
75
  paths.append(args.image)
76
  elif args.image_folder:
77
+ image_paths = glob(args.image_folder + "/*.jpg")
78
+ image_paths.extend(glob(args.image_folder + "/*.png"))
 
79
  for _ in image_paths:
80
  images.append(preprocess(_))
81
  paths.append(_)