[Updated Features] > Updated explainer and attribution with a new model, and support for better heatmap.

backend/app/ai/attribution.py

@@ -1,56 +1,83 @@
 import torch
 import cv2
 import numpy as np
-import timm
 from PIL import Image
-from torchvision import transforms
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+import logging
 
-model = timm.create_model("efficientnet_b0", pretrained=True)
-model.eval()
+logger = logging.getLogger(__name__)
 
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
-    )
-])
+model_id = "prithivMLmods/Deep-Fake-Detector-Model"
+processor = AutoImageProcessor.from_pretrained(model_id)
+model = AutoModelForImageClassification.from_pretrained(model_id)
+model.eval()
 
 def generate_attribution(image_path: str, save_path: str) -> dict:
-    img = Image.open(image_path).convert("RGB")
-    x = transform(img).unsqueeze(0)
-
-    x.requires_grad = True
-
-    features = model.forward_features(x)
-
-    score = features.mean()
-    model.zero_grad()
-    score.backward()
-
-    gradients = x.grad[0].numpy()
-    weights = np.mean(gradients, axis=(1, 2))
-
-    cam = np.zeros(gradients.shape[1:], dtype=np.float32)
-    for i, w in enumerate(weights):
-        cam += w * np.maximum(gradients[i], 0)
-    
-    cam = cv2.resize(cam, (img.width, img.height))
-    if cam.max() != 0:
-        cam = cam / cam.max()
-
-    heatmap = np.uint8(255 * cam)
-    heatmap_colored = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-    original = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-    overlay = cv2.addWeighted(original, 0.6, heatmap_colored, 0.4, 0)
-    cv2.imwrite(save_path, overlay)
-
-    y, x_coord = np.unravel_index(np.argmax(cam), cam.shape)
-
-    return {
-        "heatmap_path": save_path,
-        "top_regions": [
-            {"x": int(x_coord), "y": int(y), "importance": float(cam.max())}
-        ]
-    }
+    try:
+        img = Image.open(image_path).convert("RGB")
+        original_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+        inputs = processor(images=img, return_tensors="pt")
+        x = inputs["pixel_values"]
+
+        target_layer = model.vision_model.encoder.layers[-1].layer_norm1
+        
+        activations = []
+        gradients = []
+        
+        def forward_hook(module, args, output):
+            activations.append(output.detach())
+            
+        def backward_hook(module, grad_input, grad_output):
+            gradients.append(grad_output[0].detach())
+
+        handle_forward = target_layer.register_forward_hook(forward_hook)
+        handle_backward = target_layer.register_full_backward_hook(backward_hook)
+        
+        outputs = model(x)
+        logits = outputs.logits
+        target_class = logits.argmax(dim=-1).item()
+        
+        score = logits[0, target_class]
+        model.zero_grad()
+        score.backward()
+        
+        handle_forward.remove()
+        handle_backward.remove()
+        
+        act = activations[0][0] 
+        grad = gradients[0][0]  
+        weights = torch.mean(grad, dim=0) 
+        
+        cam = torch.zeros(act.shape[0])
+        for i in range(act.shape[1]):
+            cam += weights[i] * act[:, i]
+            
+        cam = cam.numpy()
+        cam = cam.reshape(14, 14)
+        
+        cam = np.maximum(cam, 0)
+        cam = cam - np.min(cam)
+        if np.max(cam) != 0:
+            cam = cam / np.max(cam)
+            
+        cam_resized = cv2.resize(cam, (img.width, img.height), interpolation=cv2.INTER_CUBIC)
+
+        heatmap_uint8 = np.uint8(255 * cam_resized)
+        heatmap_colored = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_INFERNO)
+        overlay = cv2.addWeighted(original_cv, 0.6, heatmap_colored, 0.4, 0)
+        
+        cv2.imwrite(save_path, overlay)
+
+        y, x_coord = np.unravel_index(np.argmax(cam_resized), cam_resized.shape)
+
+        return {
+            "heatmap_path": save_path,
+            "top_regions": [
+                {"x": int(x_coord), "y": int(y), "importance": float(np.max(cam_resized))}
+            ]
+        }
+
+    except Exception as e:
+        logger.error(f"Failed to generate attribution map: {str(e)}")
+        raise

backend/app/ai/explainer.py

@@ -2,38 +2,75 @@ import torch
 import numpy as np
 import cv2
 from PIL import Image
-import matplotlib.pyplot as plt
-import timm
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+import logging
 
-model = timm.create_model("efficientnet_b0", pretrained=True)
+logger = logging.getLogger(__name__)
+
+model_id = "prithivMLmods/Deep-Fake-Detector-Model"
+processor = AutoImageProcessor.from_pretrained(model_id)
+model = AutoModelForImageClassification.from_pretrained(model_id)
 model.eval()
 
 def generate_heatmap(image_path: str, save_path: str):
-    img = Image.open(image_path).convert("RGB")
-    img = img.resize((224, 224))
-    
-    img_np = np.array(img) / 255.0
-    input_tensor = torch.tensor(img_np.transpose(2, 0, 1)).unsqueeze(0).float()
-    
-    input_tensor.requires_grad = True
-    
-    features = model.forward_features(input_tensor)
-    
-    score = features.mean()
-    score.backward()
-    
-    gradient = input_tensor.grad[0].numpy()
-    heatmap = np.mean(np.abs(gradient), axis=0)
-    
-    heatmap = cv2.resize(heatmap, (img.width, img.height))
-    heatmap = heatmap / heatmap.max()
-    
-    heatmap = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)
-    
-    original = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-    
-    overlay = cv2.addWeighted(original, 0.6, heatmap, 0.4, 0)
-    
-    cv2.imwrite(save_path, overlay)
+    try:
+        img = Image.open(image_path).convert("RGB")
+        original_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+        inputs = processor(images=img, return_tensors="pt")
+        x = inputs["pixel_values"]
+
+        target_layer = model.vision_model.encoder.layers[-1].layer_norm1
+
+        activations = []
+        gradients = []
+
+        def forward_hook(module, args, output):
+            activations.append(output.detach())
+
+        def backward_hook(module, grad_input, grad_output):
+            gradients.append(grad_output[0].detach())
+        
+        handle_forward = target_layer.register_forward_hook(forward_hook)
+        handle_backward = target_layer.register_full_backward_hook(backward_hook)
+
+        outputs = model(x)
+        logits = outputs.logits
+        target_class = logits.argmax(dim=-1).item()
+        score = logits[0, target_class]
+
+        model.zero_grad()
+        score.backward()
+
+        handle_forward.remove()
+        handle_backward.remove()
+
+        act = activations[0][0]
+        grad = gradients[0][0]
+        weights = torch.mean(grad, dim=0)
+
+        cam = torch.zeros(act.shape[0])
+        for i in range(act.shape[1]):
+            cam += weights[i] * act[:, i]
+
+        cam = cam.numpy()
+        cam = cam.reshape(14, 14)
+
+        cam = np.maximum(cam, 0)
+        cam = cam - np.min(cam)
+        if np.max(cam) != 0:
+            cam = cam / np.max(cam)
+
+        heatmap_resized = cv2.resize(cam, (img.width, img.height), interpolation=cv2.INTER_CUBIC)
+
+        heatmap_uint8 = np.uint8(255 * heatmap_resized)
+        heatmap_colored = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_INFERNO)
+
+        overlay = cv2.addWeighted(original_cv, 0.6, heatmap_colored, 0.4, 0)
+        cv2.imwrite(save_path, overlay)
+
+        return save_path
     
-    return save_path
+    except Exception as e:
+        logger.error(f"Failed to generate heatmap: {str(e)}")
+        raise

backend/app/ai/video/frame_detector.py

@@ -1,6 +1,5 @@
 import torch
-import torch.nn as nn
-from torchvision import models, transforms
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
 import numpy as np
 import logging
@@ -12,23 +11,15 @@ class FrameDetector:
         self.device = torch.device("cpu")
 
         try:
-            self.model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1)
-
-            for param in self.model.parameters():
-                param.requires_grad = False
-
-            num_ftrs = self.model.classifier[1].in_features
-            self.model.classifier[1] = nn.Linear(num_ftrs, 1)
-
+            model_id = "prithivMLmods/Deep-Fake-Detector-Model"
+            self.processor = ViTImageProcessor.from_pretrained(model_id)
+            self.model = ViTForImageClassification.from_pretrained(model_id)
             self.model = self.model.to(self.device)
             self.model.eval()
 
-            self.transforms = transforms.Compose([
-                transforms.Resize((224, 224)),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-            ])
-            logger.info("EfficientNet-B0 CPU Inference Engine Initialized Successfully.")
+            logger.info(f"Loaded {model_id} CPU Inference Engine Successfully.")
+
+            self.fake_label_idx = 1 if "fake" in str(self.model.config.id2label.get(1, "").lower()) else 0
 
         except Exception as e:
             logger.error(f"FATAL: Could not boot AI model: {str(e)}")
@@ -45,13 +36,16 @@ class FrameDetector:
             rgb_frame = frame[:, :, ::-1]
             pil_image = Image.fromarray(rgb_frame)
 
-            input_tensor = self.transforms(pil_image).unsqueeze(0).to(self.device)
+            inputs = self.processor(images=pil_image, return_tensors="pt").to(self.device)
 
             with torch.no_grad():
-                raw_score = self.model(input_tensor)
-                probability = torch.sigmoid(raw_score).item()
+                outputs = self.model(**inputs)
+                logits = outputs.logits
+
+                probs = torch.nn.functional.softmax(logits, dim=-1)
+                ai_probability = probs[0][self.fake_label_idx].item()
             
-            return probability
+            return ai_probability
         
         except Exception as e:
             logger.error(f"AI Frame Prediction Crashed: {str(e)}")

backend/app/ai/video/noise_entropy_detector.py

@@ -0,0 +1,36 @@
+import cv2
+import numpy as np
+import logging
+from scipy.stats import entropy
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+def compute_noise_entropy_anomaly(frame: np.ndarray) -> Optional[float]:
+    try:
+        if frame is None or frame.size == 0:
+            return None
+        
+        gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        gray_small = cv2.resize(gray_frame, (512, 512), interpolation=cv2.INTER_AREA)
+
+        hist = cv2.calcHist([gray_small], [0], None, [256], [0, 256])
+        hist_prob = hist.ravel() / hist.sum()
+
+        img_entropy = entropy(hist_prob, base=2)
+
+        laplacian = cv2.Laplacian(gray_small, cv2.CV_64F)
+        noise_variance = laplacian.var()
+
+        score = 0.0
+        if img_entropy < 5.5:
+            score += 0.4
+        
+        if noise_variance < 50 or noise_variance > 3000:
+            score += 0.6
+        
+        return float(min(score, 1.0))
+
+    except Exception as e:
+        logger.error(f"Noise/Entropy Analysis Crash: {str(e)}")
+        return 0.0