File size: 1,886 Bytes
61d360d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
# ------------------------------------------------------------------------
# Copyright (c) 2022 megvii-model. All Rights Reserved.
# ------------------------------------------------------------------------
# Modified from BasicSR (https://github.com/xinntao/BasicSR)
# Copyright 2018-2020 BasicSR Authors
# ------------------------------------------------------------------------
import torch
# from basicsr.data import create_dataloader, create_dataset
from basicsr.models import create_model
from basicsr.train import parse_options
from basicsr.utils import FileClient, imfrombytes, img2tensor, padding, tensor2img, imwrite
# from basicsr.utils import (get_env_info, get_root_logger, get_time_str,
# make_exp_dirs)
# from basicsr.utils.options import dict2str
def main():
# parse options, set distributed setting, set ramdom seed
opt = parse_options(is_train=False)
opt['num_gpu'] = torch.cuda.device_count()
img_path = opt['img_path'].get('input_img')
output_path = opt['img_path'].get('output_img')
## 1. read image
file_client = FileClient('disk')
img_bytes = file_client.get(img_path, None)
try:
img = imfrombytes(img_bytes, float32=True)
except:
raise Exception("path {} not working".format(img_path))
img = img2tensor(img, bgr2rgb=True, float32=True)
## 2. run inference
opt['dist'] = False
model = create_model(opt)
model.feed_data(data={'lq': img.unsqueeze(dim=0)})
if model.opt['val'].get('grids', False):
model.grids()
model.test()
if model.opt['val'].get('grids', False):
model.grids_inverse()
visuals = model.get_current_visuals()
sr_img = tensor2img([visuals['result']])
imwrite(sr_img, output_path)
print(f'inference {img_path} .. finished. saved to {output_path}')
if __name__ == '__main__':
main()
|