prog/image.py

import torch
from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image

MODEL_NAME = "simma7/image_model"

device = "cpu"

model = SiglipForImageClassification.from_pretrained(MODEL_NAME).to(device)
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

model.eval()

def detect_image(image_path):
    try:
        image = Image.open(image_path).convert("RGB")

        inputs = processor(images=image, return_tensors="pt").to(device)

        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]

        labels = model.config.id2label

        real_score = 0.0
        fake_score = 0.0

        for i in range(len(probs)):
            label = str(labels.get(i)).lower()
            score = float(probs[i])

            if "deepfake" in label:
                fake_score += score
            elif "real" in label or "ai" in label or "artificial" in label:
                real_score += score
            else:
                fake_score += score

        total = real_score + fake_score
        if total == 0:
            return "Prediction failed"

        real_score /= total
        fake_score /= total

        prediction = "REAL" if real_score > fake_score else "FAKE"

        return f"""
IMAGE RESULT

Prediction: {prediction}

Real: {real_score:.2%}
Fake: {fake_score:.2%}
"""

    except Exception as e:
        return f"Error: {str(e)}"