Create inference.py

inference.py

@@ -0,0 +1,28 @@
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+import torchvision.models as models
+
+# Load model
+model = models.resnet18(pretrained=False, num_classes=3)
+model.load_state_dict(torch.load("pytorch_model.bin", map_location=torch.device("cpu")))
+model.eval()
+
+# Preprocessing function
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+# Labels
+labels = ["A", "B", "C", "D", "E", "F", "G"]
+
+# Required function
+def predict(image: Image.Image):
+    img_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = torch.nn.functional.softmax(outputs[0], dim=0)
+    return {labels[i]: float(probs[i]) for i in range(len(labels))}