Change to gpu inference

app.py

@@ -68,10 +68,10 @@ def main():
     
 
     with st.spinner(f'Processing {file_type}...'):
-        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)
+        processing_stdout = subprocess.run(shlex.split(f"""python extract_video.py --device cuda --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 cuda --image_folder "{base_folder}/images/{filename}" """), capture_output=True)
     st.text(f'2. Analyzing {file_type} ✅')
 
     if len(os.listdir("{}/images/{}".format(base_folder, filename.split(".")[0]))) < 1:

extract_video.py

@@ -37,7 +37,7 @@ def remove_prefix(state_dict, prefix):
     return {f(key): value for key, value in state_dict.items()}
 
 
-def load_model(model, pretrained_path, load_to_cpu):
+def load_model(model, pretrained_path, load_to_cpu=False):
     print('Loading pretrained model from {}'.format(pretrained_path))
     if load_to_cpu:
         pretrained_dict = torch.load(
@@ -226,7 +226,7 @@ if __name__ == '__main__':
 
     # net and model
     net = RetinaFace(cfg=cfg, phase='test')
-    net = load_model(net, pretrained_weights, args.device)
+    net = load_model(net, pretrained_weights, load_to_cpu=False)
     net.eval()
     print('Finished loading model!')
 

inference.py

@@ -120,9 +120,9 @@ def main():
     if args.weight and args.bin:
         raise ValueError("Only one of '--weight' or '--bin' can be set.")
     elif args.weight:
-        weights = torch.load(args.weight, map_location="cpu")
+        weights = torch.load(args.weight, map_location=device)
     elif args.bin:
-        weights = torch.load(args.bin, map_location="cpu")["model"]
+        weights = torch.load(args.bin, map_location=device)["model"]
     else:
         raise ValueError("Neither of '--weight' nor '--bin' is set. Please specify either "
                          "one of these two arguments to load model's weight properly.")