Upload 22 files

models_saved/dinov2_best.pt

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

src/data/data_loader.py

@@ -0,0 +1,83 @@
+import torch
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader, random_split
+import yaml
+import os
+
+def get_transforms(cfg):
+    """
+    DINOv2 expects ImageNet normalization.
+    We also add some light augmentation to prevent overfitting.
+    """
+    img_size = cfg['data']['image_size']
+    
+    # Training Transforms (with Augmentation)
+    train_transform = transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.RandomHorizontalFlip(p=0.5), # 50% chance to flip
+        transforms.ColorJitter(brightness=0.1, contrast=0.1), # Slight color changes
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], # DINOv2 Expected Mean
+            std=[0.229, 0.224, 0.225]   # DINOv2 Expected Std
+        )
+    ])
+
+    # Validation/Test Transforms (No Augmentation)
+    val_transform = transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+    
+    return train_transform, val_transform
+
+def create_dataloaders(config_path="configs/config.yaml"):
+    # Load config
+    with open(config_path, 'r') as f:
+        cfg = yaml.safe_load(f)
+
+    train_transform, val_transform = get_transforms(cfg)
+    data_dir = cfg['data']['train_dir'] # Should be "data/raw"
+
+    # 1. Load the Entire Dataset (REAL + FAKE)
+    full_dataset = datasets.ImageFolder(root=data_dir)
+    
+    # 2. Split: 80% Train, 20% Validation
+    total_size = len(full_dataset)
+    train_size = int(0.8 * total_size)
+    val_size = total_size - train_size
+    
+    train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])
+    
+    # Apply specific transforms
+    train_dataset.dataset.transform = train_transform
+    val_dataset.dataset.transform = val_transform
+
+    # 3. Create Loaders
+    train_loader = DataLoader(
+        train_dataset, 
+        batch_size=cfg['data']['batch_size'], 
+        shuffle=True, 
+        num_workers=cfg['data']['num_workers']
+    )
+    
+    val_loader = DataLoader(
+        val_dataset, 
+        batch_size=cfg['data']['batch_size'], 
+        shuffle=False, 
+        num_workers=cfg['data']['num_workers']
+    )
+
+    print(f"✅ Data Ready:")
+    print(f"   - Train: {len(train_dataset)} images")
+    print(f"   - Val:   {len(val_dataset)} images")
+    print(f"   - Classes: {full_dataset.class_to_idx}")
+    
+    return train_loader, val_loader
+
+if __name__ == "__main__":
+    create_dataloaders()

src/inference/predictor.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import cv2
+import os
+
+# Import GradCAM tools
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+from src.models.model import VisionGuardModel
+
+class VisionGuardPredictor:
+    def __init__(self, model_path, config_path="configs/config.yaml"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"🚀 Loading Inference Engine on: {self.device}")
+        
+        # 1. Load Model
+        self.model = VisionGuardModel(num_classes=2, pretrained=False)
+        checkpoint = torch.load(model_path, map_location=self.device)
+        self.model.load_state_dict(checkpoint)
+        self.model.to(self.device)
+        self.model.eval()
+        
+        # 2. Setup GradCAM (The Explainability Tool)
+        # We target the last normalization layer of the backbone
+        target_layers = [self.model.backbone.blocks[-1].norm1]
+        
+        # DINOv2 requires a special reshape transform because it outputs 1D sequences
+        def reshape_transform(tensor):
+            # DINOv2 small outputs: [Batch, 257, 384] (1 CLS token + 256 Patches)
+            # We discard the CLS token (index 0) and keep the 256 patches
+            result = tensor[:, 1:, :]
+            
+            # Reshape 256 -> 16x16 grid (since 224/14 = 16)
+            height = 14
+            width = 14
+            # Note: If image size is 224x224, grid is 16x16. 
+            # DINOv2-S/14 means patch size is 14. 224/14 = 16.
+            grid_size = 16
+            
+            result = result.reshape(tensor.size(0), grid_size, grid_size, tensor.size(2))
+            
+            # Bring channels first: [Batch, Channels, Height, Width]
+            result = result.transpose(2, 3).transpose(1, 2)
+            return result
+
+        self.cam = GradCAM(model=self.model, target_layers=target_layers, reshape_transform=reshape_transform)
+
+        # 3. Preprocessing
+        self.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])
+        ])
+
+        self.labels = ['FAKE', 'REAL']
+
+    def predict(self, image_path):
+        # 1. Load Image
+        image = Image.open(image_path).convert('RGB')
+        
+        # Keep a clean copy for visualization (resized to 224x224)
+        vis_image = image.resize((224, 224))
+        vis_image = np.float32(vis_image) / 255.0 # Normalize 0-1 for OpenCV
+        
+        # 2. Transform for Model
+        input_tensor = self.transform(image).unsqueeze(0).to(self.device)
+
+        # 3. Inference
+        with torch.no_grad():
+            outputs = self.model(input_tensor)
+            probs = F.softmax(outputs, dim=1)
+            confidence, predicted_class = torch.max(probs, 1)
+
+        # 4. Generate Heatmap
+        # We tell GradCAM to look for the predicted class
+        grayscale_cam = self.cam(input_tensor=input_tensor, targets=None)
+        grayscale_cam = grayscale_cam[0, :]
+        
+        # Overlay heatmap on image
+        visualization = show_cam_on_image(vis_image, grayscale_cam, use_rgb=True)
+        
+        # Convert back to PIL for Gradio
+        heatmap_pil = Image.fromarray(visualization)
+
+        # 5. Format Output
+        idx = predicted_class.item()
+        return {
+            "verdict": self.labels[idx],
+            "confidence": round(float(confidence.item()) * 100, 2),
+            "probabilities": {
+                "FAKE": round(float(probs[0][0].item()), 4),
+                "REAL": round(float(probs[0][1].item()), 4)
+            },
+            "heatmap": heatmap_pil
+        }

src/models/model.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+
+class VisionGuardModel(nn.Module):
+    def __init__(self, num_classes=2, pretrained=True):
+        super(VisionGuardModel, self).__init__()
+        # Load DINOv2 (The "Eye")
+        self.backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
+        
+        # Classification Head (The "Brain")
+        self.head = nn.Sequential(
+            nn.Linear(384, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, num_classes)
+        )
+
+    def forward(self, x):
+        features = self.backbone(x)
+        logits = self.head(features)
+        return logits

src/training/trainer.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from tqdm import tqdm
+import yaml
+import os
+
+# Import our project modules
+from src.data.data_loader import create_dataloaders
+from src.models.model import VisionGuardModel
+
+def train_one_epoch(model, loader, criterion, optimizer, device):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    loop = tqdm(loader, leave=False)
+    
+    for images, labels in loop:
+        images, labels = images.to(device), labels.to(device)
+
+        # Forward
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+
+        # Backward
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Stats
+        _, predicted = torch.max(outputs.data, 1)
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+        running_loss += loss.item()
+        
+        loop.set_description(f"Loss: {loss.item():.4f}")
+
+    return running_loss / len(loader), 100 * correct / total
+
+def validate(model, loader, criterion, device):
+    model.eval()
+    correct = 0
+    total = 0
+    running_loss = 0.0
+
+    with torch.no_grad():
+        for images, labels in loader:
+            images, labels = images.to(device), labels.to(device)
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+            running_loss += loss.item()
+
+    return running_loss / len(loader), 100 * correct / total
+
+def main():
+    # 1. Config & Device
+    with open("configs/config.yaml", "r") as f:
+        cfg = yaml.safe_load(f)
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"🚀 Training on: {device}")
+
+    # 2. Save Path (Google Drive)
+    save_dir = "/content/drive/MyDrive/VisionGuard_Models"
+    os.makedirs(save_dir, exist_ok=True)
+    save_path = f"{save_dir}/dinov2_best.pt"
+
+    # 3. Load Data & Model
+    train_loader, val_loader = create_dataloaders("configs/config.yaml")
+    model = VisionGuardModel(num_classes=2).to(device)
+
+    # 4. Optimizer (Only training the Head)
+    # Note: We only pass model.head.parameters() to optimizer because backbone is frozen!
+    optimizer = optim.AdamW(model.head.parameters(), lr=float(cfg['model']['learning_rate_head']))
+    criterion = nn.CrossEntropyLoss()
+
+    # 5. Training Loop
+    best_acc = 0.0
+    epochs = cfg['model']['epochs']
+    
+    print(f"\n🔥 Starting Training for {epochs} Epochs...")
+    print(f"💾 Best model will be saved to: {save_path}")
+    
+    for epoch in range(epochs):
+        print(f"\nEpoch {epoch+1}/{epochs}")
+        
+        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
+        val_loss, val_acc = validate(model, val_loader, criterion, device)
+        
+        print(f"   Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
+        print(f"   Val Loss:   {val_loss:.4f} | Val Acc:   {val_acc:.2f}%")
+
+        # Save Best
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save(model.state_dict(), save_path)
+            print(f"   ⭐ Saved New Best Model ({best_acc:.2f}%)")
+
+if __name__ == "__main__":
+    main()