Update app.py

app.py

@@ -7,6 +7,7 @@
 
 import sys
 import os
+
 os.system(f'pip install dlib')
 import dlib
 import argparse
@@ -24,23 +25,26 @@ import matplotlib.pyplot as plt
 from torchvision import transforms
 import traceback
 from pytorch_grad_cam import (
-    GradCAM,ScoreCAM,
+    GradCAM, ScoreCAM,
     XGradCAM, EigenCAM
 )
 from pytorch_grad_cam import GuidedBackpropReLUModel
 from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image
 
-def reshape_transform(tensor,height=14,width=14):
-    result = tensor[:, 1:, :].reshape(tensor.size(0),height,width,tensor.size(2))
-    result = result.transpose(2,3).transpose(1,2)
+
+def reshape_transform(tensor, height=14, width=14):
+    result = tensor[:, 1:, :].reshape(tensor.size(0), height, width, tensor.size(2))
+    result = result.transpose(2, 3).transpose(1, 2)
     return result
 
+
 def get_args_parser():
     parser = argparse.ArgumentParser('FSFM3C fine-tuning&Testing for image classification', add_help=False)
     parser.add_argument('--batch_size', default=64, type=int, help='Batch size per GPU')
     parser.add_argument('--epochs', default=50, type=int)
     parser.add_argument('--accum_iter', default=1, type=int, help='Accumulate gradient iterations')
-    parser.add_argument('--model', default='vit_large_patch16', type=str, metavar='MODEL', help='Name of model to train')
+    parser.add_argument('--model', default='vit_large_patch16', type=str, metavar='MODEL',
+                        help='Name of model to train')
     parser.add_argument('--input_size', default=224, type=int, help='images input size')
     parser.add_argument('--normalize_from_IMN', action='store_true', help='cal mean and std from imagenet')
     parser.set_defaults(normalize_from_IMN=True)
@@ -59,7 +63,8 @@ def get_args_parser():
     parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT', help='Random erase prob')
     parser.add_argument('--remode', type=str, default='pixel', help='Random erase mode')
     parser.add_argument('--recount', type=int, default=1, help='Random erase count')
-    parser.add_argument('--resplit', action='store_true', default=False, help='Do not random erase first augmentation split')
+    parser.add_argument('--resplit', action='store_true', default=False,
+                        help='Do not random erase first augmentation split')
     parser.add_argument('--mixup', type=float, default=0, help='mixup alpha')
     parser.add_argument('--cutmix', type=float, default=0, help='cutmix alpha')
     parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None, help='cutmix min/max ratio')
@@ -69,7 +74,8 @@ def get_args_parser():
     parser.add_argument('--finetune', default='', help='finetune from checkpoint')
     parser.add_argument('--global_pool', action='store_true')
     parser.set_defaults(global_pool=True)
-    parser.add_argument('--cls_token', action='store_false', dest='global_pool', help='Use class token for classification')
+    parser.add_argument('--cls_token', action='store_false', dest='global_pool',
+                        help='Use class token for classification')
     parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str, help='dataset path')
     parser.add_argument('--nb_classes', default=1000, type=int, help='number of the classification types')
     parser.add_argument('--output_dir', default='', help='path where to save')
@@ -116,10 +122,10 @@ def load_model(select_skpt):
     model.load_state_dict(checkpoint['model'], strict=False)
     model.eval()
     global cam
-    cam = GradCAM(model = model,
-                target_layers=[model.blocks[-1].norm1],
-                reshape_transform=reshape_transform
-                )
+    cam = GradCAM(model=model,
+                  target_layers=[model.blocks[-1].norm1],
+                  reshape_transform=reshape_transform
+                  )
     return gr.update(), f"[Loaded Model Successfully:] {args.resume}] "
 
 
@@ -168,29 +174,27 @@ def extract_face_from_fixed_num_frames(src_video, dst_path, num_frames=None):
             img.save(os.path.join(dst_path, '0', save_img_name))
     video_capture.release()
     return frame_indices
+
+
 class TargetCategory:
     def __init__(self, category_index):
         self.category_index = category_index
-    
+
     def __call__(self, output):
         return output[self.category_index]
-def preprocess_image_cam(pil_img,mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225]):
-        # 将 PIL 图像转换为 numpy 数组
+
+
+def preprocess_image_cam(pil_img, 
+                         mean=[0.5482207536697388, 0.42340534925460815, 0.3654651641845703], 
+                         std=[0.2789176106452942, 0.2438540756702423, 0.23493893444538116]):
     img_np = np.array(pil_img)
-    
-    # 归一化到 [0, 1]
     img_np = img_np.astype(np.float32) / 255.0
-    
-    # 标准化
     img_np = (img_np - mean) / std
-    
-    # 调整维度顺序以适应模型输入 (C, H, W)
     img_np = np.transpose(img_np, (2, 0, 1))
-    
-    # 添加批次维度 (B, C, H, W)
     img_np = np.expand_dims(img_np, axis=0)
-    
     return img_np
+
+
 def FSFM3C_image_detection(image):
     frame_path = os.path.join(FRAME_SAVE_PATH, str(len(os.listdir(FRAME_SAVE_PATH))))
     os.makedirs(frame_path, exist_ok=True)
@@ -204,7 +208,9 @@ def FSFM3C_image_detection(image):
     args.batch_size = 1
     dataset_val = build_dataset(is_train=False, args=args)
     sampler_val = torch.utils.data.SequentialSampler(dataset_val)
-    data_loader_val = torch.utils.data.DataLoader(dataset_val, sampler=sampler_val, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False)
+    data_loader_val = torch.utils.data.DataLoader(dataset_val, sampler=sampler_val, batch_size=args.batch_size,
+                                                  num_workers=args.num_workers, pin_memory=args.pin_mem,
+                                                  drop_last=False)
 
     if CKPT_CLASS[ckpt] > 2:
         frame_preds_list, video_pred_list = test_multi_class(data_loader_val, model, device)
@@ -217,10 +223,12 @@ def FSFM3C_image_detection(image):
 
         # Generate CAM heatmap for the detected class
         use_cuda = True
-        input_tensor = preprocess_image(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        input_tensor = preprocess_image(img, 
+                                        mean=[0.5482207536697388, 0.42340534925460815, 0.3654651641845703], 
+                                        std=[0.2789176106452942, 0.2438540756702423, 0.23493893444538116])
         if use_cuda:
             input_tensor = input_tensor.cuda()
-        
+
         # Dynamically determine the target category based on the maximum probability class
         category_names_to_index = {
             'Real or Bonafide': 0,
@@ -229,7 +237,7 @@ def FSFM3C_image_detection(image):
             'Spoofing or Presentation-attack': 3
         }
         target_category = TargetCategory(category_names_to_index[max_prob_class])
-        
+
         grayscale_cam = cam(input_tensor=input_tensor, targets=[target_category])
         grayscale_cam = 1 - grayscale_cam[0, :]
         img = np.array(img)
@@ -244,7 +252,7 @@ def FSFM3C_image_detection(image):
         # cv2.putText(visualization, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
         output_path = "./CAM_images/output_heatmap.png"
         cv2.imwrite(output_path, visualization)
-        return image_results, output_path,probabilities[max_prob_index]
+        return image_results, output_path, probabilities[max_prob_index]
 
     if CKPT_CLASS[ckpt] == 2:
         frame_preds_list, video_pred_list = test_two_class(data_loader_val, model, device)
@@ -257,7 +265,7 @@ def FSFM3C_image_detection(image):
             label = "Spoofing" if prob <= 0.5 else "Bonafide"
             prob = prob if label == "Bonafide" else 1 - prob
         image_results = f"The largest face in this image may be {label} with probability {prob * 100:.1f}%"
-        return image_results, None ,None
+        return image_results, None, None
 
 
 def FSFM3C_video_detection(video, num_frames):
@@ -270,19 +278,25 @@ def FSFM3C_video_detection(video, num_frames):
         args.batch_size = num_frames
         dataset_val = build_dataset(is_train=False, args=args)
         sampler_val = torch.utils.data.SequentialSampler(dataset_val)
-        data_loader_val = torch.utils.data.DataLoader(dataset_val, sampler=sampler_val, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False)
+        data_loader_val = torch.utils.data.DataLoader(dataset_val, sampler=sampler_val, batch_size=args.batch_size,
+                                                      num_workers=args.num_workers, pin_memory=args.pin_mem,
+                                                      drop_last=False)
 
         if CKPT_CLASS[ckpt] > 2:
             frame_preds_list, video_pred_list = test_multi_class(data_loader_val, model, device)
-            class_names = ['Real or Bonafide', 'Deepfake', 'Diffusion or AIGC generated', 'Spoofing or Presentation-attack']
+            class_names = ['Real or Bonafide', 'Deepfake', 'Diffusion or AIGC generated',
+                           'Spoofing or Presentation-attack']
             avg_video_pred = np.mean(video_pred_list, axis=0)
             max_prob_index = np.argmax(avg_video_pred)
             max_prob_class = class_names[max_prob_index]
             probabilities = [f"{class_names[i]}: {prob * 100:.1f}%" for i, prob in enumerate(avg_video_pred)]
 
-            frame_results = {f"frame_{frame_indices[i]}": [f"{class_names[j]}: {prob * 100:.1f}%" for j, prob in enumerate(frame_preds_list[i])] for i in range(len(frame_indices))}
-            video_results = (f"The largest face in this image may be {max_prob_class} with probability: \n [{', '.join(probabilities)}]\n \n"
-                             f"The frame-level detection results ['frame_index': 'probabilities']: \n{frame_results}")
+            frame_results = {f"frame_{frame_indices[i]}": [f"{class_names[j]}: {prob * 100:.1f}%" for j, prob in
+                                                           enumerate(frame_preds_list[i])] for i in
+                             range(len(frame_indices))}
+            video_results = (
+                f"The largest face in this image may be {max_prob_class} with probability: \n [{', '.join(probabilities)}]\n \n"
+                f"The frame-level detection results ['frame_index': 'probabilities']: \n{frame_results}")
             return video_results
 
         if CKPT_CLASS[ckpt] == 2:
@@ -292,23 +306,26 @@ def FSFM3C_video_detection(video, num_frames):
                 label = "Deepfake" if prob <= 0.5 else "Real"
                 prob = prob if label == "Real" else 1 - prob
                 frame_results = {f"frame_{frame_indices[i]}": f"{(frame_preds_list[i]) * 100:.1f}%" for i in
-                                 range(len(frame_indices))} if label == "Real" else {f"frame_{frame_indices[i]}": f"{(1 - frame_preds_list[i]) * 100:.1f}%" for i in
-                                 range(len(frame_indices))}
+                                 range(len(frame_indices))} if label == "Real" else {
+                    f"frame_{frame_indices[i]}": f"{(1 - frame_preds_list[i]) * 100:.1f}%" for i in
+                    range(len(frame_indices))}
 
             if ckpt == 'FAS-Checkpoint_Fine-tuned_on_MCIO':
                 prob = sum(video_pred_list) / len(video_pred_list)
                 label = "Spoofing" if prob <= 0.5 else "Bonafide"
                 prob = prob if label == "Bonafide" else 1 - prob
                 frame_results = {f"frame_{frame_indices[i]}": f"{(frame_preds_list[i]) * 100:.1f}%" for i in
-                                 range(len(frame_indices))} if label == "Bonafide" else {f"frame_{frame_indices[i]}": f"{(1 - frame_preds_list[i]) * 100:.1f}%" for i in
-                                 range(len(frame_indices))}
+                                 range(len(frame_indices))} if label == "Bonafide" else {
+                    f"frame_{frame_indices[i]}": f"{(1 - frame_preds_list[i]) * 100:.1f}%" for i in
+                    range(len(frame_indices))}
 
             video_results = (f"The largest face in this image may be {label} with probability {prob * 100:.1f}%\n \n"
-                            f"The frame-level detection results ['frame_index': 'real_face_probability']: \n{frame_results}")
+                             f"The frame-level detection results ['frame_index': 'real_face_probability']: \n{frame_results}")
             return video_results
     except Exception as e:
         return f"Error occurred. Please provide a clear face video or reduce the number of frames."
 
+
 # Paths and Constants
 P = os.path.abspath(__file__)
 FRAME_SAVE_PATH = os.path.join(os.path.dirname(P), 'frame')
@@ -337,14 +354,17 @@ CKPT_MODEL = {
 }
 
 with gr.Blocks(css=".custom-label { font-weight: bold !important; font-size: 16px !important; }") as demo:
-    gr.HTML("<h1 style='text-align: center;'>🦱 Real Facial Image&Video Detection <br> Against Face Forgery (Deepfake/Diffusion) and Spoofing (Presentation-attacks)</h1>")
-    gr.Markdown("<b>☉ Powered by the fine-tuned ViT models that is pre-trained from [FSFM-3C](https://fsfm-3c.github.io/)</b> <br> "
-                "<b>☉ We do not and cannot access or store the data you have uploaded!</b> <br> "
-                "<b>☉ Release (Continuously updating) </b> <br> <b>[V1.0] 2025/02/22-Current🎉</b>: "
-                "1) Updated <b>[✨Unified-detector_v1] for Unified Physical-Digital Face Attack&Forgery Detection, a ViT-B/16-224 (FSFM Pre-trained) detector that could identify Real&Bonafide, Deepfake, Diffusion&AIGC, Spooing&Presentation-attacks facial images or videos </b> ; 2) Provided the selection of the number of video frames (uniformly sampling 1-32 frames, more frames may time-consuming for this page without GPU acceleration); 3) Fixed some errors of V0.1 including loading and prediction. <br>"
-                "<b>[V0.1] 2024/12-2025/02/21</b>: "
-                "Create this page with basic detectors [DfD-Checkpoint_Fine-tuned_on_FF++, FAS-Checkpoint_Fine-tuned_on_MCIO] that follow the paper implementation. <br> ")
-    gr.Markdown("- Please <b>provide a facial image or video(<100s)</b>, and <b>select the model</b> for detection: <br> <b>[SUGGEST] [✨Unified-detector_v1_Fine-tuned_on_4_classes]</b> a (FSFM Pre-trained) ViT-B/16-224 for Both Real/Deepfake/Diffusion/Spoofing facial images&videos Detection <br> <b>[DfD-Checkpoint_Fine-tuned_on_FF++]</b> for deepfake detection, FSFM ViT-B/16-224 fine-tuned on the FF++_c23 train&val sets (4 manipulations, 32 frames per video) <br> <b>[FAS-Checkpoint_Fine-tuned_on_MCIO]</b> for face anti-spoofing, FSFM ViT-B/16-224 fine-tuned on the MCIO datasets (2 frames per video)")
+    gr.HTML(
+        "<h1 style='text-align: center;'>🦱 Real Facial Image&Video Detection <br> Against Face Forgery (Deepfake/Diffusion) and Spoofing (Presentation-attacks)</h1>")
+    gr.Markdown(
+        "<b>☉ Powered by the fine-tuned ViT models that is pre-trained from [FSFM-3C](https://fsfm-3c.github.io/)</b> <br> "
+        "<b>☉ We do not and cannot access or store the data you have uploaded!</b> <br> "
+        "<b>☉ Release (Continuously updating) </b> <br> <b>[V1.0] 2025/02/22-Current🎉</b>: "
+        "1) Updated <b>[✨Unified-detector_v1] for Unified Physical-Digital Face Attack&Forgery Detection, a ViT-B/16-224 (FSFM Pre-trained) detector that could identify Real&Bonafide, Deepfake, Diffusion&AIGC, Spooing&Presentation-attacks facial images or videos </b> ; 2) Provided the selection of the number of video frames (uniformly sampling 1-32 frames, more frames may time-consuming for this page without GPU acceleration); 3) Fixed some errors of V0.1 including loading and prediction. <br>"
+        "<b>[V0.1] 2024/12-2025/02/21</b>: "
+        "Create this page with basic detectors [DfD-Checkpoint_Fine-tuned_on_FF++, FAS-Checkpoint_Fine-tuned_on_MCIO] that follow the paper implementation. <br> ")
+    gr.Markdown(
+        "- Please <b>provide a facial image or video(<100s)</b>, and <b>select the model</b> for detection: <br> <b>[SUGGEST] [✨Unified-detector_v1_Fine-tuned_on_4_classes]</b> a (FSFM Pre-trained) ViT-B/16-224 for Both Real/Deepfake/Diffusion/Spoofing facial images&videos Detection <br> <b>[DfD-Checkpoint_Fine-tuned_on_FF++]</b> for deepfake detection, FSFM ViT-B/16-224 fine-tuned on the FF++_c23 train&val sets (4 manipulations, 32 frames per video) <br> <b>[FAS-Checkpoint_Fine-tuned_on_MCIO]</b> for face anti-spoofing, FSFM ViT-B/16-224 fine-tuned on the MCIO datasets (2 frames per video)")
 
     with gr.Row():
         ckpt_select_dropdown = gr.Dropdown(
@@ -358,11 +378,12 @@ with gr.Blocks(css=".custom-label { font-weight: bold !important; font-size: 16p
         model_loading_status = gr.Textbox(label="Model Loading Status")
     with gr.Row():
         with gr.Column(scale=5):
-            gr.Markdown("### Image Detection (Fast Try: copying image from [whichfaceisreal](https://whichfaceisreal.com/))")
+            gr.Markdown(
+                "### Image Detection (Fast Try: copying image from [whichfaceisreal](https://whichfaceisreal.com/))")
             image = gr.Image(label="Upload/Capture/Paste your image", type="pil")
             image_submit_btn = gr.Button("Submit")
             output_results_image = gr.Textbox(label="Detection Result")
-            
+
             with gr.Row():
                 output_heatmap = gr.Image(label="Grad_CAM")
                 output_max_prob_class = gr.Textbox(label="Detected Class")
@@ -373,15 +394,14 @@ with gr.Blocks(css=".custom-label { font-weight: bold !important; font-size: 16p
             video_submit_btn = gr.Button("Submit")
             output_results_video = gr.Textbox(label="Detection Result")
 
-    gr.HTML( 
-            '<div style="display: flex; justify-content: center; gap: 20px; margin-bottom: 20px;">' 
-            '<a href="https://mapmyvisitors.com/web/1bxvi" title="Visit tracker">' 
-            '<img src="https://mapmyvisitors.com/map.png?d=FYhBoxLDEaFAxdfRzk5TuchYOBGrnSa98Ky59EkEEpY&cl=ffffff">' 
-            '</a>' 
-            '</div>' 
-        ) 
+    gr.HTML(
+        '<div style="display: flex; justify-content: center; gap: 20px; margin-bottom: 20px;">'
+        '<a href="https://mapmyvisitors.com/web/1bxvi" title="Visit tracker">'
+        '<img src="https://mapmyvisitors.com/map.png?d=FYhBoxLDEaFAxdfRzk5TuchYOBGrnSa98Ky59EkEEpY&cl=ffffff">'
+        '</a>'
+        '</div>'
+    )
 
-            
     ckpt_select_dropdown.change(
         fn=load_model,
         inputs=[ckpt_select_dropdown],
@@ -390,7 +410,7 @@ with gr.Blocks(css=".custom-label { font-weight: bold !important; font-size: 16p
     image_submit_btn.click(
         fn=FSFM3C_image_detection,
         inputs=[image],
-        outputs=[output_results_image, output_heatmap,output_max_prob_class],
+        outputs=[output_results_image, output_heatmap, output_max_prob_class],
     )
     video_submit_btn.click(
         fn=FSFM3C_video_detection,