Update model.py

model.py

@@ -16,20 +16,29 @@ LABELS_DIR = os.path.join(BASE_DIR, "labels")
 
 @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)
+    def __init__(self, rate, **kwargs):
+        super().__init__(rate, **kwargs)
 
-# EfficientNet
 @tf.keras.utils.register_keras_serializable()
 class EfficientNetB3(tf.keras.Model):
     pass
 
+# ================= HARD INPUT SIZES =================
+
+INPUT_SIZE = {
+    "bean": 224,
+    "corn": 300,
+    "banana": 224,
+    "chilli": 224,
+    "rice": 224,
+}
+
 # ================= IMAGE PREPROCESS =================
 
-def preprocess_pytorch(img, size=224):
+def preprocess_pytorch(img):
     img = img.convert("RGB")
     transform = transforms.Compose([
-        transforms.Resize((size, size)),
+        transforms.Resize((224, 224)),
         transforms.ToTensor(),
         transforms.Normalize(
             mean=[0.485, 0.456, 0.406],
@@ -38,7 +47,8 @@ def preprocess_pytorch(img, size=224):
     ])
     return transform(img).unsqueeze(0)
 
-def preprocess_keras(img, size):
+def preprocess_keras(img, crop):
+    size = INPUT_SIZE[crop]
     img = img.convert("RGB").resize((size, size))
     arr = np.array(img).astype("float32") / 255.0
     return np.expand_dims(arr, axis=0)
@@ -56,12 +66,11 @@ def load_models():
         name, ext = os.path.splitext(file)
         crop = name.replace("_model", "").lower()
 
-        model_path = os.path.join(MODELS_DIR, file)
-        label_path = os.path.join(LABELS_DIR, f"{crop}_labels.json")
-
-        with open(label_path) as f:
+        with open(os.path.join(LABELS_DIR, f"{crop}_labels.json")) as f:
             LABELS[crop] = json.load(f)
 
+        model_path = os.path.join(MODELS_DIR, file)
+
         # ---------- PyTorch ----------
         if ext == ".pth":
             model = models.resnet18(weights=None)
@@ -73,15 +82,14 @@ def load_models():
         # ---------- Keras ----------
         elif ext in [".keras", ".h5"]:
 
-            # Bean model — load clean
             if crop == "bean":
+                # NO custom_objects
                 model = tf.keras.models.load_model(
                     model_path,
                     compile=False
                 )
 
-            # Corn model — needs EfficientNet
-            else:
+            elif crop == "corn":
                 model = tf.keras.models.load_model(
                     model_path,
                     custom_objects={
@@ -91,6 +99,8 @@ def load_models():
                     },
                     compile=False
                 )
+            else:
+                model = tf.keras.models.load_model(model_path, compile=False)
 
             KERAS_MODELS[crop] = model
 
@@ -103,20 +113,22 @@ def predict(image, crop):
 
     if crop in PYTORCH_MODELS:
         model = PYTORCH_MODELS[crop]
-        labels = LABELS[crop]
         x = preprocess_pytorch(image)
         with torch.no_grad():
             probs = torch.softmax(model(x)[0], dim=0)
         idx = probs.argmax().item()
-        return labels[idx], float(probs[idx])
+        return LABELS[crop][idx], float(probs[idx])
 
     if crop in KERAS_MODELS:
         model = KERAS_MODELS[crop]
-        labels = LABELS[crop]
-        size = 224 if crop == "bean" else 300
-        x = preprocess_keras(image, size)
+        x = preprocess_keras(image, crop)
+
+        # SAFETY ASSERT
+        expected = INPUT_SIZE[crop]
+        assert x.shape[1] == expected, f"{crop} received wrong input size!"
+
         preds = model.predict(x, verbose=0)[0]
         idx = int(np.argmax(preds))
-        return labels[idx], float(preds[idx])
+        return LABELS[crop][idx], float(preds[idx])
 
     raise ValueError(f"No model found for {crop}")