Create model.py

model.py

@@ -0,0 +1,87 @@
+import os
+import json
+import torch
+import numpy as np
+import tensorflow as tf
+from PIL import Image
+from torchvision import models, transforms
+
+BASE_DIR = os.path.dirname(__file__)
+
+MODELS_DIR = os.path.join(BASE_DIR, "models")
+LABELS_DIR = os.path.join(BASE_DIR, "labels")
+
+# ================= IMAGE PREPROCESS =================
+
+def preprocess_pytorch(img, size=224):
+    transform = transforms.Compose([
+        transforms.Resize((size, size)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+    return transform(img).unsqueeze(0)
+
+def preprocess_keras(img, size=224):
+    img = img.resize((size, size))
+    arr = np.array(img) / 255.0
+    return np.expand_dims(arr, axis=0)
+
+# ================= MODEL LOADERS =================
+
+PYTORCH_MODELS = {}
+KERAS_MODELS = {}
+LABELS = {}
+
+def load_models():
+    for file in os.listdir(MODELS_DIR):
+        name, ext = os.path.splitext(file)
+        model_path = os.path.join(MODELS_DIR, file)
+
+        # Load labels
+        with open(os.path.join(LABELS_DIR, f"{name}.json")) as f:
+            LABELS[name] = json.load(f)
+
+        if ext == ".pth":
+            num_classes = len(LABELS[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[name] = model
+
+        elif ext == ".keras":
+            KERAS_MODELS[name] = tf.keras.models.load_model(model_path)
+
+# Load once
+load_models()
+
+# ================= PREDICT =================
+
+def predict(image, crop_name):
+    crop_name = crop_name.lower()
+
+    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])
+
+    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)
+        return labels[idx], float(preds[idx])
+
+    else:
+        raise ValueError(f"No model found for crop: {crop_name}")