Update app.py

app.py

@@ -9,7 +9,8 @@ from facenet_pytorch import MTCNN, InceptionResnetV1
 import numpy as np
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-import os
+from pytorch_grad_cam.utils.image import show_cam_on_image
+import base64
 
 app = Flask(__name__)
 
@@ -27,12 +28,16 @@ DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
 mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).to(DEVICE).eval()
 
 model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=DEVICE)
-# Model Credits: https://huggingface.co/spaces/dhairyashah/deepfake-alpha-version/blob/main/CREDITS.md
 checkpoint = torch.load("resnetinceptionv1_epoch_32.pth", map_location=torch.device('cpu'))
 model.load_state_dict(checkpoint['model_state_dict'])
 model.to(DEVICE)
 model.eval()
 
+# GradCAM setup
+target_layers = [model.block8.branch1[-1]]
+cam = GradCAM(model=model, target_layers=target_layers)
+targets = [ClassifierOutputTarget(0)]
+
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
@@ -40,7 +45,7 @@ def allowed_file(filename):
 def process_frame(frame):
     face = mtcnn(frame)
     if face is None:
-        return None, None
+        return None, None, None
 
     face = face.unsqueeze(0)
     face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
@@ -53,14 +58,22 @@ def process_frame(frame):
         output = torch.sigmoid(model(face).squeeze(0))
         prediction = "fake" if output.item() >= 0.5 else "real"
 
-    return prediction, output.item()
+    # Generate GradCAM
+    grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
+    grayscale_cam = grayscale_cam[0, :]
+    
+    face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().numpy()
+    visualization = show_cam_on_image(face_image_to_plot, grayscale_cam, use_rgb=True)
+
+    return prediction, output.item(), visualization
 
 @spaces.GPU
-def analyze_video(video_path, sample_rate=30):
+def analyze_video(video_path, sample_rate=30, top_n=5):
     cap = cv2.VideoCapture(video_path)
     frame_count = 0
     fake_count = 0
     total_processed = 0
+    frames_info = []
 
     while cap.isOpened():
         ret, frame = cap.read()
@@ -69,22 +82,32 @@ def analyze_video(video_path, sample_rate=30):
 
         if frame_count % sample_rate == 0:
             rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            prediction, confidence = process_frame(rgb_frame)
+            prediction, confidence, visualization = process_frame(rgb_frame)
 
             if prediction is not None:
                 total_processed += 1
                 if prediction == "fake":
                     fake_count += 1
 
+                frames_info.append({
+                    'frame_number': frame_count,
+                    'prediction': prediction,
+                    'confidence': confidence,
+                    'visualization': visualization
+                })
+
         frame_count += 1
 
     cap.release()
 
     if total_processed > 0:
         fake_percentage = (fake_count / total_processed) * 100
-        return fake_percentage
+        frames_info.sort(key=lambda x: x['confidence'], reverse=True)
+        top_frames = frames_info[:top_n]
+
+        return fake_percentage, top_frames
     else:
-        return 0
+        return 0, []
 
 @app.route('/analyze', methods=['POST'])
 def analyze_video_api():
@@ -102,12 +125,17 @@ def analyze_video_api():
         file.save(filepath)
 
         try:
-            fake_percentage = analyze_video(filepath)
+            fake_percentage, top_frames = analyze_video(filepath)
             os.remove(filepath)  # Remove the file after analysis
 
+            # Convert numpy arrays to base64 encoded strings
+            for frame in top_frames:
+                frame['visualization'] = base64.b64encode(cv2.imencode('.png', frame['visualization'])[1]).decode('utf-8')
+
             result = {
                 'fake_percentage': round(fake_percentage, 2),
-                'is_likely_deepfake': fake_percentage >= 60
+                'is_likely_deepfake': fake_percentage >= 60,
+                'top_frames': top_frames
             }
 
             return jsonify(result), 200