LukasT9/flux-detector

📖 Translations

• 🇺🇦 Українська

Flux-Detector: AI vs Real Image Classifier

This repository hosts flux-detector, a binary image classifier implemented using the EfficientNet B2 architecture.

Trained on the Flux-1 dataset, the model is optimized for the task of differentiating between "Real" (photographic) and "AI" (Flux-generated) images.

Note: The detector is fine-tuned to work with Flux-generated images, and its performance on images from other AI generators is currently unknown.

flux-detector demonstrates 90% accuracy on a dedicated test dataset.

## Inference Code (inference.py)
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image

def load_model(model_path, device):
    """Loads the TorchScript model."""
    model = torch.jit.load(model_path, map_location=device)
    model.to(device).eval()
    return model

def preprocess_image(image_path):
    """Pre-processes the image for feeding into the model."""
    IMG_SIZE = 1024
    transform = transforms.Compose([
        transforms.Resize(IMG_SIZE + 32),
        transforms.CenterCrop(IMG_SIZE),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    img = Image.open(image_path).convert("RGB")
    return transform(img).unsqueeze(0)

def predict(model, image_tensor, device, threshold=0.5):
    """Performs model prediction."""
    with torch.no_grad():
        outputs = model(image_tensor.to(device))
        prob = torch.sigmoid(outputs).item()
    label = "Real" if prob >= threshold else "AI"
    return prob, label

if __name__ == "__main__":
    model_path = r"model.pt"  # Path to Flux-Detector
    image_path = r"test_image.jpg" # Path to test image

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = load_model(model_path, device)
    image_tensor = preprocess_image(image_path)
    prob, label = predict(model, image_tensor, device)

    print(f"Model Prediction: {prob:.4f} -> {label}")

# Example output: Model Prediction: 0.0387 -> AI

Code Features:

• TorchScript Model Loading: Efficiently loads a pre-trained TorchScript model (model.pt) for fast inference.
• Image Preprocessing: Includes necessary image preprocessing steps like resizing, center cropping to 1024x1024, converting to tensor, and normalization using ImageNet statistics.
• Binary Classification: Predicts whether an input image is "Real" or "AI-generated" based on a confidence threshold (default 0.5).
• Clear and Concise: Provides a straightforward example of how to load the model, preprocess an image, and run inference.

Requirements for Running:

• Python Libraries: Ensure you have the following Python libraries installed: torch, torchvision, and Pillow. You can install them using pip: pip install torch torchvision Pillow
• Python Environment: Requires Python 3.6 or higher.
• PyTorch: Install PyTorch and Torchvision following the instructions on the official PyTorch website. Make sure to install the version compatible with your CUDA setup if you intend to use a GPU.
• Pillow (PIL): Install Pillow for image processing: pip install Pillow (This is redundant as it's already covered in "Python Libraries", but kept for clarity if user only reads this specific point).
• Model File (model.pt): Ensure you have the model.pt file (the pre-trained flux-detector model weights) in the same directory as the inference.py script, or provide the correct path in the model_path variable.
• Test Image (test_image.png): Replace "test_image.png" with the path to your own test image, or use the test image named test_image.png in the same directory.

Security Note:

When loading this flux-detector model, you might encounter a "Detected Pickle imports" warning. This is a standard warning in PyTorch for TorchScript models that use Python's pickle module for serialization.

As the model author, I believe this model is safe to load and use. It was trained using publicly available datasets and standard PyTorch code. However, users should always exercise caution when loading and running ANY machine learning models from the internet.

Best Practices for Users:

• Review the provided inference code (inference.py) to understand how the model is loaded and used.
• Download models only from trusted sources.
• Be aware of the general security risks associated with deserializing data using pickle in Python.

Disclaimer: This model is provided "as is" without warranty of any kind. The author is not responsible for any security issues or damages arising from the use of this model. Use at your own risk.