| """Imager-1-Pro - Model architecture and inference (ConvNeXt-Base).""" |
| import os |
| import torch |
| import torch.nn as nn |
| from torchvision import models |
| from PIL import Image |
|
|
| from .preprocess import get_transform |
|
|
|
|
| class ImagerPro(nn.Module): |
| """ConvNeXt-Base backbone with a single-logit binary head.""" |
|
|
| def __init__(self, pretrained: bool = False): |
| super().__init__() |
| self.backbone = models.convnext_base( |
| weights="IMAGENET1K_V1" if pretrained else None |
| ) |
| |
| in_features = self.backbone.classifier[2].in_features |
| self.backbone.classifier[2] = nn.Linear(in_features, 1) |
|
|
| def forward(self, x): |
| return self.backbone(x) |
|
|
|
|
| def load_model(weights_path: str = None, device: str = "cpu") -> ImagerPro: |
| """Load the trained Imager-1-Pro model.""" |
| if weights_path is None: |
| weights_path = os.path.join( |
| os.path.dirname(__file__), "..", "model", "model_imager-1-pro.pth" |
| ) |
| model = ImagerPro(pretrained=False) |
| state = torch.load(weights_path, map_location=device) |
| model.load_state_dict(state) |
| model.to(device) |
| model.eval() |
| return model |
|
|
|
|
| def predict(image, model: ImagerPro = None, device: str = "cpu") -> dict: |
| """Predict whether an image is Real or AI-Generated.""" |
| if model is None: |
| model = load_model(device=device) |
|
|
| transform = get_transform(train=False) |
|
|
| if isinstance(image, str): |
| image = Image.open(image) |
| elif hasattr(image, "read"): |
| image = Image.open(image) |
| image = image.convert("RGB") |
|
|
| tensor = transform(image).unsqueeze(0).to(device) |
|
|
| with torch.no_grad(): |
| prob = torch.sigmoid(model(tensor)).item() |
|
|
| label = "AI-Generated" if prob > MODEL_CONFIG["threshold"] else "Real" |
| return { |
| "probability": round(prob, 4), |
| "label": label, |
| "confidence": round(prob if label == "AI-Generated" else 1 - prob, 4), |
| } |
|
|
|
|
| MODEL_CONFIG = { |
| "name": "imager-1-pro", |
| "architecture": "convnext_base", |
| "num_params": 89_000_000, |
| "input_size": 224, |
| "num_classes": 1, |
| "threshold": 0.5, |
| "labels": {0: "Real", 1: "AI-Generated"}, |
| "normalize": { |
| "mean": [0.485, 0.456, 0.406], |
| "std": [0.229, 0.224, 0.225], |
| }, |
| } |
|
|
