import torch
from PIL import Image
from torchvision import transforms
from model import model as Model  # import your model class

def predict_image(image_path, model_path='./models/Network_best.pth'):
    # Select device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Load model
    net = Model(pretrain=False).to(device)
    net.load_state_dict(torch.load(model_path, map_location=device))
    net.eval()

    # Define transforms (resize and normalize)
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])

    # Load and preprocess image
    img = Image.open(image_path).convert('RGB')
    x = transform(img).unsqueeze(0).to(device)

    # Make prediction
    with torch.no_grad():
        prob = torch.sigmoid(net(x)).item()

    # Show result
    result = "REAL" if prob > 0.5 else "AI-GENERATED"
    print(f"🖼️ Image: {image_path}")
    print(f"🔍 Prediction: {result}")
    print(f"📊 Confidence: {prob:.3f}")

# Example test
if __name__ == "__main__":
    # give path to any image you want to test
    image_path = "sample.jpg"   # <-- put your image file here
    predict_image(image_path)