Update streamlit_app.py

streamlit_app.py

@@ -61,13 +61,12 @@ def load_model():
         model_state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
         num_classes = len(class_names)
 
-        # Try to detect classifier shape; fallback if key missing
-        classifier_key = [k for k in model_state_dict.keys() if 'classifier.weight' in k or 'head.fc.weight' in k]
-        if not classifier_key:
-            st.error("Could not detect classifier layer in checkpoint.")
-            return None
-
-        classifier_input = model_state_dict[classifier_key[0]].shape[1]
+        # Attempt to auto-detect model from batch norm layer dimensions
+        bn2_shape = None
+        for key in model_state_dict:
+            if key.endswith("bn2.weight"):
+                bn2_shape = model_state_dict[key].shape[0]
+                break
 
         feature_map = {
             1280: 'efficientnet_b0',
@@ -79,8 +78,13 @@ def load_model():
             2304: 'efficientnet_b6',
             2560: 'efficientnet_b7'
         }
-        model_name = feature_map.get(classifier_input, 'efficientnet_b3')
-        st.info(f"Detected model architecture: {model_name}")
+
+        if bn2_shape is None:
+            st.warning("Could not detect classifier or bn2 layer in checkpoint. Defaulting to efficientnet_b3")
+            model_name = 'efficientnet_b3'
+        else:
+            model_name = feature_map.get(bn2_shape, 'efficientnet_b3')
+            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)
@@ -122,6 +126,7 @@ st.write("Tuvasta liblikaid oma kaamera abil või laadi üles pilt! / Identify b
 
 tab1, tab2 = st.tabs(["📷 Live Camera / Kaamera", "📁 Upload Image / Laadi üles"])
 
+
 with tab1:
     st.header("Kaamera jäädvustamine / Camera Capture")
     st.write("Tee pilt liblikast tuvastamiseks / Take a photo of a butterfly for identification.")