main.py

"""
@Date: 2021/07/17
@description:
"""
from cmath import nan
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import sys
import shutil
import argparse
import numpy as np
import json
import torch
import torch.nn.parallel
import torch.optim
import torch.multiprocessing as mp
import torch.utils.data
import torch.utils.data.distributed
import torch.cuda

from PIL import Image
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter
from config.defaults import get_config, get_rank_config
from models.other.criterion import calc_criterion
from models.build import build_model
from models.other.init_env import init_env
from utils.logger import build_logger
from utils.misc import tensor2np_d, tensor2np
from dataset.build import build_loader
from evaluation.accuracy import calc_accuracy, show_heat_map, calc_ce, calc_pe, calc_rmse_delta_1, \
    show_depth_normal_grad, calc_f1_score
from postprocessing.post_process import post_process

try:
    from apex import amp
except ImportError:
    amp = None


def parse_option():
    debug = True if sys.gettrace() else False
    parser = argparse.ArgumentParser(description='Panorama Layout Transformer training and evaluation script')
    parser.add_argument('--cfg',
                        type=str,
                        metavar='FILE',
                        help='path to config file')

    parser.add_argument('--mode',
                        type=str,
                        default='train',
                        choices=['train', 'val', 'test'],
                        help='train/val/test mode')

    parser.add_argument('--val_name',
                        type=str,
                        choices=['val', 'test'],
                        help='val name')

    parser.add_argument('--bs', type=int,
                        help='batch size')

    parser.add_argument('--save_eval', action='store_true',
                        help='save eval result')

    parser.add_argument('--post_processing', type=str,
                        choices=['manhattan', 'atalanta', 'manhattan_old'],
                        help='type of postprocessing ')

    parser.add_argument('--need_cpe', action='store_true',
                        help='need to evaluate corner error and pixel error')

    parser.add_argument('--need_f1', action='store_true',
                        help='need to evaluate f1-score of corners')

    parser.add_argument('--force_cube', action='store_true',
                        help='force cube shape when eval')
                        
    parser.add_argument('--need_rmse', action='store_true',
                        help='need to evaluate root mean squared error and delta error')

    parser.add_argument('--wall_num', type=int,
                        help='wall number')

    args = parser.parse_args()
    args.debug = debug
    print("arguments:")
    for arg in vars(args):
        print(arg, ":", getattr(args, arg))
    print("-" * 50)
    return args


def main():
    args = parse_option()
    config = get_config(args)

    os.makedirs(config.CKPT.DIR, exist_ok=True)
    os.makedirs(config.CKPT.RESULT_DIR, exist_ok=True)
    os.makedirs(config.LOGGER.DIR, exist_ok=True)

    if config.MODE == 'train':
        with open(os.path.join(config.CKPT.DIR, "config.yaml"), "w") as f:
            f.write(config.dump(allow_unicode=True))

    print(f"Use single process, device:{config.TRAIN.DEVICE}")
    main_worker(0, config, 1)


def main_worker(local_rank, cfg, world_size):
    config = get_rank_config(cfg, local_rank, world_size)
    logger = build_logger(config)
    writer = SummaryWriter(config.CKPT.DIR)
    logger.info(f"Comment: {config.COMMENT}")
    cur_pid = os.getpid()
    logger.info(f"Current process id: {cur_pid}")
    init_env(config.SEED, config.TRAIN.DETERMINISTIC, config.DATA.NUM_WORKERS)

    model, optimizer, criterion, scheduler = build_model(config, logger)
    train_data_loader, val_data_loader = build_loader(config, logger)

    if 'cuda' in config.TRAIN.DEVICE:
        torch.cuda.set_device(config.TRAIN.DEVICE)

    if config.MODE == 'train':
        train(model, train_data_loader, val_data_loader, optimizer, criterion, config, logger, writer, scheduler)
    else:
        iou_results, other_results = val_an_epoch(model, val_data_loader,
                                                  criterion, config, logger, writer=None,
                                                  epoch=config.TRAIN.START_EPOCH)
        results = dict(iou_results, **other_results)
        if config.SAVE_EVAL:
            save_path = os.path.join(config.CKPT.RESULT_DIR, f"result.json")
            with open(save_path, 'w+') as f:
                json.dump(results, f, indent=4)


def save(model, optimizer, epoch, iou_d, logger, writer, config):
    model.save(optimizer, epoch, accuracy=iou_d['full_3d'], logger=logger, acc_d=iou_d, config=config)
    for k in model.acc_d:
        writer.add_scalar(f"BestACC/{k}", model.acc_d[k]['acc'], epoch)


def train(model, train_data_loader, val_data_loader, optimizer, criterion, config, logger, writer, scheduler):
    for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS):
        logger.info("=" * 100)
        train_an_epoch(model, train_data_loader, optimizer, criterion, config, logger, writer, epoch)

        epoch_iou_d, _ = val_an_epoch(model, val_data_loader, criterion, config, logger, writer, epoch)

        save(model, optimizer, epoch, epoch_iou_d, logger, writer, config)

        if scheduler is not None:
            if scheduler.min_lr is not None and optimizer.param_groups[0]['lr'] <= scheduler.min_lr:
                continue
            scheduler.step()
    writer.close()


def train_an_epoch(model, train_data_loader, optimizer, criterion, config, logger, writer, epoch=0):
    logger.info(f'Start Train Epoch {epoch}/{config.TRAIN.EPOCHS - 1}')
    model.train()

    optimizer.zero_grad()

    data_len = len(train_data_loader)
    start_i = data_len * epoch * config.WORLD_SIZE
    bar = enumerate(train_data_loader)

    bar = tqdm(bar, total=data_len, ncols=200)

    device = config.TRAIN.DEVICE
    epoch_loss_d = {}
    for i, gt in bar:
        imgs = gt['image'].to(device, non_blocking=True)
        gt['depth'] = gt['depth'].to(device, non_blocking=True)
        gt['ratio'] = gt['ratio'].to(device, non_blocking=True)
        gt['segmentation'] = gt['segmentation'].to(device, non_blocking=True)

        dt = model(imgs)
        loss, batch_loss_d, epoch_loss_d = calc_criterion(criterion, gt, dt, epoch_loss_d)

        bar.set_postfix(batch_loss_d)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        global_step = start_i + i * config.WORLD_SIZE + config.LOCAL_RANK
        for key, val in batch_loss_d.items():
            writer.add_scalar(f'TrainBatchLoss/{key}', val, global_step)

    epoch_loss_d = dict(zip(epoch_loss_d.keys(), [np.array(epoch_loss_d[k]).mean() for k in epoch_loss_d.keys()]))
    s = 'TrainEpochLoss: '
    for key, val in epoch_loss_d.items():
        writer.add_scalar(f'TrainEpochLoss/{key}', val, epoch)
        s += f" {key}={val}"
    logger.info(s)
    writer.add_scalar('LearningRate', optimizer.param_groups[0]['lr'], epoch)
    logger.info(f"LearningRate: {optimizer.param_groups[0]['lr']}")


@torch.no_grad()
def val_an_epoch(model, val_data_loader, criterion, config, logger, writer, epoch=0):
    model.eval()
    logger.info(f'Start Validate Epoch {epoch}/{config.TRAIN.EPOCHS - 1}')
    data_len = len(val_data_loader)
    start_i = data_len * epoch * config.WORLD_SIZE
    bar = enumerate(val_data_loader)
    bar = tqdm(bar, total=data_len, ncols=200)
    device = config.TRAIN.DEVICE
    epoch_loss_d = {}
    epoch_iou_d = {
        'visible_2d': [],
        'visible_3d': [],
        'full_2d': [],
        'full_3d': [],
        'height': []
    }

    epoch_other_d = {
        'ce': [],
        'pe': [],
        'f1': [],
        'precision': [],
        'recall': [],
        'rmse': [],
        'delta_1': []
    }

    for i, gt in bar:
        imgs = gt['image'].to(device, non_blocking=True)
        gt['ratio'] = gt['ratio'].to(device, non_blocking=True)
        gt['depth'] = gt['depth'].to(device, non_blocking=True)
        gt['segmentation'] = gt['segmentation'].to(device, non_blocking=True)

        dt = model(imgs)

        vis_w = config.TRAIN.VIS_WEIGHT  # 1024
        visualization = True

        loss, batch_loss_d, epoch_loss_d = calc_criterion(criterion, gt, dt, epoch_loss_d)

        if config.EVAL.POST_PROCESSING is not None:
            depth = tensor2np(dt['depth'])
            dt['processed_xyz'] = post_process(depth, type_name=config.EVAL.POST_PROCESSING,
                                               need_cube=config.EVAL.FORCE_CUBE)

            if config.EVAL.FORCE_CUBE and config.EVAL.NEED_CPE:
                ce = calc_ce(tensor2np_d(dt), tensor2np_d(gt))
                pe = calc_pe(tensor2np_d(dt), tensor2np_d(gt))

                epoch_other_d['ce'].append(ce)
                epoch_other_d['pe'].append(pe)

            if config.EVAL.NEED_F1:
                f1, precision, recall = calc_f1_score(tensor2np_d(dt), tensor2np_d(gt))
                epoch_other_d['f1'].append(f1)
                epoch_other_d['precision'].append(precision)
                epoch_other_d['recall'].append(recall)
                
        if config.EVAL.NEED_RMSE:
            rmse, delta_1 = calc_rmse_delta_1(tensor2np_d(dt), tensor2np_d(gt))
            epoch_other_d['rmse'].append(rmse)
            epoch_other_d['delta_1'].append(delta_1)

        visb_iou, full_iou, iou_height, pano_bds, full_iou_2ds = calc_accuracy(tensor2np_d(dt), tensor2np_d(gt),
                                                                               visualization, h=vis_w // 2)
        epoch_iou_d['visible_2d'].append(visb_iou[0])
        epoch_iou_d['visible_3d'].append(visb_iou[1])
        epoch_iou_d['full_2d'].append(full_iou[0])
        epoch_iou_d['full_3d'].append(full_iou[1])
        epoch_iou_d['height'].append(iou_height)

        bar.set_postfix(batch_loss_d)

        global_step = start_i + i * config.WORLD_SIZE + config.LOCAL_RANK

        if writer:
            for key, val in batch_loss_d.items():
                writer.add_scalar(f'ValBatchLoss/{key}', val, global_step)

        if not visualization:
            continue

        gt_grad_imgs, dt_grad_imgs = show_depth_normal_grad(dt, gt, device, vis_w)

        dt_heat_map_imgs = None
        gt_heat_map_imgs = None
        
        if config.TRAIN.VIS_MERGE or config.SAVE_EVAL:
            imgs = []
            for j in range(len(pano_bds)):
                # floorplan = np.concatenate([visb_iou[2][j], full_iou[2][j]], axis=-1)
                floorplan = full_iou[2][j]
                margin_w = int(floorplan.shape[-1] * (60/512))
                floorplan = floorplan[:, :, margin_w:-margin_w]
                floorplan = np.concatenate([floorplan, floorplan[:, 452:, :]], axis=-2)
                
                grad_h = dt_grad_imgs[0].shape[1]
                vis_merge = [
                    pano_bds[j][:, :],
                    dt_grad_imgs[j]
                ]
                if 'corner_heat_map' in gt:
                    vis_merge = [dt_heat_map_imgs[j], gt_heat_map_imgs[j]] + vis_merge
                img = np.concatenate(vis_merge, axis=-2)

                img = np.concatenate([img, floorplan], axis=-1)
                # img = gt_grad_imgs[j]
                imgs.append(img)
            if writer:
                writer.add_images('VIS/Merge', np.array(imgs), global_step)

            if config.SAVE_EVAL:
                for k in range(len(imgs)):
                    img = imgs[k] * 255.0
                    save_path = os.path.join(config.CKPT.RESULT_DIR, f"{gt['id'][k]}_{full_iou_2ds[k]:.5f}.png")
                    Image.fromarray(img.transpose(1, 2, 0).astype(np.uint8)).save(save_path)

        elif writer:
            writer.add_images('IoU/Visible_Floorplan', visb_iou[2], global_step)
            writer.add_images('IoU/Full_Floorplan', full_iou[2], global_step)
            writer.add_images('IoU/Boundary', pano_bds, global_step)
            writer.add_images('Grad/gt', gt_grad_imgs, global_step)
            writer.add_images('Grad/dt', dt_grad_imgs, global_step)

    epoch_loss_d = dict(zip(epoch_loss_d.keys(), [np.array(epoch_loss_d[k]).mean() for k in epoch_loss_d.keys()]))
    s = 'ValEpochLoss: '
    for key, val in epoch_loss_d.items():
        if writer:
            writer.add_scalar(f'ValEpochLoss/{key}', val, epoch)
        s += f" {key}={val}"
    logger.info(s)

    epoch_iou_d = dict(zip(epoch_iou_d.keys(), [np.array(epoch_iou_d[k]).mean() for k in epoch_iou_d.keys()]))
    s = 'ValEpochIoU: '
    for key, val in epoch_iou_d.items():
        if writer:
            writer.add_scalar(f'ValEpochIoU/{key}', val, epoch)
        s += f" {key}={val}"
    logger.info(s)
    epoch_other_d = dict(zip(epoch_other_d.keys(),
                             [np.array(epoch_other_d[k]).mean() if len(epoch_other_d[k]) > 0 else 0 for k in
                              epoch_other_d.keys()]))

    logger.info(f'other acc: {epoch_other_d}')
    return epoch_iou_d, epoch_other_d


if __name__ == '__main__':
    main()