init commit

.gitattributes

@@ -29,3 +29,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model_weights/** filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,8 +1,8 @@
 ---
 title: OAGen
-emoji: 😻
-colorFrom: pink
-colorTo: gray
+emoji: 💻
+colorFrom: gray
+colorTo: red
 sdk: streamlit
 sdk_version: 1.10.0
 app_file: app.py

app.py

@@ -0,0 +1,52 @@
+import os, glob, sys
+import pickle
+import streamlit as st
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+sys.path.append('stylegan3')
+
+class SampleFromGAN:
+    def __init__(self, G, z_shp, in_gpu=False) -> None:
+        self.G = G
+        self.in_gpu = in_gpu
+        self.z_shp = z_shp #[#images, z_dim]
+    def __call__(self,):
+        z = torch.randn(self.z_shp)
+        if self.in_gpu:
+            z = z.cuda()
+        ims = G(z, c=None)
+        ims = ims[:,0,...]
+        return ims
+
+class Plot:
+    def __init__(self, im_gen) -> None:
+        self.im_gen = im_gen
+        assert callable(im_gen)
+    def __call__(self):
+        ims = self.im_gen()
+        # plot first image
+        im = ims[0,...]
+        fig, ax = plt.subplots(1, figsize=(12,12))
+        fig.subplots_adjust(left=0,right=1,bottom=0,top=1)
+        ax.imshow(im, cmap='gray')
+        ax.axis('tight')
+        ax.axis('off')
+        st.pyplot(fig)
+
+
+# path_ckpt = "/home/firat/saved_models/DLBIRHOUI/stylegan/00001-stylegan2-SWFD-linear_BP-linear_BP-gpus2-batch32-gamma8.2-augnoaug"
+path_ckpt = "./model_weights"
+fname_pkl = os.path.join(path_ckpt, 'network-snapshot-005000.pkl')
+in_gpu = False
+num_images = 1
+with open(fname_pkl, 'rb') as f:
+    G = pickle.load(f)['G_ema']  # torch.nn.ModuleDict
+if in_gpu:
+    G = G.cuda()
+
+sampler = SampleFromGAN(G=G, z_shp=[num_images, G.z_dim], in_gpu=in_gpu)
+
+button_on_click = Plot(im_gen=sampler)
+button_gen_clicked = st.button(label='Generate an image', key='n', on_click=button_on_click)
+

requirements.txt

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