Create utils.py

utils.py

@@ -0,0 +1,37 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+from torchvision.models import resnet18
+
+# Same as your original custom classifier
+class ResNet18Classifier(torch.nn.Module):
+    def __init__(self, num_classes=3, pretrained=False):
+        super(ResNet18Classifier, self).__init__()
+        self.model = resnet18(pretrained=pretrained)
+        self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+def load_model(model_path="model/best_classification_model.pth", num_classes=3):
+    model = ResNet18Classifier(num_classes=num_classes, pretrained=False)
+    state_dict = torch.load(model_path, map_location='cpu')
+    model.load_state_dict(state_dict)
+    model.eval()
+    return model
+
+def predict_image(image_path, model, class_names):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225])
+    ])
+    image = Image.open(image_path).convert('RGB')
+    image_tensor = transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        _, predicted = torch.max(outputs, 1)
+
+    return class_names[predicted.item()]