import os import streamlit as st import torch import torch.nn as nn from torchvision import transforms as tr from PIL import Image import numpy as np import random class ResidualBlock(nn.Module): def __init__(self, in_channels): super(ResidualBlock, self).__init__() self.block = nn.Sequential( nn.ReflectionPad2d(1), nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=0), nn.InstanceNorm2d(in_channels), nn.ReLU(inplace=True), nn.ReflectionPad2d(1), nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=0), nn.InstanceNorm2d(in_channels), ) def forward(self, x): return x + self.block(x) class Generator(nn.Module): def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=9): super(Generator, self).__init__() model = [ nn.ReflectionPad2d(3), nn.Conv2d(in_channels, 64, kernel_size=7, padding=0), nn.InstanceNorm2d(64), nn.ReLU(inplace=True) ] in_features = 64 out_features = in_features * 2 for _ in range(2): model += [ nn.Conv2d(in_features, out_features, kernel_size=3, stride=2, padding=1), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True) ] in_features, out_features = out_features, out_features * 2 for _ in range(n_residual_blocks): model += [ResidualBlock(in_features)] out_features = in_features // 2 for _ in range(2): model += [ nn.Upsample(scale_factor=2, mode='nearest'), nn.Conv2d(in_features, out_features, kernel_size=3, stride=1, padding=1), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True) ] in_features, out_features = out_features, out_features // 2 model += [ nn.ReflectionPad2d(3), nn.Conv2d(64, out_channels, kernel_size=7, padding=0), nn.Tanh() ] self.model = nn.Sequential(*model) def forward(self, x): return self.model(x) class Discriminator(nn.Module): def __init__(self, in_channels=3): super(Discriminator, self).__init__() def conv_block(in_f, out_f, norm=True, stride=2): layers = [nn.Conv2d(in_f, out_f, kernel_size=4, stride=stride, padding=1)] if norm: layers.append(nn.InstanceNorm2d(out_f)) layers.append(nn.LeakyReLU(0.2, inplace=True)) return layers self.model = nn.Sequential( *conv_block(in_channels, 64, norm=False, stride=2), *conv_block(64, 128, stride=2), *conv_block(128, 256, stride=2), *conv_block(256, 512, stride=1), nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=1) ) def forward(self, x): return self.model(x) class CycleGAN(nn.Module): def __init__(self, mean_a, std_a, mean_b, std_b, in_channels=3, out_channels=3, n_residual_blocks=9): super(CycleGAN, self).__init__() self.generators = nn.ModuleDict({ "a_to_b": Generator(in_channels, out_channels, n_residual_blocks), "b_to_a": Generator(out_channels, in_channels, n_residual_blocks), }) self.discriminators = nn.ModuleDict({ "a": Discriminator(in_channels), "b": Discriminator(out_channels), }) self.register_buffer('mean_a', torch.tensor(mean_a).view(1, in_channels, 1, 1)) self.register_buffer('std_a', torch.tensor(std_a).view(1, in_channels, 1, 1)) self.register_buffer('mean_b', torch.tensor(mean_b).view(1, out_channels, 1, 1)) self.register_buffer('std_b', torch.tensor(std_b).view(1, out_channels, 1, 1)) def forward(self, x, direction="a_to_b"): if direction == "a_to_b": return self.generators["a_to_b"](x) else: return self.generators["b_to_a"](x) def get_transforms(mean, std, crop_size=256): transform = tr.Compose([ tr.Resize(crop_size), tr.CenterCrop(crop_size), tr.ToTensor(), tr.Normalize(mean=mean, std=std), ]) def de_normalize(tensor): device = tensor.device mean_t = torch.tensor(mean, device=device).view(-1, 1, 1) std_t = torch.tensor(std, device=device).view(-1, 1, 1) tensor = tensor * std_t + mean_t tensor = tensor.clamp(0, 1) return tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() return transform, de_normalize DEFAULT_MEAN = [0.5, 0.5, 0.5] DEFAULT_STD = [0.5, 0.5, 0.5] @st.cache_resource(show_spinner=False) def load_model(checkpoint_path, mean_a, std_a, mean_b, std_b): try: model = CycleGAN( mean_a=mean_a, std_a=std_a, mean_b=mean_b, std_b=std_b, n_residual_blocks=9 ) if os.path.exists(checkpoint_path): checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False) model.load_state_dict(checkpoint['model_state_dict']) model.eval() return model, True else: return None, False except Exception as e: st.error(f"Error while loading model: {e}") return None, False def process_image(model, image, direction): transform, de_norm = get_transforms(DEFAULT_MEAN, DEFAULT_STD) input_tensor = transform(image).unsqueeze(0) with torch.no_grad(): output = model(input_tensor, direction=direction) output = de_norm(output) return output st.title(":material/auto_awesome: CycleGAN Studio") with st.expander(":material/info: About the Model & Training", expanded=False): st.markdown(""" ### Architecture This application utilizes **CycleGAN** (Cycle-Consistent Adversarial Networks). * **Generators:** Built on a ResNet architecture (9 blocks). Using `InstanceNorm` instead of `BatchNorm` helps preserve the specific style of an individual image better. * **Discriminators:** PatchGAN, which evaluates local patches rather than the full image to maintain high-frequency details and fine textures. ### Training Process The training process minimizes three key loss functions: 1. **Adversarial Loss:** Ensures the generated images look realistic enough to fool the discriminator. 2. **Cycle Consistency Loss:** Guarantees that translating an image from domain A to B, and then back to A, reconstructs the original image. This enables training on *unpaired* datasets. 3. **Identity Loss:** Encourages the generator to preserve the overall color composition of the input image. """) if 'random_image_path' not in st.session_state: st.session_state['random_image_path'] = None def reset_image(): st.session_state['random_image_path'] = None with st.sidebar: st.header(":material/settings: Settings") model_choice = st.radio( "Select Model", ["Summer <-> Winter", "Sketch <-> Art"], index=1, on_change=reset_image ) paths = { "Summer <-> Winter": { "chkp": "src/chkp/summer2winter.pt", "dirs": ["a_to_b", "b_to_a"], "folder_a": "src/dataset/summer2winter/testA", "folder_b": "src/dataset/summer2winter/testB", "labels": ["To Winter", "To Summer"]}, "Sketch <-> Art": { "chkp": "src/chkp/sketch2art.pt", "dirs": ["a_to_b", "b_to_a"], "folder_a": "src/dataset/sketch2art/testA", "folder_b": "src/dataset/sketch2art/testB", "labels": ["To Art", "To Sketch"]} } current_config = paths[model_choice] direction_label = st.selectbox( "Translation Direction", current_config["labels"] ) dir_idx = current_config["labels"].index(direction_label) active_direction = current_config["dirs"][dir_idx] st.divider() st.markdown("### :material/history: Dataset Examples") example_path = f"examples/{model_choice.lower().replace(' <-> ', '2')}/" if st.button("Load Random Example"): target_folder = current_config["folder_a"] if active_direction == "a_to_b" else current_config["folder_b"] if os.path.exists(target_folder): valid_extensions = ('.png', '.jpg', '.jpeg') files = [f for f in os.listdir(target_folder) if f.lower().endswith(valid_extensions)] if files: random_file = random.choice(files) st.session_state['random_image_path'] = os.path.join(target_folder, random_file) st.toast(f"Loaded: {random_file}") else: st.error("No images found in the folder.") else: st.error(f"Path not found: {target_folder}") st.subheader(f":material/swap_horiz: {model_choice}: {direction_label}") uploaded_file = st.file_uploader( "Upload an image for processing", type=['png', 'jpg', 'jpeg'], label_visibility="collapsed" ) model, is_loaded = load_model(current_config["chkp"], DEFAULT_MEAN, DEFAULT_STD, DEFAULT_MEAN, DEFAULT_STD) if model is None: st.error(f":material/error: Weights file `{current_config['chkp']}` not found.") else: img = None if uploaded_file is not None: img = Image.open(uploaded_file).convert("RGB") st.session_state['random_image_path'] = None elif st.session_state['random_image_path'] is not None: img = Image.open(st.session_state['random_image_path']).convert("RGB") if img is not None: run_button = st.button(":material/magic_button: Run Processing", use_container_width=True) col1, col2 = st.columns(2) with col1: st.image(img, caption="Original Image", use_container_width=True) with col2: if run_button: with st.spinner("Generating..."): result = process_image(model, img, active_direction) st.image(result, caption="CycleGAN Result", use_container_width=True) st.toast("Done!") else: st.info(":material/image: Please upload an image or select an example from the sidebar.")