Update app.py

app.py

@@ -10,8 +10,13 @@ from VAE_model_tablets_class import VAE
 
 import gradio as gr
 
+num_classes = len(TabletPeriodDataset.PERIOD_INDICES)
 
-model = VAE()
+class_weights = torch.load("data/class_weights_period.pt")
+
+checkpoint_path = 'epoch=22-step=213621.ckpt'
+vae_model = VAE.load_from_checkpoint(checkpoint_path,image_channels=1,z_dim=12, lr =0.0001, use_classification_loss=True, num_classes=num_classes,
+            loss_type="weighted", class_weights=class_weights, device = device)
 model.load_state_dict(torch.load('epoch=22-step=213621.ckpt'))
 model.eval()