File size: 2,669 Bytes
1afdd99 3e99287 1afdd99 6e80436 1afdd99 cde0e53 1afdd99 6e80436 f27a3bb 1afdd99 6e80436 1afdd99 6e80436 c814c49 1afdd99 6e80436 1afdd99 6e80436 c814c49 1afdd99 161d5d5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import streamlit as st
import torch
import DCGAN
import SRGAN
from utils import color_histogram_mapping, denormalize_images
import torch.nn as nn
import random
device = torch.device("cpu")
if torch.cuda.is_available():
device = torch.device("cuda")
latent_size = 100
checkpoint_path = "Checkpoints/150epochs.chkpt"
st.title("Generating Abstract Art")
st.sidebar.subheader("Configurations")
seed = st.sidebar.slider('Seed', -100000, 100000, 0)
num_images = st.sidebar.slider('Number of Images', 1, 10, 4)
use_srgan = st.sidebar.selectbox(
'Apply image enhancement',
('Yes', 'No')
)
generate = st.sidebar.button("Generate")
st.write("Get started using the left side bar :sunglasses:")
# caching the expensive model loading
@st.cache(allow_output_mutation=True)
def load_dcgan():
model = torch.jit.load('Checkpoints/dcgan.pt', map_location=device)
return model
@st.cache(allow_output_mutation=True)
def load_esrgan():
model_state_dict = torch.load("Checkpoints/esrgan.pt", map_location=device)
return model_state_dict
# if the user wants to generate something new
if generate:
torch.manual_seed(seed)
random.seed(seed)
sampled_noise = torch.randn(num_images, latent_size, 1, 1, device=device)
generator = load_dcgan()
generator.eval()
with torch.no_grad():
fakes = generator(sampled_noise).detach()
# use srgan for super resolution
if use_srgan == "Yes":
# restore to the checkpoint
esrgan_generator = SRGAN.GeneratorRRDB(channels=3, filters=64, num_res_blocks=23).to(device)
esrgan_checkpoint = load_esrgan()
esrgan_generator.load_state_dict(esrgan_checkpoint)
esrgan_generator.eval()
with torch.no_grad():
enhanced_fakes = esrgan_generator(fakes).detach().cpu()
color_match = color_histogram_mapping(enhanced_fakes, fakes.cpu())
cols = st.columns(num_images)
for i in range(len(color_match)):
# denormalize and permute to correct color channel
cols[i].image(denormalize_images(color_match[i]).permute(1, 2, 0).numpy(), use_column_width=True)
st.image("pointing.jpg", use_column_width=True, caption="https://knowyourmeme.com/memes/two-soyjaks-pointing")
# default setting -> vanilla dcgan generation
if use_srgan == "No":
fakes = fakes.cpu()
cols = st.columns(num_images)
for i in range(len(fakes)):
cols[i].image(denormalize_images(fakes[i]).permute(1, 2, 0).numpy(), use_column_width=True)
st.image("pointing.jpg", use_column_width=True, caption="https://knowyourmeme.com/memes/two-soyjaks-pointing")
|