binarize

app.py

@@ -10,7 +10,7 @@ models = {
     "Deep ConvAE": torch.load("deep_convae.th", map_location="cpu"),
     "Dense K-Sparse": torch.load("fc_sparse.th", map_location="cpu"),
 }
-def gen(md, model_name, seed, nb_iter, nb_samples, width, height, nb_active, only_last, black_bg):
+def gen(md, model_name, seed, nb_iter, nb_samples, width, height, nb_active, only_last, black_bg, binarize, binarize_threshold):
     torch.manual_seed(int(seed))
     bs = 64
     model = models[model_name]
@@ -22,6 +22,7 @@ def gen(md, model_name, seed, nb_iter, nb_samples, width, height, nb_active, onl
         nb_examples=int(nb_samples), 
         w=int(width), h=int(height), c=1, 
         batch_size=bs,
+        binarize_threshold=binarize_threshold if binarize else None,
     )
     if not black_bg:
         samples = 1 - samples
@@ -48,7 +49,7 @@ iface = gr.Interface(
     fn=gen,
     inputs=[
         gr.Markdown(text),
-        gr.Dropdown(list(models.keys()), value="Deep ConvAE"), gr.Number(value=0), gr.Number(value=25), gr.Number(value=1), gr.Number(value=28), gr.Number(value=28),gr.Slider(minimum=0,maximum=800, value=800, step=1), gr.Checkbox(value=False, label="Only show last iteration"), gr.Checkbox(value=True, label="Black background")
+        gr.Dropdown(list(models.keys()), value="Deep ConvAE"), gr.Number(value=0), gr.Number(value=25), gr.Number(value=1), gr.Number(value=28), gr.Number(value=28),gr.Slider(minimum=0,maximum=800, value=800, step=1), gr.Checkbox(value=False, label="Only show last iteration"), gr.Checkbox(value=True, label="Black background"), gr.Checkbox(value=False, label="binarize"), gr.Number(value=0.5)
     ],
     outputs="image"
 )

cli.py

@@ -91,7 +91,7 @@ def save_weights(m, folder='.'):
             imsave('{}/feat.png'.format(folder), gr)
 
 @torch.no_grad()
-def iterative_refinement(ae, nb_examples=1, nb_iter=10, w=28, h=28, c=1, batch_size=None):
+def iterative_refinement(ae, nb_examples=1, nb_iter=10, w=28, h=28, c=1, batch_size=None, binarize_threshold=None):
     if batch_size is None:
         batch_size = nb_examples
     x = torch.rand(nb_iter, nb_examples, c, w, h)
@@ -99,6 +99,8 @@ def iterative_refinement(ae, nb_examples=1, nb_iter=10, w=28, h=28, c=1, batch_s
         for j in range(0, nb_examples, batch_size):
             oldv = x[i-1][j:j + batch_size].to(device)
             newv = ae(oldv)
+            if binarize_threshold is not None:
+                newv = (newv>binarize_threshold).float()
             newv = newv.data.cpu()
             x[i][j:j + batch_size] = newv
     return x