import argparse
import os
import warnings
from transformers import SwinForImageClassification, AutoImageProcessor
from PIL import Image
import torch

warnings.filterwarnings("ignore")
os.environ["TRANSFORMERS_VERBOSITY"] = "error"

def classify_image(image_path):
    model = SwinForImageClassification.from_pretrained("./")
    processor = AutoImageProcessor.from_pretrained("./", use_fast=False)

    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class = torch.argmax(logits, dim=1).item()

    label = model.config.id2label[predicted_class]
    return "Real" if label == "human" else "Fake"

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Detecta se uma imagem é real ou gerada por IA.")
    parser.add_argument("--image", type=str, required=True, help="Caminho para a imagem a ser analisada.")
    
    args = parser.parse_args()
    result = classify_image(args.image)
    print(result)