Spaces:
Runtime error
Runtime error
File size: 2,466 Bytes
e4eca40 |
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 |
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
display_width = 450
checkpoint_path = "Checkpoints/150epochs.chkpt"
st.title("Generating Abstract Art")
st.text("start generating (left side bar)")
st.text("Made by Xingyu B.")
st.sidebar.subheader("Configurations")
seed = st.sidebar.slider('Seed', -10000, 10000, 0)
num_images = st.sidebar.slider('Number of Images', 1, 10, 1)
use_srgan = st.sidebar.selectbox(
'Apply image enhancement',
('Yes', 'No')
)
generate = st.sidebar.button("Generate")
# 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
st.write("Using DCGAN then ESRGAN upscale...")
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())
for i in range(len(color_match)):
# denormalize and permute to correct color channel
st.image(denormalize_images(color_match[i]).permute(1, 2, 0).numpy(), width=display_width)
# default setting -> vanilla dcgan generation
if use_srgan == "No":
fakes = fakes.cpu()
st.write("Using DCGAN Model...")
for i in range(len(fakes)):
st.image(denormalize_images(fakes[i]).permute(1, 2, 0).numpy(), width=display_width)
|