Deploy Backend (No Frontend)

.dockerignore

@@ -0,0 +1,10 @@
+.git
+.gitignore
+__pycache__
+*.pyc
+.DS_Store
+model/test_images
+venv
+env
+node_modules
+frontend/tests

Dockerfile

@@ -0,0 +1,31 @@
+FROM python:3.9
+
+WORKDIR /app
+
+# Install system dependencies for OpenCV and GLib
+RUN apt-get update && apt-get install -y \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first to leverage Docker cache
+COPY backend/requirements_web.txt .
+RUN pip install --no-cache-dir -r requirements_web.txt
+
+# Copy the rest of the application
+# Copy the rest of the application
+COPY backend/ backend/
+COPY model/ model/
+
+# Set working directory to backend
+WORKDIR /app/backend
+
+# Create necessary directories
+RUN mkdir -p uploads
+RUN mkdir -p history_uploads
+
+# Expose Hugging Face default port
+EXPOSE 7860
+
+# Run the application
+CMD ["python", "app.py"]

backend/app.py

@@ -0,0 +1,368 @@
+from flask import Flask, request, jsonify, send_from_directory
+from flask_cors import CORS
+import sys
+import os
+
+# Add model directory to path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'model')))
+import datetime
+import torch
+import cv2
+import os
+import numpy as np
+import ssl
+import base64
+from werkzeug.utils import secure_filename
+import io
+from PIL import Image
+from src import video_inference
+
+# Disable SSL verification
+ssl._create_default_https_context = ssl._create_unverified_context
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from albumentations.pytorch import ToTensorV2
+from src.models import DeepfakeDetector
+from src.config import Config
+import database
+
+try:
+    from safetensors.torch import load_file
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+
+app = Flask(__name__, static_folder='../frontend', static_url_path='')
+CORS(app)
+
+# Configuration
+UPLOAD_FOLDER = os.path.join(os.path.dirname(__file__), 'uploads')
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp', 'mp4', 'avi', 'mov', 'webm'}
+HISTORY_FOLDER = os.path.join(os.path.dirname(__file__), 'history_uploads')
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(HISTORY_FOLDER, exist_ok=True)
+
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
+
+# Global model and transform
+device = torch.device(Config.DEVICE)
+model = None
+transform = None
+
+def get_transform():
+    return A.Compose([
+        A.Resize(Config.IMAGE_SIZE, Config.IMAGE_SIZE),
+        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ToTensorV2(),
+    ])
+
+def load_model():
+    """Load the trained deepfake detection model"""
+    global model, transform
+    
+    checkpoint_dir = Config.CHECKPOINT_DIR
+    # Explicitly target the model requested by the user
+    target_model_name = "best_model.safetensors"
+    checkpoint_path = os.path.join(checkpoint_dir, target_model_name)
+    
+    print(f"Using device: {device}")
+    
+    # Initialize with pretrained=True to ensure missing keys (frozen layers) have valid ImageNet weights
+    # instead of random noise. This fixes the "random prediction" issue when the checkpoint 
+    # only contains finetuned layers.
+    model = DeepfakeDetector(pretrained=True)
+    model.to(device)
+    model.eval()
+    
+    # Check if file exists first
+    if not os.path.exists(checkpoint_path):
+        print(f"❌ CRITICAL ERROR: Model file not found at: {checkpoint_path}")
+        print(f"Please ensure '{target_model_name}' exists in '{checkpoint_dir}'")
+        model = None
+        transform = get_transform()
+        return model, transform
+
+    try:
+        print(f"Loading checkpoint: {checkpoint_path}")
+        if checkpoint_path.endswith(".safetensors") and SAFETENSORS_AVAILABLE:
+            state_dict = load_file(checkpoint_path)
+        else:
+            state_dict = torch.load(checkpoint_path, map_location=device)
+            
+        # Use strict=False because the checkpoint might be a partial save (e.g. only finetuned layers)
+        # or there might be minor architecture mismatches.
+        # Since we use pretrained=True, the missing keys will remain as ImageNet weights (valid features).
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+        
+        print(f"✅ Model loaded successfully!")
+        if missing_keys:
+            print(f"ℹ️  {len(missing_keys)} keys missing from checkpoint (using pretrained defaults).")
+        if unexpected_keys:
+            print(f"ℹ️  {len(unexpected_keys)} unexpected keys in checkpoint.")
+        
+    except Exception as e:
+        print(f"❌ Error loading checkpoint: {e}")
+        print("Predictions will fail until this is resolved.")
+        model = None
+    
+    transform = get_transform()
+    return model, transform
+
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+def predict_image(image_path):
+    """Make prediction on a single image"""
+    if model is None:
+        return None, "Error: Model not loaded. Check backend logs for 'best_finetuned_datasetB.safetensors' error."
+
+    try:
+        # Read and preprocess image
+        image = cv2.imread(image_path)
+        if image is None:
+            return None, "Error: Could not read image"
+        
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        augmented = transform(image=image)
+        image_tensor = augmented['image'].unsqueeze(0).to(device)
+        
+        # Make prediction
+        logits = model(image_tensor)
+        prob = torch.sigmoid(logits).item()
+        
+        # Generate Heatmap
+        heatmap = model.get_heatmap(image_tensor)
+        
+        # Process Heatmap for Visualization
+        # Resize to original image size
+        heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))
+        heatmap = np.uint8(255 * heatmap)
+        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+        
+        # Superimpose
+        # Heatmap is BGR (from cv2), Image is RGB. Convert Image to BGR.
+        image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        superimposed_img = heatmap * 0.4 + image_bgr * 0.6
+        superimposed_img = np.clip(superimposed_img, 0, 255).astype(np.uint8)
+        
+        # Encode to Base64
+        _, buffer = cv2.imencode('.jpg', superimposed_img)
+        heatmap_b64 = base64.b64encode(buffer).decode('utf-8')
+        
+        is_fake = prob > 0.5
+        label = "FAKE" if is_fake else "REAL"
+        confidence = prob if is_fake else 1 - prob
+        
+        return {
+            'prediction': label,
+            'confidence': float(confidence),
+            'fake_probability': float(prob),
+            'real_probability': float(1 - prob),
+            'heatmap': heatmap_b64
+        }, None
+    except Exception as e:
+        return None, str(e)
+
+
+@app.route('/')
+def index():
+    """Backend Root"""
+    return jsonify({
+        "status": "online",
+        "message": "DeepGuard Backend is Running",
+        "endpoints": ["/api/predict", "/api/history", "/api/health"]
+    })
+
+@app.route('/history_uploads/<path:filename>')
+def serve_history_image(filename):
+    """Serve history images"""
+    return send_from_directory(HISTORY_FOLDER, filename)
+
+@app.route('/api/health', methods=['GET'])
+def health_check():
+    """Health check endpoint"""
+    return jsonify({
+        'status': 'healthy',
+        'model_loaded': model is not None,
+        'device': str(device)
+    })
+
+@app.route('/api/predict', methods=['POST'])
+def predict():
+    """Handle image upload and prediction"""
+    try:
+        # Check if file is present
+        if 'file' not in request.files:
+            return jsonify({'error': 'No file provided'}), 400
+        
+        file = request.files['file']
+        
+        if file.filename == '':
+            return jsonify({'error': 'No file selected'}), 400
+        
+        if not allowed_file(file.filename):
+            return jsonify({'error': 'Invalid file type. Allowed types: png, jpg, jpeg, webp'}), 400
+        
+        # Save file
+        filename = secure_filename(file.filename)
+        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+        file.save(filepath)
+        
+        # Make prediction
+        result, error = predict_image(filepath)
+        
+        # Save to History
+        import shutil
+        history_filename = f"scan_{int(datetime.datetime.now().timestamp())}_{filename}"
+        history_path = os.path.join(HISTORY_FOLDER, history_filename)
+        
+        # Copy original file to history folder
+        # We need to read the file again or just copy if we haven't deleted it?
+        # We read via cv2, the file is still at filepath.
+        shutil.copy(filepath, history_path)
+        
+        # Relative path for frontend
+        relative_path = f"history_uploads/{history_filename}"
+
+        database.add_scan(
+            filename=filename,
+            prediction=result['prediction'],
+            confidence=result['confidence'],
+            fake_prob=result['fake_probability'],
+            real_prob=result['real_probability'],
+            image_path=relative_path
+        )
+        
+        # Clean up uploaded file
+        try:
+            os.remove(filepath)
+        except:
+            pass
+        
+        return jsonify(result)
+    
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/api/predict_video', methods=['POST'])
+def predict_video():
+    """Handle video upload and prediction"""
+    try:
+        if 'file' not in request.files:
+            return jsonify({'error': 'No file provided'}), 400
+        
+        file = request.files['file']
+        
+        if file.filename == '':
+            return jsonify({'error': 'No file selected'}), 400
+            
+        if not allowed_file(file.filename):
+             return jsonify({'error': 'Invalid file type'}), 400
+             
+        # Save file
+        filename = secure_filename(file.filename)
+        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+        file.save(filepath)
+        
+        # Process Video
+        # Note: process_video needs sys.path to be correct to import models inside it if it was standalone,
+        # but here we pass the already loaded 'model' object.
+        if model is None:
+             return jsonify({'error': 'Model not loaded'}), 500
+             
+        result = video_inference.process_video(filepath, model, transform, device)
+        
+        if "error" in result:
+             return jsonify(result), 500
+             
+        # Save to History (Using the first frame or a placeholder icon for now?)
+        # For video, we might want to save the video file itself to history_uploads
+        # or just a thumbnail. Let's save the video for now.
+        import shutil
+        history_filename = f"scan_{int(datetime.datetime.now().timestamp())}_{filename}"
+        history_path = os.path.join(HISTORY_FOLDER, history_filename)
+        shutil.copy(filepath, history_path)
+        
+        relative_path = f"history_uploads/{history_filename}"
+        
+        # Add to database
+        # Note: The database 'add_scan' might expect image-specific fields.
+        # We'll re-use 'fake_prob' as 'avg_fake_prob'
+        database.add_scan(
+            filename=filename,
+            prediction=result['prediction'],
+            confidence=result['confidence'],
+            fake_prob=result['avg_fake_prob'],
+            real_prob=1 - result['avg_fake_prob'],
+            image_path=relative_path 
+        )
+        
+        # Clean up
+        try:
+            os.remove(filepath)
+        except:
+            pass
+            
+        # Add video URL for frontend playback
+        result['video_url'] = relative_path
+        
+        return jsonify(result)
+
+    except Exception as e:
+        print(f"Video Error: {e}")
+        return jsonify({'error': str(e)}), 500
+
+
+@app.route('/api/history', methods=['GET'])
+def get_history():
+    """Get all past scans"""
+    history = database.get_history()
+    history = database.get_history()
+    return jsonify(history)
+
+@app.route('/api/history/<int:scan_id>', methods=['DELETE'])
+def delete_scan(scan_id):
+    """Delete a specific scan"""
+    if database.delete_scan(scan_id):
+        return jsonify({'message': 'Scan deleted'})
+    return jsonify({'error': 'Failed to delete scan'}), 500
+
+@app.route('/api/history', methods=['DELETE'])
+def clear_history():
+    """Clear all history"""
+    if database.clear_history():
+        return jsonify({'message': 'History cleared'})
+    return jsonify({'error': 'Failed to clear history'}), 500
+
+@app.route('/api/model-info', methods=['GET'])
+def model_info():
+    """Return model information"""
+    return jsonify({
+        'model_name': 'DeepGuard: Advanced Deepfake Detector',
+        'architecture': 'Hybrid CNN-ViT',
+        'components': {
+            'RGB Analysis': Config.USE_RGB,
+            'Frequency Domain': Config.USE_FREQ,
+            'Patch-based Detection': Config.USE_PATCH,
+            'Vision Transformer': Config.USE_VIT
+        },
+        'image_size': Config.IMAGE_SIZE,
+        'device': str(device),
+        'threshold': 0.5
+    })
+
+if __name__ == '__main__':
+    print("=" * 60)
+    print("🚀 DeepGuard - Deepfake Detection System")
+    print("=" * 60)
+    
+    # Load model
+    load_model()
+    
+    print("=" * 60)
+    port = int(os.environ.get("PORT", 7860))
+    print(f"🌐 Starting server on http://0.0.0.0:{port}")
+    print("=" * 60)
+    
+    app.run(debug=False, host='0.0.0.0', port=port)

backend/database.py

@@ -0,0 +1,106 @@
+import sqlite3
+import datetime
+import os
+
+DB_NAME = os.path.join(os.path.dirname(__file__), 'database.db')
+
+def get_db_connection():
+    try:
+        conn = sqlite3.connect(DB_NAME)
+        conn.row_factory = sqlite3.Row
+        return conn
+    except sqlite3.Error as e:
+        print(f"Database error: {e}")
+        return None
+
+def init_db():
+    conn = get_db_connection()
+    if conn:
+        try:
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS history (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    filename TEXT NOT NULL,
+                    prediction TEXT NOT NULL,
+                    confidence REAL NOT NULL,
+                    fake_probability REAL NOT NULL,
+                    real_probability REAL NOT NULL,
+                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
+                )
+            ''')
+            conn.commit()
+            print("✅ Database initialized successfully.")
+        except sqlite3.Error as e:
+            print(f"Error initializing database: {e}")
+        
+        # Migration: Add image_path if not exists
+        try:
+            conn.execute('ALTER TABLE history ADD COLUMN image_path TEXT')
+            print("✅ Added image_path column.")
+        except sqlite3.Error:
+            pass # Column likely exists
+            
+        finally:
+            conn.close()
+
+def add_scan(filename, prediction, confidence, fake_prob, real_prob, image_path=""):
+    conn = get_db_connection()
+    if conn:
+        try:
+            conn.execute('''
+                INSERT INTO history (filename, prediction, confidence, fake_probability, real_probability, image_path)
+                VALUES (?, ?, ?, ?, ?, ?)
+            ''', (filename, prediction, confidence, fake_prob, real_prob, image_path))
+            conn.commit()
+            return True
+        except sqlite3.Error as e:
+            print(f"Error adding scan: {e}")
+            return False
+        finally:
+            conn.close()
+    return False
+
+def get_history():
+    conn = get_db_connection()
+    if conn:
+        try:
+            cursor = conn.execute('SELECT * FROM history ORDER BY timestamp DESC')
+            history = [dict(row) for row in cursor.fetchall()]
+            return history
+        except sqlite3.Error as e:
+            print(f"Error retrieving history: {e}")
+            return []
+        finally:
+            conn.close()
+    return []
+
+def clear_history():
+    conn = get_db_connection()
+    if conn:
+        try:
+            conn.execute('DELETE FROM history')
+            conn.commit()
+            return True
+        except sqlite3.Error as e:
+            print(f"Error clearing history: {e}")
+            return False
+        finally:
+            conn.close()
+    return False
+
+def delete_scan(scan_id):
+    conn = get_db_connection()
+    if conn:
+        try:
+            conn.execute('DELETE FROM history WHERE id = ?', (scan_id,))
+            conn.commit()
+            return True
+        except sqlite3.Error as e:
+            print(f"Error deleting scan: {e}")
+            return False
+        finally:
+            conn.close()
+    return False
+
+# Initialize DB on module load
+init_db()

backend/requirements_web.txt

@@ -0,0 +1,9 @@
+flask==3.0.0
+flask-cors==4.0.0
+torch
+torchvision
+opencv-python
+albumentations
+Pillow
+numpy
+safetensors

model/results/checkpoints/best_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a58fd840e8ebab964a3021acb9e365fe445a86dbcb8af93d808f68eb3a254ad4
+size 202457588

model/src/config.py

@@ -0,0 +1,50 @@
+import os
+import torch
+import platform
+
+class Config:
+    # System
+    PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    DATA_DIR = os.path.join(PROJECT_ROOT, "data")
+    RESULTS_DIR = os.path.join(PROJECT_ROOT, "results")
+    
+    # Model Architecture
+    IMAGE_SIZE = 256
+    NUM_CLASSES = 1  # Logic: 0=Real, 1=Fake (Sigmoid output)
+    
+    # Component Flags
+    USE_RGB = True
+    USE_FREQ = True
+    USE_PATCH = True
+    USE_VIT = True
+    
+    # Training Hyperparameters
+    BATCH_SIZE = 32  # Optimized for Mac M4 (Unified Memory)
+    EPOCHS = 3
+    LEARNING_RATE = 1e-4
+    WEIGHT_DECAY = 1e-5
+    NUM_WORKERS = 8  # Leverage M4 Performance Cores
+    
+    # Hardware
+    DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    
+    # Paths
+    # Docker Deployment: Use relative paths
+    DATA_DIR = os.path.join(PROJECT_ROOT, "data")
+    
+    # Since we are using the root folder, the script will recursively find ALL images
+    # in all sub-datasets and split them 80/20 for training/validation.
+    TRAIN_DATA_PATH = DATA_DIR 
+    TEST_DATA_PATH = DATA_DIR 
+    CHECKPOINT_DIR = os.path.join(RESULTS_DIR, "checkpoints")
+
+    @classmethod
+    def setup(cls):
+        os.makedirs(cls.RESULTS_DIR, exist_ok=True)
+        os.makedirs(cls.CHECKPOINT_DIR, exist_ok=True)
+        os.makedirs(cls.DATA_DIR, exist_ok=True)
+        print(f"Project initialized at {cls.PROJECT_ROOT}")
+        print(f"Using device: {cls.DEVICE}")
+
+if __name__ == "__main__":
+    Config.setup()

model/src/dataset.py

@@ -0,0 +1,109 @@
+import os
+import cv2
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from src.config import Config
+
+class DeepfakeDataset(Dataset):
+    def __init__(self, root_dir=None, file_paths=None, labels=None, phase='train', max_samples=None):
+        """
+        Args:
+            root_dir (str): Directory with subfolders containing images. (Optional if file_paths provided)
+            file_paths (list): List of absolute paths to images.
+            labels (list): List of labels corresponding to file_paths.
+            phase (str): 'train' or 'val'.
+            max_samples (int): Optional limit for quick debugging.
+        """
+        self.phase = phase
+        
+        if file_paths is not None and labels is not None:
+            self.image_paths = file_paths
+            self.labels = labels
+        elif root_dir is not None:
+            self.image_paths, self.labels = self.scan_directory(root_dir)
+        else:
+            raise ValueError("Either root_dir or (file_paths, labels) must be provided.")
+            
+        if max_samples:
+            self.image_paths = self.image_paths[:max_samples]
+            self.labels = self.labels[:max_samples]
+            
+        self.transform = self._get_transforms()
+        
+        print(f"Initialized {self.phase} dataset with {len(self.image_paths)} samples.")
+
+    @staticmethod
+    def scan_directory(root_dir):
+        image_paths = []
+        labels = []
+        print(f"Scanning dataset at {root_dir}...")
+        
+        # Valid extensions
+        exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tif')
+        
+        for root, dirs, files in os.walk(root_dir):
+            for file in files:
+                if file.lower().endswith(exts):
+                    path = os.path.join(root, file)
+                    # Label inference based on full path
+                    path_lower = path.lower()
+                    
+                    label = None
+                    # Prioritize explicit folder names
+                    if "real" in path_lower:
+                        label = 0.0
+                    elif any(x in path_lower for x in ["fake", "df", "synthesis", "generated", "ai"]):
+                        label = 1.0
+                    
+                    if label is not None:
+                        image_paths.append(path)
+                        labels.append(label)
+        
+        return image_paths, labels
+
+    def _get_transforms(self):
+        size = Config.IMAGE_SIZE
+        if self.phase == 'train':
+            return A.Compose([
+                A.Resize(size, size),
+                A.HorizontalFlip(p=0.5),
+                A.RandomBrightnessContrast(p=0.2),
+                A.GaussNoise(p=0.2),
+                # A.GaussianBlur(p=0.1),
+                # Fixed for newer albumentations versions
+                A.ImageCompression(quality_lower=60, quality_upper=100, p=0.3),
+                A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+                ToTensorV2(),
+            ])
+        else:
+            return A.Compose([
+                A.Resize(size, size),
+                A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+                ToTensorV2(),
+            ])
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        path = self.image_paths[idx]
+        label = self.labels[idx]
+        
+        try:
+            image = cv2.imread(path)
+            if image is None:
+                raise ValueError("Image not found or corrupt")
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        except Exception as e:
+            # print(f"Error loading {path}: {e}")
+            # Fallback to next image
+            return self.__getitem__((idx + 1) % len(self))
+        
+        if self.transform:
+            augmented = self.transform(image=image)
+            image = augmented['image']
+            
+        return image, torch.tensor(label, dtype=torch.float32)

model/src/finetune.py

@@ -0,0 +1,182 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import random
+import ssl
+# Disable SSL verification for downloading pretrained weights
+ssl._create_default_https_context = ssl._create_unverified_context
+
+from src.config import Config
+from src.models import DeepfakeDetector
+from src.dataset import DeepfakeDataset
+
+try:
+    from safetensors.torch import save_file, load_model
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+    print("Warning: safetensors not installed. Checkpoints will be saved as .pt")
+
+def finetune():
+    # Setup
+    Config.setup()
+    device = torch.device(Config.DEVICE)
+    
+    # Fine-tuning dataset path
+    FINETUNE_DATA_PATH = "/Users/harshvardhan/Developer/dataset/Dataset c"
+    
+    print(f"\n{'='*80}")
+    print("FINE-TUNING ON DATASET C")
+    print(f"{'='*80}\n")
+    
+    # --- Data Loading ---
+    print(f"Loading data from: {FINETUNE_DATA_PATH}")
+    all_paths, all_labels = DeepfakeDataset.scan_directory(FINETUNE_DATA_PATH)
+    
+    if len(all_paths) == 0:
+        print(f"No images found in {FINETUNE_DATA_PATH}")
+        return
+
+    # Shuffle and split
+    combined = list(zip(all_paths, all_labels))
+    random.shuffle(combined)
+    
+    split_idx = int(len(combined) * 0.8)
+    train_data = combined[:split_idx]
+    val_data = combined[split_idx:]
+    
+    train_paths, train_labels = zip(*train_data)
+    val_paths, val_labels = zip(*val_data)
+    
+    train_dataset = DeepfakeDataset(file_paths=list(train_paths), labels=list(train_labels), phase='train')
+    val_dataset = DeepfakeDataset(file_paths=list(val_paths), labels=list(val_labels), phase='val')
+    
+    # Dataloaders
+    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True,
+                              num_workers=Config.NUM_WORKERS, 
+                              pin_memory=True if device.type=='cuda' else False,
+                              persistent_workers=True if Config.NUM_WORKERS > 0 else False)
+    val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE, shuffle=False,
+                            num_workers=Config.NUM_WORKERS, 
+                            pin_memory=True if device.type=='cuda' else False,
+                            persistent_workers=True if Config.NUM_WORKERS > 0 else False)
+    
+    # Load pre-trained model from Dataset A
+    print("\n🔄 Loading pre-trained model from Dataset A...")
+    model = DeepfakeDetector(pretrained=False).to(device)
+    
+    checkpoint_path = "results/checkpoints/best_model.safetensors"
+    if os.path.exists(checkpoint_path):
+        load_model(model, checkpoint_path, strict=False)
+        print(f"✅ Loaded checkpoint: {checkpoint_path}")
+    else:
+        print("⚠️ No checkpoint found! Starting from random weights.")
+    
+    model.to(device)
+    
+    # Optimization with LOWER learning rate for fine-tuning
+    FINETUNE_LR = 1e-5  # 10x lower than original training
+    FINETUNE_EPOCHS = 2
+    
+    print(f"\n📝 Fine-tuning settings:")
+    print(f"   Learning Rate: {FINETUNE_LR} (10x lower for fine-tuning)")
+    print(f"   Epochs: {FINETUNE_EPOCHS}")
+    print(f"   Batch Size: {Config.BATCH_SIZE}")
+    
+    criterion = nn.BCEWithLogitsLoss()
+    optimizer = optim.AdamW(model.parameters(), lr=FINETUNE_LR, weight_decay=Config.WEIGHT_DECAY)
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
+    
+    # Loop
+    best_acc = 0.0
+    
+    for epoch in range(FINETUNE_EPOCHS):
+        model.train()
+        train_loss = 0.0
+        train_correct = 0
+        train_total = 0
+        
+        loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{FINETUNE_EPOCHS}")
+        for images, labels in loop:
+            images = images.to(device)
+            labels = labels.to(device).unsqueeze(1)
+            
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            
+            train_loss += loss.item()
+            preds = (torch.sigmoid(outputs) > 0.5).float()
+            correct = (preds == labels).sum().item()
+            train_correct += correct
+            train_total += labels.size(0)
+            
+            loop.set_postfix(loss=loss.item(), acc=correct/labels.size(0))
+            
+        train_acc = train_correct / train_total if train_total > 0 else 0
+        print(f"Epoch {epoch+1} Train Loss: {train_loss/len(train_loader):.4f} Acc: {train_acc:.4f}")
+        
+        # Save checkpoint after every epoch
+        save_checkpoint(model, epoch+1, train_acc, name=f"finetuned_datasetC_ep{epoch+1}")
+        
+        # Validation
+        if len(val_dataset) > 0:
+            val_loss, val_acc = validate(model, val_loader, criterion, device)
+            print(f"Epoch {epoch+1} Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")
+            
+            # Save best model if validation accuracy improved
+            if val_acc > best_acc:
+                best_acc = val_acc
+                print(f"⭐ New best model! Validation Accuracy: {val_acc:.4f}")
+                save_checkpoint(model, epoch+1, val_acc, name="best_finetuned_datasetC")
+        
+        scheduler.step()
+    
+    print(f"\n🎉 Fine-tuning Complete!")
+    print(f"Best Validation Accuracy: {best_acc:.4f}")
+    print(f"\n💾 Checkpoints saved in: results/checkpoints/")
+
+def validate(model, loader, criterion, device):
+    model.eval()
+    val_loss = 0.0
+    correct = 0
+    total = 0
+    
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            labels = labels.to(device).unsqueeze(1)
+            
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            
+            val_loss += loss.item()
+            preds = (torch.sigmoid(outputs) > 0.5).float()
+            correct += (preds == labels).sum().item()
+            total += labels.size(0)
+            
+    return val_loss / len(loader), correct / total
+
+def save_checkpoint(model, epoch, acc, name="checkpoint"):
+    state_dict = model.state_dict()
+    filename = f"{name}.safetensors"
+    path = os.path.join(Config.CHECKPOINT_DIR, filename)
+    
+    if SAFETENSORS_AVAILABLE:
+        try:
+            from safetensors.torch import save_model
+            save_model(model, path)
+            print(f"✅ Saved: {filename}")
+        except Exception as e:
+            print(f"SafeTensors save failed, falling back to .pth: {e}")
+            torch.save(state_dict, path.replace(".safetensors", ".pth"))
+    else:
+        torch.save(state_dict, path.replace(".safetensors", ".pth"))
+
+if __name__ == "__main__":
+    finetune()

model/src/finetune_dataset_a.py

@@ -0,0 +1,204 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import random
+import ssl
+import platform
+
+# Disable SSL verification for downloading pretrained weights
+ssl._create_default_https_context = ssl._create_unverified_context
+
+from src.config import Config
+from src.models import DeepfakeDetector
+from src.dataset import DeepfakeDataset
+
+try:
+    from safetensors.torch import save_file, load_model
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+    print("Warning: safetensors not installed. Checkpoints will be saved as .pt")
+
+def finetune():
+    # Setup
+    Config.setup()
+    device = torch.device(Config.DEVICE)
+    
+    # Fine-tuning dataset path - Dataset A
+    if platform.system() == "Windows":
+        FINETUNE_DATA_PATH = r"C:\Users\kanna\Downloads\Dataset\Dataset A\Dataset A"
+    else:
+        FINETUNE_DATA_PATH = "/Users/harshvardhan/Developer/dataset/Dataset A"
+    
+    print(f"\n{'='*80}")
+    print("FINE-TUNING ON DATASET A")
+    print(f"{'='*80}\n")
+    
+    # --- Data Loading ---
+    print(f"Loading data from: {FINETUNE_DATA_PATH}")
+    if not os.path.exists(FINETUNE_DATA_PATH):
+        print(f"❌ Error: Dataset path not found: {FINETUNE_DATA_PATH}")
+        return
+
+    all_paths, all_labels = DeepfakeDataset.scan_directory(FINETUNE_DATA_PATH)
+    
+    if len(all_paths) == 0:
+        print(f"No images found in {FINETUNE_DATA_PATH}")
+        return
+
+    # Shuffle and split
+    combined = list(zip(all_paths, all_labels))
+    random.shuffle(combined)
+    
+    # Use 80/20 split for fine-tuning dataset
+    split_idx = int(len(combined) * 0.8)
+    train_data = combined[:split_idx]
+    val_data = combined[split_idx:]
+    
+    train_paths, train_labels = zip(*train_data)
+    val_paths, val_labels = zip(*val_data)
+    
+    train_dataset = DeepfakeDataset(file_paths=list(train_paths), labels=list(train_labels), phase='train')
+    val_dataset = DeepfakeDataset(file_paths=list(val_paths), labels=list(val_labels), phase='val')
+    
+    # Dataloaders - Use Config.BATCH_SIZE but ensure it fits GPU
+    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True,
+                              num_workers=Config.NUM_WORKERS, 
+                              pin_memory=True if device.type=='cuda' else False,
+                              persistent_workers=True if Config.NUM_WORKERS > 0 else False)
+    val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE, shuffle=False,
+                            num_workers=Config.NUM_WORKERS, 
+                            pin_memory=True if device.type=='cuda' else False,
+                            persistent_workers=True if Config.NUM_WORKERS > 0 else False)
+    
+    # Load pre-trained model
+    print("\n🔄 Loading pre-trained model (best_model)...")
+    model = DeepfakeDetector(pretrained=False).to(device)
+    
+    # Try to load the best model found so far
+    checkpoint_path = os.path.join(Config.CHECKPOINT_DIR, "best_model.safetensors")
+    if not os.path.exists(checkpoint_path):
+        # Fallback to .pth if safetensors logic above failed or not used previously
+        checkpoint_path = os.path.join(Config.CHECKPOINT_DIR, "best_model.pth")
+    
+    if os.path.exists(checkpoint_path):
+        try:
+            if checkpoint_path.endswith(".safetensors"):
+                load_model(model, checkpoint_path, strict=False)
+            else:
+                model.load_state_dict(torch.load(checkpoint_path, map_location=device))
+            print(f"✅ Loaded checkpoint: {checkpoint_path}")
+        except Exception as e:
+            print(f"⚠️ Error loading checkpoint: {e}")
+            print("Starting from random weights (not ideal for fine-tuning!)")
+    else:
+        print("⚠️ No checkpoint found! Starting from random weights.")
+    
+    model.to(device)
+    
+    # Optimization with LOWER learning rate for fine-tuning
+    FINETUNE_LR = 1e-5  # 10x lower than original training
+    FINETUNE_EPOCHS = 5 # Give it a few epochs to adapt
+    
+    print(f"\n📝 Fine-tuning settings:")
+    print(f"   Learning Rate: {FINETUNE_LR} (Low LR for fine-tuning)")
+    print(f"   Epochs: {FINETUNE_EPOCHS}")
+    print(f"   Batch Size: {Config.BATCH_SIZE}")
+    
+    criterion = nn.BCEWithLogitsLoss()
+    optimizer = optim.AdamW(model.parameters(), lr=FINETUNE_LR, weight_decay=Config.WEIGHT_DECAY)
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=2, verbose=True)
+    
+    # Loop
+    best_acc = 0.0
+    
+    for epoch in range(FINETUNE_EPOCHS):
+        model.train()
+        train_loss = 0.0
+        train_correct = 0
+        train_total = 0
+        
+        loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{FINETUNE_EPOCHS}")
+        for images, labels in loop:
+            images = images.to(device)
+            labels = labels.to(device).unsqueeze(1)
+            
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            
+            train_loss += loss.item()
+            preds = (torch.sigmoid(outputs) > 0.5).float()
+            correct = (preds == labels).sum().item()
+            train_correct += correct
+            train_total += labels.size(0)
+            
+            loop.set_postfix(loss=loss.item(), acc=correct/labels.size(0) if labels.size(0) > 0 else 0)
+            
+        train_acc = train_correct / train_total if train_total > 0 else 0
+        print(f"Epoch {epoch+1} Train Loss: {train_loss/len(train_loader):.4f} Acc: {train_acc:.4f}")
+        
+        # Save checkpoint after every epoch
+        save_checkpoint(model, epoch+1, train_acc, name=f"finetuned_datasetA_ep{epoch+1}")
+        
+        # Validation
+        if len(val_dataset) > 0:
+            val_loss, val_acc = validate(model, val_loader, criterion, device)
+            print(f"Epoch {epoch+1} Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")
+            
+            scheduler.step(val_acc)
+
+            # Save best model if validation accuracy improved
+            if val_acc > best_acc:
+                best_acc = val_acc
+                print(f"⭐ New best model! Validation Accuracy: {val_acc:.4f}")
+                save_checkpoint(model, epoch+1, val_acc, name="best_finetuned_datasetA")
+        
+    print(f"\n🎉 Fine-tuning Complete!")
+    print(f"Best Validation Accuracy: {best_acc:.4f}")
+    print(f"\n💾 Checkpoints saved in: {Config.CHECKPOINT_DIR}")
+
+def validate(model, loader, criterion, device):
+    model.eval()
+    val_loss = 0.0
+    correct = 0
+    total = 0
+    
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            labels = labels.to(device).unsqueeze(1)
+            
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            
+            val_loss += loss.item()
+            preds = (torch.sigmoid(outputs) > 0.5).float()
+            correct += (preds == labels).sum().item()
+            total += labels.size(0)
+            
+    return val_loss / len(loader), correct / total
+
+def save_checkpoint(model, epoch, acc, name="checkpoint"):
+    state_dict = model.state_dict()
+    filename = f"{name}.safetensors"
+    path = os.path.join(Config.CHECKPOINT_DIR, filename)
+    
+    if SAFETENSORS_AVAILABLE:
+        try:
+            from safetensors.torch import save_model
+            save_model(model, path)
+            print(f"✅ Saved: {filename}")
+        except Exception as e:
+            print(f"SafeTensors save failed, falling back to .pth: {e}")
+            torch.save(state_dict, path.replace(".safetensors", ".pth"))
+    else:
+        torch.save(state_dict, path.replace(".safetensors", ".pth"))
+
+if __name__ == "__main__":
+    finetune()

model/src/inference.py

@@ -0,0 +1,139 @@
+import argparse
+import torch
+import cv2
+import os
+import glob
+import numpy as np
+import ssl
+# Disable SSL verification
+ssl._create_default_https_context = ssl._create_unverified_context
+
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from src.models import DeepfakeDetector
+from src.config import Config
+
+try:
+    from safetensors.torch import load_file
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+
+def get_transform():
+    return A.Compose([
+        A.Resize(Config.IMAGE_SIZE, Config.IMAGE_SIZE),
+        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ToTensorV2(),
+    ])
+
+def load_models(checkpoints_arg, device):
+    """
+    Load one or multiple models for ensemble inference.
+    checkpoints_arg: Comma-separated list of paths, or single path, or directory.
+    """
+    paths = []
+    if os.path.isdir(checkpoints_arg):
+        paths = glob.glob(os.path.join(checkpoints_arg, "*.safetensors"))
+        if not paths:
+            paths = glob.glob(os.path.join(checkpoints_arg, "*.pth"))
+    else:
+        paths = checkpoints_arg.split(',')
+    
+    models = []
+    print(f"Loading {len(paths)} model(s) for ensemble inference...")
+    
+    for path in paths:
+        path = path.strip()
+        if not path: continue
+        
+        print(f"Loading: {path}")
+        model = DeepfakeDetector(pretrained=False) # Structure only
+        model.to(device)
+        model.eval()
+        
+        try:
+            if path.endswith(".safetensors") and SAFETENSORS_AVAILABLE:
+                state_dict = load_file(path)
+            else:
+                state_dict = torch.load(path, map_location=device)
+            model.load_state_dict(state_dict)
+            models.append(model)
+            print(f"✅ Successfully loaded: {os.path.basename(path)}")
+        except Exception as e:
+            print(f"❌ Failed to load {path}: {e}")
+            import traceback
+            traceback.print_exc()
+            
+    if not models:
+        # Fallback for testing if no checkpoint exists yet
+        print("Warning: No valid checkoints loaded. Using random initialization for testing flow.")
+        model = DeepfakeDetector(pretrained=False).to(device)
+        model.eval()
+        models.append(model)
+        
+    return models
+
+def predict_ensemble(models, image_path, device, transform):
+    try:
+        image = cv2.imread(image_path)
+        if image is None:
+            return None, "Error: Could not read image"
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    except Exception as e:
+        return None, str(e)
+
+    augmented = transform(image=image)
+    image_tensor = augmented['image'].unsqueeze(0).to(device)
+    
+    probs = []
+    with torch.no_grad():
+        for model in models:
+            logits = model(image_tensor)
+            prob = torch.sigmoid(logits).item()
+            probs.append(prob)
+            
+    # Ensemble Strategy: Average Probability
+    avg_prob = sum(probs) / len(probs)
+    return avg_prob, None
+
+def main():
+    parser = argparse.ArgumentParser(description="Deepfake Detection Inference (Ensemble Support)")
+    parser.add_argument("--source", type=str, required=True, help="Path to image or directory")
+    parser.add_argument("--checkpoints", type=str, default="results/checkpoints", help="Path to checkpoint file, list of files (comma-separated), or directory")
+    parser.add_argument("--device", type=str, default=Config.DEVICE, help="Device to use (cuda/mps/cpu)")
+    args = parser.parse_args()
+    
+    device = torch.device(args.device)
+    print(f"Using device: {device}")
+    
+    # Load Models
+    models = load_models(args.checkpoints, device)
+    transform = get_transform()
+    
+    # Process Source
+    if os.path.isdir(args.source):
+        files = glob.glob(os.path.join(args.source, "*.*"))
+        # Filter images
+        files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+    else:
+        files = [args.source]
+        
+    print(f"Processing {len(files)} images with {len(models)} model(s)...")
+    print("-" * 65)
+    print(f"{'Image Name':<40} | {'Prediction':<10} | {'Confidence':<10}")
+    print("-" * 65)
+    
+    for file_path in files:
+        prob, error = predict_ensemble(models, file_path, device, transform)
+        if error:
+            print(f"{os.path.basename(file_path):<40} | ERROR: {error}")
+            continue
+            
+        is_fake = prob > 0.5
+        label = "FAKE" if is_fake else "REAL"
+        confidence = prob if is_fake else 1 - prob
+        
+        print(f"{os.path.basename(file_path):<40} | {label:<10} | {confidence:.2%}")
+
+if __name__ == "__main__":
+    main()

model/src/models.py

@@ -0,0 +1,197 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import numpy as np
+from src.utils import get_fft_feature
+
+class RGBBranch(nn.Module):
+    def __init__(self, pretrained=True):
+        super().__init__()
+        # EfficientNet V2 Small: Robust and efficient spatial features
+        weights = models.EfficientNet_V2_S_Weights.DEFAULT if pretrained else None
+        self.net = models.efficientnet_v2_s(weights=weights)
+        # Extract features before classification head
+        self.features = self.net.features
+        self.avgpool = self.net.avgpool
+        self.out_dim = 1280 
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return x
+
+class FreqBranch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Simple CNN to analyze frequency domain patterns
+        self.net = nn.Sequential(
+            nn.Conv2d(3, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((1,1))
+        )
+        self.out_dim = 128
+    
+    def forward(self, x):
+        return torch.flatten(self.net(x), 1)
+
+class PatchBranch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Analyzes local patches for inconsistencies
+        # Shared lightweight CNN for each patch
+        self.patch_encoder = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2), # 64 -> 32
+            nn.Conv2d(16, 32, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2), # 32 -> 16
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((1,1))
+        )
+        self.out_dim = 64
+
+    def forward(self, x):
+        # x: (B, 3, 256, 256)
+        # Create 4x4=16 patches of size 64x64
+        # Unfold logic: kernel_size=64, stride=64
+        patches = x.unfold(2, 64, 64).unfold(3, 64, 64)
+        # patches shape: (B, 3, 4, 4, 64, 64)
+        B, C, H_grid, W_grid, H_patch, W_patch = patches.shape
+        
+        # Merge batch and grid dimensions for parallel processing
+        patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()
+        patches = patches.view(B * H_grid * W_grid, C, H_patch, W_patch)
+        
+        # Encode
+        feats = self.patch_encoder(patches) # (B*16, 64, 1, 1)
+        feats = torch.flatten(feats, 1) # (B*16, 64)
+        
+        # Aggregate back to B
+        feats = feats.view(B, H_grid * W_grid, -1) # (B, 16, 64)
+        
+        # Max pool over patches to capture the "most fake" patch signal
+        feats_max, _ = torch.max(feats, dim=1) # (B, 64)
+        
+        return feats_max
+
+class ViTBranch(nn.Module):
+    def __init__(self, pretrained=True):
+        super().__init__()
+        # Swin Transformer Tiny: Capture long-range dependencies
+        weights = models.Swin_V2_T_Weights.DEFAULT if pretrained else None
+        self.net = models.swin_v2_t(weights=weights)
+        
+        # Replace head with Identity to get features
+        self.out_dim = self.net.head.in_features
+        self.net.head = nn.Identity()
+    
+    def forward(self, x):
+        return self.net(x)
+
+class DeepfakeDetector(nn.Module):
+    def __init__(self, pretrained=True):
+        super().__init__()
+        self.rgb_branch = RGBBranch(pretrained)
+        self.freq_branch = FreqBranch()
+        self.patch_branch = PatchBranch()
+        self.vit_branch = ViTBranch(pretrained)
+        
+        input_dim = (self.rgb_branch.out_dim + 
+                     self.freq_branch.out_dim + 
+                     self.patch_branch.out_dim + 
+                     self.vit_branch.out_dim)
+        
+        # Confidence-based fusion head
+        self.classifier = nn.Sequential(
+            nn.Linear(input_dim, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(512, 1)
+        )
+        
+    def forward(self, x):
+        # 1. Spatial Analysis
+        rgb_feat = self.rgb_branch(x)
+        
+        # 2. Frequency Analysis
+        freq_img = get_fft_feature(x)
+        freq_feat = self.freq_branch(freq_img)
+        
+        # 3. Patch Analysis (Local Inconsistencies)
+        patch_feat = self.patch_branch(x)
+        
+        # 4. Global Consistency (ViT)
+        vit_feat = self.vit_branch(x)
+        
+        # 5. Feature Fusion
+        combined = torch.cat([rgb_feat, freq_feat, patch_feat, vit_feat], dim=1)
+        
+        return self.classifier(combined)
+
+    def get_heatmap(self, x):
+        """Generate Grad-CAM heatmap for the input image"""
+        # We'll use the RGB branch for visualization as it contains spatial features
+        # Enable gradients for the input if needed, though typically we hook into layers
+        
+        # 1. Forward pass through RGB branch
+        # We need to register a hook on the last conv layer of the efficientnet features
+        # Target layer: self.rgb_branch.features[-1] (the last block)
+        
+        gradients = []
+        activations = []
+        
+        def backward_hook(module, grad_input, grad_output):
+            gradients.append(grad_output[0])
+            
+        def forward_hook(module, input, output):
+            activations.append(output)
+            
+        # Register hooks on the last convolutional layer of RGB branch
+        target_layer = self.rgb_branch.features[-1]
+        hook_b = target_layer.register_full_backward_hook(backward_hook)
+        hook_f = target_layer.register_forward_hook(forward_hook)
+        
+        # Forward pass
+        logits = self(x)
+        pred_idx = 0 # Binary classification, output is scalar logic
+        
+        # Backward pass
+        self.zero_grad()
+        logits.backward(retain_graph=True)
+        
+        # Get gradients and activations
+        pooled_gradients = torch.mean(gradients[0], dim=[0, 2, 3])
+        activation = activations[0][0]
+        
+        # Weight activations by gradients (Grad-CAM)
+        for i in range(activation.shape[0]):
+            activation[i, :, :] *= pooled_gradients[i]
+            
+        heatmap = torch.mean(activation, dim=0).cpu().detach().numpy()
+        heatmap = np.maximum(heatmap, 0) # ReLU
+        
+        # Normalize
+        if np.max(heatmap) != 0:
+            heatmap /= np.max(heatmap)
+            
+        # Remove hooks
+        hook_b.remove()
+        hook_f.remove()
+        
+        return heatmap

model/src/test_dataloading.py

@@ -0,0 +1,47 @@
+import os
+import random
+from src.config import Config
+from src.dataset import DeepfakeDataset
+
+def test_dataloading():
+    print("Testing Data Loading & Splitting Logic...")
+    Config.setup()
+    
+    print(f"Data Path: {Config.TRAIN_DATA_PATH}")
+    
+    # 1. Test Scan
+    paths, labels = DeepfakeDataset.scan_directory(Config.TRAIN_DATA_PATH)
+    total_files = len(paths)
+    print(f"Total images found: {total_files}")
+    
+    if total_files == 0:
+        print("[FAIL] No images found! Check path.")
+        return
+
+    # 2. Simulate Split Logic
+    combined = list(zip(paths, labels))
+    random.shuffle(combined)
+    split_idx = int(len(combined) * 0.8)
+    train_data = combined[:split_idx]
+    val_data = combined[split_idx:]
+    
+    print(f"Train Split: {len(train_data)} images")
+    print(f"Val Split:   {len(val_data)} images")
+    
+    # 3. Test Dataset Initialization
+    try:
+        train_paths, train_labels = zip(*train_data)
+        ds = DeepfakeDataset(file_paths=list(train_paths), labels=list(train_labels), phase='train')
+        print(f"[Pass] Dataset initialized with {len(ds)} samples.")
+        
+        # Test Get Item
+        img, lbl = ds[0]
+        print(f"[Pass] Loaded sample image. Shape: {img.shape}, Label: {lbl}")
+    except Exception as e:
+        print(f"[FAIL] Dataset initialization or loading error: {e}")
+        return
+
+    print("\nSUCCESS: Data loading verification passed!")
+
+if __name__ == "__main__":
+    test_dataloading()

model/src/test_dryrun.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+from src.models import DeepfakeDetector
+from src.config import Config
+
+def test_model_architecture():
+    print("Testing DeepfakeDetector Architecture...")
+    
+    # Check device
+    device = torch.device("cpu") # Test on CPU for simplicity or Config.DEVICE
+    print(f"Device: {device}")
+    
+    # Initialize Model
+    try:
+        model = DeepfakeDetector(pretrained=False).to(device)
+        print("[Pass] Model Initialization")
+    except Exception as e:
+        print(f"[Fail] Model Initialization: {e}")
+        return
+
+    # Create dummy input
+    batch_size = 2
+    x = torch.randn(batch_size, 3, Config.IMAGE_SIZE, Config.IMAGE_SIZE).to(device)
+    print(f"Input Shape: {x.shape}")
+    
+    # Forward Pass
+    try:
+        out = model(x)
+        print(f"Output Shape: {out.shape}")
+        
+        if out.shape == (batch_size, 1):
+            print("[Pass] Output Shape Correct")
+        else:
+            print(f"[Fail] Output Shape Incorrect. Expected ({batch_size}, 1), got {out.shape}")
+    except Exception as e:
+        print(f"[Fail] Forward Pass: {e}")
+        # Debug trace
+        import traceback
+        traceback.print_exc()
+        return
+
+    # Loss and Backward
+    try:
+        criterion = nn.BCEWithLogitsLoss()
+        target = torch.ones(batch_size, 1).to(device)
+        loss = criterion(out, target)
+        loss.backward()
+        print(f"[Pass] Backward Pass (Loss: {loss.item():.4f})")
+    except Exception as e:
+        print(f"[Fail] Backward Pass: {e}")
+        return
+
+    print("\nSUCCESS: Model architecture verification passed!")
+
+if __name__ == "__main__":
+    test_model_architecture()

model/src/train.py

@@ -0,0 +1,271 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import random
+import ssl
+# Disable SSL verification for downloading pretrained weights
+ssl._create_default_https_context = ssl._create_unverified_context
+from torch.cuda.amp import GradScaler, autocast
+
+from src.config import Config
+from src.models import DeepfakeDetector
+from src.dataset import DeepfakeDataset
+
+try:
+    from safetensors.torch import save_file, load_file
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+    print("Warning: safetensors not installed. Checkpoints will be saved as .pt")
+
+def train():
+    # Setup
+    Config.setup()
+    device = torch.device(Config.DEVICE)
+    
+    # --- Data Loading with Automatic Split ---
+    if Config.TRAIN_DATA_PATH == Config.TEST_DATA_PATH:
+        print("Train and Test paths are identical. Performing automatic 80/20 shuffle split...")
+        all_paths, all_labels = DeepfakeDataset.scan_directory(Config.TRAIN_DATA_PATH)
+        
+        if len(all_paths) == 0:
+            print(f"No images found in {Config.TRAIN_DATA_PATH}")
+            return
+
+        # Combine and shuffle
+        combined = list(zip(all_paths, all_labels))
+        random.shuffle(combined)
+        
+        split_idx = int(len(combined) * 0.8)
+        train_data = combined[:split_idx]
+        val_data = combined[split_idx:]
+        
+        train_paths, train_labels = zip(*train_data)
+        val_paths, val_labels = zip(*val_data)
+        
+        train_dataset = DeepfakeDataset(file_paths=list(train_paths), labels=list(train_labels), phase='train')
+        val_dataset = DeepfakeDataset(file_paths=list(val_paths), labels=list(val_labels), phase='val')
+    else:
+        # Standard folder-based loading
+        train_dataset = DeepfakeDataset(root_dir=Config.TRAIN_DATA_PATH, phase='train')
+        val_dataset = DeepfakeDataset(root_dir=Config.TEST_DATA_PATH, phase='val')
+    
+    # Dataloaders
+    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True,
+                              num_workers=Config.NUM_WORKERS, 
+                              pin_memory=True if device.type=='cuda' else False,
+                              persistent_workers=True if Config.NUM_WORKERS > 0 else False)
+    val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE, shuffle=False,
+                            num_workers=Config.NUM_WORKERS, 
+                            pin_memory=True if device.type=='cuda' else False,
+                            persistent_workers=True if Config.NUM_WORKERS > 0 else False)
+    
+    # Model
+    print("Initializing Multi-Branch DeepfakeDetector...")
+    model = DeepfakeDetector(pretrained=True).to(device)
+    
+    # Optimization
+    criterion = nn.BCEWithLogitsLoss()
+    optimizer = optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE, weight_decay=Config.WEIGHT_DECAY)
+    # Optimization
+    criterion = nn.BCEWithLogitsLoss()
+    optimizer = optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE, weight_decay=Config.WEIGHT_DECAY)
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
+    
+    # Enable AMP only for CUDA (Windows NVIDIA)
+    use_amp = (Config.DEVICE == 'cuda')
+    scaler = GradScaler() if use_amp else None
+    if use_amp:
+        print("🚀 Mixed Precision (AMP) Enabled for RTX GPU")
+    else:
+        print("🐌 Standard Precision (No AMP) for CPU/MPS")
+    
+    # Resume from checkpoint if exists
+    start_epoch = 0
+    best_acc = 0.0
+    
+    # Priority:
+    # 1. best_model.safetensors (if we crashed mid-training)
+    # 2. patched_model.safetensors (the model we want to improve)
+    
+    resume_path = os.path.join(Config.CHECKPOINT_DIR, "best_model.safetensors")
+    if not os.path.exists(resume_path):
+        # Look for latest epoch checkpoint
+        import glob
+        import re
+        checkpoints = glob.glob(os.path.join(Config.CHECKPOINT_DIR, "checkpoint_ep*.safetensors"))
+        if checkpoints:
+            # Sort by epoch number
+            def get_epoch(p):
+                match = re.search(r"checkpoint_ep(\d+)", p)
+                return int(match.group(1)) if match else 0
+            
+            latest_ckpt = max(checkpoints, key=get_epoch)
+            resume_path = latest_ckpt
+            start_epoch = get_epoch(latest_ckpt)
+            print(f"🔄 Auto-Resuming from latest epoch: {start_epoch}")
+        else:
+            resume_path = os.path.join(Config.CHECKPOINT_DIR, "patched_model.safetensors")
+    
+    if os.path.exists(resume_path):
+        print(f"\n🔄 Found existing checkpoint: {resume_path}")
+        print("Auto-resuming to FINETUNE this model...")
+        
+        try:
+            if resume_path.endswith(".safetensors") and SAFETENSORS_AVAILABLE:
+                state_dict = load_file(resume_path)
+            else:
+                state_dict = torch.load(resume_path, map_location=device)
+            
+            # Use strict=False to allow for minor architecture changes or missing keys
+            model.load_state_dict(state_dict, strict=False)
+            print("✅ Weights loaded. Starting Fine-Tuning.")
+        except Exception as e:
+            print(f"⚠ Failed to load checkpoint: {e}")
+            print("Starting from ImageNet weights.")
+    
+    # Loop
+    
+    for epoch in range(start_epoch, Config.EPOCHS):
+        model.train()
+        train_loss = 0.0
+        train_correct = 0
+        train_total = 0
+        
+        loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{Config.EPOCHS}")
+        for images, labels in loop:
+            images = images.to(device)
+            labels = labels.to(device).unsqueeze(1)
+            
+            optimizer.zero_grad()
+            
+            if use_amp:
+                with autocast():
+                    outputs = model(images)
+                    loss = criterion(outputs, labels)
+                
+                scaler.scale(loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                # Standard training for Mac/CPU
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+            
+            train_loss += loss.item()
+            preds = (torch.sigmoid(outputs) > 0.5).float()
+            correct = (preds == labels).sum().item()
+            train_correct += correct
+            train_total += labels.size(0)
+            
+            loop.set_postfix(loss=loss.item(), acc=correct/labels.size(0))
+            
+        train_acc = train_correct / train_total if train_total > 0 else 0
+        print(f"Epoch {epoch+1} Train Loss: {train_loss/len(train_loader):.4f} Acc: {train_acc:.4f}")
+        
+        # Save checkpoint after every epoch
+        save_checkpoint(model, epoch+1, train_acc, best=False)
+        
+        # Validation
+        if len(val_dataset) > 0:
+            val_loss, val_acc = validate(model, val_loader, criterion, device)
+            print(f"Epoch {epoch+1} Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")
+            
+            # Save best model if validation accuracy improved
+            if val_acc > best_acc:
+                best_acc = val_acc
+                print(f"⭐ New best model! Validation Accuracy: {val_acc:.4f}")
+                save_checkpoint(model, epoch+1, val_acc, best=True)
+        
+        scheduler.step()
+    
+    print(f"\n🎉 Training Complete!")
+    print(f"Best Validation Accuracy: {best_acc:.4f}")
+
+def validate(model, loader, criterion, device):
+    model.eval()
+    val_loss = 0.0
+    correct = 0
+    total = 0
+    
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            labels = labels.to(device).unsqueeze(1)
+            
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            
+            val_loss += loss.item()
+            preds = (torch.sigmoid(outputs) > 0.5).float()
+            correct += (preds == labels).sum().item()
+            total += labels.size(0)
+            
+    return val_loss / len(loader), correct / total
+
+def save_checkpoint(model, epoch, acc, best=False):
+    state_dict = model.state_dict()
+    name = "best_model.safetensors" if best else f"checkpoint_ep{epoch}.safetensors"
+    path = os.path.join(Config.CHECKPOINT_DIR, name)
+    
+    if SAFETENSORS_AVAILABLE:
+        try:
+            # Try with shared tensors support
+            from safetensors.torch import save_model
+            save_model(model, path)
+            print(f"Saved Checkpoint: {path}")
+            
+            # 📝 Auto-Log to History
+            try:
+                from datetime import datetime
+                log_path = os.path.join(Config.PROJECT_ROOT, "TRAINING_HISTORY.md")
+                timestamp = datetime.now().strftime("%Y-%m-%d | %I:%M %p")
+                
+                # Create file with header if doesn't exist
+                if not os.path.exists(log_path):
+                    with open(log_path, "w", encoding="utf-8") as f:
+                        f.write("# 📜 Training History Log\n\n")
+                        f.write("| Date | Time | Model Name | Dataset | Epochs | Accuracy | Loss | Status |\n")
+                        f.write("| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n")
+                
+                # Append Entry to Summary Log
+                with open(log_path, "a", encoding="utf-8") as f:
+                    # Format: Date | Time | Name | Dataset | Epoch | Acc | Loss | Status
+                    dataset_name = os.path.basename(Config.DATA_DIR)
+                    entry = f"| **{timestamp.split(' | ')[0]}** | {timestamp.split(' | ')[1]} | {name} | {dataset_name} | {epoch} | {acc*100:.2f}% | N/A | ✅ Saved |\n"
+                    f.write(entry)
+                    print(f"📝 Logged to TRAINING_HISTORY.md")
+
+                # 📝 Detailed Lab Notebook Logging
+                detail_path = os.path.join(Config.PROJECT_ROOT, "DETAILED_HISTORY.md")
+                with open(detail_path, "a", encoding="utf-8") as f:
+                    f.write(f"\n## Model: {name} (Epoch {epoch})\n")
+                    f.write(f"| Feature | Detail |\n| :--- | :--- |\n")
+                    f.write(f"| **Date** | {timestamp} |\n")
+                    f.write(f"| **Training Accuracy** | {acc*100:.2f}% |\n")
+                    f.write(f"| **Dataset** | {Config.DATA_DIR} |\n")
+                    f.write(f"| **Batch Size** | {Config.BATCH_SIZE} |\n")
+                    f.write(f"| **Optimizer** | AdamW (lr={Config.LEARNING_RATE}) |\n")
+                    f.write(f"| **Device** | {Config.DEVICE.upper()} |\n")
+                    f.write("\n---\n")
+                    print(f"📘 Detailed log written to DETAILED_HISTORY.md")
+
+            except Exception as e:
+                print(f"⚠️ Failed to write log: {e}")
+
+        except Exception as e:
+            # Fallback to regular torch save if safetensors fails
+            print(f"SafeTensors save failed ({e}), falling back to .pth format")
+            torch.save(state_dict, path.replace(".safetensors", ".pth"))
+            print(f"Saved Checkpoint (Legacy): {path.replace('.safetensors', '.pth')}")
+    else:
+        torch.save(state_dict, path.replace(".safetensors", ".pth"))
+        print(f"Saved Checkpoint (Legacy): {path}")
+
+if __name__ == "__main__":
+    train()

model/src/utils.py

@@ -0,0 +1,37 @@
+import torch
+import numpy as np
+import cv2
+
+def get_fft_feature(x):
+    """
+    Computes the Log-Magnitude Spectrum of the input images.
+    Args:
+        x (torch.Tensor): Input images of shape (B, C, H, W)
+    Returns:
+        torch.Tensor: Log-magnitude spectrum of shape (B, C, H, W)
+    """
+    if x.dim() == 3:
+        x = x.unsqueeze(0)
+        
+    # Compute 2D FFT
+    fft = torch.fft.fft2(x, norm='ortho')
+    
+    # Compute magnitude
+    mag = torch.abs(fft)
+    
+    # Apply log scale (add epsilon for stability)
+    mag = torch.log(mag + 1e-6)
+    
+    # Shift zero-frequency component to the center of the spectrum
+    mag = torch.fft.fftshift(mag, dim=(-2, -1))
+    
+    return mag
+
+def min_max_normalize(tensor):
+    """
+    Min-max normalization for visualization or stable training provided tensor.
+    """
+    min_val = tensor.min()
+    max_val = tensor.max()
+    return (tensor - min_val) / (max_val - min_val + 1e-8)
+

model/src/video_inference.py

@@ -0,0 +1,172 @@
+import cv2
+import torch
+import numpy as np
+import os
+from PIL import Image
+
+def process_video(video_path, model, transform, device, frames_per_second=1):
+    """
+    Process a video file frame-by-frame using the deepfake detection model.
+    
+    Args:
+        video_path (str): Path to the video file.
+        model (torch.nn.Module): Loaded PyTorch model.
+        transform (callable): Albumentations transform pipeline.
+        device (torch.device): Device to run inference on.
+        frames_per_second (int): Number of frames to sample per second of video. 
+                                 Default is 1 to keep processing fast.
+    
+    Returns:
+        dict: Aggregated results including verdict, average confidence, and frame-level details.
+    """
+    if model is None:
+        return {"error": "Model not loaded"}
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return {"error": "Could not open video file"}
+
+    # specific video properties
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    if fps <= 0: fps = 30 # Fallback
+    
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    duration = total_frames / fps
+    
+    # Calculate sampling interval (step size)
+    # If we want 1 frame per second, we step by 'fps' frames
+    step = int(fps / frames_per_second)
+    if step < 1: step = 1
+
+    frame_indices = []
+    probs = []
+    
+    print(f"Processing video: {video_path}")
+    print(f"Duration: {duration:.2f}s, FPS: {fps}, Total Frames: {total_frames}")
+    print(f"Sampling every {step} frames...")
+
+    count = 0
+    processed_count = 0
+    
+    suspicious_frames = [] # Store frames with high fake probability
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+            
+        if count % step == 0:
+            # Process this frame
+            try:
+                # Convert BGR (OpenCV) to RGB
+                image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                
+                # --- Face Extraction ---
+                # Load Haar Cascade (lazy load)
+                if not hasattr(process_video, "face_cascade"):
+                    try:
+                        cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
+                        process_video.face_cascade = cv2.CascadeClassifier(cascade_path)
+                    except:
+                        process_video.face_cascade = None
+
+                face_crop = None
+                if process_video.face_cascade:
+                    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+                    faces = process_video.face_cascade.detectMultiScale(
+                        gray, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60)
+                    )
+                    
+                    if len(faces) > 0:
+                        # Find largest face
+                        largest_face = max(faces, key=lambda rect: rect[2] * rect[3])
+                        x, y, w, h = largest_face
+                        
+                        # Add margin (20%)
+                        margin = int(max(w, h) * 0.2)
+                        x_start = max(x - margin, 0)
+                        y_start = max(y - margin, 0)
+                        x_end = min(x + w + margin, frame.shape[1])
+                        y_end = min(y + h + margin, frame.shape[0])
+                        
+                        face_crop = image[y_start:y_end, x_start:x_end]
+                
+                # Use face crop if found, otherwise use full image
+                input_image = face_crop if face_crop is not None else image
+                
+                # Apply transforms
+                augmented = transform(image=input_image)
+                image_tensor = augmented['image'].unsqueeze(0).to(device)
+                
+                # Inference
+                with torch.no_grad():
+                    logits = model(image_tensor)
+                    prob = torch.sigmoid(logits).item()
+                
+                probs.append(prob)
+                frame_indices.append(count)
+                processed_count += 1
+                
+                # If highly fake, store metadata (timestamp)
+                if prob > 0.5:
+                    timestamp = count / fps
+                    suspicious_frames.append({
+                        "timestamp": round(timestamp, 2),
+                        "frame_index": count,
+                        "fake_prob": round(prob, 4)
+                    })
+                    
+            except Exception as e:
+                print(f"Error processing frame {count}: {e}")
+        
+        count += 1
+
+    cap.release()
+
+    if processed_count == 0:
+        return {"error": "No frames processed"}
+
+    # Aggregation
+    avg_prob = sum(probs) / len(probs)
+    max_prob = max(probs)
+    fake_frame_count = len([p for p in probs if p > 0.6]) # Stricter frame threshold
+    fake_ratio = fake_frame_count / processed_count
+    
+    # Verdict Logic (Tuned for High Efficiency Model)
+    # The new model is detecting everything as fake, so we need stricter rules.
+    
+    # 1. Standard Average Check (shifted)
+    cond1 = avg_prob > 0.65
+    
+    # 2. Density Check: Require at least 15% of frames to be strictly fake
+    # Was 5%, which is too low for a sensitive model
+    cond2 = fake_ratio > 0.15 and max_prob > 0.7
+    
+    # 3. Peak Check: Only flag single-frame anomalies if EXTREMELY suspicious
+    cond3 = max_prob > 0.95
+    
+    is_fake = cond1 or cond2 or cond3
+    
+    verdict = "FAKE" if is_fake else "REAL"
+    
+    # Confidence Calculation
+    if is_fake:
+        confidence = max(max_prob, 0.6)
+    else:
+        confidence = 1 - avg_prob
+    
+    return {
+        "type": "video",
+        "prediction": verdict,
+        "confidence": float(confidence),
+        "avg_fake_prob": float(avg_prob),
+        "max_fake_prob": float(max_prob),
+        "fake_frame_ratio": float(fake_ratio),
+        "processed_frames": processed_count,
+        "duration": float(duration),
+        "timeline": [
+            {"time": round(i / fps, 2), "prob": round(p, 3)} 
+            for i, p in zip(frame_indices, probs)
+        ],
+        "suspicious_frames": suspicious_frames[:10] # Top 10 suspicious moments
+    }