Update model.py

model.py

@@ -19,11 +19,16 @@ 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)
 
-# Dummy class to satisfy EfficientNet deserialization
+# EfficientNetB3
 @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
+
 # ================= INPUT SIZE PER MODEL =================
 
 KERAS_INPUT_SIZES = {
@@ -34,6 +39,7 @@ KERAS_INPUT_SIZES = {
 # ================= IMAGE PREPROCESS =================
 
 def preprocess_pytorch(img, size=224):
+    img = img.convert("RGB")
     transform = transforms.Compose([
         transforms.Resize((size, size)),
         transforms.ToTensor(),
@@ -73,7 +79,7 @@ def load_models():
         with open(label_path, "r") as f:
             LABELS[crop_name] = json.load(f)
 
-        # ---------- PyTorch ----------
+        # ---------- PyTorch (.pth) ----------
         if ext == ".pth":
             num_classes = len(LABELS[crop_name])
             model = models.resnet18(weights=None)
@@ -82,7 +88,7 @@ def load_models():
             model.eval()
             PYTORCH_MODELS[crop_name] = model
 
-        # ---------- Keras ----------
+        # ---------- Keras (.keras / .h5) ----------
         elif ext in [".keras", ".h5"]:
             KERAS_MODELS[crop_name] = tf.keras.models.load_model(
                 model_path,
@@ -90,10 +96,12 @@ def load_models():
                     "swish": tf.keras.activations.swish,
                     "FixedDropout": FixedDropout,
                     "EfficientNetB3": EfficientNetB3,
+                    "MobileNetV2": MobileNetV2,
                 },
                 compile=False
             )
 
+# Load once at startup
 load_models()
 
 # ================= PREDICTION =================
@@ -101,22 +109,28 @@ 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: