import os import argparse import albumentations from albumentations import Resize import torch import torch.backends.cudnn as cudnn import torchvision.transforms as transforms from torch.utils.data import DataLoader from model.build_model import build_model from datasets.build_dataset import dataset_generator from utils import misc, metrics def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('--workers', type=int, default=1, metavar='N', help='Dataloader threads.') parser.add_argument('--batch_size', type=int, default=1, help='You can override model batch size by specify positive number.') parser.add_argument('--device', type=str, default='cuda', help="Whether use cuda, 'cuda' or 'cpu'.") parser.add_argument('--save_path', type=str, default="./logs", help='Where to save logs and checkpoints.') parser.add_argument('--dataset_path', type=str, default=r".\iHarmony4", help='Dataset path.') parser.add_argument('--base_size', type=int, default=256, help='Base size. Resolution of the image input into the Encoder') parser.add_argument('--input_size', type=int, default=256, help='Input size. Resolution of the image that want to be generated by the Decoder') parser.add_argument('--INR_input_size', type=int, default=256, help='INR input size. Resolution of the image that want to be generated by the Decoder. ' 'Should be the same as `input_size`') parser.add_argument('--INR_MLP_dim', type=int, default=32, help='Number of channels for INR linear layer.') parser.add_argument('--LUT_dim', type=int, default=7, help='Dim of the output LUT. Refer to https://ieeexplore.ieee.org/abstract/document/9206076') parser.add_argument('--activation', type=str, default='leakyrelu_pe', help='INR activation layer type: leakyrelu_pe, sine') parser.add_argument('--pretrained', type=str, default=r'.\pretrained_models\Resolution_RAW_iHarmony4.pth', help='Pretrained weight path') parser.add_argument('--param_factorize_dim', type=int, default=10, help='The intermediate dimensions of the factorization of the predicted MLP parameters. ' 'Refer to https://arxiv.org/abs/2011.12026') parser.add_argument('--embedding_type', type=str, default="CIPS_embed", help='Which embedding_type to use.') parser.add_argument('--optim', type=str, default='adamw', help='Which optimizer to use.') parser.add_argument('--INRDecode', action="store_false", help='Whether INR decoder. Set it to False if you want to test the baseline ' '(https://github.com/SamsungLabs/image_harmonization)') parser.add_argument('--isMoreINRInput', action="store_false", help='Whether to cat RGB and mask. See Section 3.4 in the paper.') parser.add_argument('--hr_train', action="store_true", help='Whether use hr_train. See section 3.4 in the paper.') parser.add_argument('--isFullRes', action="store_true", help='Whether for original resolution. See section 3.4 in the paper.') opt = parser.parse_args() opt.save_path = misc.increment_path(os.path.join(opt.save_path, "test1")) return opt def inference(val_loader, model, logger, opt): current_process = 10 model.eval() metric_log = { 'HAdobe5k': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'HCOCO': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'Hday2night': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'HFlickr': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'All': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, } lut_metric_log = { 'HAdobe5k': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'HCOCO': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'Hday2night': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'HFlickr': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, 'All': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0}, } for step, batch in enumerate(val_loader): composite_image = batch['composite_image'].to(opt.device) real_image = batch['real_image'].to(opt.device) mask = batch['mask'].to(opt.device) category = batch['category'] fg_INR_coordinates = batch['fg_INR_coordinates'].to(opt.device) with torch.no_grad(): fg_content_bg_appearance_construct, _, lut_transform_image = model( composite_image, mask, fg_INR_coordinates, ) if opt.INRDecode: pred_fg_image = fg_content_bg_appearance_construct[-1] else: pred_fg_image = misc.lin2img(fg_content_bg_appearance_construct, val_loader.dataset.INR_dataset.size) if fg_content_bg_appearance_construct is not None else None if not opt.INRDecode: pred_harmonized_image = None else: pred_harmonized_image = pred_fg_image * (mask > 100 / 255.) + real_image * (~(mask > 100 / 255.)) lut_transform_image = lut_transform_image * (mask > 100 / 255.) + real_image * (~(mask > 100 / 255.)) misc.visualize(real_image, composite_image, mask, pred_fg_image, pred_harmonized_image, lut_transform_image, opt, -1, show=False, wandb=False, isAll=True, step=step) if opt.INRDecode: mse, fmse, psnr, ssim = metrics.calc_metrics(misc.normalize(pred_harmonized_image, opt, 'inv'), misc.normalize(real_image, opt, 'inv'), mask) lut_mse, lut_fmse, lut_psnr, lut_ssim = metrics.calc_metrics(misc.normalize(lut_transform_image, opt, 'inv'), misc.normalize(real_image, opt, 'inv'), mask) for idx in range(len(category)): if opt.INRDecode: metric_log[category[idx]]['Samples'] += 1 metric_log[category[idx]]['MSE'] += mse[idx] metric_log[category[idx]]['fMSE'] += fmse[idx] metric_log[category[idx]]['PSNR'] += psnr[idx] metric_log[category[idx]]['SSIM'] += ssim[idx] metric_log['All']['Samples'] += 1 metric_log['All']['MSE'] += mse[idx] metric_log['All']['fMSE'] += fmse[idx] metric_log['All']['PSNR'] += psnr[idx] metric_log['All']['SSIM'] += ssim[idx] lut_metric_log[category[idx]]['Samples'] += 1 lut_metric_log[category[idx]]['MSE'] += lut_mse[idx] lut_metric_log[category[idx]]['fMSE'] += lut_fmse[idx] lut_metric_log[category[idx]]['PSNR'] += lut_psnr[idx] lut_metric_log[category[idx]]['SSIM'] += lut_ssim[idx] lut_metric_log['All']['Samples'] += 1 lut_metric_log['All']['MSE'] += lut_mse[idx] lut_metric_log['All']['fMSE'] += lut_fmse[idx] lut_metric_log['All']['PSNR'] += lut_psnr[idx] lut_metric_log['All']['SSIM'] += lut_ssim[idx] if (step + 1) / len(val_loader) * 100 >= current_process: logger.info(f'Processing: {current_process}') current_process += 10 logger.info('=========================') for key in metric_log.keys(): if opt.INRDecode: msg = f"{key}-'MSE': {metric_log[key]['MSE'] / metric_log[key]['Samples']:.2f}\n" \ f"{key}-'fMSE': {metric_log[key]['fMSE'] / metric_log[key]['Samples']:.2f}\n" \ f"{key}-'PSNR': {metric_log[key]['PSNR'] / metric_log[key]['Samples']:.2f}\n" \ f"{key}-'SSIM': {metric_log[key]['SSIM'] / metric_log[key]['Samples']:.4f}\n" \ f"{key}-'LUT_MSE': {lut_metric_log[key]['MSE'] / lut_metric_log[key]['Samples']:.2f}\n" \ f"{key}-'LUT_fMSE': {lut_metric_log[key]['fMSE'] / lut_metric_log[key]['Samples']:.2f}\n" \ f"{key}-'LUT_PSNR': {lut_metric_log[key]['PSNR'] / lut_metric_log[key]['Samples']:.2f}\n" \ f"{key}-'LUT_SSIM': {lut_metric_log[key]['SSIM'] / lut_metric_log[key]['Samples']:.4f}\n" else: msg = f"{key}-'LUT_MSE': {lut_metric_log[key]['MSE'] / lut_metric_log[key]['Samples']:.2f}\n" \ f"{key}-'LUT_fMSE': {lut_metric_log[key]['fMSE'] / lut_metric_log[key]['Samples']:.2f}\n" \ f"{key}-'LUT_PSNR': {lut_metric_log[key]['PSNR'] / lut_metric_log[key]['Samples']:.2f}\n" \ f"{key}-'LUT_SSIM': {lut_metric_log[key]['SSIM'] / lut_metric_log[key]['Samples']:.4f}\n" logger.info(msg) logger.info('=========================') def main_process(opt): logger = misc.create_logger(os.path.join(opt.save_path, "log.txt")) cudnn.benchmark = True valset_path = os.path.join(opt.dataset_path, "IHD_test.txt") opt.transform_mean = [.5, .5, .5] opt.transform_var = [.5, .5, .5] torch_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(opt.transform_mean, opt.transform_var)]) valset_alb_transform = albumentations.Compose([Resize(opt.input_size, opt.input_size)], additional_targets={'real_image': 'image', 'object_mask': 'image'}) valset = dataset_generator(valset_path, valset_alb_transform, torch_transform, opt, mode='Val') val_loader = DataLoader(valset, opt.batch_size, shuffle=False, drop_last=False, pin_memory=True, num_workers=opt.workers, persistent_workers=True) model = build_model(opt).to(opt.device) logger.info(f"Load pretrained weight from {opt.pretrained}") load_dict = torch.load(opt.pretrained)['model'] for k in load_dict.keys(): if k not in model.state_dict().keys(): print(f"Skip {k}") model.load_state_dict(load_dict, strict=False) inference(val_loader, model, logger, opt) if __name__ == '__main__': opt = parse_args() os.makedirs(opt.save_path, exist_ok=True) main_process(opt)