Grad_CAM

app.py

@@ -7,7 +7,6 @@
 
 import sys
 import os
-os.system(f'pip install grad-cam')
 os.system(f'pip install dlib')
 import dlib
 import argparse
@@ -21,21 +20,6 @@ import gradio as gr
 import models_vit
 from util.datasets import build_dataset
 from engine_finetune import test_two_class, test_multi_class
-import matplotlib.pyplot as plt
-from torchvision import transforms
-import traceback
-from pytorch_grad_cam import (
-    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)
-    return result
 
 
 def get_args_parser():
@@ -43,8 +27,7 @@ def get_args_parser():
     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)
@@ -63,8 +46,7 @@ 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')
@@ -74,8 +56,7 @@ 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')
@@ -121,11 +102,6 @@ def load_model(select_skpt):
     checkpoint = torch.load(args.resume, map_location=device)
     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
-                  )
     return gr.update(), f"[Loaded Model Successfully:] {args.resume}] "
 
 
@@ -176,25 +152,6 @@ def extract_face_from_fixed_num_frames(src_video, dst_path, num_frames=None):
     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.5482207536697388, 0.42340534925460815, 0.3654651641845703],
-                         std=[0.2789176106452942, 0.2438540756702423, 0.23493893444538116]):
-    img_np = np.array(pil_img)
-    img_np = img_np.astype(np.float32) / 255.0
-    img_np = (img_np - mean) / std
-    img_np = np.transpose(img_np, (2, 0, 1))
-    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)
@@ -208,9 +165,7 @@ 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)
@@ -220,46 +175,7 @@ def FSFM3C_image_detection(image):
         max_prob_class = class_names[max_prob_index]
         probabilities = [f"{class_names[i]}: {prob * 100:.1f}%" for i, prob in enumerate(avg_video_pred)]
         image_results = f"The largest face in this image may be {max_prob_class} with probability: \n [{', '.join(probabilities)}]"
-
-        # Generate CAM heatmap for the detected class
-        use_cuda = torch.cuda.is_available()
-        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,
-            'Deepfake': 1,
-            'Diffusion or AIGC generated': 2,
-            'Spoofing or Presentation-attack': 3
-        }
-        target_category = TargetCategory(category_names_to_index[max_prob_class])
-
-        cam = GradCAM(model=model,
-                      target_layers=[model.blocks[-1].norm1],
-                      reshape_transform=reshape_transform
-                      )
-        grayscale_cam = cam(input_tensor=input_tensor, targets=[target_category], aug_smooth=False, eigen_smooth=True)
-        grayscale_cam = 1 - grayscale_cam[0, :]
-        img = np.array(img)
-        if img.shape[2] == 4:
-            img = img[:, :, :3]
-        img = img.astype(np.float32) / 255.0
-        visualization = show_cam_on_image(img, grayscale_cam)
-        visualization = cv2.cvtColor(visualization, cv2.COLOR_RGB2BGR)
-
-        # Add text overlay to the heatmap
-        # text = f"Detected: {max_prob_class}"
-        # cv2.putText(visualization, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-        cam_path = os.path.join(CAM_SAVE_PATH, str(len(os.listdir(CAM_SAVE_PATH))))
-        os.makedirs(cam_path, exist_ok=True)
-        os.makedirs(os.path.join(cam_path, '0'), exist_ok=True)
-        output_path = os.path.join(cam_path, "output_heatmap.png")
-        cv2.imwrite(output_path, visualization)
-        return image_results, output_path, probabilities[max_prob_index]
+        return image_results
 
     if CKPT_CLASS[ckpt] == 2:
         frame_preds_list, video_pred_list = test_two_class(data_loader_val, model, device)
@@ -272,7 +188,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
 
 
 def FSFM3C_video_detection(video, num_frames):
@@ -285,25 +201,19 @@ 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:
@@ -313,33 +223,28 @@ 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')
-CAM_SAVE_PATH = os.path.join(os.path.dirname(P), 'cam')
 CKPT_SAVE_PATH = os.path.join(os.path.dirname(P), 'checkpoints')
 os.makedirs(FRAME_SAVE_PATH, exist_ok=True)
-os.makedirs(CAM_SAVE_PATH, exist_ok=True)
 os.makedirs(CKPT_SAVE_PATH, exist_ok=True)
 CKPT_NAME = [
     '✨Unified-detector_v1_Fine-tuned_on_4_classes',
@@ -362,18 +267,17 @@ CKPT_MODEL = {
     'FAS-Checkpoint_Fine-tuned_on_MCIO': 'vit_base_patch16',
 }
 
+
 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 [by [Gaojian Wang/汪高健](https://scholar.google.com/citations?user=tpP4cFQAAAAJ&hl=zh-CN&oi=ao), [Tong Wu/吴桐](https://github.com/Coco-T-T), [Xingtang Luo/罗兴塘](https://github.com/Rox-C)]) </b> <br> <b>[V1.0] 2025/02/22-Current🎉</b>: "
-        "1) Updated <b>[✨Unified-detector_v1] for 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. <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(
@@ -387,29 +291,16 @@ 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")
         with gr.Column(scale=5):
             gr.Markdown("### Video Detection")
             video = gr.Video(label="Upload/Capture your video")
             frame_slider = gr.Slider(minimum=1, maximum=32, step=1, value=32, label="Number of Frames for Detection")
             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" title="Visit tracker">'
-        '<img src="https://mapmyvisitors.com/map.png?cl=ffffff&w=a&t=tt&d=FYhBoxLDEaFAxdfRzk5TuchYOBGrnSa98Ky59EkEEpY">'
-        '</a>'
-        '</div>'
-    )
 
     ckpt_select_dropdown.change(
         fn=load_model,
@@ -419,7 +310,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],
     )
     video_submit_btn.click(
         fn=FSFM3C_video_detection,
@@ -427,6 +318,7 @@ with gr.Blocks(css=".custom-label { font-weight: bold !important; font-size: 16p
         outputs=[output_results_video],
     )
 
+
 if __name__ == "__main__":
     args = get_args_parser()
     args = args.parse_args()