Update model.py

model.py

@@ -6,24 +6,25 @@ 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__)
 MODELS_DIR = os.path.join(BASE_DIR, "models")
 LABELS_DIR = os.path.join(BASE_DIR, "labels")
 
-# ================= FIX: EfficientNet Custom Layer =================
+# ================= TF CUSTOM OBJECTS =================
 
 @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)
 
+# Dummy class to satisfy EfficientNet deserialization
+@tf.keras.utils.register_keras_serializable()
+class EfficientNetB3(tf.keras.Model):
+    pass
+
 # ================= INPUT SIZE PER MODEL =================
-# Only corn differs — others remain 224
 
 KERAS_INPUT_SIZES = {
     "corn": 300,
@@ -55,13 +56,11 @@ PYTORCH_MODELS = {}
 KERAS_MODELS = {}
 LABELS = {}
 
-# ================= LOAD MODELS ONCE =================
+# ================= LOAD MODELS =================
 
 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)
@@ -73,31 +72,27 @@ 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)
             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) ----------
+        # ---------- 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,
-                    # Needed ONLY for corn, harmless for others
                     "EfficientNetB3": EfficientNetB3,
                 },
-                compile=False  # IMPORTANT for HF
+                compile=False
             )
 
-# Load models at startup
 load_models()
 
 # ================= PREDICTION =================
@@ -105,28 +100,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: