Update app.py

Fix fatal error

app.py

@@ -169,19 +169,20 @@ model = models_vit.__dict__['vit_base_patch16'](
     num_classes=args.nb_classes,
     drop_path_rate=args.drop_path,
     global_pool=args.global_pool,
-)
+).to(device)
 
 
 def load_model(ckpt):
     if ckpt == 'choose from here' or 'continuously updating...':
         return gr.update()
-    args.resume = os.path.join(CKPT_SAVE_PATH, ckpt)
+    args.resume = os.path.join(CKPT_SAVE_PATH, CKPT_NAME[ckpt])
     if os.path.isfile(args.resume) == False:
         hf_hub_download(local_dir=CKPT_SAVE_PATH,
                         repo_id='Wolowolo/fsfm-3c/' + CKPT_NAME[ckpt],
                         filename=ckpt)
     checkpoint = torch.load(args.resume, map_location='cpu')
     model.load_state_dict(checkpoint['model'])
+    model.eval()
     return gr.update()
 
 
@@ -276,9 +277,7 @@ def extract_face_from_fixed_num_frames(src_video, dst_path, num_frames=None, dev
     return frame_indices
 
 
-def FSFM3C_video_detection(video):
-    model.to(device)
-
+def FSFM3C_video_detection(video, ckpt_select_dropdown):
     # extract frames
     num_frames = 32
 
@@ -308,21 +307,18 @@ def FSFM3C_video_detection(video):
 
     real_prob_video = int(round(1. - (sum(video_pred_list) / len(video_pred_list)), 2) * 100)
     if real_prob_video > 50:
-        result_message = "real"
+        result_message = "real" if 'FAS' not in ckpt_select_dropdown else 'spoof'
     else:
-        result_message = "fake"
+        result_message = "fake" if 'FAS' not in ckpt_select_dropdown else 'real'
+    prob = 1 - real_prob_image if real_prob_video <= 50 else real_prob_video
+    image_results = (f"The face in this image may be {result_message} with probability is {real_prob_image}%")
 
-    video_results = (f"The face in this video may be {result_message}, "
-                     f"and the video-level real_face_probability is {real_prob_video}% \n"
-                     f"The frame-level detection results ['sampled_frame_index': 'real_face_probability']: \n"
-                     f"{frame_results} \n")
+    video_results = (f"The face in this video may be {result_message} with probability {prob}")
 
     return video_results
 
 
-def FSFM3C_image_detection(image):
-    model.to(device)
-
+def FSFM3C_image_detection(image, ckpt_select_dropdown):
     files = os.listdir(FRAME_SAVE_PATH)
     num_files = len(files)
     frame_path = os.path.join(FRAME_SAVE_PATH, str(num_files))
@@ -352,12 +348,11 @@ def FSFM3C_image_detection(image):
 
     real_prob_image = int(round(1. - (sum(video_pred_list) / len(video_pred_list)), 2) * 100)
     if real_prob_image > 50:
-        result_message = "real"
+        result_message = "real" if 'FAS' not in ckpt_select_dropdown else 'spoof'
     else:
-        result_message = "fake"
-
-    image_results = (f"The face in this image may be {result_message},"
-                     f"and the real_face_probability is {real_prob_image}%")
+        result_message = "fake" if 'FAS' not in ckpt_select_dropdown else 'real'
+    prob = 1 - real_prob_image if real_prob_image <= 50 else real_prob_image
+    image_results = (f"The face in this image may be {result_message} with probability is {real_prob_image}%")
 
     return image_results
 
@@ -406,12 +401,12 @@ with gr.Blocks() as demo:
 
     image_submit_btn.click(
         fn=FSFM3C_image_detection,
-        inputs=[image],
+        inputs=[image, ckpt_select_dropdown],
         outputs=[output_results_image],
     )
     video_submit_btn.click(
         fn=FSFM3C_video_detection,
-        inputs=[video],
+        inputs=[video, ckpt_select_dropdown],
         outputs=[output_results_video],
     )
     ckpt_select_dropdown.change(