Modify inference script args

app.py

@@ -18,7 +18,7 @@ from file_picker import st_file_selector
 
 import os
 
-DEBUG = False
+DEBUG = True
 
 def main():
     st.markdown("###")    
@@ -36,6 +36,7 @@ def main():
         filename = "{}.{}".format(random_id, uploaded_file.type.split("/")[-1])
         file_type = uploaded_file.type.split("/")[0]
         filepath = f"{base_folder}/{filename}"
+        faces_folder = f"{base_folder}/images/{random_id}"
 
         if uploaded_file.type == 'video/mp4':
             with open(f"temps/{filename}", mode='wb') as f:
@@ -52,6 +53,7 @@ def main():
         file_type = selected_file.split(".")[-1]
         filename = selected_file.split("/")[-1]
         filepath = f"{base_folder}/{selected_file}"
+        faces_folder = f"{base_folder}/images/" + selected_file.split(".")[0]
 
         if file_type == 'mp4':
             video_file = open(filepath, 'rb')
@@ -71,10 +73,10 @@ def main():
         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)
     st.text(f'1. Processing {file_type} ✅')
     with st.spinner(f'Analyzing {file_type}...'):
-        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)
+        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)
     st.text(f'2. Analyzing {file_type} ✅')
 
-    if len(os.listdir("{}/images/{}".format(base_folder, filename.split(".")[0]))) < 1:
+    if len(os.listdir(faces_folder)) < 1:
         st.text("No faces could be detected! 🚨")
         return
     
@@ -87,8 +89,7 @@ def main():
         st.text("fake probability {:.2f}".format(fake_probability))
 
         # os.remove(f"{base_folder}/{filename}")
-        folder_name = ".".join(filename.split(".")[:-1])
-        shutil.rmtree(f"{base_folder}/images/{folder_name}")
+        shutil.rmtree(faces_folder)
     except Exception as e:
         if DEBUG:
             st.text(processing_stdout.stdout.decode('utf-8'))

extract_video.py

@@ -187,7 +187,6 @@ def extract_method_videos(data_path, interval):
 
     image_folder = video.split('.')[0]
     try:
-        print(data_path)
         image_list = extract_frames(data_path, interval)
         detect(image_list, join(images_path, image_folder))
     except Exception as ex:
@@ -225,6 +224,7 @@ if __name__ == '__main__':
     print(device)
 
     # net and model
+    print('loading the model...')
     net = RetinaFace(cfg=cfg, phase='test')
     net = load_model(net, pretrained_weights, load_to_cpu=False if device=='cuda' else True)
     net.eval()

inference.py

@@ -74,9 +74,8 @@ def prepare_data():
         images.append(preprocess(args.image))
         paths.append(args.image)
     elif args.image_folder:
-        image_folder = '.'.join(args.image_folder.split('.')[:-1])
-        image_paths = glob(image_folder + "/*.jpg")
-        image_paths.extend(glob(image_folder + "/*.png"))
+        image_paths = glob(args.image_folder + "/*.jpg")
+        image_paths.extend(glob(args.image_folder + "/*.png"))
         for _ in image_paths:
             images.append(preprocess(_))
             paths.append(_)