|
import pandas as pd |
|
import numpy as np |
|
import streamlit as st |
|
|
|
from models import Generator, Discriminrator |
|
from StyleMix import style_mix |
|
import torch |
|
import torchvision.transforms as T |
|
from torchvision.utils import make_grid |
|
from PIL import Image |
|
|
|
from streamlit_lottie import st_lottie |
|
from streamlit_option_menu import option_menu |
|
import requests |
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
|
|
model_name = { |
|
"aurora": 'huggan/fastgan-few-shot-aurora', |
|
"painting": 'huggan/fastgan-few-shot-painting', |
|
"shell": 'huggan/fastgan-few-shot-shells', |
|
"fauvism": 'huggan/fastgan-few-shot-fauvism-still-life', |
|
"universe": 'huggan/fastgan-few-shot-universe', |
|
"grumpy cat": 'huggan/fastgan-few-shot-grumpy-cat', |
|
"anime": 'huggan/fastgan-few-shot-anime-face', |
|
"moon gate": 'huggan/fastgan-few-shot-moongate', |
|
} |
|
|
|
|
|
def load_generator(model_name_or_path): |
|
generator = Generator(in_channels=256, out_channels=3) |
|
generator = generator.from_pretrained(model_name_or_path, in_channels=256, out_channels=3) |
|
_ = generator.to(device) |
|
_ = generator.eval() |
|
|
|
return generator |
|
|
|
def _denormalize(input: torch.Tensor) -> torch.Tensor: |
|
return (input * 127.5) + 127.5 |
|
|
|
|
|
def generate_images(generator, number_imgs): |
|
noise = torch.zeros(number_imgs, 256, 1, 1, device=device).normal_(0.0, 1.0) |
|
with torch.no_grad(): |
|
gan_images, _ = generator(noise) |
|
|
|
gan_images = _denormalize(gan_images.detach()).cpu() |
|
gan_images = [i for i in gan_images] |
|
gan_images = [make_grid(i, nrow=1, normalize=True) for i in gan_images] |
|
gan_images = [i.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() for i in gan_images] |
|
gan_images = [Image.fromarray(i) for i in gan_images] |
|
return gan_images |
|
|
|
def load_lottieurl(url: str): |
|
r = requests.get(url) |
|
if r.status_code != 200: |
|
return None |
|
return r.json() |
|
|
|
def show_model_summary(expanded): |
|
st.subheader("Model gallery") |
|
with st.expander('Image gallery', expanded=expanded): |
|
col1, col2, col3, col4 = st.columns(4) |
|
with col1: |
|
st.markdown('Fauvism GAN [model](https://huggingface.co/huggan/fastgan-few-shot-fauvism-still-life)', unsafe_allow_html=True) |
|
st.image('assets/image/fauvism.png', width=200) |
|
st.markdown('Painting GAN [model](https://huggingface.co/huggan/fastgan-few-shot-painting)', unsafe_allow_html=True) |
|
st.image('assets/image/painting.png', width=200) |
|
|
|
with col2: |
|
st.markdown('Aurora GAN [model](https://huggingface.co/huggan/fastgan-few-shot-aurora)', unsafe_allow_html=True) |
|
st.image('assets/image/aurora.png', width=200) |
|
st.markdown('Universe GAN [model](https://huggingface.co/huggan/fastgan-few-shot-universe)', unsafe_allow_html=True) |
|
st.image('assets/image/universe.png', width=200) |
|
|
|
with col3: |
|
st.markdown('Anime GAN [model](https://huggingface.co/huggan/fastgan-few-shot-anime-face)', unsafe_allow_html=True) |
|
st.image('assets/image/anime.png', width=200) |
|
st.markdown('Shell GAN [model](https://huggingface.co/huggan/fastgan-few-shot-shells)', unsafe_allow_html=True) |
|
st.image('assets/image/shell.png', width=200) |
|
|
|
with col4: |
|
st.markdown('Grumpy cat GAN [model](https://huggingface.co/huggan/fastgan-few-shot-grumpy-cat)', unsafe_allow_html=True) |
|
st.image('assets/image/grumpy_cat.png', width=200) |
|
st.markdown('Moon gate GAN [model](https://huggingface.co/huggan/fastgan-few-shot-moongate)', unsafe_allow_html=True) |
|
st.image('assets/image/moon_gate.png', width=200) |
|
|
|
with st.expander('Video gallery', expanded=True): |
|
cols=st.columns(4) |
|
|
|
cols[0].write("Universe GAN") |
|
cols[0].video('assets/video/universe.mp4') |
|
cols[0].write("Fauvism still life GAN") |
|
cols[0].video('assets/video/fauvism.mp4') |
|
|
|
cols[1].write("Aurora GAN") |
|
cols[1].video('assets/video/aurora.mp4') |
|
cols[1].write("Moon gate GAN") |
|
cols[1].video('assets/video/moongate.mp4') |
|
|
|
cols[2].write("Anime GAN") |
|
cols[2].video('assets/video/anime.mp4') |
|
cols[2].write("Painting GAN") |
|
cols[2].video('assets/video/painting.mp4') |
|
|
|
cols[3].write("Grumpy cat GAN") |
|
cols[3].video('assets/video/grumpy.mp4') |
|
|
|
|
|
def main(): |
|
|
|
st.set_page_config( |
|
page_title="FastGAN Generator", |
|
page_icon="🖥️", |
|
layout="wide", |
|
initial_sidebar_state="expanded" |
|
) |
|
|
|
lottie_penguin = load_lottieurl('https://assets7.lottiefiles.com/packages/lf20_mm4bsl3l.json') |
|
|
|
with st.sidebar: |
|
st_lottie(lottie_penguin, height=200) |
|
choose = option_menu("FastGAN", ["Model Gallery", "Generate images", "Mix style"], |
|
icons=['collection', 'file-plus', 'intersect'], |
|
menu_icon="infinity", default_index=0, |
|
styles={ |
|
"container": {"padding": ".0rem", "font-size": "14px"}, |
|
"nav-link-selected": {"color": "#000000", "font-size": "16px"}, |
|
} |
|
) |
|
st.sidebar.markdown( |
|
""" |
|
___ |
|
<p style='text-align: center'> |
|
FastGAN is a few-shot GAN model trained on high-fidelity images which requires less computation resource and samples for training. |
|
<br/> |
|
<a href="https://arxiv.org/abs/2101.04775" target="_blank">Article</a> |
|
</p> |
|
<p style='text-align: center; font-size: 14px;'> |
|
Model training and Spaces creating by |
|
<br/> |
|
<a href="https://www.linkedin.com/in/vumichien/" target="_blank">Chien Vu</a> | <a href="https://www.linkedin.com/in/nhu-hoang/" target="_blank">Nhu Hoang</a> |
|
<br/> |
|
</p> |
|
""", |
|
unsafe_allow_html=True, |
|
) |
|
|
|
if choose == 'Model Gallery': |
|
st.header("Welcome to FastGAN") |
|
show_model_summary(True) |
|
elif choose == 'Generate images': |
|
st.header("Generate images") |
|
col11, col12, col13 = st.columns([3,3.5,3.5]) |
|
with col11: |
|
img_type = st.selectbox("Choose type of image to generate", index=0, |
|
options=["aurora", "anime", "painting", "fauvism", "shell", "universe", "grumpy cat", "moon gate"]) |
|
|
|
number_imgs = st.slider('How many images you want to generate ?', min_value=1, max_value=5) |
|
if number_imgs is None: |
|
st.write('Invalid number ! Please insert number of images to generate !') |
|
raise ValueError('Invalid number ! Please insert number of images to generate !') |
|
|
|
generate_button = st.button('Get Image') |
|
if generate_button: |
|
st.markdown(""" |
|
<small><i>Predictions may take up to 1 minute under high load. Please stand by.</i></small> |
|
""", |
|
unsafe_allow_html=True,) |
|
|
|
if generate_button: |
|
with col11: |
|
with st.spinner(text=f"Loading selected model..."): |
|
generator = load_generator(model_name[img_type]) |
|
with st.spinner(text=f"Generating images..."): |
|
gan_images = generate_images(generator, number_imgs) |
|
with col12: |
|
st.image(gan_images[0], width=300) |
|
if len(gan_images) > 1: |
|
with col13: |
|
if len(gan_images) <= 2: |
|
st.image(gan_images[1], width=300) |
|
else: |
|
st.image(gan_images[1:], width=150) |
|
|
|
elif choose == 'Mix style': |
|
st.header("Mix style") |
|
st.markdown( |
|
""" |
|
<p style='text-align: left'> |
|
Get the style representations of 2 images generated from the model to create a new one that mixes the style of two. |
|
</p> |
|
""", |
|
unsafe_allow_html=True, |
|
) |
|
st.markdown("""___""") |
|
col21, col22 = st.columns([3, 6]) |
|
with col21: |
|
img_type = st.selectbox("Choose type of image to mix", index=0, |
|
options=["aurora", "anime", "painting", "fauvism", "shell", "universe", "grumpy cat", "moon gate"]) |
|
number_imgs = st.slider('How many images you want to generate ?', min_value=1, max_value=3) |
|
generate_button = st.button('Mix style') |
|
|
|
if generate_button: |
|
with col21: |
|
with st.spinner(text=f"Mixing styles..."): |
|
mix_imgs = style_mix(model_name[img_type], number_imgs, device) |
|
mix_imgs = make_grid(mix_imgs, nrow=number_imgs+1, normalize=True) |
|
mix_imgs = mix_imgs.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() |
|
mix_imgs = Image.fromarray(mix_imgs) |
|
with col22: |
|
st.image(mix_imgs, width=600) |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
