Update model.py

model.py

@@ -38,22 +38,31 @@ LABELS = {}
 def load_models():
     for file in os.listdir(MODELS_DIR):
         name, ext = os.path.splitext(file)
+
+        # Normalize crop key
+        crop_name = name.replace("_model", "").lower()
+
         model_path = os.path.join(MODELS_DIR, file)
+        label_path = os.path.join(LABELS_DIR, f"{crop_name}_labels.json")
+
+        if not os.path.exists(label_path):
+            raise FileNotFoundError(
+                f"Label file missing for {crop_name}: {label_path}"
+            )
 
-        # Load labels
-        with open(os.path.join(LABELS_DIR, f"{name}.json")) as f:
-            LABELS[name] = json.load(f)
+        with open(label_path) as f:
+            LABELS[crop_name] = json.load(f)
 
         if ext == ".pth":
-            num_classes = len(LABELS[name])
+            num_classes = len(LABELS[crop_name])
             model = models.resnet18(weights=None)
             model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
             model.load_state_dict(torch.load(model_path, map_location="cpu"))
             model.eval()
-            PYTORCH_MODELS[name] = model
+            PYTORCH_MODELS[crop_name] = model
 
-        elif ext == ".keras":
-            KERAS_MODELS[name] = tf.keras.models.load_model(model_path)
+        elif ext in [".keras", ".h5"]:
+            KERAS_MODELS[crop_name] = tf.keras.models.load_model(model_path)
 
 # Load once
 load_models()