built space

MnistVAEmodel.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d6ab1a824858a37b3dbeffce09cd2de481906e689b4817e505cb2550e992d3d
+size 4796991

app.py

@@ -0,0 +1,35 @@
+import torch
+from model import VariationalAutoEncoder
+from torchvision import transforms
+from PIL import Image
+import gradio as gr
+
+
+INPUT_DIM = 784
+H_DIM = 512
+Z_DIM = 256
+
+model = VariationalAutoEncoder(INPUT_DIM, H_DIM, Z_DIM)
+model.load_state_dict(torch.load("MnistVAEmodel.pth"))
+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):
+        epsilon = torch.randn_like(sigma)
+        z = mu + sigma * epsilon
+        out = model.decode(z)
+        out = out.view(-1,1,28,28)
+        res.append(transforms.ToPILImage()(out[0]))
+    return res
+
+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.outputs.Gallery(),
+             examples=examples, title=title, description=description).launch(inline=False)

model.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+
+
+class VariationalAutoEncoder(nn.Module):
+    # Input image -> hidden dim -> mean, std -> parametirazation trick -> Decoder -> output image
+    def __init__(self, inpud_dim, h_dim=200, z_dim=20):
+        super().__init__()
+
+        # encoder
+        self.img_2hid = nn.Linear(inpud_dim, h_dim)
+        self.hid_2mu = nn.Linear(h_dim, z_dim)
+        self.hid_2sigma = nn.Linear(h_dim, z_dim)
+
+        # decoder
+        self.z_2hi = nn.Linear(z_dim, h_dim)
+        self.hid_2img = nn.Linear(h_dim, inpud_dim)
+
+        self.relu = nn.ReLU()
+
+    def encode(self, x):
+        # q_phi(z/x)
+        h = self.relu(self.img_2hid(x))
+        mu, sigma = self.hid_2mu(h), self.hid_2sigma(h)
+
+        return mu, sigma
+
+    def decode(self, z):
+        # p_theta(x/z)
+        h = self.relu(self.z_2hi(z))
+        x = self.hid_2img(h)
+        return torch.sigmoid(x) # image values should be between zero and one.
+
+    def forward(self, x):
+        mu, sigma = self.encode(x)
+        # parametirazation trick
+        epsilon = torch.randn_like(sigma) # Returns a tensor with the same size as input that is filled with random numbers from a normal distribution with mean 0 and variance 1
+        z_reparametrized = mu + sigma * epsilon
+        x_reconstructed = self.decode(z_reparametrized)
+        return x_reconstructed, mu, sigma  # 2 parts of loss: 1- mu, sigma pushed to normal distribution. 2 the x_reconstructed should be same as x
+
+if __name__ == "__main__":
+
+    x = torch.randn(4,28*28)
+    vae = VariationalAutoEncoder(inpud_dim=784)
+    x_reconstructed, mu, sigma = vae(x)
+    print(x_reconstructed.shape)
+    print(mu.shape)
+    print(sigma.shape)
+    print(torch.mean(mu))

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+torch
+torchvision