almost final

app.py

@@ -1,20 +1,17 @@
-from pydoc import ModuleScanner
-import re
 import streamlit as st # HF spaces at v1.2.0
 from demo import load_model,generate,get_dataset,embed,make_meme
 from PIL import Image
 import numpy as np
+import io
 # TODOs
 # Add markdown short readme project intro
+# Add link to wandb logs
 
 
 st.sidebar.subheader("This butterfly does not exist! ")
 st.sidebar.image("assets/logo.png", width=200)
 
-st.header("ButterflyGAN")
-
-st.write("Demo prep still in progress!! Come back later")
-
+st.title("ButterflyGAN")
 
 @st.experimental_singleton
 def load_model_intocache(model_name,model_version):
@@ -27,16 +24,29 @@ def load_dataset():
     dataset=get_dataset()
     return dataset
 
+@st.experimental_singleton
+def load_variables():# Don't want to open read files over and over. not sure if it makes a diff
+    st.session_state['latent_walk_code']=open("assets/code_snippets/latent_walk.py").read()
+    st.session_state['latent_walk_code_music']=open("assets/code_snippets/latent_walk_music.py").read()
+
+def img2download(image):
+        imgByteArr = io.BytesIO()
+        image.save(imgByteArr, format="JPEG")
+        imgByteArr = imgByteArr.getvalue()
+        return imgByteArr
+
 model_name='ceyda/butterfly_cropped_uniq1K_512'
 # model_version='0edac54b81958b82ce9fd5c1f688c33ac8e4f223'
 model_version=None ##TBD
 model=load_model_intocache(model_name,model_version)
 dataset=load_dataset()
+load_variables()
 
 generate_menu="🦋 Make butterflies"
 latent_walk_menu="🎧 Take a latent walk"
 make_meme_menu="🐦 Make a meme"
 mosaic_menu="👀 See the mosaic"
+fun_menu="Release the butterflies"
 
 screen = st.sidebar.radio("Pick a destination",[generate_menu,latent_walk_menu,make_meme_menu,mosaic_menu])
 
@@ -52,8 +62,10 @@ if screen == generate_menu:
         st.session_state['ims'] = None
         run()
     ims=st.session_state["ims"]
-    
-    runb=st.button("Generate", on_click=run)
+    st.write("Light-GAN model trained on 1000 butterfly images taken from the Smithsonian Museum collection. \n \
+Based on [paper:](https://openreview.net/forum?id=1Fqg133qRaI) *Towards Faster and Stabilized GAN Training for High-fidelity Few-shot Image Synthesis*")
+
+    runb=st.button("Generate", on_click=run ,help="generated on the fly maybe slow")
     if ims is not None:
         cols=st.columns(col_num)
         picks=[False]*batch_size
@@ -79,34 +91,65 @@ if screen == generate_menu:
                     scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(ims[i]), k=5)
                     for r in retrieved_examples["image"]:
                         cols[i].image(r)
-
-    st.write(f"Latent dimension: {model.latent_dim}, Image size:{model.image_size}")
+    st.write("Nearest neighbors found in the training set according to L2 distance on 'microsoft/beit-base-patch16-224' embeddings")
+    st.write(f"Latent dimension: {model.latent_dim}, image size:{model.image_size}")
 
 elif screen == latent_walk_menu:
-    st.write("Take a latent walk :musical_note:")
+
+    latent_walk_code=open("assets/code_snippets/latent_walk.py").read()
+    latent_walk_music_code=open("assets/code_snippets/latent_walk_music.py").read()
+    st.write("Take a latent walk :musical_note: with cute butterflies")
 
     cols=st.columns(3)
 
+    cols[0].caption("A regular walk (no music)")
     cols[0].video("assets/latent_walks/regular_walk.mp4")
-    cols[0].caption("Regular walk")
+    
+    cols[1].caption("Walk with music :butterfly:")
     cols[1].video("assets/latent_walks/walk_happyrock.mp4")
-    cols[1].caption("walk with music :butterfly:")
+    cols[2].caption("Walk with music :butterfly:")
     cols[2].video("assets/latent_walks/walk_cute.mp4")
-    cols[2].caption(":musical_note: walk with cute butterflies")
-    cols[1].caption("Royalty Free Music from Bensound")
+    
+    st.caption("Royalty Free Music from Bensound")
+    st.write("🎧Did those butterflies seem to be dancing to the music?!Here is the secret:")
+    with st.expander("See the Code Snippets"):
+        st.write("A regular latent walk:")
+        st.code(st.session_state['latent_walk_code'], language='python')
+        st.write(":musical_note: latent walk with music:")
+        st.code(st.session_state['latent_walk_code_music'], language='python')
 
 
 elif screen == make_meme_menu:
-    im = generate(model,1)[0]
-    no_bg=st.checkbox("Remove background?",True)
-    meme_text=st.text_input("Meme text","Is this a pigeon?")
-    meme=make_meme(im,text=meme_text,show_text=True,remove_background=no_bg)
+    if "pigeon" not in st.session_state:
+        st.session_state['pigeon'] = generate(model,1)[0]
+
+    def get_pigeon():
+        st.session_state['pigeon'] = generate(model,1)[0]
+    
+    cols= st.columns(2)
+    cols[0].button("change pigeon",on_click=get_pigeon)
+    no_bg=cols[1].checkbox("Remove background?",True,help="Remove the background from pigeon")
+    show_text=cols[1].checkbox("Show text?",True)
+    
+    meme_text=st.text_input("Enter text","Is this a pigeon?")
+   
+    
+    meme=make_meme(st.session_state['pigeon'],text=meme_text,show_text=show_text,remove_background=no_bg)
     st.image(meme)
+    coly=st.columns(2)
+    coly[0].download_button("Download", img2download(meme),mime="image/jpeg")
+    coly[1].write("Made a cool one? [Share](https://twitter.com/intent/tweet?text=Check%20out%20the%20demo%20for%20Butterfly%20GAN%20%F0%9F%A6%8Bhttps%3A//huggingface.co/spaces/huggan/butterfly-gan%0Amade%20by%20%40ceyda_cinarel%20%26%20%40johnowhitaker%20) on Twitter")
 
 
 elif screen == mosaic_menu:
-    st.markdown("Todo add explanation about data")
-    st.image("assets/training_data_lowres.png")
+    cols=st.columns(2)
+    cols[0].markdown("These are all the butterflies in our [training set](https://huggingface.co/huggan/smithsonian_butterflies_subset)")
+    cols[0].image("assets/train_data_mosaic_lowres.jpg")
+    cols[0].write("🔎 view the high-res version [here](https://www.easyzoom.com/imageaccess/0c77e0e716f14ea7bc235447e5a4c397)")
+
+    cols[1].markdown("These are the butterflies our model generated.")
+    cols[1].image("assets/gen_mosaic_lowres.jpg")
+    cols[1].write("🔎 view the high-res version [here](https://www.easyzoom.com/imageaccess/cbb04e81106c4c54a9d9f9dbfb236eab)")
 
 
 # footer stuff
@@ -116,6 +159,6 @@ st.sidebar.caption(f"[Model](https://huggingface.co/ceyda/butterfly_cropped_uniq
 # Credits
 st.sidebar.caption(f"Made during the [huggan](https://github.com/huggingface/community-events) hackathon")
 st.sidebar.caption(f"Contributors:")
-st.sidebar.caption(f"[Ceyda Cinarel](https://huggingface.co/ceyda) & [Jonathan Whitaker](https://datasciencecastnet.home.blog/)")
+st.sidebar.caption(f"[Ceyda Cinarel](https://github.com/cceyda) & [Jonathan Whitaker](https://datasciencecastnet.home.blog/)")
 
 ## Feel free to add more & change stuff ^

assets/code_snippets/latent_walk.py

@@ -0,0 +1,15 @@
+# Some parameters
+n_points = 6 #@param
+n_steps = 300 #@param
+latents = torch.randn(n_points, 256)
+
+# Loop through generating the frames
+frames = []
+for i in tqdm(range(n_steps)):
+  p1 = max(0, int(n_points*i/n_steps))
+  p2 = min(n_points, int(n_points*i/n_steps)+1)%n_points # so it wraps back to 0
+  frac = (i-(p1*(n_steps/n_points))) / (n_steps/n_points)
+  l = latents[p1]*(1-frac) + latents[p2]*frac
+  im = model.G(l.unsqueeze(0)).clamp_(0., 1.)
+  frame=(im[0].permute(1, 2, 0).detach().cpu().numpy()*255).astype(np.uint8)
+  frames.append(frame)

assets/code_snippets/latent_walk_music.py

@@ -0,0 +1,55 @@
+#Code Author: Jonathan Whitaker 😎
+
+import librosa
+import soundfile as sf
+from scipy.signal import savgol_filter
+
+# The driving audio file
+audio_file = './sounds/bensound-cute.wav' #@param
+
+# How many points in the base latent walk loop
+n_points = 6 #@param
+
+# Smooths the animation effect, smaller=jerkier, must be odd
+filter_window_size=301 #@param 
+
+# How much should we scale position based on music vs the base path?
+chr_scale = 0.5 #@param
+base_scale = 0.3 #@param
+
+# Load the file
+X, sample_rate = sf.read(audio_file, dtype='float32')
+
+X= X[:int(len(X)*0.5)]
+
+# Remove percussive elements
+harmonic = librosa.effects.harmonic(X[:,0])
+
+# Get chroma_stft (power in different notes)
+chroma = librosa.feature.chroma_stft(harmonic) # Just one channel
+
+# Smooth these out
+chroma = savgol_filter(chroma, filter_window_size, 3)
+
+# Calculate how many frames we want
+fps = 25
+duration = X.shape[0] / sample_rate
+print('Duration:', duration)
+n_steps = int(fps * duration)
+print('N frames:', n_steps, fps * duration)
+
+latents = torch.randn(n_points, 256)*base_scale
+chroma_latents = torch.randn(12, 256)*chr_scale
+
+frames=[]
+for i in tqdm(range(n_steps)):
+  p1 = max(0, int(n_points*i/n_steps))
+  p2 = min(n_points, int(n_points*i/n_steps)+1)%n_points # so it wraps back to 0
+  frac = (i-(p1*(n_steps/n_points))) / (n_steps/n_points)
+  l = latents[p1]*(1-frac) + latents[p2]*frac
+  for c in range(12): # HERE adding the music influence to the latent 
+    scale_factor = chroma[c, int(i*chroma.shape[1]/n_steps)]
+    l += chroma_latents[c]*chr_scale*scale_factor
+  im = model.G(l.unsqueeze(0)).clamp_(0., 1.)
+  frame=(im[0].permute(1, 2, 0).detach().cpu().numpy()*255).astype(np.uint8)
+  frames.append(frame)