Spaces:
Running
Running
| import numpy as np | |
| import cv2 | |
| import os, argparse, json | |
| from os.path import join | |
| from glob import glob | |
| import torch | |
| import torch.nn.functional as F | |
| from model.model import ResHalf | |
| from model.model import Quantize | |
| from model.loss import l1_loss | |
| from utils import util | |
| from utils.dct import DCT_Lowfrequency | |
| from utils.filters_tensor import bgr2gray | |
| from collections import OrderedDict | |
| class Inferencer: | |
| def __init__(self, checkpoint_path, model, use_cuda=True, multi_gpu=True): | |
| self.checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu')) | |
| self.use_cuda = use_cuda | |
| self.model = model.eval() | |
| if multi_gpu: | |
| self.model = torch.nn.DataParallel(self.model) | |
| state_dict = self.checkpoint['state_dict'] | |
| else: | |
| ## remove keyword "module" in the state_dict | |
| state_dict = OrderedDict() | |
| for k, v in self.checkpoint['state_dict'].items(): | |
| name = k[7:] | |
| state_dict[name] = v | |
| if self.use_cuda: | |
| self.model = self.model.cuda() | |
| self.model.load_state_dict(state_dict) | |
| def __call__(self, input_img, decoding_only=False): | |
| with torch.no_grad(): | |
| scale = 8 | |
| _, _, H, W = input_img.shape | |
| if H % scale != 0 or W % scale != 0: | |
| input_img = F.pad(input_img, [0, scale - W % scale, 0, scale - H % scale], mode='reflect') | |
| if self.use_cuda: | |
| input_img = input_img.cuda() | |
| if decoding_only: | |
| resColor = self.model(input_img, decoding_only) | |
| if H % scale != 0 or W % scale != 0: | |
| resColor = resColor[:, :, :H, :W] | |
| return resColor | |
| else: | |
| resHalftone, resColor = self.model(input_img, decoding_only) | |
| resHalftone = Quantize.apply((resHalftone + 1.0) * 0.5) * 2.0 - 1. | |
| if H % scale != 0 or W % scale != 0: | |
| resHalftone = resHalftone[:, :, :H, :W] | |
| resColor = resColor[:, :, :H, :W] | |
| return resHalftone, resColor | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser(description='invHalf') | |
| parser.add_argument('--model', default=None, type=str, | |
| help='model weight file path') | |
| parser.add_argument('--decoding', action='store_true', default=False, help='restoration from halftone input') | |
| parser.add_argument('--data_dir', default=None, type=str, | |
| help='where to load input data (RGB images)') | |
| parser.add_argument('--save_dir', default=None, type=str, | |
| help='where to save the result') | |
| args = parser.parse_args() | |
| invhalfer = Inferencer( | |
| checkpoint_path=args.model, | |
| model=ResHalf(train=False) | |
| ) | |
| save_dir = os.path.join(args.save_dir) | |
| util.ensure_dir(save_dir) | |
| test_imgs = glob(join(args.data_dir, '*.*g')) | |
| print('------loaded %d images.' % len(test_imgs) ) | |
| for img in test_imgs: | |
| print('[*] processing %s ...' % img) | |
| if args.decoding: | |
| input_img = cv2.imread(img, flags=cv2.IMREAD_GRAYSCALE) / 127.5 - 1. | |
| c = invhalfer(util.img2tensor(input_img), decoding_only=True) | |
| c = util.tensor2img(c / 2. + 0.5) * 255. | |
| cv2.imwrite(join(save_dir, 'restored_' + img.split('/')[-1].split('.')[0] + '.png'), c) | |
| else: | |
| input_img = cv2.imread(img, flags=cv2.IMREAD_COLOR) / 127.5 - 1. | |
| h, c = invhalfer(util.img2tensor(input_img), decoding_only=False) | |
| h = util.tensor2img(h / 2. + 0.5) * 255. | |
| c = util.tensor2img(c / 2. + 0.5) * 255. | |
| cv2.imwrite(join(save_dir, 'halftone_' + img.split('/')[-1].split('.')[0] + '.png'), h) | |
| cv2.imwrite(join(save_dir, 'restored_' + img.split('/')[-1].split('.')[0] + '.png'), c) | |