Update streamlit_app.py

streamlit_app.py

@@ -61,7 +61,7 @@ def load_model():
         model_state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
         num_classes = len(class_names)
 
-        # Detect model size via known feature dimensions
+        # Detect model from classifier shape
         classifier_input = model_state_dict['classifier.weight'].shape[1]
         feature_map = {
             1280: 'efficientnet_b0',
@@ -74,7 +74,7 @@ def load_model():
             2560: 'efficientnet_b7'
         }
         model_name = feature_map.get(classifier_input, 'efficientnet_b3')
-        print(f"Auto-detected model architecture: {model_name}")
+        st.info(f"Detected model architecture: {model_name}")
 
         model = timm.create_model(model_name, pretrained=False, num_classes=num_classes, drop_rate=0.4, drop_path_rate=0.3)
         model.load_state_dict(model_state_dict, strict=False)
@@ -84,9 +84,12 @@ def load_model():
         st.error(f"Error loading model: {str(e)}")
         return None
 
+model = load_model()
 
 def predict_butterfly(image, threshold=0.5):
     try:
+        if model is None:
+            raise ValueError("Model is not loaded.")
         if image is None:
             return None, None
         if isinstance(image, np.ndarray):