import streamlit as st
import requests
import os.path
import torch
from PIL import Image
from utils import utils_model
from utils import utils_image as util
import numpy as np
import cv2

st.set_page_config(layout="wide", page_title="Image Denoising Demo")

st.title("Image Real Denoising Demo")
st.write("The model removes the noise from real world images.")
st.sidebar.write("## Upload and download :gear:")

# Create the columns
col1, col2 = st.columns(2)

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

if upfile is not None:
    img = Image.open(upfile)
    img.save('test1.png')

    n_channels = 3
    img_U = util.imread_uint('test1.png', n_channels=n_channels)
    #img = Image.open(upfile)
    col1.write("Original Noisy Image")
    col1.image(upfile)
    
    
    model_name = 'team15_SAKDNNet'
   
    model_path = os.path.join('model_zoo', model_name+'.pth')

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    from models.team15_SAKDNNet import SAKDNNet as net
    model = net(in_nc=n_channels,config=[4,4,4,4,4,4,4],dim=64)

    model.load_state_dict(torch.load(model_path), strict=True)
    model.eval()
    for k, v in model.named_parameters():
        v.requires_grad = False
        
    model = model.to(device)
    img_N = util.uint2tensor4(img_U)
    img_N = img_N.to(device)
    
    img_DN = utils_model.inference(model, img_N, refield=64, min_size=512, mode=2)
    img_DN = model(img_N)
    img_DN = util.tensor2uint(img_DN)
    
    col2.write("Denoised Image")
    col2.image(img_DN)
    
    #st.sidebar.markdown("\\n")

    st.write("The method Dense Residual Swin Transformer is included in NTIRE 2023 Image Denoising Challenge at CVPR 2023")