File size: 2,600 Bytes
0212735 | 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 | """
OurNet Model Definition with ConvNeXt backbone
This model is designed for image forgery detection using a ConvNeXt backbone
with dual projection heads and a detection head.
"""
import torch
import torch.nn as nn
import timm
class OurNet(nn.Module):
def __init__(self, config=None):
super().__init__()
# Load config if provided
if config is None:
backbone_name = "convnext_base"
n_features = 1024
else:
backbone_name = config.get("backbone", {}).get("name", "convnext_base")
n_features = config.get("backbone", {}).get("n_features", 1024)
self.backbone = timm.create_model(backbone_name, pretrained=False)
# Remove classification head
if hasattr(self.backbone, "head"):
self.n_features = self.backbone.head.in_features
self.backbone.head = nn.Identity()
elif hasattr(self.backbone, "fc"):
self.n_features = self.backbone.fc.in_features
self.backbone.fc = nn.Identity()
elif hasattr(self.backbone, "classifier"):
self.n_features = self.backbone.classifier.in_features
self.backbone.classifier = nn.Identity()
else:
raise ValueError("Unsupported backbone architecture")
# Projection heads
self.aux_fc1 = nn.Sequential(
nn.Linear(self.n_features, self.n_features),
nn.ReLU(),
nn.Linear(self.n_features, 128),
)
self.aux_fc2 = nn.Sequential(
nn.Linear(self.n_features, self.n_features),
nn.ReLU(),
nn.Linear(self.n_features, 128),
)
# Detection heads
self.det_fc1 = nn.Sequential(
nn.Linear(self.n_features, self.n_features),
nn.ReLU(),
nn.Linear(self.n_features, 128),
)
self.det_fc2 = nn.Sequential(
nn.Linear(self.n_features, 256),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(256, 1),
)
def forward_det(self, x):
"""Detection forward pass"""
feats = self.backbone.forward_features(x)
feats = feats.mean([-2, -1])
homo_head = self.det_fc1(feats)
det_head = self.det_fc2(feats)
return homo_head, det_head
def forward_proj(self, x):
"""Projection forward pass"""
feats = self.backbone.forward_features(x)
feats = feats.mean([-2, -1])
heter_head = self.aux_fc1(feats)
homo_head = self.aux_fc2(feats)
return heter_head, homo_head
|