Update app.py

app.py

@@ -22,7 +22,30 @@ dinov2_vits14.to(device)
 # Define the transformations: convert to tensor, resize, and normalize
 transform_image = transforms.Compose([transforms.ToTensor(), transforms.Resize(224), transforms.Normalize([0.5], [0.5])])
 
-model = torch.load("dress_model.pth")
+#model = torch.load("dress_model.pth")
+class Dress_Class(nn.Module):
+    def __init__(self, num_dim, hidden_layers, num_classes):
+        super(Dress_Class, self).__init__()
+        layers = []
+        current_num = num_dim
+        
+        for num_neurons in hidden_layers:
+            layers.append(nn.Linear(current_num, num_neurons))
+            layers.append(nn.ReLU())
+            current_num = num_neurons
+        
+        layers.append(nn.Linear(current_num, num_classes))
+        layers.append(nn.Softmax(dim=1))
+        self.model = nn.Sequential(*layers)
+        
+    def forward(self, x):
+        x = self.model(x)
+        return x
+        
+best = torch.load("best_mo.pth")       
+hidden_layers = [192, 187, 331, 186, 256]
+model = Dress_Class(384, hidden_layers, 10)
+model.state_dict(best)
 model.eval()
 
 with open('saved_dress_morph.pkl', 'rb') as f: