Update model.py

model.py

@@ -6,11 +6,19 @@ import tensorflow as tf
 from PIL import Image
 from torchvision import models, transforms
 
-BASE_DIR = os.path.dirname(__file__)
+# ================= PATHS =================
 
+BASE_DIR = os.path.dirname(__file__)
 MODELS_DIR = os.path.join(BASE_DIR, "models")
 LABELS_DIR = os.path.join(BASE_DIR, "labels")
 
+# ================= FIX: EfficientNet Custom Layer =================
+
+@tf.keras.utils.register_keras_serializable()
+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)
+
 # ================= IMAGE PREPROCESS =================
 
 def preprocess_pytorch(img, size=224):
@@ -26,52 +34,64 @@ def preprocess_pytorch(img, size=224):
 
 def preprocess_keras(img, size=224):
     img = img.resize((size, size))
-    arr = np.array(img) / 255.0
+    arr = np.array(img).astype("float32") / 255.0
     return np.expand_dims(arr, axis=0)
 
-# ================= MODEL LOADERS =================
+# ================= MODEL REGISTRIES =================
 
 PYTORCH_MODELS = {}
 KERAS_MODELS = {}
 LABELS = {}
 
+# ================= LOAD MODELS ONCE =================
+
 def load_models():
     for file in os.listdir(MODELS_DIR):
         name, ext = os.path.splitext(file)
+
+        # Normalize crop key (banana_model -> banana)
         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"Missing labels for {crop_name}")
+            raise FileNotFoundError(f"❌ Missing label file: {label_path}")
 
-        with open(label_path) as f:
+        with open(label_path, "r") as f:
             LABELS[crop_name] = json.load(f)
 
+        # ---------- PyTorch (.pth) ----------
         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.load_state_dict(torch.load(model_path, map_location="cpu"))
             model.eval()
+
             PYTORCH_MODELS[crop_name] = model
 
+        # ---------- Keras (.keras / .h5) ----------
         elif ext in [".keras", ".h5"]:
             KERAS_MODELS[crop_name] = tf.keras.models.load_model(
                 model_path,
-                custom_objects={"swish": tf.keras.activations.swish},
-                compile=False
+                custom_objects={
+                    "swish": tf.keras.activations.swish,
+                    "FixedDropout": FixedDropout,
+                },
+                compile=False  # IMPORTANT for HF
             )
 
-# Load once
+# Load models at startup
 load_models()
 
-# ================= PREDICT =================
+# ================= PREDICTION =================
 
 def predict(image, crop_name):
-    crop_name = crop_name.lower()
+    crop_name = crop_name.strip().lower()
 
+    # -------- PyTorch --------
     if crop_name in PYTORCH_MODELS:
         model = PYTORCH_MODELS[crop_name]
         labels = LABELS[crop_name]
@@ -81,16 +101,19 @@ def predict(image, crop_name):
             output = model(tensor)
             probs = torch.softmax(output[0], dim=0)
             idx = probs.argmax().item()
-            return labels[idx], float(probs[idx])
 
+        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)
-        preds = model.predict(arr)[0]
-        idx = np.argmax(preds)
+        preds = model.predict(arr, 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: {crop_name}")