Fine-Tuning MantraNet for Image Forgery Detection

Cracking Image Forgeries with ManTraNet: Clear Fine‑Tuning Paths, Practical Code, and Smarter Results

ManTraNet is a deep learning model for image forgery detection and localization. It predicts a pixel-wise “forgery heatmap” that highlights manipulated regions.

What is ManTraNet?

ManTraNet stands for Manipulation Tracing Network. It has two main components:

• IMTFE: Image Manipulation Trace Feature Extractor (the backbone)
• AnomalyDetector: Head that maps features to a binary forgery heatmap

Idea: Image manipulations (splicing, copy-move, inpainting) leave subtle traces — like JPEG grid inconsistencies, noise pattern irregularities, and edge artifacts. IMTFE is trained to pick up these generic manipulation traces, while the AnomalyDetector turns them into a per-pixel prediction map.

Fine‑Tuning Strategies: Full vs. Partial

You can fine-tune ManTraNet in two ways:

Full fine-tuning:

• Update both IMTFE and AnomalyDetector.
• Use when switching to a very different domain where manipulation traces or acquisition pipelines shift significantly.
• Examples: medical imaging forensics, legal/scanned documents, satellite imagery.

Head-only fine-tuning (parameter-efficient):

• Freeze IMTFE, train only AnomalyDetector.
• Use for standard natural-image forensics tasks or when data is limited.
• Faster, less risk of catastrophic forgetting, and works well if your data distribution isn’t drastically different.

Rule of thumb:

• Big domain shift → full fine-tuning.
• Same or similar domain, small to medium dataset → train only the anomaly head.

Datasets for Forgery Detection (Quick Picks)

• CASIA v2.0: ~12,000 images (authentic + tampered). Good for splicing, copy-move, general tampering.
• COVERAGE: ~100 base images with ~400 copy‑move tampered versions — great for quick experiments.
• Columbia Uncompressed Splicing (CUISD): ~180 images — simple and small.
• Korus: ~220 forged/original pairs — useful for quick runs.

For this tutorial, we’ll use the COVERAGE dataset because it’s small and fast to iterate on.

Colab Setup

Using a free T4 GPU on Google Colab

Install dependencies:

# Install dependencies
!pip install pytorch-lightning==1.9.5
!pip install h5py scipy opencv-python matplotlib

Clone the ManTraNet (PyTorch) repo, then enter the folder that contains the model code:

!git clone https://github.com/RonyAbecidan/ManTraNet-pytorch.git
%cd ManTraNet-pytorch/MantraNet

Note:

• The repo folder structure places the core code in MantraNet/ with mantranet.py. After cd, Python imports should work as expected.
• If Colab restarts, rerun the cells and re‑cd into the folder.

Load the Pretrained Model

import torch
import matplotlib.pyplot as plt
from mantranet import pre_trained_model

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load pretrained ManTraNet
model = pre_trained_model(weight_path="./MantraNetv4.pt", device=device)
model.eval()

Why model.eval()?

eval() switches the model to inference behavior:

• Disables dropout randomness
• Uses running means/variances for BatchNorm

This ensures deterministic, stable outputs during validation/inference and when using the backbone as a frozen feature extractor.

Build a Minimal Dataset Class for COVERAGE

import os
import cv2
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import numpy as np
from PIL import Image

class CoverageDataset(Dataset):
    def __init__(self, img_dir, mask_dir, transform=None):
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.img_files = [f for f in os.listdir(img_dir) if f.lower().endswith((".tif", ".png", ".jpg", ".jpeg"))]

    def __len__(self):
        return len(self.img_files)

    def __getitem__(self, idx):
        img_name = self.img_files[idx]
        img_path = os.path.join(self.img_dir, img_name)

        image = Image.open(img_path).convert("RGB")

        base = os.path.splitext(img_name)[0]
        # Find mask by prefix match
        mask_candidates = [f for f in os.listdir(self.mask_dir) if os.path.splitext(f)[0].startswith(base)]
        if mask_candidates:
            mask_path = os.path.join(self.mask_dir, mask_candidates[0])
            mask = Image.open(mask_path).convert("L")
        else:
            mask = Image.new("L", image.size, 0)  # empty mask

        if self.transform:
            image = self.transform(image)
            # Important: for masks, use nearest-neighbor resize and map to {0,1}
            mask_np = np.array(mask, dtype=np.uint8)
            mask_pil = Image.fromarray(mask_np)
            mask_pil = transforms.functional.resize(mask_pil, (256, 256), interpolation=transforms.InterpolationMode.NEAREST)
            mask_tensor = torch.from_numpy(np.array(mask_pil, dtype=np.uint8)).float().unsqueeze(0) / 255.0
            return image, mask_tensor

        # Fallback if no transform
        image = transforms.ToTensor()(image)
        mask = transforms.ToTensor()(mask)
        return image, mask

import torch
transform = transforms.Compose([
    transforms.Resize((256, 256), antialias=True),
    transforms.ToTensor(),
])

train_dataset = CoverageDataset(
    img_dir="/content/drive/MyDrive/coverage/image",
    mask_dir="/content/drive/MyDrive/coverage/mask",
    transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=2, pin_memory=True)

Head‑Only Training: Freeze IMTFE, Train AnomalyDetector

from mantranet import AnomalyDetector, IMTFE

# Load pretrained feature extractor (IMTFE)
feature_extractor = IMTFE()
feature_extractor.load_state_dict(torch.load("./IMTFEv4.pt", map_location=device))
feature_extractor.eval().to(device)
for p in feature_extractor.parameters():
    p.requires_grad = False  # Freeze

# Initialize anomaly detector head
anomaly_detector = AnomalyDetector().to(device)

# Loss & Optimizer
criterion = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(anomaly_detector.parameters(), lr=1e-4)

Why BCEWithLogitsLoss?

• The AnomalyDetector outputs raw logits (unbounded scores).
• BCEWithLogitsLoss combines a sigmoid layer with binary cross-entropy in a numerically stable way.
• Tip : If the model already applied sigmoid inside forward(), switch to BCELoss instead. In ManTraNet, the head typically outputs logits, so BCEWithLogitsLoss is appropriate.

Training (demo: 3 epochs):

feature_extractor.eval()  # stays frozen

for epoch in range(3):  # demo; use more epochs for real runs
    anomaly_detector.train()
    running_loss = 0.0

    for imgs, masks in train_loader:
        imgs, masks = imgs.to(device), masks.to(device)

        with torch.no_grad():
            feats = feature_extractor(imgs)

        outputs = anomaly_detector(feats)  # logits, shape [B,1,H,W]
        # Ensure masks shape matches outputs
        if masks.shape != outputs.shape:
            masks = torch.nn.functional.interpolate(masks, size=outputs.shape[-2:], mode="nearest")

        loss = criterion(outputs, masks)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * imgs.size(0)

    epoch_loss = running_loss / len(train_dataset)
    print(f"Epoch {epoch+1}: loss={epoch_loss:.4f}")

Visualize Predictions

def visualize_result(img, mask, pred):
    img = img.permute(1, 2, 0).cpu().numpy()
    mask = mask.squeeze().cpu().numpy()
    pred = torch.sigmoid(pred).squeeze().detach().cpu().numpy()

    fig, axs = plt.subplots(1, 3, figsize=(12,4))
    axs[0].imshow(img)
    axs[0].set_title("Original Image")
    axs[0].axis("off")
    axs[1].imshow(mask, cmap="gray", vmin=0, vmax=1)
    axs[1].set_title("Ground Truth Mask")
    axs[1].axis("off")
    axs[2].imshow(pred, cmap="hot")
    axs[2].set_title("Predicted Forgery Heatmap")
    axs[2].axis("off")
    plt.tight_layout()
    plt.show()

# Test on one mini-batch
anomaly_detector.eval()
imgs, masks = next(iter(train_loader))
with torch.no_grad():
    feats = feature_extractor(imgs.to(device))
    preds = anomaly_detector(feats)

visualize_result(imgs[0], masks[0], preds[0])

Output:

As it was for 3 epochs only you can try to train for more better results

When to Do Full Fine‑Tuning

Switch to full fine‑tuning (unfreeze IMTFE) when:

• The imagery differs a lot from the pretraining domain: medical, documents, satellite, compression settings, sensors/scanners.
• The manipulations have domain-specific signatures (e.g., resave pipelines, scanner streaks).

How to unfreeze:

for p in feature_extractor.parameters():
    p.requires_grad = True

optimizer = torch.optim.Adam(
    list(feature_extractor.parameters()) + list(anomaly_detector.parameters()),
    lr=1e-5  # lower LR for backbone
)

Tip: Use discriminative learning rates (lower for backbone, higher for head) or separate optimizers/param groups.

Quality Boosters and Practical Tips

Data preprocessing:

• Ensure masks are binary {0,1}.
• Use nearest-neighbor resizing for masks.
• Match output resolution: interpolate masks to the head’s output size.

Data augmentation (strongly recommended):

• Photometric: brightness/contrast, JPEG re-encoding, Gaussian noise.
• Geometric: flips, small rotations, random crops.
• Important: Apply identical geometric transforms to both image and mask; do not apply photometric transforms to masks.

Loss improvements:

• Combine BCEWithLogitsLoss with Dice or Focal to handle class imbalance:
  total_loss = bce + λ * (1 — dice)
• Use pos_weight in BCEWithLogitsLoss to upweight tampered pixels if masks are sparse.

Thresholding:

• Convert heatmaps to binary by threshold (default 0.5). Tune threshold on a validation set (PR curve, F1 max).

Evaluation metrics:

• Pixel F1, IoU (Jaccard), AUC-ROC/PR on validation split.
• Report both per-image metrics and dataset averages.

Learning rate & schedule:

• Start with 1e-4 (head), reduce on plateau by factor 0.5–0.2.
• Warmup for a few epochs if fully fine‑tuning.

Mixed precision:

• Use torch.cuda.amp for speed on T4.

scaler = torch.cuda.amp.GradScaler()
for imgs, masks in train_loader:
    optimizer.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        with torch.no_grad():
            feats = feature_extractor(imgs.to(device))
        outputs = anomaly_detector(feats)
        loss = criterion(outputs, masks.to(device))
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Checkpoints:

• Save best model by validation IoU/F1, not just loss.

torch.save(anomaly_detector.state_dict(), "anomaly_detector_coverage.pt")

Class imbalance handling:

• If manipulations are small regions, set pos_weight in BCEWithLogitsLoss:

pos_weight = torch.tensor([5.0], device=device)  # tune
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)

Post-processing:

• Apply CRF or simple morphological ops (opening/closing) to smooth masks.
• Small connected components removal can reduce false positives.

Reproducibility:

• Set seeds and ensure deterministic behavior where possible.

import random, numpy as np, torch
def seed_all(seed=42):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
seed_all(42)

Common Gotchas (and Fixes)

“My mask looks gray, not binary”

• Ensure mask is 0/255 and then normalized to 0/1. Avoid bilinear interpolation for masks.

“Training loss won’t go down”

• Verify that the mask aligns spatially with outputs; interpolate masks to match model output size.
• Check that anomaly head outputs logits (no extra sigmoid inside the model when using BCEWithLogitsLoss).

“Overfitting fast”

• Add more augmentation, reduce LR, use weight decay, or early stopping.
• Consider a validation split and monitor generalization metrics, not just training loss.

“Predictions too blurry”

• Train longer; add Dice/Focal; try deeper/stronger augmentation; use a CRF/morphology post‑process.