Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # -*- coding:utf-8 -*- | |
| # Power by Zongsheng Yue 2022-12-16 16:17:14 | |
| import os | |
| import torch | |
| import argparse | |
| import numpy as np | |
| import gradio as gr | |
| from pathlib import Path | |
| from einops import rearrange | |
| from omegaconf import OmegaConf | |
| from skimage import img_as_ubyte | |
| from utils import util_opts | |
| from utils import util_image | |
| from utils import util_common | |
| from sampler import DifIRSampler | |
| from ResizeRight.resize_right import resize | |
| from basicsr.utils.download_util import load_file_from_url | |
| # setting configurations | |
| cfg_path = 'configs/sample/iddpm_ffhq512_swinir.yaml' | |
| configs = OmegaConf.load(cfg_path) | |
| configs.aligned = False | |
| configs.diffusion.timestep_respacing = '200' | |
| # build the sampler for diffusion | |
| sampler_dist = DifIRSampler(configs) | |
| def predict(im_path, background_enhance, face_upsample, upscale, started_timesteps): | |
| assert isinstance(im_path, str) | |
| print(f'Processing image: {im_path}...') | |
| configs.background_enhance = background_enhance | |
| configs.face_upsample = face_upsample | |
| started_timesteps = int(started_timesteps) | |
| assert started_timesteps < int(configs.diffusion.params.timestep_respacing) | |
| # prepare the checkpoint | |
| if not Path(configs.model.ckpt_path).exists(): | |
| load_file_from_url( | |
| url="https://github.com/zsyOAOA/DifFace/releases/download/V1.0/iddpm_ffhq512_ema500000.pth", | |
| model_dir=str(Path(configs.model.ckpt_path).parent), | |
| progress=True, | |
| file_name=Path(configs.model.ckpt_path).name, | |
| ) | |
| if not Path(configs.model_ir.ckpt_path).exists(): | |
| load_file_from_url( | |
| url="https://github.com/zsyOAOA/DifFace/releases/download/V1.0/General_Face_ffhq512.pth", | |
| model_dir=str(Path(configs.model_ir.ckpt_path).parent), | |
| progress=True, | |
| file_name=Path(configs.model_ir.ckpt_path).name, | |
| ) | |
| # Load image | |
| im_lq = util_image.imread(im_path, chn='bgr', dtype='uint8') | |
| if upscale > 4: | |
| upscale = 4 # avoid momory exceeded due to too large upscale | |
| if upscale > 2 and min(im_lq.shape[:2])>1280: | |
| upscale = 2 # avoid momory exceeded due to too large img resolution | |
| configs.detection.upscale = int(upscale) | |
| if background_enhance: | |
| image_restored, face_restored, face_cropped = sampler_dist.sample_func_bfr_unaligned( | |
| y0=im_lq, | |
| start_timesteps=started_timesteps, | |
| need_restoration=True, | |
| draw_box=False, | |
| ) # h x w x c, numpy array, [0, 255], uint8, BGR | |
| image_restored = util_image.bgr2rgb(image_restored) | |
| else: | |
| image_restored = sampler_dist.sample_func_ir_aligned( | |
| y0=im_lq, | |
| start_timesteps=started_timesteps, | |
| need_restoration=True, | |
| )[0] # b x c x h x w, [0, 1], torch tensor, RGB | |
| image_restored = util_image.tensor2img( | |
| image_restored.cpu(), | |
| rgb2bgr=False, | |
| out_type=np.uint8, | |
| min_max=(0, 1), | |
| ) # h x w x c, [0, 255], uint8, RGB, numpy array | |
| restored_image_dir = Path('restored_output') | |
| if not restored_image_dir.exists(): | |
| restored_image_dir.mkdir() | |
| # save the whole image | |
| save_path = restored_image_dir / Path(im_path).name | |
| util_image.imwrite(image_restored, save_path, chn='rgb', dtype_in='uint8') | |
| return image_restored, str(save_path) | |
| # title = "DifFace: Blind Face Restoration with Diffused Error Contraction" | |
| # description = r""" | |
| # <b>Official Gradio demo</b> for <a href='https://github.com/zsyOAOA/DifFace' target='_blank'><b>DifFace: Blind Face Restoration with Diffused Error Contraction</b></a>.<br> | |
| # 🔥 DifFace is a robust face restoration algorithm for old or corrupted photos.<br> | |
| # """ | |
| # article = r""" | |
| # If DifFace is helpful for your work, please help to ⭐ the <a href='https://github.com/zsyOAOA/DifFace' target='_blank'>Github Repo</a>. Thanks! | |
| # [](https://github.com/zsyOAOA/DifFace) | |
| # --- | |
| # 📝 **Citation** | |
| # If our work is useful for your research, please consider citing: | |
| # ```bibtex | |
| # @article{yue2022difface, | |
| # title={DifFace: Blind Face Restoration with Diffused Error Contraction}, | |
| # author={Yue, Zongsheng and Loy, Chen Change}, | |
| # journal={arXiv preprint arXiv:2212.06512}, | |
| # year={2022} | |
| # } | |
| # ``` | |
| # 📋 **License** | |
| # This project is licensed under <a rel="license" href="https://github.com/zsyOAOA/DifFace/blob/master/LICENSE">S-Lab License 1.0</a>. | |
| # Redistribution and use for non-commercial purposes should follow this license. | |
| # 📧 **Contact** | |
| # If you have any questions, please feel free to contact me via <b>zsyzam@gmail.com</b>. | |
| #  | |
| # """ | |
| # demo = gr.Interface( | |
| # predict, | |
| # inputs=[ | |
| # gr.Image(type="filepath", label="Input"), | |
| # gr.Checkbox(value=True, label="Background_Enhance"), | |
| # gr.Checkbox(value=True, label="Face_Upsample"), | |
| # gr.Number(value=2, label="Rescaling_Factor (up to 4)"), | |
| # gr.Slider(1, 160, value=80, step=10, label='Realism-Fidelity Trade-off') | |
| # ], | |
| # outputs=[ | |
| # gr.Image(type="numpy", label="Output"), | |
| # gr.outputs.File(label="Download the output") | |
| # ], | |
| # title=title, | |
| # description=description, | |
| # article=article, | |
| # examples=[ | |
| # ['./testdata/whole_imgs/00.jpg', True, True, 2, 80], | |
| # ['./testdata/whole_imgs/01.jpg', True, True, 2, 80], | |
| # ['./testdata/whole_imgs/04.jpg', True, True, 2, 80], | |
| # ['./testdata/whole_imgs/05.jpg', True, True, 2, 80], | |
| # ] | |
| # ) | |
| # demo.queue(concurrency_count=4) | |
| # demo.launch() | |