update to use model v13_1_e9 for top 30 morphs prediction

app.py

@@ -1,43 +1,56 @@
+import os
 import torch
 from torch import nn
 from torchvision import models
 from torchvision.transforms import v2
-import os
 from huggingface_hub import hf_hub_download
 import gradio as gr
 
-labels = ['Pastel',
-            'Yellow Belly',
-            'Enchi',
-            'Clown',
-            'Leopard',
-            'Piebald',
-            'Orange Dream',
-            'Fire',
-            'Mojave',
-            'Pinstripe',
-            'Banana',
-            'Normal',
-            'Black Pastel',
-            'Lesser',
-            'Spotnose',
-            'Cinnamon',
-            'GHI',
-            'Hypo',
-            'Spider',
-            'Super Pastel']
+labels = [
+    'Pastel',
+    'Yellow Belly',
+    'Enchi',
+    'Clown',
+    'Leopard',
+    'Piebald',
+    'Orange Dream',
+    'Fire',
+    'Mojave',
+    'Pinstripe',
+    'Banana',
+    'Normal',
+    'Black Pastel',
+    'Lesser',
+    'Spotnose',
+    'Cinnamon',
+    'GHI',
+    'Hypo',
+    'Spider',
+    'Super Pastel',
+    'Desert Ghost',
+    'Black Head',
+    'Vanilla',
+    'Red Stripe',
+    'Asphalt',
+    'Gravel',
+    'Butter',
+    'Calico',
+    'Albino',
+    'Chocolate'
+    ]
+
 num_labels = len(labels)
 
 # 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="samfhy/morphmarket_model", filename="model_v9_epoch8.pt", token=hf_token)
+model_path = hf_hub_download(repo_id="samfhy/morphmarket_model", filename="model_v13_1_epoch9.pt", token=hf_token)
 checkpoint = torch.load(model_path, map_location=device)
 
 new_layers = nn.Sequential(
-nn.LazyLinear(1280),  
-nn.BatchNorm1d(1280),  
+nn.LazyLinear(2048),  
+nn.BatchNorm1d(2048),  
 nn.ReLU(),            
 nn.Dropout(0.5),      
 nn.LazyLinear(num_labels)
@@ -47,7 +60,7 @@ IMAGE_SIZE = 480
 transform = v2.Compose([
 v2.ToImage(),
 v2.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-v2.ToDtype(torch.float32, scale=True),
+v2.ToDtype(torch.float32),
 v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])