ash12321/sdxl-detector-resnet50

1#!/usr/bin/env python32"""3Example: Using SDXL Detector from HuggingFace4==============================================5 6Simple example showing how to use the SDXL detector7to classify images as real or SDXL-generated.8"""9 10import torch11from torchvision import transforms12from PIL import Image13from huggingface_hub import hf_hub_download14import torch.nn as nn15import torchvision.models as models16 17# ============================================================================18# MODEL DEFINITION19# ============================================================================20 21class SDXLDetector(nn.Module):22    """ResNet-50 based SDXL detector"""23    24    def __init__(self):25        super().__init__()26        self.backbone = models.resnet50(pretrained=False)27        num_features = self.backbone.fc.in_features28        29        self.backbone.fc = nn.Sequential(30            nn.Dropout(p=0.3),31            nn.Linear(num_features, 512),32            nn.BatchNorm1d(512),33            nn.ReLU(inplace=True),34            nn.Dropout(p=0.15),35            nn.Linear(512, 2)36        )37    38    def forward(self, x):39        return self.backbone(x)40 41# ============================================================================42# LOAD MODEL43# ============================================================================44 45def load_model(device='cpu'):46    """Load model from HuggingFace Hub"""47    48    # Download checkpoint49    model_path = hf_hub_download(50        repo_id="ash12321/sdxl-detector-resnet50",51        filename="best.pth"52    )53    54    # Load checkpoint55    checkpoint = torch.load(model_path, map_location=device)56    57    # Create model and load weights58    model = SDXLDetector()59    model.load_state_dict(checkpoint['model_state_dict'])60    model.to(device)61    model.eval()62    63    print(f"✅ Model loaded from {model_path}")64    print(f"   Trained for {checkpoint['epoch'] + 1} epochs")65    print(f"   Best validation accuracy: {checkpoint['best_val_acc']:.2f}%")66    67    return model68 69# ============================================================================70# PREPROCESSING71# ============================================================================72 73def get_transform():74    """Get image preprocessing transform"""75    return transforms.Compose([76        transforms.Resize(256),77        transforms.CenterCrop(224),78        transforms.ToTensor(),79        transforms.Normalize(80            mean=[0.485, 0.456, 0.406],81            std=[0.229, 0.224, 0.225]82        )83    ])84 85# ============================================================================86# PREDICTION87# ============================================================================88 89def predict_image(model, image_path, device='cpu'):90    """91    Predict if an image is real or SDXL-generated92    93    Args:94        model: Loaded SDXLDetector model95        image_path: Path to image file96        device: Device to run inference on97    98    Returns:99        dict with prediction, confidence, and probabilities100    """101    102    # Load and preprocess image103    image = Image.open(image_path).convert('RGB')104    transform = get_transform()105    input_tensor = transform(image).unsqueeze(0).to(device)106    107    # Predict108    with torch.no_grad():109        outputs = model(input_tensor)110        probs = torch.softmax(outputs, dim=1)111        prediction = torch.argmax(probs, dim=1).item()112        confidence = probs[0][prediction].item()113    114    # Format results115    labels = ['Real', 'SDXL-generated']116    117    return {118        'prediction': labels[prediction],119        'confidence': confidence,120        'probabilities': {121            'real': probs[0][0].item(),122            'sdxl': probs[0][1].item()123        }124    }125 126# ============================================================================127# MAIN128# ============================================================================129 130def main():131    """Example usage"""132    133    # Setup134    device = 'cuda' if torch.cuda.is_available() else 'cpu'135    print(f"Using device: {device}")136    137    # Load model138    model = load_model(device)139    140    # Example prediction141    image_path = "test_image.jpg"  # Replace with your image142    143    result = predict_image(model, image_path, device)144    145    print(f"\n📊 Results for {image_path}:")146    print(f"   Prediction: {result['prediction']}")147    print(f"   Confidence: {result['confidence']*100:.2f}%")148    print(f"   \nProbabilities:")149    print(f"      Real: {result['probabilities']['real']*100:.2f}%")150    print(f"      SDXL: {result['probabilities']['sdxl']*100:.2f}%")151 152if __name__ == "__main__":153    main()154