Update app.py

app.py

@@ -10,39 +10,47 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # 🔹 Load EfficientNet-B0 and modify classifier for 100 classes
 model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)
-in_features = model.classifier[1].in_features
-model.classifier[1] = nn.Linear(in_features, 100)  # ✅ 100 classes now
+in_features = effecientnet.classifier[1].in_features
+
+# Replace the classifier with a new linear layer (for 100 classes, for example)
+effecientnet.classifier = nn.Sequential(
+    nn.Linear(in_features, 512),
+    nn.ReLU(),
+    nn.Dropout(0.5),
+    nn.Linear(512, 101))
 
 # 🔹 Load trained weights (make sure the model was trained for 100 classes!)
-model.load_state_dict(torch.load("best_model (1).pth", map_location=device))
+model.load_state_dict(torch.load("best_model (10).pth", map_location=device))
 model.to(device)
 model.eval()
 
 # 🔹 Image preprocessing (should match validation transforms)
 val_transforms = transforms.Compose([
-    transforms.Resize((224, 224)),
+    transforms.Lambda(lambda x: x.convert('RGB')),
+    transforms.Resize((224,224)), # Resize to a larger size first
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225])
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 ])
 
 # 🔹 Correct class names for 100 fruits
 class_names = [
-    "oil_palm", "hog_plum", "apricot", "sugar_apple", "rambutan", "yellow_plum", "pineapple", "dewberry",
-    "plumcot", "ambarella", "grenadilla", "elderberry", "barbadine", "horned_melon", "chico", "mock_strawberry",
-    "cashew", "finger_lime", "apple", "eggplant", "pawpaw", "chokeberry", "jamaica_cherry", "grapefruit",
-    "santol", "chenet", "salak", "indian_strawberry", "cherimoya", "otaheite_apple", "dragonfruit", "raspberry",
-    "mangosteen", "yali_pear", "taxus_baccata", "coconut", "guava", "black_mullberry", "durian", "ackee",
-    "olive", "mandarine", "black_berry", "acerola", "jaboticaba", "fig", "langsat", "redcurrant", "gooseberry",
-    "camu_camu", "barberry", "rose_hip", "jalapeno", "brazil_nut", "damson", "acai", "prikly_pear", "morinda",
-    "sea_buckthorn", "avocado", "strawberry_guava", "jackfruit", "greengage", "cupuacu", "longan",
-    "passion_fruit", "feijoa", "betel_nut", "kaffir_lime", "sapodilla", "cempedak", "hawthorn", "mango",
-    "malay_apple", "cranberry", "jocote", "cluster_fig", "corn_kernel", "kumquat", "rose_leaf_bramble",
-    "jujube", "grape", "pea", "papaya", "bitter_gourd", "ugli_fruit", "jambul", "mabolo", "abiu", "quince",
-    "custard_apple", "medlar", "mountain_soursop", "banana", "goumi", "hard_kiwi", "pomegranate",
-    "white_currant", "lablab", "emblic"
+    "abiu", "acai", "acerola", "ackee", "ambarella", "apple", "apricot", "avocado", "banana",
+    "barbadine", "barberry", "betel_nut", "bitter_gourd", "black_berry", "black_mullberry",
+    "brazil_nut", "camu_camu", "cashew", "cempedak", "chenet", "cherimoya", "chico", "chokeberry",
+    "cluster_fig", "coconut", "corn_kernel", "cranberry", "cupuaçu", "custard_apple", "damson",
+    "dewberry", "dragonfruit", "durian", "eggplant", "elderberry", "emblic", "feijoa", "fig",
+    "finger_lime", "gooseberry", "goumi", "grape", "grapefruit", "greengage", "grenadilla", "guava",
+    "hard_kiwi", "hawthorn", "hog_plum", "horned_melon", "indian_strawberry", "jaboticaba",
+    "jackfruit", "jalapeno", "jamaica_cherry", "jambul", "jocote", "jujube", "kaffir_lime",
+    "kumquat", "lablab", "langsat", "longan", "mabolo", "malay_apple", "mandarine", "mango",
+    "mangosteen", "medlar", "mock_strawberry", "morinda", "mountain_soursop", "oil_palm", "olive",
+    "otahiete_apple", "papaya", "passion_fruit", "pawpaw", "pea", "pineapple", "plumcot",
+    "pomegranate", "prickly_pear", "quince", "rambutan", "raspberry", "redcurrant", "rose_hip",
+    "rose_leaf_bramble", "salak", "santol", "sapodilla", "sea_buckthorn", "strawberry_guava",
+    "sugar_apple", "taxus_baccata", "ugli_fruit", "white_currant", "yali_pear", "yellow_plum"
 ]
 
+
 # 🔹 Prediction function
 def classify_image(img):
     img = val_transforms(img).unsqueeze(0).to(device)