update app.py (fruitsvegs0 -> freshvision0 b2)

app.py

@@ -16,28 +16,35 @@ manual_transform = transforms.Compose([
     normalize
 ])
 
-class_names = ['Fresh Banana',
-  'Fresh Lemon',
-  'Fresh Lulo',
-  'Fresh Mango',
-  'Fresh Orange',
-  'Fresh Strawberry',
-  'Fresh Tamarillo',
-  'Fresh Tomato',
-  'Spoiled Banana',
-  'Spoiled Lemon',
-  'Spoiled Lulo',
-  'Spoiled Mango',
-  'Spoiled Orange',
-  'Spoiled Strawberry',
-  'Spoiled Tamarillo',
-  'Spoiled Tomato']
-
-model_0 = mb.create_model_baseline_effnetb0(out_feats=len(class_names), device=device)
-model_0.load_state_dict(torch.load(f="models/effnetb0_fruitsvegs0_5_epochs.pt", map_location="cpu"))
+# class_names = ['Fresh Banana',
+#   'Fresh Lemon',
+#   'Fresh Lulo',
+#   'Fresh Mango',
+#   'Fresh Orange',
+#   'Fresh Strawberry',
+#   'Fresh Tamarillo',
+#   'Fresh Tomato',
+#   'Spoiled Banana',
+#   'Spoiled Lemon',
+#   'Spoiled Lulo',
+#   'Spoiled Mango',
+#   'Spoiled Orange',
+#   'Spoiled Strawberry',
+#   'Spoiled Tamarillo',
+#   'Spoiled Tomato']
+
+class_names = ['Fresh Apple',
+ 'Fresh Banana',
+ 'Fresh Orange',
+ 'Rotten Apple',
+ 'Rotten Banana',
+ 'Rotten Orange']
+
+model = mb.create_model_baseline_effnetb2(out_feats=len(class_names), device=device)
+model.load_state_dict(torch.load(f="models/effnetb2_freshvisionv0_10_epochs.pt", map_location="cpu"))
 
 def pred(img):
-    model_0.eval()
+    model.eval()
     transformed = manual_transform(img).to(device)
     with torch.inference_mode():
         logits = model_0(transformed.unsqueeze(dim=0))