make demo prettier half way there

.gitattributes

@@ -27,3 +27,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.faiss filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,22 +1,25 @@
+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
-
+from demo import load_model,generate,get_dataset,embed,make_meme
+from PIL import Image
+import numpy as np
 # TODOs
 # Add markdown short readme project intro
 
 
-st.sidebar.image("assets/logo.png", use_column_width=True)
+st.sidebar.subheader("This butterfly does not exist! ")
+st.sidebar.image("assets/logo.png", width=200)
+
 st.header("ButterflyGAN")
-st.caption("This butterfly does not exist! ")
-st.write("Demo prep still in progress!!")
+
+st.write("Demo prep still in progress!! Come back later")
 
 
 @st.experimental_singleton
-def load_model_intocache(model_name):
-    
+def load_model_intocache(model_name,model_version):
     # model_name='ceyda/butterfly_512_base'
-    gan = load_model(model_name)
+    gan = load_model(model_name,model_version)
     return gan
 
 @st.experimental_singleton
@@ -25,33 +28,46 @@ def load_dataset():
     return dataset
 
 model_name='ceyda/butterfly_cropped_uniq1K_512'
-model=load_model_intocache(model_name)
+# model_version='0edac54b81958b82ce9fd5c1f688c33ac8e4f223'
+model_version=None ##TBD
+model=load_model_intocache(model_name,model_version)
 dataset=load_dataset()
 
-screen = st.sidebar.radio("Pick a destination",["Make butterflies","Take a latent walk", "See the data mosaic"])
+generate_menu="🦋 Make butterflies"
+latent_walk_menu="🎧 Take a latent walk"
+make_meme_menu="🐦 Make a meme"
+mosaic_menu="👀 See the mosaic"
 
-if screen == "Make butterflies":
+screen = st.sidebar.radio("Pick a destination",[generate_menu,latent_walk_menu,make_meme_menu,mosaic_menu])
 
+if screen == generate_menu:
     
-    
-    if 'ims' not in st.session_state:
-        st.session_state['ims'] = None
-
-    ims=st.session_state["ims"]
     batch_size=4 #generate 4 butterflies 
+    col_num=4
     def run():
         with st.spinner("Generating..."):
             ims=generate(model,batch_size)
             st.session_state['ims'] = ims
-
+    if 'ims' not in st.session_state:
+        st.session_state['ims'] = None
+        run()
+    ims=st.session_state["ims"]
+    
     runb=st.button("Generate", on_click=run)
     if ims is not None:
-        cols=st.columns(batch_size)
+        cols=st.columns(col_num)
         picks=[False]*batch_size
-        for i,im in enumerate(ims):
+        for j,im in enumerate(ims):
+            i=j%col_num
             cols[i].image(im)
-            picks[i]=cols[i].button("Find Nearest",key="pick_"+str(i))
-            # if picks[i]:
+            picks[j]=cols[i].button("Find Nearest",key="pick_"+str(j))
+            # meme_it=cols[i].button("What is this?",key="meme_"+str(j))
+            # if meme_it:
+            #     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)
+            #     st.image(meme)
+            # if picks[j]:
             #     scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(im), k=5)
             #     for r in retrieved_examples["image"]:
             #         st.image(r)
@@ -66,13 +82,40 @@ if screen == "Make butterflies":
 
     st.write(f"Latent dimension: {model.latent_dim}, Image size:{model.image_size}")
 
-elif screen == "Take a latent walk":
-    st.write("Take a latent walk")
+elif screen == latent_walk_menu:
+    st.write("Take a latent walk :musical_note:")
+
+    cols=st.columns(3)
+
+    cols[0].video("assets/latent_walks/regular_walk.mp4")
+    cols[0].caption("Regular walk")
+    cols[1].video("assets/latent_walks/walk_happyrock.mp4")
+    cols[1].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")
 
-elif screen == "Input data mosaic":
+
+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)
+    st.image(meme)
+
+
+elif screen == mosaic_menu:
     st.markdown("Todo add explanation about data")
     st.image("assets/training_data_lowres.png")
 
 
 # footer stuff
-st.sidebar.info(f"Model {model_name} is loaded")
+st.sidebar.caption(f"[Model](https://huggingface.co/ceyda/butterfly_cropped_uniq1K_512) & [Dataset](https://huggingface.co/huggan/smithsonian_butterflies_subset) used")
+# Link project repo( scripts etc )
+
+# 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/)")
+
+## Feel free to add more & change stuff ^

assets/latent_walks/regular_walk.mp4

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

assets/latent_walks/walk_cute.mp4

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

assets/latent_walks/walk_happyrock.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5219340b5fe3e509f02e83a0f0c972bd8f0ecd76df4019cedb0abe373b0fb5e8
+size 6594393

demo.py

@@ -1,6 +1,67 @@
 import torch
 from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
 from datasets import load_dataset
+from PIL import Image
+import numpy as np
+import paddlehub as hub
+import random
+from PIL import ImageDraw,ImageFont
+
+import streamlit as st
+
+@st.experimental_singleton
+def load_bg_model():
+    bg_model = hub.Module(name='U2NetP', directory='assets/models/')
+    return bg_model
+
+
+bg_model = load_bg_model()
+def remove_bg(img):
+    result = bg_model.Segmentation(
+        images=[np.array(img)[:,:,::-1]],
+        paths=None,
+        batch_size=1,
+        input_size=320,
+        output_dir=None,
+        visualization=False)
+    output = result[0]
+    mask=Image.fromarray(output['mask'])
+    front=Image.fromarray(output['front'][:,:,::-1]).convert("RGBA")
+    front.putalpha(mask)
+    return front
+
+meme_template=Image.open("./assets/pigeon_meme.jpg").convert("RGBA")
+def make_meme(pigeon,text="Is this a pigeon?",show_text=True,remove_background=True):
+    
+    meme=meme_template.copy()
+    approx_butterfly_center=(850,30)
+
+    if remove_background:
+        pigeon=remove_bg(pigeon)
+        meme=meme.convert("RGBA")
+
+    random_rotate=random.randint(-30,30)
+    random_size=random.randint(150,200)
+    pigeon=pigeon.resize((random_size,random_size)).rotate(random_rotate,expand=True)
+
+    meme.alpha_composite(pigeon, approx_butterfly_center)
+
+    #ref: https://blog.lipsumarium.com/caption-memes-in-python/
+    def drawTextWithOutline(text, x, y):
+        draw.text((x-2, y-2), text,(0,0,0),font=font)
+        draw.text((x+2, y-2), text,(0,0,0),font=font)
+        draw.text((x+2, y+2), text,(0,0,0),font=font)
+        draw.text((x-2, y+2), text,(0,0,0),font=font)
+        draw.text((x, y), text, (255,255,255), font=font)
+
+    if show_text:
+        draw = ImageDraw.Draw(meme)
+        font_size=52
+        font = ImageFont.truetype("assets/impact.ttf", font_size)
+        w, h = draw.textsize(text, font) # measure the size the text will take
+        drawTextWithOutline(text, meme.width/2 - w/2, meme.height - font_size*2)
+    meme = meme.convert("RGB")    
+    return meme
 
 def get_train_data(dataset_name="huggan/smithsonian_butterflies_subset"):
     dataset=load_dataset(dataset_name)
@@ -8,13 +69,13 @@ def get_train_data(dataset_name="huggan/smithsonian_butterflies_subset"):
     return dataset["train"]
 
 from transformers import BeitFeatureExtractor, BeitForImageClassification
-feature_extractor = BeitFeatureExtractor.from_pretrained('microsoft/beit-base-patch16-224')
-model = BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224')
+emb_feature_extractor = BeitFeatureExtractor.from_pretrained('microsoft/beit-base-patch16-224')
+emb_model = BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224')
 def embed(images):
-    inputs = feature_extractor(images=images, return_tensors="pt")
-    outputs = model(**inputs,output_hidden_states= True)
+    inputs = emb_feature_extractor(images=images, return_tensors="pt")
+    outputs = emb_model(**inputs,output_hidden_states= True)
     last_hidden=outputs.hidden_states[-1]
-    pooler=model.base_model.pooler
+    pooler=emb_model.base_model.pooler
     final_emb=pooler(last_hidden).detach().numpy()
     return final_emb    
     
@@ -29,15 +90,15 @@ def get_dataset():
     dataset.load_faiss_index('beit_embeddings', 'beit_index.faiss')
     return dataset
 
-def load_model(model_name='ceyda/butterfly_cropped_uniq1K_512'):
-    gan = LightweightGAN.from_pretrained(model_name)
+def load_model(model_name='ceyda/butterfly_cropped_uniq1K_512',model_version="95a9596a1e47e2419c9bd5252d809eecb14fdcf4"):
+    gan = LightweightGAN.from_pretrained(model_name,version=model_version)
     gan.eval()
     return gan
     
 def generate(gan,batch_size=1):
     with torch.no_grad():
-        ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0., 1.)
-        ims = ims.permute(0,2,3,1).detach().cpu().numpy()
+        ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0., 1.)*255
+        ims = ims.permute(0,2,3,1).detach().cpu().numpy().astype(np.uint8)
     return ims
 
 def interpolate():

requirements.txt

@@ -1,3 +1,5 @@
 git+https://github.com/huggingface/community-events.git@3fea10c5d5a50c69f509e34cd580fe9139905d04#egg=huggan
 transformers
-faiss-cpu
+faiss-cpu
+paddlehub
+paddlepaddle