test_image.py

import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


model = models.resnet18(pretrained=False)
model.fc = nn.Linear(model.fc.in_features, 2)
model = model.to(device)


model.load_state_dict(torch.load("../models/image_model.pth", map_location=device))
model.eval()


transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])


image_path = r"C:\Users\Isha\.ipython\ai_generated_content_detector\data\images\ai\IMG-20240917-WA0010.jpg"

image = Image.open(image_path).convert("RGB")
image = transform(image).unsqueeze(0).to(device)


import torch.nn.functional as F

with torch.no_grad():
    output = model(image)
    probabilities = F.softmax(output, dim=1)
    confidence, predicted = torch.max(probabilities, 1)

classes = ["AI Generated", "Real"]

print("Prediction:", classes[predicted.item()])
print("Confidence:", round(confidence.item() * 100, 2), "%")