Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| from torchvision import models, transforms | |
| from safetensors.torch import load_file | |
| from PIL import Image | |
| # 1. Rebuild the blank ResNet-50 architecture | |
| model = models.resnet50(weights=None) | |
| model.fc = nn.Linear(model.fc.in_features, 2) # 2 classes: Fake (0) and Real (1) | |
| # 2. Load YOUR trained weights from the safetensors file | |
| # Using map_location='cpu' ensures it works on Hugging Face's free CPU tier! | |
| state_dict = load_file("model.safetensors", device="cpu") | |
| # (Optional safety check) If you trained with DataParallel, keys might have "module." in front. | |
| # This removes it so the weights load perfectly no matter what. | |
| state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} | |
| # Inject the weights into the skeleton and set to test mode | |
| model.load_state_dict(state_dict) | |
| model.eval() | |
| # 3. Define the strict patch transform (No randomness!) | |
| test_transform = transforms.Compose([ | |
| # transforms.Resize(256), # Standard practice to resize slightly before center cropping | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| classes = ['FAKE', 'REAL'] | |
| # 4. Create the prediction function that Gradio will call | |
| def predict_image(img): | |
| if img is None: | |
| return None | |
| # Ensure image is strictly RGB (drops alpha channels from PNGs) | |
| img = img.convert('RGB') | |
| # Apply transforms and add the batch dimension | |
| img_tensor = test_transform(img).unsqueeze(0) | |
| with torch.no_grad(): | |
| # Pass through model | |
| preds = model(img_tensor) | |
| # Convert raw numbers to percentages (0.0 to 1.0) | |
| probs = torch.nn.functional.softmax(preds[0], dim=0) | |
| # Gradio expects a dictionary of { "Class Name": probability } | |
| return {classes[i]: float(probs[i]) for i in range(2)} | |
| # 5. Build and launch the Web App! | |
| description_text = """ | |
| ### How it works: | |
| This model analyzes patches of an image to determine if an image is real or AI-generated. | |
| ### Limitations for best results: | |
| * **Resolution Sweet Spot:** Works flawlessly on standard AI resolutions and mid-sized images (from **512x512 up to around 1000x1500 pixels**, like 640x832 or 880x1320). | |
| * **The 4K Danger Zone:** Ultra-high-resolution (like **3840x2160 / 4K**) images will cause the model to fail. Because the model's 'magnifying glass' is strictly fixed to a 224x224 pixel crop, it ends up looking through a pinhole at less than 0.6% of a 4K image, causing it to lose context and guess randomly. | |
| * **Centered Subjects:** The model strictly scans the dead-center of the image. If the AI artifacts or mistakes (like extra fingers or warped backgrounds) are on the far edges, the model won't see them! | |
| * **No Screenshots:** Heavy compression (like taking a screenshot or downloading from messaging apps) destroys the microscopic forensic evidence. Please upload the raw, original files. | |
| """ | |
| # Inside your interface: | |
| interface = gr.Interface( | |
| fn=predict_image, | |
| inputs=gr.Image(type="pil", label="Upload an Image"), | |
| outputs=gr.Label(num_top_classes=2, label="Prediction"), | |
| title="FakeOut: AI Image Detector", | |
| description=description_text, | |
| flagging_mode="never" | |
| ) | |
| interface.launch() |