renamed files

app.py

@@ -0,0 +1,86 @@
+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)
+
+
+
+

utils.py

@@ -0,0 +1,40 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import torchvision.utils as vutils
+import torchvision.transforms as transforms
+from skimage.exposure import match_histograms
+import torch
+
+# contains utility functions that we need in the main program
+
+# matches the color histogram of original and the super resolution output
+def color_histogram_mapping(images, references):
+    matched_list = []
+    for i in range(len(images)):
+        matched = match_histograms(images[i].permute(1, 2, 0).numpy(), references[i].permute(1, 2, 0).numpy(),
+                                   channel_axis=-1)
+        matched_list.append(matched)
+    return torch.tensor(np.array(matched_list)).permute(0, 3, 1, 2)
+
+
+def visualize_generations(seed, images):
+    plt.figure(figsize=(16, 16))
+    plt.title(f"Seed: {seed}")
+    plt.axis("off")
+    plt.imshow(np.transpose(vutils.make_grid(images, padding=2, nrow=5, normalize=True), (2, 1, 0)))
+    plt.show()
+
+
+# denormalize the images for proper display
+def denormalize_images(images):
+    mean= [0.5, 0.5, 0.5]
+    std= [0.5, 0.5, 0.5]
+    inv_normalize = transforms.Normalize(
+        mean=[-m / s for m, s in zip(mean, std)],
+        std=[1 / s for s in std]
+    )
+    return inv_normalize(images)
+
+
+
+