File size: 5,182 Bytes
fc666d4 26d751c 151f8a5 26d751c 151f8a5 26d751c 151f8a5 fc666d4 26d751c 151f8a5 26d751c 151f8a5 26d751c 151f8a5 b6e5890 151f8a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | ---
license: apache-2.0
tags:
- vision
- image-classification
- dinov2
- moire-detection
- screen-detection
- anti-spoofing
- digital-forensics
datasets:
- soumikrakshit/uhdm-dataset
pipeline_tag: image-classification
base_model:
- facebook/dinov2-with-registers-base
---
# 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:
```bash
pip install torch torchvision transformers pillow huggingface-hub requests
```
### 2. Inference Script
```python
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[int(pred.item())]} (Confidence: {conf.item()*100:.1f}%)")
``` |