Update model.py

model.py

@@ -6,6 +6,9 @@ import tensorflow as tf
 from PIL import Image
 from torchvision import models, transforms
 
+# EfficientNet (needed ONLY for corn)
+from efficientnet.tfkeras import EfficientNetB3
+
 # ================= PATHS =================
 
 BASE_DIR = os.path.dirname(__file__)
@@ -19,6 +22,14 @@ 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)
 
+# ================= INPUT SIZE PER MODEL =================
+# Only corn differs — others remain 224
+
+KERAS_INPUT_SIZES = {
+    "corn": 300,
+    "bean": 224,
+}
+
 # ================= IMAGE PREPROCESS =================
 
 def preprocess_pytorch(img, size=224):
@@ -32,7 +43,8 @@ def preprocess_pytorch(img, size=224):
     ])
     return transform(img).unsqueeze(0)
 
-def preprocess_keras(img, size=224):
+def preprocess_keras(img, crop_name):
+    size = KERAS_INPUT_SIZES.get(crop_name, 224)
     img = img.resize((size, size))
     arr = np.array(img).astype("float32") / 255.0
     return np.expand_dims(arr, axis=0)
@@ -56,7 +68,7 @@ def load_models():
         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}")
+            raise FileNotFoundError(f"Missing label file: {label_path}")
 
         with open(label_path, "r") as f:
             LABELS[crop_name] = json.load(f)
@@ -79,6 +91,8 @@ def load_models():
                 custom_objects={
                     "swish": tf.keras.activations.swish,
                     "FixedDropout": FixedDropout,
+                    # Needed ONLY for corn, harmless for others
+                    "EfficientNetB3": EfficientNetB3,
                 },
                 compile=False  # IMPORTANT for HF
             )
@@ -109,11 +123,11 @@ def predict(image, crop_name):
         model = KERAS_MODELS[crop_name]
         labels = LABELS[crop_name]
 
-        arr = preprocess_keras(image)
+        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:
-        raise ValueError(f"❌ No model found for crop: {crop_name}")
+        raise ValueError(f"No model found for crop: {crop_name}")