switched model to effnetv2 L

app.py

@@ -31,14 +31,14 @@ num_labels = len(labels)
 def predict(img, confidence):
 
     new_layers = nn.Sequential(
-    nn.Linear(1920, 1000),  # Reduce dimension from 1024 to 500
-    nn.BatchNorm1d(1000),   # Normalize the activations from the previous layer
-    nn.ReLU(),             # Non-linear activation function
-    nn.Dropout(0.5),       # Dropout for regularization (50% probability)
-    nn.Linear(1000, num_labels)  # Final layer for class predictions
+    nn.LazyLinear(1280),  
+    nn.BatchNorm1d(1280),  
+    nn.ReLU(),            
+    nn.Dropout(0.5),      
+    nn.LazyLinear(num_labels)
     )
 
-    IMAGE_SIZE = 512
+    IMAGE_SIZE = 480
     transform = v2.Compose([
     v2.ToImage(),
     v2.Resize((IMAGE_SIZE, IMAGE_SIZE)),
@@ -47,26 +47,26 @@ def predict(img, confidence):
     ])
 
 
-    densenet = models.densenet201(weights='DenseNet201_Weights.DEFAULT')
-    densenet.classifier = new_layers
+    efficientnet = models.efficientnet_v2_l(weights='EfficientNet_V2_L_Weights.DEFAULT')
+    efficientnet.classifier = new_layers
 
     # If using GPU
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     hf_token = os.getenv('HF_token')
-    model_path = hf_hub_download(repo_id="fohy24/morphmarket_model", filename="model_v8_epoch9.pt", token=hf_token)
+    model_path = hf_hub_download(repo_id="fohy24/morphmarket_model", filename="model_v9_epoch8.pt", token=hf_token)
 
     checkpoint = torch.load(model_path, map_location=device)
-    densenet.load_state_dict(checkpoint['model_state_dict'])
+    efficientnet.load_state_dict(checkpoint['model_state_dict'])
 
-    densenet.eval()
+    efficientnet.eval()
 
     input_img = transform(img)
     input_img = input_img.unsqueeze(0)
 
 
     with torch.no_grad():
-        output = densenet(input_img)
+        output = efficientnet(input_img)
 
     predicted_probs = torch.sigmoid(output).to('cpu').flatten().tolist()
     prediction_dict = {labels[i]: predicted_probs[i] for i in range(len(labels)) if predicted_probs[i] > confidence}