Update model.py

model.py

@@ -24,16 +24,6 @@ class FixedDropout(tf.keras.layers.Dropout):
 class EfficientNetB3(tf.keras.Model):
     pass
 
-# MobileNetV2
-from tensorflow.keras.applications import MobileNetV2
-
-# ================= INPUT SIZE PER MODEL =================
-
-KERAS_INPUT_SIZES = {
-    "corn": 300,
-    "bean": 224,
-}
-
 # ================= IMAGE PREPROCESS =================
 
 def preprocess_pytorch(img, size=224):
@@ -48,14 +38,12 @@ def preprocess_pytorch(img, size=224):
     ])
     return transform(img).unsqueeze(0)
 
-def preprocess_keras(img, crop_name):
-    img = img.convert("RGB")
-    size = KERAS_INPUT_SIZES.get(crop_name, 224)
-    img = img.resize((size, size))
+def preprocess_keras(img, size):
+    img = img.convert("RGB").resize((size, size))
     arr = np.array(img).astype("float32") / 255.0
     return np.expand_dims(arr, axis=0)
 
-# ================= MODEL REGISTRIES =================
+# ================= REGISTRIES =================
 
 PYTORCH_MODELS = {}
 KERAS_MODELS = {}
@@ -66,63 +54,69 @@ LABELS = {}
 def load_models():
     for file in os.listdir(MODELS_DIR):
         name, ext = os.path.splitext(file)
-        crop_name = name.replace("_model", "").lower()
+        crop = 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"Missing label file: {label_path}")
+        label_path = os.path.join(LABELS_DIR, f"{crop}_labels.json")
 
-        with open(label_path, "r") as f:
-            LABELS[crop_name] = json.load(f)
+        with open(label_path) as f:
+            LABELS[crop] = json.load(f)
 
         # ---------- PyTorch ----------
         if ext == ".pth":
-            num_classes = len(LABELS[crop_name])
             model = models.resnet18(weights=None)
-            model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
+            model.fc = torch.nn.Linear(model.fc.in_features, len(LABELS[crop]))
             model.load_state_dict(torch.load(model_path, map_location="cpu"))
             model.eval()
-            PYTORCH_MODELS[crop_name] = model
+            PYTORCH_MODELS[crop] = model
 
         # ---------- Keras ----------
         elif ext in [".keras", ".h5"]:
-            KERAS_MODELS[crop_name] = tf.keras.models.load_model(
-                model_path,
-                custom_objects={
-                    "swish": tf.keras.activations.swish,
-                    "FixedDropout": FixedDropout,
-                    "EfficientNetB3": EfficientNetB3,  # corn
-                    "MobileNetV2": MobileNetV2,        # bean (REAL)
-                },
-                compile=False
-            )
+
+            # Bean model — load clean
+            if crop == "bean":
+                model = tf.keras.models.load_model(
+                    model_path,
+                    compile=False
+                )
+
+            # Corn model — needs EfficientNet
+            else:
+                model = tf.keras.models.load_model(
+                    model_path,
+                    custom_objects={
+                        "FixedDropout": FixedDropout,
+                        "EfficientNetB3": EfficientNetB3,
+                        "swish": tf.keras.activations.swish,
+                    },
+                    compile=False
+                )
+
+            KERAS_MODELS[crop] = model
 
 load_models()
 
-# ================= PREDICTION =================
+# ================= PREDICT =================
 
-def predict(image, crop_name):
-    crop_name = crop_name.strip().lower()
+def predict(image, crop):
+    crop = crop.lower()
 
-    if crop_name in PYTORCH_MODELS:
-        model = PYTORCH_MODELS[crop_name]
-        labels = LABELS[crop_name]
-        tensor = preprocess_pytorch(image)
+    if crop in PYTORCH_MODELS:
+        model = PYTORCH_MODELS[crop]
+        labels = LABELS[crop]
+        x = preprocess_pytorch(image)
         with torch.no_grad():
-            output = model(tensor)
-            probs = torch.softmax(output[0], dim=0)
-            idx = probs.argmax().item()
+            probs = torch.softmax(model(x)[0], dim=0)
+        idx = probs.argmax().item()
         return labels[idx], float(probs[idx])
 
-    elif crop_name in KERAS_MODELS:
-        model = KERAS_MODELS[crop_name]
-        labels = LABELS[crop_name]
-        arr = preprocess_keras(image, crop_name)
-        preds = model.predict(arr, verbose=0)[0]
+    if crop in KERAS_MODELS:
+        model = KERAS_MODELS[crop]
+        labels = LABELS[crop]
+        size = 224 if crop == "bean" else 300
+        x = preprocess_keras(image, size)
+        preds = model.predict(x, verbose=0)[0]
         idx = int(np.argmax(preds))
         return labels[idx], float(preds[idx])
 
-    else:
-        raise ValueError(f"No model found for crop: {crop_name}")
+    raise ValueError(f"No model found for {crop}")