Spaces:
Running
Running
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") | |