update app.py

edit by chatGPT

app.py

@@ -9,24 +9,13 @@ import torchvision
 
 from torch import nn
 
+from torchvision.models import densenet121
 
-def create_effnetb2_model(num_classes: int = 1,
-                          seed: int = 42):
-    """Creates an EfficientNetB2 feature extractor model and transforms.
-
-    Args:
-        num_classes (int, optional): number of classes in the classifier head. 
-            Defaults to 3.
-        seed (int, optional): random seed value. Defaults to 42.
-
-    Returns:
-        model (torch.nn.Module): EffNetB2 feature extractor model. 
-        transforms (torchvision.transforms): EffNetB2 image transforms.
-    """
-    # Create EffNetB2 pretrained weights, transforms and model
-    weights = torchvision.models.AlexNet_Weights.DEFAULT
-    transforms = weights.transforms()
-    model = torchvision.models.alexnet(weights=weights)
+def create_densenet121_model(num_classes: int = 1, seed: int = 42):
+    """Creates a DenseNet121 model and transforms."""
+    
+    # Create DenseNet121 model
+    model = densenet121(pretrained=False)  # Set to False since we will be loading our own weights
 
     # Freeze all layers in base model
     for param in model.parameters():
@@ -34,31 +23,23 @@ def create_effnetb2_model(num_classes: int = 1,
 
     # Change classifier head with random seed for reproducibility
     torch.manual_seed(seed)
-    model.classifier = nn.Sequential(
-        nn.Dropout(p=0.2,),
-        nn.Linear(in_features=9216, out_features=1),
-    )
-
-    return model, transforms
-
+    model.classifier = nn.Linear(model.classifier.in_features, num_classes)
 
-# Setup class names
-class_names = ["Normal", "Pneumonia"]
+    # You might want to use the appropriate transforms for densenet121 here
+    transforms = torchvision.transforms.Compose([
+        torchvision.transforms.Resize(224),
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
 
-### 2. Model and transforms preparation ###
+    return model, transforms
 
-# Create EffNetB2 model
-effnetb2, effnetb2_transforms = create_effnetb2_model(
-    num_classes=1,  # len(class_names) would also work
-)
+# Create densenet121 model
+densenet, densenet_transforms = create_densenet121_model(num_classes=1)
 
 # Load saved weights
-effnetb2.load_state_dict(
-    torch.load(
-        f="FL_global_model.pt",
-        map_location=torch.device("cpu"),  # load to CPU
-    )
-)
+densenet.load_state_dict(torch.load("FL_global_model.pt", map_location=torch.device("cpu")))
+
 
 
 def predict(img) -> Tuple[Dict, float]: