from config import Config
from preprocessing import get_transforms, tensor_to_image
import model as base_model
import monet_model

import os
import torch
import streamlit as st
from PIL import Image, ImageEnhance


def main():
    artist = st.selectbox("Select an artist", ["Monet", "van Gogh", "Cezanne"])

    artist = artist.lower().replace(" ", "")

    mode = st.radio(
        "Select conversion mode", ("Painting to Photo", "Photo to Painting")
    )

    transforms, de_normalize = get_transforms(artist)

    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        input_image = Image.open(uploaded_file).convert("RGB")
         
        width, height = input_image.size
        new_height = 300
        aspect_ratio = width / height
        new_width = int(new_height * aspect_ratio)
        input_image = input_image.resize((new_width, new_height), Image.LANCZOS)
        st.image(input_image, caption="Input Image")

        input_tensor = transforms(input_image)

        if artist == "monet":
            model = monet_model.get_model(artist)
        else:
            model = base_model.get_model(artist)

        model.eval()

        with torch.no_grad():
            if mode == "Painting to Photo":
                output_tensor = model.generator_A2B(input_tensor)
            elif mode == "Photo to Painting":
                output_tensor = model.generator_B2A(input_tensor)

        output_image = tensor_to_image(output_tensor, de_normalize)
        output_image = Image.fromarray(output_image)

        output_image = output_image.resize((new_width, new_height), Image.LANCZOS)
        
        st.image(output_image, caption="Converted Image")


if __name__ == "__main__":
    main()