| import streamlit as st
|
| import torch
|
| import numpy as np
|
| from PIL import Image
|
| import torchvision.transforms as T
|
| from model import load_generator
|
| from io import BytesIO
|
|
|
| st.set_page_config(page_title="Summer ↔ Winter CycleGAN", page_icon="🏔️", layout="centered")
|
| st.title("🏔️ Summer ↔ Winter Translation")
|
| st.markdown("Upload a landscape photo and convert it between **summer** and **winter**.")
|
|
|
| @st.cache_resource
|
| def get_generators():
|
| device = "cpu"
|
| gen_a2b = load_generator("gen_a2b_fp16.pth", device)
|
| gen_b2a = load_generator("gen_b2a_fp16.pth", device)
|
| return gen_a2b, gen_b2a, device
|
|
|
| gen_a2b, gen_b2a, device = get_generators()
|
| st.success(f"Model loaded on **{device}**", icon="✅")
|
|
|
| MEAN = (0.5, 0.5, 0.5)
|
| STD = (0.5, 0.5, 0.5)
|
|
|
| to_tensor = T.Compose([
|
| T.Resize((256, 256)),
|
| T.ToTensor(),
|
| T.Normalize(MEAN, STD),
|
| ])
|
|
|
| def to_pil(tensor, original_size=None):
|
| img = tensor.squeeze(0).cpu().float()
|
| for i, (m, s) in enumerate(zip(MEAN, STD)):
|
| img[i] = img[i] * s + m
|
| img = torch.clamp(img, 0, 1)
|
|
|
| pil_img = T.ToPILImage()(img)
|
|
|
| if original_size:
|
| pil_img = pil_img.resize(original_size, Image.LANCZOS)
|
| return pil_img
|
|
|
| direction = st.radio(
|
| "Translation direction",
|
| ["☀️ Summer → ❄️ Winter", "❄️ Winter → ☀️ Summer"],
|
| horizontal=True,
|
| )
|
|
|
| uploaded = st.file_uploader("Upload landscape photo (JPG/PNG)", type=["jpg", "jpeg", "png"])
|
|
|
| if uploaded is not None:
|
| try:
|
| raw = uploaded.read()
|
| input_img = Image.open(BytesIO(raw)).convert("RGB")
|
| original_size = input_img.size
|
|
|
| col1, col2 = st.columns(2)
|
| with col1:
|
| st.subheader("Input")
|
| st.image(input_img, use_container_width=True)
|
|
|
| with st.spinner("Translating..."):
|
| tensor = to_tensor(input_img).unsqueeze(0).to(device)
|
| generator = gen_a2b if "Summer" in direction.split("→")[0] else gen_b2a
|
| with torch.no_grad():
|
| output_tensor = generator(tensor)
|
| output_img = to_pil(output_tensor, original_size=original_size)
|
|
|
| with col2:
|
| st.subheader("Output")
|
| st.image(output_img, use_container_width=True)
|
|
|
| buf = BytesIO()
|
| output_img.save(buf, format="PNG")
|
| st.download_button("⬇️ Download result", buf.getvalue(), "translated.png", "image/png")
|
|
|
| except Exception as error:
|
| st.error(f"Error: {error}")
|
| st.exception(error)
|
|
|
| st.markdown("---")
|
| st.markdown("**Model:** CycleGAN ResNet-9 blocks (64 channels) · **Train / Test:** Yosemite / Alpine (Unsplash)") |
