Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import numpy as np | |
| import torch | |
| from models.network_swinir import SwinIR | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print("device: %s" % device) | |
| default_models = { | |
| "sr": "weights/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth", | |
| "denoise": "weights/005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth" | |
| } | |
| torch.backends.cudnn.enabled = True | |
| torch.backends.cudnn.benchmark = True | |
| denoise_model = SwinIR(upscale=1, in_chans=3, img_size=128, window_size=8, | |
| img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6], | |
| mlp_ratio=2, upsampler='', resi_connection='1conv').to(device) | |
| param_key_g = 'params' | |
| try: | |
| pretrained_model = torch.load(default_models["denoise"]) | |
| denoise_model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True) | |
| except: print("Loading model failed") | |
| denoise_model.eval() | |
| sr_model = SwinIR(upscale=4, in_chans=3, img_size=64, window_size=8, | |
| img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6], | |
| mlp_ratio=2, upsampler='nearest+conv', resi_connection='1conv').to(device) | |
| param_key_g = 'params_ema' | |
| try: | |
| pretrained_model = torch.load(default_models["sr"]) | |
| sr_model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True) | |
| except: print("Loading model failed") | |
| sr_model.eval() | |
| def sr(input_img): | |
| window_size = 8 | |
| # read image | |
| img_lq = input_img.astype(np.float32) / 255. | |
| img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1)) # HCW-BGR to CHW-RGB | |
| img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device) # CHW-RGB to NCHW-RGB | |
| # inference | |
| with torch.no_grad(): | |
| # pad input image to be a multiple of window_size | |
| _, _, h_old, w_old = img_lq.size() | |
| h_pad = (h_old // window_size + 1) * window_size - h_old | |
| w_pad = (w_old // window_size + 1) * window_size - w_old | |
| img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :] | |
| img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad] | |
| output = sr_model(img_lq) | |
| output = output[..., :h_old * 4, :w_old * 4] | |
| # save image | |
| output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy() | |
| if output.ndim == 3: | |
| output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR | |
| output = (output * 255.0).round().astype(np.uint8) # float32 to uint8 | |
| return output | |
| def denoise(input_img): | |
| window_size = 8 | |
| # read image | |
| img_lq = input_img.astype(np.float32) / 255. | |
| img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1)) # HCW-BGR to CHW-RGB | |
| img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device) # CHW-RGB to NCHW-RGB | |
| # inference | |
| with torch.no_grad(): | |
| # pad input image to be a multiple of window_size | |
| _, _, h_old, w_old = img_lq.size() | |
| h_pad = (h_old // window_size + 1) * window_size - h_old | |
| w_pad = (w_old // window_size + 1) * window_size - w_old | |
| img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :] | |
| img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad] | |
| output = denoise_model(img_lq) | |
| output = output[..., :h_old * 4, :w_old * 4] | |
| # save image | |
| output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy() | |
| if output.ndim == 3: | |
| output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR | |
| output = (output * 255.0).round().astype(np.uint8) # float32 to uint8 | |
| return output | |
| title = " AISeed AI Application Demo " | |
| description = "# A Demo of Deep Learning for Image Restoration" | |
| example_list = [["examples/" + example] for example in os.listdir("examples")] | |
| with gr.Blocks() as demo: | |
| demo.title = title | |
| gr.Markdown(description) | |
| with gr.Row(): | |
| with gr.Column(): | |
| im = gr.Image(label="Input Image") | |
| im_2 = gr.Image(label="Enhanced Image") | |
| with gr.Column(): | |
| btn1 = gr.Button(value="Enhance Resolution") | |
| btn1.click(sr, inputs=[im], outputs=[im_2]) | |
| btn2 = gr.Button(value="Denoise") | |
| btn2.click(denoise, inputs=[im], outputs=[im_2]) | |
| gr.Examples(examples=example_list, | |
| inputs=[im], | |
| outputs=[im_2]) | |
| if __name__ == "__main__": | |
| demo.launch() |