Spaces:
Runtime error
Runtime error
File size: 2,754 Bytes
e0f2d0e |
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 |
import torch
import timm
from torchvision import transforms
from PIL import Image
from typing import Any
from schemas.vision_schemas import AIImageDetector
class ENetAIImageDetector(AIImageDetector):
"""
EfficientNet-B3 AI Image Detector that classifies whether an image
is AI-generated or real using a pre-trained PyTorch model.
Attributes:
model_path (str): Path to the trained model file (.pt).
model (Any): Loaded PyTorch model.
device (str): Device to run inference on ('cuda' or 'cpu').
"""
def __init__(self, model_path: str = "./models/efficientnet_b3_full_ai_image_classifier.pt"):
"""
Initialize the ENetAIImageDetector.
Args:
model_path (str, optional): Path to the trained EfficientNet model.
"""
self.model_path = model_path
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = self._load_model()
self.transform = self._build_transform()
def _load_model(self) -> Any:
"""Load the trained EfficientNet-B3 model."""
if self.model_path.endswith(".pt"):
model = torch.load(self.model_path, map_location=self.device, weights_only=False)
else:
model = timm.create_model("efficientnet_b3", pretrained=False, num_classes=1)
model.load_state_dict(torch.load(self.model_path, map_location=self.device))
model.to(self.device)
model.eval()
return model
def _build_transform(self) -> Any:
"""Return preprocessing pipeline for input images."""
return transforms.Compose([
transforms.Resize((300, 300)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
def _preprocess_image(self, image: Image.Image) -> torch.Tensor:
"""Convert a PIL Image to a normalized tensor."""
return self.transform(image).unsqueeze(0).to(self.device)
def detect(self, image: Image.Image) -> bool:
"""
Detect whether a given PIL image is AI-generated.
Args:
image (PIL.Image.Image): The input image.
Returns:
bool: True if AI-generated, False if real.
"""
if not isinstance(image, Image.Image):
raise TypeError("Input must be a PIL.Image.Image object.")
img_tensor = self._preprocess_image(image)
with torch.no_grad():
outputs = self.model(img_tensor)
prob = torch.sigmoid(outputs).item()
is_ai_generated = prob >= 0.001
return is_ai_generated
|