Update model.py

model.py

@@ -16,29 +16,25 @@ LABELS_DIR = os.path.join(BASE_DIR, "labels")
 
 @tf.keras.utils.register_keras_serializable()
 class FixedDropout(tf.keras.layers.Dropout):
-    def __init__(self, rate, **kwargs):
-        super().__init__(rate, **kwargs)
+    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
 
-# ================= HARD INPUT SIZES =================
+# ================= INPUT SIZE PER MODEL =================
 
-INPUT_SIZE = {
-    "bean": 224,
+KERAS_INPUT_SIZES = {
     "corn": 300,
-    "banana": 224,
-    "chilli": 224,
-    "rice": 224,
 }
 
 # ================= IMAGE PREPROCESS =================
 
-def preprocess_pytorch(img):
-    img = img.convert("RGB")
+def preprocess_pytorch(img, size=224):
     transform = transforms.Compose([
-        transforms.Resize((224, 224)),
+        transforms.Resize((size, size)),
         transforms.ToTensor(),
         transforms.Normalize(
             mean=[0.485, 0.456, 0.406],
@@ -47,13 +43,14 @@ def preprocess_pytorch(img):
     ])
     return transform(img).unsqueeze(0)
 
-def preprocess_keras(img, crop):
-    size = INPUT_SIZE[crop]
-    img = img.convert("RGB").resize((size, size))
+def preprocess_keras(img, crop_name):
+    img = img.convert("RGB")
+    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)
 
-# ================= REGISTRIES =================
+# ================= MODEL REGISTRIES =================
 
 PYTORCH_MODELS = {}
 KERAS_MODELS = {}
@@ -64,71 +61,63 @@ LABELS = {}
 def load_models():
     for file in os.listdir(MODELS_DIR):
         name, ext = os.path.splitext(file)
-        crop = name.replace("_model", "").lower()
-
-        with open(os.path.join(LABELS_DIR, f"{crop}_labels.json")) as f:
-            LABELS[crop] = json.load(f)
+        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 label file: {label_path}")
+
+        with open(label_path, "r") as f:
+            LABELS[crop_name] = json.load(f)
 
         # ---------- PyTorch ----------
         if ext == ".pth":
+            num_classes = len(LABELS[crop_name])
             model = models.resnet18(weights=None)
-            model.fc = torch.nn.Linear(model.fc.in_features, len(LABELS[crop]))
+            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] = model
+            PYTORCH_MODELS[crop_name] = model
 
         # ---------- Keras ----------
         elif ext in [".keras", ".h5"]:
-
-            if crop == "bean":
-                # NO custom_objects
-                model = tf.keras.models.load_model(
-                    model_path,
-                    compile=False
-                )
-
-            elif crop == "corn":
-                model = tf.keras.models.load_model(
-                    model_path,
-                    custom_objects={
-                        "FixedDropout": FixedDropout,
-                        "EfficientNetB3": EfficientNetB3,
-                        "swish": tf.keras.activations.swish,
-                    },
-                    compile=False
-                )
-            else:
-                model = tf.keras.models.load_model(model_path, compile=False)
-
-            KERAS_MODELS[crop] = model
-
+            KERAS_MODELS[crop_name] = tf.keras.models.load_model(
+                model_path,
+                custom_objects={
+                    "swish": tf.keras.activations.swish,
+                    "FixedDropout": FixedDropout,
+                    "EfficientNetB3": EfficientNetB3,
+                },
+                compile=False
+            )
+
+# Load models at startup
 load_models()
 
-# ================= PREDICT =================
+# ================= PREDICTION =================
 
-def predict(image, crop):
-    crop = crop.lower()
+def predict(image, crop_name):
+    crop_name = crop_name.strip().lower()
 
-    if crop in PYTORCH_MODELS:
-        model = PYTORCH_MODELS[crop]
-        x = preprocess_pytorch(image)
+    if crop_name in PYTORCH_MODELS:
+        model = PYTORCH_MODELS[crop_name]
+        labels = LABELS[crop_name]
+        tensor = preprocess_pytorch(image)
         with torch.no_grad():
-            probs = torch.softmax(model(x)[0], dim=0)
-        idx = probs.argmax().item()
-        return LABELS[crop][idx], float(probs[idx])
-
-    if crop in KERAS_MODELS:
-        model = KERAS_MODELS[crop]
-        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]
+            output = model(tensor)
+            probs = torch.softmax(output[0], dim=0)
+            idx = probs.argmax().item()
+        return labels[idx], float(probs[idx])
+
+    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[crop][idx], float(preds[idx])
+        return labels[idx], float(preds[idx])
 
-    raise ValueError(f"No model found for {crop}")
+    else:
+        raise ValueError(f"No model found for crop: {crop_name}")