| """Imager-1-Tiny - Model architecture and inference.""" |
| import os |
| import torch |
| import torch.nn as nn |
| from torchvision import models |
| from PIL import Image |
|
|
| from .preprocess import get_transform |
|
|
|
|
| class ImagerTiny(nn.Module): |
| """ResNet-50 backbone with a single-logit binary head.""" |
|
|
| def __init__(self, pretrained: bool = False): |
| super().__init__() |
| self.backbone = models.resnet50( |
| weights="IMAGENET1K_V1" if pretrained else None |
| ) |
| in_features = self.backbone.fc.in_features |
| self.backbone.fc = nn.Linear(in_features, 1) |
|
|
| def forward(self, x): |
| return self.backbone(x) |
|
|
|
|
| def load_model(weights_path: str = None, device: str = "cpu") -> ImagerTiny: |
| """Load the trained Imager-1-Tiny model.""" |
| if weights_path is None: |
| weights_path = os.path.join( |
| os.path.dirname(__file__), "..", "model", "model_imager-1-tiny.pth" |
| ) |
| model = ImagerTiny(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: ImagerTiny = None, device: str = "cpu") -> dict: |
| """ |
| Predict whether an image is Real or AI-Generated. |
| |
| Args: |
| image: PIL.Image, str path, or file-like object. |
| model: Pre-loaded ImagerTiny (optional). |
| device: "cpu" or "cuda". |
| |
| Returns: |
| dict with probability, label, confidence. |
| """ |
| 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-tiny", |
| "architecture": "resnet50", |
| "num_params": 25_000_000, |
| "input_size": 224, |
| "num_classes": 1, |
| "threshold": 0.385, |
| "labels": {0: "Real", 1: "AI-Generated"}, |
| "normalize": { |
| "mean": [0.485, 0.456, 0.406], |
| "std": [0.229, 0.224, 0.225], |
| }, |
| } |
|
|
