Update app.py

app.py

@@ -21,15 +21,13 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 num_classes = 6
 
 # Load the pre-trained ResNet model
-model = models.resnet152(pretrained=True)
+model = models.resnet18(pretrained=True)
 for param in model.parameters():
     param.requires_grad = False  # Freeze feature extractor
 
 # Modify the classifier for 6 classes with an additional hidden layer
 model.fc = nn.Sequential(
-    nn.Linear(model.fc.in_features, 512),
-    nn.ReLU(),
-    nn.Linear(512, num_classes)
+    nn.Linear(model.fc.in_features, num_classes)
 )
 
 # Load trained weights
@@ -43,7 +41,7 @@ class_labels = ['bird', 'cat', 'deer', 'dog', 'frog', 'horse']
 def transform_image(image):
     """Preprocess the input image."""
     transform = transforms.Compose([
-        transforms.Resize((32, 32)),
+        transforms.Resize(224),
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ])