added better functionality

README.md

@@ -1,8 +1,8 @@
 ---
 title: Sefa
-emoji: 💻
+emoji: 🔮
 colorFrom: blue
-colorTo: green
+colorTo: red
 sdk: streamlit
 sdk_version: 1.2.0
 app_file: app.py

app.py

@@ -54,11 +54,16 @@ def synthesize(model, gan_type, code):
     image = postprocess(image)[0]
     return image
 
+def _update_slider():
+    num_semantics = st.session_state["num_semantics"]
+    for sem_idx in range(num_semantics):
+        st.session_state[f"semantic_slider_{sem_idx}"] = 0
+
 
 """Main function (loop for StreamLit)."""
 st.title('Closed-Form Factorization of Latent Semantics in GANs')
 st.sidebar.title('Options')
-reset = st.sidebar.button('Reset')
+st.sidebar.button('Reset', on_click=_update_slider, kwargs={})
 
 model_name = st.sidebar.selectbox(
     'Model to Interpret',
@@ -72,7 +77,7 @@ layer_idx = st.sidebar.selectbox(
 layers, boundaries, eigen_values = factorize_model(model, layer_idx)
 
 num_semantics = st.sidebar.number_input(
-    'Number of semantics', value=5, min_value=0, max_value=None, step=1)
+    'Number of semantics', value=5, min_value=0, max_value=None, step=1, key="num_semantics")
 steps = {sem_idx: 0 for sem_idx in range(num_semantics)}
 if gan_type == 'pggan':
     max_step = 5.0
@@ -87,10 +92,12 @@ for sem_idx in steps:
         value=0.0,
         min_value=-max_step,
         max_value=max_step,
-        step=0.04 * max_step if not reset else 0.0)
+        step=0.04 * max_step,
+        key=f"semantic_slider_{sem_idx}")
 
 image_placeholder = st.empty()
 button_placeholder = st.empty()
+button_totally_random = st.empty()
 
 try:
     base_codes = np.load(f'latent_codes/{model_name}_latents.npy')
@@ -105,13 +112,16 @@ if state.model_name != model_name:
     state.code_idx = 0
     state.codes = base_codes[0:1]
 
-if button_placeholder.button('Random', key=0):
+if button_placeholder.button('Next Sample'):
     state.code_idx += 1
     if state.code_idx < base_codes.shape[0]:
         state.codes = base_codes[state.code_idx][np.newaxis]
     else:
         state.codes = sample(model, gan_type)
 
+if button_totally_random.button('Totally Random'):
+    state.codes = sample(model, gan_type)
+
 code = state.codes.copy()
 for sem_idx, step in steps.items():
     if gan_type == 'pggan':

interface.py

@@ -1,5 +1,6 @@
 # python 3.7
 """Demo."""
+import random
 
 import numpy as np
 import torch
@@ -55,11 +56,18 @@ def synthesize(model, gan_type, code):
     return image
 
 
+def _update_slider():
+    num_semantics = st.session_state["num_semantics"]
+    for sem_idx in range(num_semantics):
+        st.session_state[f"semantic_slider_{sem_idx}"] = 0
+
+
 def main():
     """Main function (loop for StreamLit)."""
+
     st.title('Closed-Form Factorization of Latent Semantics in GANs')
     st.sidebar.title('Options')
-    reset = st.sidebar.button('Reset')
+    st.sidebar.button('Reset', on_click=_update_slider, kwargs={})
 
     model_name = st.sidebar.selectbox(
         'Model to Interpret',
@@ -73,7 +81,7 @@ def main():
     layers, boundaries, eigen_values = factorize_model(model, layer_idx)
 
     num_semantics = st.sidebar.number_input(
-        'Number of semantics', value=5, min_value=0, max_value=None, step=1)
+        'Number of semantics', value=5, min_value=0, max_value=None, step=1, key="num_semantics")
     steps = {sem_idx: 0 for sem_idx in range(num_semantics)}
     if gan_type == 'pggan':
         max_step = 5.0
@@ -88,10 +96,12 @@ def main():
             value=0.0,
             min_value=-max_step,
             max_value=max_step,
-            step=0.04 * max_step if not reset else 0.0)
+            step=0.04 * max_step,
+            key=f"semantic_slider_{sem_idx}")
 
     image_placeholder = st.empty()
     button_placeholder = st.empty()
+    button_totally_random = st.empty()
 
     try:
         base_codes = np.load(f'latent_codes/{model_name}_latents.npy')
@@ -106,13 +116,16 @@ def main():
         state.code_idx = 0
         state.codes = base_codes[0:1]
 
-    if button_placeholder.button('Random', key=0):
+    if button_placeholder.button('Next Sample'):
         state.code_idx += 1
         if state.code_idx < base_codes.shape[0]:
             state.codes = base_codes[state.code_idx][np.newaxis]
         else:
             state.codes = sample(model, gan_type)
 
+    if button_totally_random.button('Totally Random'):
+        state.codes = sample(model, gan_type)
+
     code = state.codes.copy()
     for sem_idx, step in steps.items():
         if gan_type == 'pggan':