|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
|
|
|
|
|
|
from basicsr.models import create_model |
|
|
from basicsr.utils import FileClient, imfrombytes, img2tensor, padding, tensor2img, imwrite, set_random_seed |
|
|
|
|
|
import argparse |
|
|
from basicsr.utils.options import dict2str, parse |
|
|
from basicsr.utils.dist_util import get_dist_info, init_dist |
|
|
import random |
|
|
|
|
|
def parse_options(is_train=True): |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument( |
|
|
'-opt', type=str, required=True, help='Path to option YAML file.') |
|
|
parser.add_argument( |
|
|
'--launcher', |
|
|
choices=['none', 'pytorch', 'slurm'], |
|
|
default='none', |
|
|
help='job launcher') |
|
|
parser.add_argument('--local_rank', type=int, default=0) |
|
|
|
|
|
parser.add_argument('--input_l_path', type=str, required=True, help='The path to the input left image. For stereo image inference only.') |
|
|
parser.add_argument('--input_r_path', type=str, required=True, help='The path to the input right image. For stereo image inference only.') |
|
|
parser.add_argument('--output_l_path', type=str, required=True, help='The path to the output left image. For stereo image inference only.') |
|
|
parser.add_argument('--output_r_path', type=str, required=True, help='The path to the output right image. For stereo image inference only.') |
|
|
|
|
|
args = parser.parse_args() |
|
|
opt = parse(args.opt, is_train=is_train) |
|
|
|
|
|
|
|
|
if args.launcher == 'none': |
|
|
opt['dist'] = False |
|
|
print('Disable distributed.', flush=True) |
|
|
else: |
|
|
opt['dist'] = True |
|
|
if args.launcher == 'slurm' and 'dist_params' in opt: |
|
|
init_dist(args.launcher, **opt['dist_params']) |
|
|
else: |
|
|
init_dist(args.launcher) |
|
|
print('init dist .. ', args.launcher) |
|
|
|
|
|
opt['rank'], opt['world_size'] = get_dist_info() |
|
|
|
|
|
|
|
|
seed = opt.get('manual_seed') |
|
|
if seed is None: |
|
|
seed = random.randint(1, 10000) |
|
|
opt['manual_seed'] = seed |
|
|
set_random_seed(seed + opt['rank']) |
|
|
|
|
|
opt['img_path'] = { |
|
|
'input_l': args.input_l_path, |
|
|
'input_r': args.input_r_path, |
|
|
'output_l': args.output_l_path, |
|
|
'output_r': args.output_r_path |
|
|
} |
|
|
|
|
|
return opt |
|
|
|
|
|
def imread(img_path): |
|
|
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) |
|
|
return img |
|
|
|
|
|
def main(): |
|
|
|
|
|
opt = parse_options(is_train=False) |
|
|
opt['num_gpu'] = torch.cuda.device_count() |
|
|
|
|
|
img_l_path = opt['img_path'].get('input_l') |
|
|
img_r_path = opt['img_path'].get('input_r') |
|
|
output_l_path = opt['img_path'].get('output_l') |
|
|
output_r_path = opt['img_path'].get('output_r') |
|
|
|
|
|
|
|
|
img_l = imread(img_l_path) |
|
|
img_r = imread(img_r_path) |
|
|
img = torch.cat([img_l, img_r], dim=0) |
|
|
|
|
|
|
|
|
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_l = visuals['result'][:,:3] |
|
|
sr_img_r = visuals['result'][:,3:] |
|
|
sr_img_l, sr_img_r = tensor2img([sr_img_l, sr_img_r]) |
|
|
imwrite(sr_img_l, output_l_path) |
|
|
imwrite(sr_img_r, output_r_path) |
|
|
|
|
|
print(f'inference {img_l_path} .. finished. saved to {output_l_path}') |
|
|
print(f'inference {img_r_path} .. finished. saved to {output_r_path}') |
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
|
|
|
|
|
|
