# IMPORT LIBRARIES & FUNCTIONS
# External Libraries
import streamlit as st

# Project Functions
# - Model
from src.model.model import load_model, stylize_content_image
from src.model.utils import use_low_resource_settings, suppress_warnings

# - Data
from src.data.data import upload_image
from src.data.utils import remove_source_images

use_low_resource_settings("Yes")

suppress_warnings("Yes")

st.write("# 🖼️Neural🎨Style🖌️Transfer🖼️")

# LOAD MODEL
with st.spinner("🖌️Loading Model"):
    model = load_model()
st.success("🖌️Model Loaded")

content_image_def, style_image_def = st.columns(2)
with content_image_def:
    st.write("🖼️ Content Image: Image to style")
with style_image_def:
    st.write("🎨 Style Image: Style to transfer to content image")

content_image, style_image = st.columns(2)

# UPLOAD IMAGES
with content_image:
    ContentColumnTitle = "## 🖼️ Content Image 🖼️"
    ContentImageSelectionPrompt = "Pick a Content image"
    content_image_file = upload_image(
        ContentColumnTitle, ContentImageSelectionPrompt, "Content", "content"
    )

with style_image:
    StyleColumnTitle = "## 🎨 Style Image 🎨"
    StyleImageSelectionPrompt = "Pick a Style image"
    style_image_file = upload_image(
        StyleColumnTitle, StyleImageSelectionPrompt, "Style", "style"
    )

# STYLIZE CONTENT IMAGE
stylize_image = st.button("🖼️🖌️🎨 Start Neural Style Transfer 🖼️🖌️🎨")

if stylize_image:
    final_image = stylize_content_image(model, content_image_file, style_image_file)
    st.write("# Styled Image:")
    st.image(final_image)
    try:
        remove_source_images()
    except:
        pass