bug fix

app.py

@@ -16,7 +16,7 @@ with open(food_vision_class_names_path, "r") as f:
     class_names = f.read().splitlines()
 
 # Specify number of classes
-num_classes = len(class_names) - 1 # "unknown class to be removed to properly load the model"
+NUM_CLASSES = len(class_names) - 1 # "unknown class to be removed to properly load the model"
 
 # Load the food description file
 food_descriptions_json = "food_descriptions.json"
@@ -28,7 +28,7 @@ classification_model_name_path = "effnetb0_classif_epoch13.pth"
 effnetb0_model, _ = create_effnetb0(
     model_weights_dir=".",
     model_weights_name=classification_model_name_path,
-    num_classes=num_classes
+    num_classes=2
     )
 
 # Load the ViT-Base transformer
@@ -38,7 +38,7 @@ vitbase_model = create_vitbase_model(
     model_weights_dir=".",
     model_weights_name=food_vision_model_name_path,
     img_size=IMG_SIZE,
-    num_classes=num_classes
+    num_classes=NUM_CLASSES
 )
 
 # Specify manual transforms
@@ -76,7 +76,7 @@ def predict(img) -> Tuple[Dict, float]:
         else:
 
             # Set all probabilites to zero
-            pred_probs = torch.tensor([[0.0] * num_classes])
+            pred_probs = torch.tensor([[0.0] * NUM_CLASSES])
     
     # Calculate entropy
     entropy = -torch.sum(pred_probs * torch.log(pred_probs), dim=1).item()

model.py

@@ -83,7 +83,7 @@ def create_vitbase_model(
 def create_effnetb0(
         model_weights_dir: Path,
         model_weights_name: str,
-        num_classes: int=101,
+        num_classes: int=2,
         dropout: float=0.2
         ):
     """Creates an EfficientNetB0 feature extractor model and transforms.