UserPollo/moire-pattern-detector

Dual-Branch DINOv2 Moiré & Screen Recapture Detector

A robust digital forensics classifier built on DINOv2 (with registers) designed to detect screen recaptures and Moiré patterns.

Optimized for anti-spoofing pipelines and automated asset valuation platforms, this model overcomes the traditional scale dilemma in Moiré detection by combining an un-resized native crop (capturing high-frequency pixel interference) with a global thumbnail (capturing screen-wide periodic banding).

---

📁 Repository Files

• best_screen_detector_backbone.pt: Weights for the fine-tuned top transformer blocks of the DINOv2 backbone.
• best_screen_detector_mlp.pt: Weights for the 2-layer classification MLP head.
• classes.json: Class index mapping (0: "gt", 1: "moire").

---

🚀 Quick Start & Inference

1. Requirements

If you are running this in Google Colab, you do not need to run pip install for most of these packages, as PyTorch, Transformers, and Pillow are pre-installed. You only need to ensure huggingface-hub is up to date.

For local environments, install the dependencies:

pip install torch torchvision transformers pillow huggingface-hub requests

2. Inference Script

import torch
import torch.nn as nn
import json
import requests
from PIL import Image
from torchvision import transforms
from transformers import AutoModel
from huggingface_hub import hf_hub_download

# 1. Download weights and classes from the Hub
repo_id = "UserPollo/moire-pattern-detector"
backbone_ckpt = hf_hub_download(repo_id=repo_id, filename="best_screen_detector_backbone.pt")
mlp_ckpt = hf_hub_download(repo_id=repo_id, filename="best_screen_detector_mlp.pt")
classes_file = hf_hub_download(repo_id=repo_id, filename="classes.json")

with open(classes_file, "r") as f:
    classes = json.load(f)

# 2. Define the custom MLP Head
class ScreenDetectorMLP(nn.Module):
    def __init__(self, input_size=3072, hidden_size=256, num_classes=2, dropout=0.3):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.GELU(),
            nn.BatchNorm1d(hidden_size),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, num_classes),
        )

    def forward(self, x):
        return self.mlp(x)

# 3. Load Models and apply weights
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

backbone = AutoModel.from_pretrained("facebook/dinov2-with-registers-base").to(device)
backbone.eval()
for p in backbone.parameters():
    p.requires_grad_(False)

# Load fine-tuned weights into the last 2 blocks of the backbone
total_layers = len(backbone.encoder.layer)
unfrozen_state = torch.load(backbone_ckpt, map_location=device, weights_only=True)
for i, layer in enumerate(backbone.encoder.layer[total_layers - 2:]):
    layer.load_state_dict(unfrozen_state[f"layer.{total_layers - 2 + i}"])

head = ScreenDetectorMLP(input_size=3072).to(device)
head.load_state_dict(torch.load(mlp_ckpt, map_location=device, weights_only=True))
head.eval()

# 4. Prepare Dual-Branch Image Transforms
IMAGENET_MEAN, IMAGENET_STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]

local_transform = transforms.Compose([
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

global_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

# 5. Load Image
url = "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcS2pnVr5QT5OjVf8H4YtrJfRoPvGBFAG5pBG8F-LfPGB0sEcyF1JkT6Okly&s=10"
img = Image.open(requests.get(url, stream=True).raw).convert("RGB")

local_tensor = local_transform(img).unsqueeze(0).to(device)
global_tensor = global_transform(img).unsqueeze(0).to(device)

# Concatenate for a single forward pass
batch = torch.cat([local_tensor, global_tensor], dim=0)

# 6. Run Inference
with torch.no_grad():
    out = backbone(pixel_values=batch)
    hidden = out.last_hidden_state.float()

    # Extract CLS and patch mean (ignoring register tokens)
    n_reg = getattr(backbone.config, "num_register_tokens", 0)
    cls_tok = hidden[:, 0, :]
    patch_mean = hidden[:, 1 + n_reg:, :].mean(dim=1)
    feat = torch.cat([cls_tok, patch_mean], dim=-1)

    # Split back into local and global, then concatenate horizontally
    local_feat, global_feat = feat[0:1], feat[1:2]
    combined_feat = torch.cat([local_feat, global_feat], dim=-1)

    # Classify
    logits = head(combined_feat)
    probs = torch.softmax(logits, dim=1)
    conf, pred = probs.max(dim=1)

print(f"Prediction: {classes[str(pred.item())]} (Confidence: {conf.item()*100:.1f}%)")