"""
Vision Transformer-Based CMFD System
Author: AI Assistant (2023)
License: MIT
"""
import torch
import [Link] as nn
import [Link] as optim
import numpy as np
import cv2
import albumentations as A
from [Link] import Dataset, DataLoader
from [Link] import f1_score, roc_auc_score
import timm
import gradio as gr
# --- Synthetic Dataset Generation ---
class SyntheticCMFDGenerator:
def __init__(self, img_size=224):
self.img_size = img_size
[Link] = [[Link](224,224,3) for _ in range(100)] # Mock data
def create_forgery(self):
# 1. Random background
bg = [Link][[Link](0,100)]
# 2. Copy-move operation
obj = bg[50:150, 50:150].copy()
# 3. Apply transformations
transform = [Link]([
[Link](limit=45, p=0.7),
[Link](scale_limit=0.2, p=0.5),
[Link](p=0.3)
])
transformed = transform(image=obj)['image']
# 4. Paste back
mask = [Link]((224,224), dtype=np.float32)
x, y = [Link](0,100), [Link](0,100)
bg[y:y+100, x:x+100] = transformed
mask[y:y+100, x:x+100] = 1.0
return (bg*255).astype(np.uint8), mask
# --- Dataset Class ---
class CMFDDataset(Dataset):
def __init__(self, generator, num_samples=1000, transform=None):
[Link] = generator
self.num_samples = num_samples
[Link] = transform
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
# Generate synthetic sample
img, mask = [Link].create_forgery()
# Augmentations
if [Link]:
augmented = [Link](image=img, mask=mask)
img, mask = augmented['image'], augmented['mask']
# Convert to tensors
img_rgb = [Link](img).permute(2,0,1)/255.0
img_lab = [Link]([Link](img,
cv2.COLOR_RGB2LAB)).permute(2,0,1)/255.0
edges = [Link]([Link](img, 100, 200))[None,:,:]/255.0
mask = [Link](mask)[None,:,:]
return {'rgb': img_rgb, 'lab': img_lab, 'edges': edges, 'mask': mask}
# --- Model Architecture ---
class MultiModalAttentionCMFD([Link]):
def __init__(self):
super().__init__()
# Vision Transformer Branches
self.vit_rgb = timm.create_model('vit_base_patch16_224', pretrained=True,
num_classes=0)
self.vit_lab = timm.create_model('vit_base_patch16_224', pretrained=True,
num_classes=0)
# Edge Feature Extractor
self.edge_net = [Link](
nn.Conv2d(1, 64, 3, padding=1),
[Link](),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
[Link](),
nn.MaxPool2d(2)
)
# Cross-Modal Attention
self.cross_attn = [Link](embed_dim=768, num_heads=8, batch_first=True)
# Decoder
[Link] = [Link](
nn.ConvTranspose2d(768, 256, 4, stride=2),
[Link](),
nn.ConvTranspose2d(256, 128, 4, stride=2),
[Link](),
nn.Conv2d(128, 1, 1),
[Link]()
)
def forward(self, rgb, lab, edges):
# Extract features
rgb_feats = self.vit_rgb(rgb)
lab_feats = self.vit_lab(lab)
edge_feats = self.edge_net(edges)
# Reshape for attention (B, N, C)
edge_feats = edge_feats.flatten(2).permute(0,2,1)
# Cross-modal attention
attn_out, _ = self.cross_attn(
query=rgb_feats,
key=lab_feats,
value=edge_feats
)
# Decode
B, N, C = attn_out.shape
h = w = int(N**0.5)
attn_out = attn_out.permute(0,2,1).view(B, C, h, w)
return [Link](attn_out)
# --- Training Utilities ---
class CMFDTrainer:
def __init__(self):
# Initialize components
[Link] = SyntheticCMFDGenerator()
[Link] = [Link]([
[Link](p=0.5),
[Link](p=0.5),
[Link](p=0.3)
])
# Dataset & Loaders
self.train_set = CMFDDataset([Link], 1000, [Link])
self.val_set = CMFDDataset([Link], 200)
self.train_loader = DataLoader(self.train_set, batch_size=8, shuffle=True)
self.val_loader = DataLoader(self.val_set, batch_size=8)
# Model & Optimizer
[Link] = MultiModalAttentionCMFD()
[Link] = [Link]([Link](), lr=3e-5)
[Link] = [Link]()
def dice_loss(self, pred, target):
smooth = 1e-5
intersection = (pred * target).sum()
return 1 - (2. * intersection + smooth) / ([Link]() + [Link]() + smooth)
def train_epoch(self):
[Link]()
total_loss = 0.0
for batch in self.train_loader:
[Link].zero_grad()
outputs = [Link](
batch['rgb'],
batch['lab'],
batch['edges']
)
loss = 0.7*self.dice_loss(outputs, batch['mask']) + 0.3*[Link](outputs,
batch['mask'])
[Link]()
[Link]()
total_loss += [Link]()
return total_loss/len(self.train_loader)
def validate(self):
[Link]()
preds, targets = [], []
with torch.no_grad():
for batch in self.val_loader:
outputs = [Link](
batch['rgb'],
batch['lab'],
batch['edges']
)
[Link]([Link]().numpy().flatten())
[Link](batch['mask'].cpu().numpy().flatten())
return {
'f1': f1_score(targets, [Link](preds)),
'auc': roc_auc_score(targets, preds)
}
# --- Gradio Interface ---
def create_demo(model):
def predict(image):
# Preprocess
image = [Link](image, (224, 224))
lab = [Link](image, cv2.COLOR_RGB2LAB)
edges = [Link](image, 100, 200)
# To tensor
rgb_tensor = [Link](image).permute(2,0,1)[None]/255.0
lab_tensor = [Link](lab).permute(2,0,1)[None]/255.0
edges_tensor = [Link](edges)[None,None]/255.0
# Predict
with torch.no_grad():
mask = model(rgb_tensor, lab_tensor, edges_tensor)
return (mask[0,0].numpy() > 0.5).astype(np.uint8)*255
return [Link](
fn=predict,
inputs=[Link](label="Input Image"),
outputs=[Link](label="Forgery Mask"),
title="CMFD Detection Demo",
examples=[["[Link]"], ["[Link]"]] # Add real examples
)
# --- Main Execution ---
if __name__ == "__main__":
# Initialize system
trainer = CMFDTrainer()
# Training loop
for epoch in range(10):
train_loss = trainer.train_epoch()
val_metrics = [Link]()
print(f"Epoch {epoch+1}:")
print(f" Train Loss: {train_loss:.4f}")
print(f" Val F1: {val_metrics['f1']:.4f}, AUC: {val_metrics['auc']:.4f}")
# Launch demo
demo = create_demo([Link])
[Link]()