Update app.py

app.py

@@ -17,19 +17,19 @@ effecientnet.classifier = nn.Sequential(
     nn.Linear(in_features, 512),
     nn.ReLU(),
     nn.Dropout(0.5),
-    nn.Linear(512, 100))
+    nn.Linear(512, 100)
+)
 
 # 🔹 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.to(device)
-model.eval()
+effecientnet.load_state_dict(torch.load("best_model (1).pth", map_location=device))
+effecientnet.to(device)
+effecientnet.eval()
 
 # 🔹 Image preprocessing (should match validation transforms)
 val_transforms = transforms.Compose([
-    transforms.Lambda(lambda x: x.convert('RGB')),
-    transforms.Resize((224,224)), # Resize to a larger size first
+    transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.225, 0.224])
 ])
 
 # 🔹 Correct class names for 100 fruits
@@ -50,12 +50,11 @@ class_names = [
     "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)
     with torch.no_grad():
-        outputs = model(img)
+        outputs = effecientnet(img)
         probs = torch.nn.functional.softmax(outputs, dim=1)
         top5 = torch.topk(probs[0], 5)
         return {class_names[i]: float(top5.values[j]) for j, i in enumerate(top5.indices)}