Update utils.py

utils.py

@@ -1,20 +1,20 @@
 import torch
+import torch.nn as nn
 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)
+class ResNet18Classifier(nn.Module):
+    def __init__(self, num_classes=3):
+        super().__init__()
+        self.resnet = resnet18(weights=None)  # modern way
+        self.resnet.fc = nn.Linear(self.resnet.fc.in_features, num_classes)
 
     def forward(self, x):
-        return self.model(x)
+        return self.resnet(x)
 
 def load_model(model_path="model/best_classification_model.pth", num_classes=3):
-    model = ResNet18Classifier(num_classes=num_classes, pretrained=False)
+    model = ResNet18Classifier(num_classes=num_classes)
     state_dict = torch.load(model_path, map_location='cpu')
     model.load_state_dict(state_dict)
     model.eval()
@@ -34,4 +34,4 @@ def predict_image(image_path, model, class_names):
         outputs = model(image_tensor)
         _, predicted = torch.max(outputs, 1)
 
-    return class_names[predicted.item()]
+    return class_names[predicted.item()]