Update app.py

app.py

@@ -13,14 +13,11 @@ model = VariationalAutoEncoder(INPUT_DIM, H_DIM, Z_DIM)
 model.load_state_dict(torch.load("MnistVAEmodel.pt"))
 model.eval()
 def predict(img):
-    img = img.convert('1')
     img = transforms.ToTensor()(img)
-    img = transforms.CenterCrop(size=28)(img)
-    print(type(img), img.shape)
     mu, sigma = model.encode(img.view(1, INPUT_DIM))
 
     res = []
-    for example in range(10):
+    for example in range(5):
         epsilon = torch.randn_like(sigma)
         z = mu + sigma * epsilon
         out = model.decode(z)
@@ -30,6 +27,6 @@ def predict(img):
 
 title = "Variational-Autoencoder-on-MNIST "
 description = "TO DO"
-examples = ["original_5.png", "original_8.png"]
-gr.Interface(fn=predict, inputs = gr.inputs.Image(), outputs= gr.Gallery(),
+examples = ["original_8.png"]
+gr.Interface(fn=predict, inputs = gr.inputs.Image(shape=(28,28), image_mode="L"), outputs= gr.Gallery(),
              examples=examples, title=title, description=description).launch(inline=False)