processing.py

import os
import time
import datetime

import torch
import torchvision

from utils import misc, metrics

best_psnr = 0


def train(train_loader, val_loader, model, optimizer, scheduler, loss_fn, logger, opt):
    total_step = opt.epochs * len(train_loader)

    step_time_log = misc.AverageMeter()
    loss_log = misc.AverageMeter(':6f')
    loss_fg_content_bg_appearance_construct_log = misc.AverageMeter(':6f')
    loss_lut_transform_image_log = misc.AverageMeter(':6f')
    loss_lut_regularize_log = misc.AverageMeter(':6f')

    start_epoch = 0

    "Load pretrained checkpoints"
    if opt.pretrained is not None:
        logger.info(f"Load pretrained weight from {opt.pretrained}")
        load_state = torch.load(opt.pretrained)
        model = model.cpu()
        model.load_state_dict(load_state['model'], strict=False)
        model = model.to(opt.device)
        optimizer.load_state_dict(load_state['optimizer'])
        scheduler.load_state_dict(load_state['scheduler'])
        start_epoch = load_state['last_epoch'] + 1

    for epoch in range(start_epoch, opt.epochs):
        model.train()
        time_ckp = time.time()
        for step, batch in enumerate(train_loader):
            current_step = epoch * len(train_loader) + step + 1

            if opt.INRDecode and opt.hr_train:
                "List with 4 elements: [Input to Encoder, three different resolutions' crop to INR Decoder]"
                composite_image = [batch[f'composite_image{name}'].to(opt.device) for name in range(4)]
                real_image = [batch[f'real_image{name}'].to(opt.device) for name in range(4)]
                mask = [batch[f'mask{name}'].to(opt.device) for name in range(4)]
                coordinate_map = [batch[f'coordinate_map{name}'].to(opt.device) for name in range(4)]

                fg_INR_coordinates = coordinate_map[1:]

            else:
                composite_image = batch['composite_image'].to(opt.device)
                real_image = batch['real_image'].to(opt.device)
                mask = batch['mask'].to(opt.device)

                fg_INR_coordinates = batch['fg_INR_coordinates'].to(opt.device)

            fg_content_bg_appearance_construct, fit_lut3d, lut_transform_image = model(
                composite_image, mask, fg_INR_coordinates)

            if opt.INRDecode:
                loss_fg_content_bg_appearance_construct = 0
                """
                    Our LRIP module requires three different resolution layers, thus here 
                    `loss_fg_content_bg_appearance_construct` is calculated in multiple layers. 
                    Besides, when leverage `hr_train`, i.e. use RSC strategy (See Section 3.4), the `real_image`
                    and `mask` are list type, corresponding different resolutions' crop.
                """
                if opt.hr_train:
                    for n in range(3):
                        loss_fg_content_bg_appearance_construct += loss_fn['masked_mse'] \
                            (fg_content_bg_appearance_construct[n], real_image[3 - n], mask[3 - n])
                    loss_fg_content_bg_appearance_construct /= 3
                    loss_lut_transform_image = loss_fn['masked_mse'](lut_transform_image, real_image[1], mask[1])
                else:
                    for n in range(3):
                        loss_fg_content_bg_appearance_construct += loss_fn['MaskWeightedMSE'] \
                            (fg_content_bg_appearance_construct[n],
                             torchvision.transforms.Resize(opt.INR_input_size // 2 ** (3 - n - 1))(real_image),
                             torchvision.transforms.Resize(opt.INR_input_size // 2 ** (3 - n - 1))(mask))
                    loss_fg_content_bg_appearance_construct /= 3
                    loss_lut_transform_image = loss_fn['masked_mse'](lut_transform_image, real_image, mask)
                loss_lut_regularize = loss_fn['regularize_LUT'](fit_lut3d)

            else:
                loss_fg_content_bg_appearance_construct = 0
                loss_lut_transform_image = loss_fn['masked_mse'](lut_transform_image, real_image, mask)
                loss_lut_regularize = 0

            loss = loss_fg_content_bg_appearance_construct + loss_lut_transform_image + loss_lut_regularize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()

            step_time_log.update(time.time() - time_ckp)

            loss_fg_content_bg_appearance_construct_log.update(0 if isinstance(loss_fg_content_bg_appearance_construct,
                                                                               int) else loss_fg_content_bg_appearance_construct.item())
            loss_lut_transform_image_log.update(
                0 if isinstance(loss_lut_transform_image, int) else loss_lut_transform_image.item())
            loss_lut_regularize_log.update(0 if isinstance(loss_lut_regularize, int) else loss_lut_regularize.item())
            loss_log.update(loss.item())

            if current_step % opt.print_freq == 0:
                remain_secs = (total_step - current_step) * step_time_log.avg
                remain_time = datetime.timedelta(seconds=round(remain_secs))
                finish_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time() + remain_secs))

                log_msg = f'Epoch: [{epoch}/{opt.epochs}]\t' \
                          f'Step: [{step}/{len(train_loader)}]\t' \
                          f'StepTime {step_time_log.val:.3f} ({step_time_log.avg:.3f})\t' \
                          f'lr {optimizer.param_groups[0]["lr"]}\t' \
                          f'Loss {loss_log.val:.4f} ({loss_log.avg:.4f})\t' \
                          f'Loss_fg_bg_cons {loss_fg_content_bg_appearance_construct_log.val:.4f} ({loss_fg_content_bg_appearance_construct_log.avg:.4f})\t' \
                          f'Loss_lut_trans {loss_lut_transform_image_log.val:.4f} ({loss_lut_transform_image_log.avg:.4f})\t' \
                          f'Loss_lut_reg {loss_lut_regularize_log.val:.4f} ({loss_lut_regularize_log.avg:.4f})\t' \
                          f'Remaining Time {remain_time} ({finish_time})'
                logger.info(log_msg)

                if opt.wandb:
                    import wandb
                    wandb.log(
                        {'Train/Epoch': epoch, 'Train/lr': optimizer.param_groups[0]['lr'], 'Train/Step': current_step,
                         'Train/Loss': loss_log.val,
                         'Train/Loss_fg_bg_cons': loss_fg_content_bg_appearance_construct_log.val,
                         'Train/Loss_lut_trans': loss_lut_transform_image_log.val,
                         'Train/Loss_lut_reg': loss_lut_regularize_log.val,
                         })

            time_ckp = time.time()

        state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'last_epoch': epoch,
                 'scheduler': scheduler.state_dict()}

        """
            As the validation of original resolution Harmonization will have no consistent resolution among images 
            (so fail to form a batch) and also may lead to out-of-memory problem when combined with training phase,
            we here only save the model when `opt.isFullRes` is True, leaving the evaluation in `inference.py`.
        """
        if opt.isFullRes and opt.hr_train:
            if epoch % 5 == 0:
                torch.save(state, os.path.join(opt.save_path, f"epoch{epoch}.pth"))
            else:
                torch.save(state, os.path.join(opt.save_path, "last.pth"))
        else:
            val(val_loader, model, logger, opt, state)


def val(val_loader, model, logger, opt, state):
    global best_psnr
    current_process = 10
    model.eval()

    metric_log = {
        'HYouTube': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0},
        'All': {'Samples': 0, 'MSE': 0, 'fMSE': 0, 'PSNR': 0, 'SSIM': 0},
    }

    lut_metric_log = {
        'HYouTube': {'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)
        bg_INR_coordinates = batch['bg_INR_coordinates'].to(opt.device)
        fg_transfer_INR_RGB = batch['fg_transfer_INR_RGB'].to(opt.device)

        with torch.no_grad():
            fg_content_bg_appearance_construct, _, lut_transform_image = model(
                composite_image,
                mask,
                fg_INR_coordinates,
                bg_INR_coordinates)
        if opt.INRDecode:
            pred_fg_image = fg_content_bg_appearance_construct[-1]
        else:
            pred_fg_image = None
        fg_transfer_INR_RGB = misc.lin2img(fg_transfer_INR_RGB,
                                           val_loader.dataset.INR_dataset.size) if fg_transfer_INR_RGB is not None else None

        "For INR"
        mask_INR = torchvision.transforms.Resize(opt.INR_input_size)(mask)

        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.))

        "Save the output images. For every 10 epochs, save more results, otherwise, save little. Thus save storage."
        if state['last_epoch'] % 10 == 0:
            misc.visualize(real_image, composite_image, mask, pred_fg_image,
                           pred_harmonized_image, lut_transform_image, opt, state['last_epoch'], show=False,
                           wandb=opt.wandb, isAll=True, step=step)
        elif step == 0:
            misc.visualize(real_image, composite_image, mask, pred_fg_image,
                           pred_harmonized_image, lut_transform_image, opt, state['last_epoch'], show=False,
                           wandb=opt.wandb, step=step)

        if opt.INRDecode:
            mse, fmse, psnr, ssim = metrics.calc_metrics(misc.normalize(pred_harmonized_image, opt, 'inv'),
                                                         misc.normalize(fg_transfer_INR_RGB, opt, 'inv'), mask_INR)

        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)

        if opt.wandb:
            import wandb
            if opt.INRDecode:
                wandb.log(
                    {f'Val/{key}/Epoch': state['last_epoch'],
                     f'Val/{key}/MSE': metric_log[key]['MSE'] / metric_log[key]['Samples'],
                     f'Val/{key}/fMSE': metric_log[key]['fMSE'] / metric_log[key]['Samples'],
                     f'Val/{key}/PSNR': metric_log[key]['PSNR'] / metric_log[key]['Samples'],
                     f'Val/{key}/SSIM': metric_log[key]['SSIM'] / metric_log[key]['Samples'],
                     f'Val/{key}/LUT_MSE': lut_metric_log[key]['MSE'] / lut_metric_log[key]['Samples'],
                     f'Val/{key}/LUT_fMSE': lut_metric_log[key]['fMSE'] / lut_metric_log[key]['Samples'],
                     f'Val/{key}/LUT_PSNR': lut_metric_log[key]['PSNR'] / lut_metric_log[key]['Samples'],
                     f'Val/{key}/LUT_SSIM': lut_metric_log[key]['SSIM'] / lut_metric_log[key]['Samples']
                     })
            else:
                wandb.log(
                    {f'Val/{key}/Epoch': state['last_epoch'],
                     f'Val/{key}/LUT_MSE': lut_metric_log[key]['MSE'] / lut_metric_log[key]['Samples'],
                     f'Val/{key}/LUT_fMSE': lut_metric_log[key]['fMSE'] / lut_metric_log[key]['Samples'],
                     f'Val/{key}/LUT_PSNR': lut_metric_log[key]['PSNR'] / lut_metric_log[key]['Samples'],
                     f'Val/{key}/LUT_SSIM': lut_metric_log[key]['SSIM'] / lut_metric_log[key]['Samples']
                     })

    logger.info('=========================')

    if not opt.INRDecode:
        if lut_metric_log['All']['PSNR'] / lut_metric_log['All']['Samples'] > best_psnr:
            logger.info("Best Save!")
            best_psnr = lut_metric_log['All']['PSNR'] / lut_metric_log['All']['Samples']
            torch.save(state, os.path.join(opt.save_path, "best.pth"))
        else:
            logger.info("Last Save!")
            torch.save(state, os.path.join(opt.save_path, "last.pth"))
    else:
        if metric_log['All']['PSNR'] / metric_log['All']['Samples'] > best_psnr:
            logger.info("Best Save!")
            best_psnr = metric_log['All']['PSNR'] / metric_log['All']['Samples']
            torch.save(state, os.path.join(opt.save_path, "best.pth"))
        else:
            logger.info("Last Save!")
            torch.save(state, os.path.join(opt.save_path, "last.pth"))