Update utils/model_loader.py

utils/model_loader.py

@@ -2,19 +2,21 @@ import torch
 import json
 import os
 from models.model import PlantCNN
+from utils.config import load_config 
 
 def load_model_and_config():
-    """Load the trained model and all configuration files"""
     
     # Paths
     MODEL_PATH = "saved_models/plant_cnn.pt"
     CLASS_NAMES_PATH = "ui_text/class_names.json"
-    
-    # Config
+
+    # Load config
+    config = load_config()
     DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    CHANNELS = [96, 192, 384, 768]
-    DROPOUT = 0.4
-    NUM_CLASSES = 39
+
+    CHANNELS = config["channels"]          
+    DROPOUT = config["dropout"]             
+    NUM_CLASSES = config["num_classes"]    
     
     # Load class names
     with open(CLASS_NAMES_PATH, "r") as f:
@@ -28,9 +30,11 @@ def load_model_and_config():
     model = PlantCNN(num_classes=NUM_CLASSES, channels=CHANNELS, dropout=DROPOUT).to(DEVICE)
     if os.path.exists(MODEL_PATH):
         print("Loading trained model weights...")
-        model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
+        state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
+        model.load_state_dict(state_dict)   
         model.eval()
     else:
+        print("Model file not found at", MODEL_PATH)
         exit()
     
     return {
@@ -39,13 +43,3 @@ def load_model_and_config():
         'disease_db': disease_db,
         'device': DEVICE
     }
-
-def load_ui_text():
-    """Load intro and about markdown files"""
-    with open("ui_text/intro.md", "r", encoding="utf-8") as f:
-        intro_md = f.read()
-    
-    with open("ui_text/about.md", "r", encoding="utf-8") as f:
-        about_md = f.read()
-    
-    return intro_md, about_md