Update streamlit_app.py

streamlit_app.py

@@ -52,12 +52,12 @@ inference_transform = A.Compose([
 # Enhanced model loading function
 @st.cache_resource
 def load_model():
-    """Enhanced model loading with architecture detection and fallback options"""
+    """Enhanced model loading with better architecture detection"""
     
     # Try different model file names
     model_files = [
         "butterfly_classifier.pth",
-        "best_butterfly_model_v3.pth",
+        "best_butterfly_model_v3.pth", 
         "best_butterfly_model.pth"
     ]
     
@@ -68,7 +68,7 @@ def load_model():
             break
     
     if MODEL_PATH is None:
-        st.error("No model file found! Please ensure one of these files exists: " + ", ".join(model_files))
+        st.error("No model file found!")
         return None
     
     st.info(f"Loading model from: {MODEL_PATH}")
@@ -80,87 +80,42 @@ def load_model():
         # Extract model state dict
         if 'model_state_dict' in checkpoint:
             model_state_dict = checkpoint['model_state_dict']
-            if 'class_names' in checkpoint:
-                st.info(f"Model trained on {len(checkpoint['class_names'])} classes")
         else:
             model_state_dict = checkpoint
-        
+            
         num_classes = len(class_names)
         
-        # Architecture detection based on model state dict
-        def detect_model_architecture(state_dict):
-            """Detect model architecture from state dict"""
-            
-            # Check for EfficientNet variants by looking at key layer dimensions
-            architecture_indicators = {
-                'conv_head.weight': 'efficientnet',
-                'head.weight': 'efficientnet_v2',
-                'classifier.weight': 'other'
-            }
-            
-            # Look for specific layer patterns
-            for key in state_dict.keys():
-                if 'conv_head.weight' in key:
-                    shape = state_dict[key].shape
-                    if len(shape) >= 2:
-                        feature_dim = shape[1]
-                        # EfficientNet feature dimensions
-                        efficientnet_map = {
-                            1280: 'efficientnet_b0',
-                            1408: 'efficientnet_b1', 
-                            1536: 'efficientnet_b2',
-                            1792: 'efficientnet_b3',
-                            1920: 'efficientnet_b4',
-                            2048: 'efficientnet_b5',
-                            2304: 'efficientnet_b6',
-                            2560: 'efficientnet_b7'
-                        }
-                        return efficientnet_map.get(feature_dim, 'efficientnet_b3')
-                
-                if 'head.weight' in key:
-                    shape = state_dict[key].shape
-                    if len(shape) >= 2:
-                        feature_dim = shape[1]
-                        # EfficientNetV2 feature dimensions
-                        efficientnetv2_map = {
-                            1280: 'tf_efficientnetv2_s',
-                            1408: 'tf_efficientnetv2_m',
-                            1792: 'tf_efficientnetv2_l'
-                        }
-                        return efficientnetv2_map.get(feature_dim, 'tf_efficientnetv2_s')
-            
-            # Fallback: check bn2 layer for EfficientNet variants
-            for key in state_dict.keys():
-                if key.endswith("bn2.weight"):
-                    bn2_shape = state_dict[key].shape[0]
-                    feature_map = {
-                        1280: 'efficientnet_b0',
-                        1408: 'efficientnet_b1',
-                        1536: 'efficientnet_b2',
-                        1792: 'efficientnet_b3',
-                        1920: 'efficientnet_b4',
-                        2048: 'efficientnet_b5',
-                        2304: 'efficientnet_b6',
-                        2560: 'efficientnet_b7'
+        # Better architecture detection based on conv_stem channels
+        def detect_architecture_by_channels(state_dict):
+            """Detect architecture by examining conv_stem channels"""
+            for key, tensor in state_dict.items():
+                if key.endswith('conv_stem.weight'):
+                    channels = tensor.shape[0]  # Output channels
+                    # Map channels to likely architectures
+                    channel_map = {
+                        24: ['tf_efficientnetv2_s', 'efficientnet_b0'],
+                        32: ['tf_efficientnetv2_s', 'efficientnet_b1'], 
+                        40: ['efficientnet_b3', 'efficientnet_b2'],
+                        48: ['efficientnet_b4', 'tf_efficientnetv2_m'],
+                        56: ['efficientnet_b5'],
+                        64: ['efficientnet_b6', 'tf_efficientnetv2_l'],
+                        72: ['efficientnet_b7']
                     }
-                    return feature_map.get(bn2_shape, 'efficientnet_b3')
-            
-            return 'efficientnet_b3'  # Default fallback
+                    return channel_map.get(channels, ['tf_efficientnetv2_s'])
+            return ['tf_efficientnetv2_s']
         
-        # Detect architecture
-        detected_arch = detect_model_architecture(model_state_dict)
-        st.info(f"Detected model architecture: {detected_arch}")
+        # Get likely architectures based on channels
+        likely_architectures = detect_architecture_by_channels(model_state_dict)
         
-        # List of architectures to try in order
-        architectures_to_try = [
-            detected_arch,
-            'efficientnet_b3',
-            'efficientnet_b2',
-            'efficientnet_b0',
-            'efficientnet_b1',
-            'efficientnet_b4',
+        # Expanded list of architectures to try
+        architectures_to_try = likely_architectures + [
             'tf_efficientnetv2_s',
-            'tf_efficientnetv2_m'
+            'efficientnet_b0', 
+            'efficientnet_b1',
+            'efficientnet_b2',
+            'efficientnet_b3',
+            'tf_efficientnetv2_m',
+            'efficientnet_b4'
         ]
         
         # Remove duplicates while preserving order
@@ -175,61 +130,30 @@ def load_model():
             try:
                 st.info(f"Trying architecture: {arch}")
                 
-                # Create model with the detected/guessed architecture
+                # Create model
                 model = timm.create_model(
                     arch,
                     pretrained=False,
                     num_classes=num_classes,
-                    drop_rate=0.4,
-                    drop_path_rate=0.3
+                    drop_rate=0.0,  # Set to 0 for inference
+                    drop_path_rate=0.0  # Set to 0 for inference
                 )
                 
-                # Check if the model has a custom head/classifier in the checkpoint
-                if any('head.' in key for key in model_state_dict.keys()):
-                    # Model has custom head - try to load it
+                # Try to load the state dict
+                try:
+                    model.load_state_dict(model_state_dict, strict=True)
+                    st.success(f"✅ Successfully loaded model with architecture: {arch}")
+                    successful_arch = arch
+                    break
+                except Exception as e:
+                    # Try with strict=False
                     try:
                         model.load_state_dict(model_state_dict, strict=False)
-                        st.success(f"✅ Successfully loaded model with architecture: {arch}")
+                        st.warning(f"⚠️ Loaded {arch} with some mismatched weights")
                         successful_arch = arch
                         break
-                    except Exception as e:
-                        st.warning(f"Failed to load custom head for {arch}: {str(e)}")
-                        continue
-                
-                elif any('classifier.' in key for key in model_state_dict.keys()):
-                    # Model has custom classifier - try to load it
-                    try:
-                        model.load_state_dict(model_state_dict, strict=False)
-                        st.success(f"✅ Successfully loaded model with architecture: {arch}")
-                        successful_arch = arch
-                        break
-                    except Exception as e:
-                        st.warning(f"Failed to load custom classifier for {arch}: {str(e)}")
-                        continue
-                
-                else:
-                    # Try to create custom head/classifier and load backbone
-                    try:
-                        # Load backbone weights (ignore head/classifier mismatches)
-                        backbone_dict = {k: v for k, v in model_state_dict.items() 
-                                       if not (k.startswith('head.') or k.startswith('classifier.'))}
-                        
-                        model.load_state_dict(backbone_dict, strict=False)
-                        
-                        # Create new head/classifier
-                        if hasattr(model, 'classifier'):
-                            in_features = model.classifier.in_features
-                            model.classifier = torch.nn.Linear(in_features, num_classes)
-                        elif hasattr(model, 'head'):
-                            in_features = model.head.in_features
-                            model.head = torch.nn.Linear(in_features, num_classes)
-                        
-                        st.warning(f"⚠️ Loaded {arch} with new head/classifier (backbone weights only)")
-                        successful_arch = arch
-                        break
-                        
-                    except Exception as e:
-                        st.warning(f"Failed to load backbone for {arch}: {str(e)}")
+                    except Exception as e2:
+                        st.warning(f"Failed to load {arch}: {str(e2)}")
                         continue
                         
             except Exception as e:
@@ -254,7 +178,6 @@ def load_model():
         
     except Exception as e:
         st.error(f"❌ Error loading model: {str(e)}")
-        st.error("Please check your model file and ensure it's compatible")
         return None
 
 # Load model