Update model.py

model.py

@@ -19,15 +19,13 @@ class FixedDropout(tf.keras.layers.Dropout):
     def __init__(self, rate, noise_shape=None, seed=None, **kwargs):
         super().__init__(rate, noise_shape=noise_shape, seed=seed, **kwargs)
 
-# EfficientNetB3
+# EfficientNet
 @tf.keras.utils.register_keras_serializable()
 class EfficientNetB3(tf.keras.Model):
     pass
 
 # MobileNetV2
-@tf.keras.utils.register_keras_serializable()
-class MobileNetV2(tf.keras.Model):
-    pass
+from tensorflow.keras.applications import MobileNetV2
 
 # ================= INPUT SIZE PER MODEL =================
 
@@ -79,7 +77,7 @@ def load_models():
         with open(label_path, "r") as f:
             LABELS[crop_name] = json.load(f)
 
-        # ---------- PyTorch (.pth) ----------
+        # ---------- PyTorch ----------
         if ext == ".pth":
             num_classes = len(LABELS[crop_name])
             model = models.resnet18(weights=None)
@@ -88,20 +86,19 @@ def load_models():
             model.eval()
             PYTORCH_MODELS[crop_name] = model
 
-        # ---------- Keras (.keras / .h5) ----------
+        # ---------- 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,
-                    "MobileNetV2": MobileNetV2,
+                    "EfficientNetB3": EfficientNetB3,  # corn
+                    "MobileNetV2": MobileNetV2,        # bean (REAL)
                 },
                 compile=False
             )
 
-# Load once at startup
 load_models()
 
 # ================= PREDICTION =================
@@ -109,28 +106,22 @@ load_models()
 def predict(image, crop_name):
     crop_name = crop_name.strip().lower()
 
-    # -------- PyTorch --------
     if crop_name in PYTORCH_MODELS:
         model = PYTORCH_MODELS[crop_name]
         labels = LABELS[crop_name]
-
         tensor = preprocess_pytorch(image)
         with torch.no_grad():
             output = model(tensor)
             probs = torch.softmax(output[0], dim=0)
             idx = probs.argmax().item()
-
         return labels[idx], float(probs[idx])
 
-    # -------- Keras --------
     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]
         idx = int(np.argmax(preds))
-
         return labels[idx], float(preds[idx])
 
     else: