src/streamlit_app.py

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.")