Spaces:
Running on T4
Running on T4
| import argparse | |
| import os | |
| import warnings | |
| import torch | |
| from torchvision import transforms | |
| import gradio as gr | |
| import PIL | |
| import json | |
| from tqdm import tqdm | |
| means = [0.5] | |
| stds = [0.5] | |
| device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') | |
| t_stds = torch.tensor(stds).to(device)[:, None, None] | |
| t_means = torch.tensor(means).to(device)[:, None, None] | |
| def tensor2im(var): | |
| return var.mul(t_stds).add(t_means).mul(255.).clamp(0, 255) | |
| def process(input_image, step=None): | |
| progressbar = gr.Progress(track_tqdm=True) | |
| img = tfs(input_image)[None, ...].to(device) | |
| b, *_ = img.shape | |
| if step == None: | |
| timesteps = 200 | |
| else: | |
| timesteps = step | |
| with torch.no_grad(): | |
| model.eval() | |
| y_t = torch.randn_like(img) | |
| count = timesteps | |
| for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps): | |
| t = torch.full((b,), count, device=img.device, dtype=torch.long) | |
| y_t = model(img, y_t, t) | |
| count -= 1 | |
| # out, _ = model.netG.restoration(img, sample_num=8) | |
| out = y_t | |
| out = out[0, :, :, :] | |
| out = tensor2im(out)[0, :, :] | |
| output_image = out.detach().cpu().numpy().astype('uint8') | |
| output_image = PIL.Image.fromarray(output_image) | |
| return output_image | |
| if __name__ == '__main__': | |
| torch.backends.cudnn.enabled = False | |
| model = torch.jit.load('./emdiffuse-n.pt', map_location=torch.device('cpu')).to(device) | |
| tfs = transforms.Compose([ | |
| transforms.Resize((768, 768)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=means, std=stds) | |
| ]) | |
| title = "EMDiffuse-n" | |
| description = "EMDiffuse-n demo for electron micrscope image denoising. \n" \ | |
| "The model is trained on mouse brain cortex scanning electron microscope dataset. To use it, simply upload your image, or click one of the examples to load them.\n" \ | |
| "Gradio does not currently support the TIFF image format, so the uploaded image may not appear in the input box, but it still functions when you click submit. Read more at the links below." | |
| article = "<div style='text-align: center;'>Chixiang et al. <a href='https://www.biorxiv.org/content/10.1101/2023.07.12.548636v1' target='_blank'> EMDiffuse S</a> | <a href='https://github.com/Luchixiang/EMDiffuse' target='_blank'>Github Repo</a> | <center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_arcanegan' alt='visitor badge'></center></div>" | |
| gr.Interface( | |
| process, | |
| inputs=[gr.Image(type="pil", label="Input", image_mode='L'), gr.Slider(200, 1000, label='Step numbers', | |
| info="The diffusion model's step count directly impacts image quality: more steps lead to higher quality images.")], | |
| outputs=[gr.Image(type="pil", label="Output", image_mode='L')], | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=[['./demo/22.png', None], ['./demo/23.png', None], ['./demo/24.png', None]], | |
| allow_flagging=False, | |
| ).queue(max_size=10).launch() | |