modify model

app.py

@@ -6,19 +6,19 @@ from torchvision import models, transforms
 from torch.utils.data import DataLoader
 from torchvision.datasets import ImageFolder
 
-vgg16 = models.vgg16(pretrained=True)
-
-# Freeze the convolutional base to prevent updating weights during training
-for param in vgg16.features.parameters():
+#define the model architecture
+model_resnet = models.resnet18(weights='IMAGENET1K_V1')
+for param in model_resnet.parameters():
     param.requires_grad = False
 
-num_features = vgg16.classifier[6].in_features
-num_classes = 3  
-vgg16.classifier[6] = torch.nn.Linear(num_features, num_classes)
+# Parameters of newly constructed modules have requires_grad=True by default
+num_ftrs = model_resnet.fc.in_features
+model_resnet.fc = nn.Linear(num_ftrs, 15)  #mengganti jumlah classifier sesuai output kelas
+
 
 # Load the model
-model = vgg16
-state_dict = torch.load('vgg16_transfer_learning.pth', map_location=torch.device('cpu'))
+model = model_resnet
+state_dict = torch.load('transfer_learning_resnet_15class.pth', map_location=torch.device('cpu'))
 model.load_state_dict(state_dict)
 model.eval()
 
@@ -26,10 +26,10 @@ model.eval()
 transform = transforms.Compose([
     transforms.Resize((224, 224)),  
     transforms.ToTensor(),  
-    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))        
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])        
 ])
 
-classes = ('broccoli', 'cabbage', 'cauliflower')
+classes = ('Bean', 'Bitter_Gourd', 'Bottle_Gourd', 'Brinjal', 'Broccoli', 'Cabbage', 'Capsicum', 'Carrot', 'Cauliflower', 'Cucumber', 'Papaya', 'Potato', 'Pumpkin', 'Radish', 'Tomato')
 
 def predict(image):
     input_tensor = transform(image)
@@ -42,7 +42,7 @@ def predict(image):
     max_value, predicted_class = torch.max(probabilities, 0)
     return classes[predicted_class.item()], max_value.item() * 100
 
-st.title('Vegetable Classification')
+st.title('Vegetable Classification for learning')
 st.write('you can upload your image of veggies below')
 
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])