Create inference.py

inference.py

@@ -0,0 +1,41 @@
+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 = "model.pt"  # Path to Flux-Detector
+    image_path = "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}")