update app.py

app.py

@@ -9,7 +9,9 @@ import torchvision.transforms as transforms
 import matplotlib.pyplot as plt
 import timm
 
+
 class BaseModel(nn.Module):
+    
     def predict(self, x: torch.Tensor) -> torch.Tensor:
         with torch.no_grad():
             logits = self(x)
@@ -20,6 +22,7 @@ class BaseModel(nn.Module):
 
 
 class CNNModel(BaseModel):
+    
     def __init__(self, num_classes: int, input_size: int = 224):
         super(CNNModel, self).__init__()
         
@@ -55,6 +58,21 @@ class CNNModel(BaseModel):
             nn.Linear(256, num_classes)
         )
 
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.constant_(m.bias, 0)
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.conv_layers(x)
         return self.classifier(x)
@@ -64,6 +82,7 @@ class CNNModel(BaseModel):
 
 
 class EfficientNetModel(BaseModel):
+    
     def __init__(
         self,
         num_classes: int,
@@ -98,6 +117,7 @@ class EfficientNetModel(BaseModel):
 
 class AnimalClassifierApp:
     def __init__(self):
+        """Initialize the application."""
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.labels = ["bird", "cat", "dog", "horse"]
         
@@ -115,12 +135,12 @@ class AnimalClassifierApp:
             print("Warning: No models found in checkpoints directory!")
 
     def load_models(self):
+        """Load both trained models."""
         models = {}
         
-        # Load EfficientNet
         try:
             efficientnet = EfficientNetModel(num_classes=len(self.labels))
-            efficientnet_path = "efficientnet_best_model.pth"
+            efficientnet_path = os.path.join("checkpoints", "efficientnet", "efficientnet_best_model.pth")
             if os.path.exists(efficientnet_path):
                 checkpoint = torch.load(efficientnet_path, map_location=self.device, weights_only=True)
                 state_dict = checkpoint.get('model_state_dict', checkpoint)
@@ -131,10 +151,9 @@ class AnimalClassifierApp:
         except Exception as e:
             print(f"Error loading EfficientNet model: {str(e)}")
         
-        # Load CNN
         try:
             cnn = CNNModel(num_classes=len(self.labels))
-            cnn_path = "cnn_best_model.pth"
+            cnn_path = os.path.join("checkpoints", "cnn", "cnn_best_model.pth")
             if os.path.exists(cnn_path):
                 checkpoint = torch.load(cnn_path, map_location=self.device, weights_only=True)
                 state_dict = checkpoint.get('model_state_dict', checkpoint)
@@ -149,8 +168,9 @@ class AnimalClassifierApp:
 
     def predict(self, image: Image.Image):
         if not self.models:
-            return "No trained models found. Please train the models first."
+            return ["No trained models found. Please train the models first.", ""]
         
+        # Preprocess image
         img_tensor = self.transform(image).unsqueeze(0).to(self.device)
         
         results = {}
@@ -209,10 +229,12 @@ class AnimalClassifierApp:
             description="Upload an image of an animal to see predictions from both EfficientNet and CNN models."
         )
 
+
 def main():
     app = AnimalClassifierApp()
     interface = app.create_interface()
     interface.launch()
 
+
 if __name__ == "__main__":
     main()