models/parsing_gen_model.py

import logging
import math
from collections import OrderedDict

import mmcv
import numpy as np
import torch
from torchvision.utils import save_image

from models.archs.fcn_arch import FCNHead
from models.archs.shape_attr_embedding_arch import ShapeAttrEmbedding
from models.archs.unet_arch import ShapeUNet
from models.losses.accuracy import accuracy
from models.losses.cross_entropy_loss import CrossEntropyLoss

logger = logging.getLogger('base')


class ParsingGenModel():
    """Paring Generation model.
    """

    def __init__(self, opt):
        self.opt = opt
        self.device = torch.device('cuda')
        self.is_train = opt['is_train']

        self.attr_embedder = ShapeAttrEmbedding(
            dim=opt['embedder_dim'],
            out_dim=opt['embedder_out_dim'],
            cls_num_list=opt['attr_class_num']).to(self.device)
        self.parsing_encoder = ShapeUNet(
            in_channels=opt['encoder_in_channels']).to(self.device)
        self.parsing_decoder = FCNHead(
            in_channels=opt['fc_in_channels'],
            in_index=opt['fc_in_index'],
            channels=opt['fc_channels'],
            num_convs=opt['fc_num_convs'],
            concat_input=opt['fc_concat_input'],
            dropout_ratio=opt['fc_dropout_ratio'],
            num_classes=opt['fc_num_classes'],
            align_corners=opt['fc_align_corners'],
        ).to(self.device)

        self.init_training_settings()

        self.palette = [[0, 0, 0], [255, 250, 250], [220, 220, 220],
                        [250, 235, 215], [255, 250, 205], [211, 211, 211],
                        [70, 130, 180], [127, 255, 212], [0, 100, 0],
                        [50, 205, 50], [255, 255, 0], [245, 222, 179],
                        [255, 140, 0], [255, 0, 0], [16, 78, 139],
                        [144, 238, 144], [50, 205, 174], [50, 155, 250],
                        [160, 140, 88], [213, 140, 88], [90, 140, 90],
                        [185, 210, 205], [130, 165, 180], [225, 141, 151]]

    def init_training_settings(self):
        optim_params = []
        for v in self.attr_embedder.parameters():
            if v.requires_grad:
                optim_params.append(v)
        for v in self.parsing_encoder.parameters():
            if v.requires_grad:
                optim_params.append(v)
        for v in self.parsing_decoder.parameters():
            if v.requires_grad:
                optim_params.append(v)
        # set up optimizers
        self.optimizer = torch.optim.Adam(
            optim_params,
            self.opt['lr'],
            weight_decay=self.opt['weight_decay'])
        self.log_dict = OrderedDict()
        self.entropy_loss = CrossEntropyLoss().to(self.device)

    def feed_data(self, data):
        self.pose = data['densepose'].to(self.device)
        self.attr = data['attr'].to(self.device)
        self.segm = data['segm'].to(self.device)

    def optimize_parameters(self):
        self.attr_embedder.train()
        self.parsing_encoder.train()
        self.parsing_decoder.train()

        self.attr_embedding = self.attr_embedder(self.attr)
        self.pose_enc = self.parsing_encoder(self.pose, self.attr_embedding)
        self.seg_logits = self.parsing_decoder(self.pose_enc)

        loss = self.entropy_loss(self.seg_logits, self.segm)

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

        self.log_dict['loss_total'] = loss

    def get_vis(self, save_path):
        img_cat = torch.cat([
            self.pose,
            self.segm,
        ], dim=3).detach()
        img_cat = ((img_cat + 1) / 2)

        img_cat = img_cat.clamp_(0, 1)

        save_image(img_cat, save_path, nrow=1, padding=4)

    def inference(self, data_loader, save_dir):
        self.attr_embedder.eval()
        self.parsing_encoder.eval()
        self.parsing_decoder.eval()

        acc = 0
        num = 0

        for _, data in enumerate(data_loader):
            pose = data['densepose'].to(self.device)
            attr = data['attr'].to(self.device)
            segm = data['segm'].to(self.device)
            img_name = data['img_name']

            num += pose.size(0)
            with torch.no_grad():
                attr_embedding = self.attr_embedder(attr)
                pose_enc = self.parsing_encoder(pose, attr_embedding)
                seg_logits = self.parsing_decoder(pose_enc)
            seg_pred = seg_logits.argmax(dim=1)
            acc += accuracy(seg_logits, segm)
            palette_label = self.palette_result(segm.cpu().numpy())
            palette_pred = self.palette_result(seg_pred.cpu().numpy())
            pose_numpy = ((pose[0] + 1) / 2. * 255.).expand(
                3,
                pose[0].size(1),
                pose[0].size(2),
            ).cpu().numpy().clip(0, 255).astype(np.uint8).transpose(1, 2, 0)
            concat_result = np.concatenate(
                (pose_numpy, palette_pred, palette_label), axis=1)
            mmcv.imwrite(concat_result, f'{save_dir}/{img_name[0]}')

        self.attr_embedder.train()
        self.parsing_encoder.train()
        self.parsing_decoder.train()
        return (acc / num).item()

    def get_current_log(self):
        return self.log_dict

    def update_learning_rate(self, epoch):
        """Update learning rate.

        Args:
            current_iter (int): Current iteration.
            warmup_iter (int): Warmup iter numbers. -1 for no warmup.
                Default: -1.
        """
        lr = self.optimizer.param_groups[0]['lr']

        if self.opt['lr_decay'] == 'step':
            lr = self.opt['lr'] * (
                self.opt['gamma']**(epoch // self.opt['step']))
        elif self.opt['lr_decay'] == 'cos':
            lr = self.opt['lr'] * (
                1 + math.cos(math.pi * epoch / self.opt['num_epochs'])) / 2
        elif self.opt['lr_decay'] == 'linear':
            lr = self.opt['lr'] * (1 - epoch / self.opt['num_epochs'])
        elif self.opt['lr_decay'] == 'linear2exp':
            if epoch < self.opt['turning_point'] + 1:
                # learning rate decay as 95%
                # at the turning point (1 / 95% = 1.0526)
                lr = self.opt['lr'] * (
                    1 - epoch / int(self.opt['turning_point'] * 1.0526))
            else:
                lr *= self.opt['gamma']
        elif self.opt['lr_decay'] == 'schedule':
            if epoch in self.opt['schedule']:
                lr *= self.opt['gamma']
        else:
            raise ValueError('Unknown lr mode {}'.format(self.opt['lr_decay']))
        # set learning rate
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

        return lr

    def save_network(self, save_path):
        """Save networks.
        """

        save_dict = {}
        save_dict['embedder'] = self.attr_embedder.state_dict()
        save_dict['encoder'] = self.parsing_encoder.state_dict()
        save_dict['decoder'] = self.parsing_decoder.state_dict()

        torch.save(save_dict, save_path)

    def load_network(self):
        checkpoint = torch.load(self.opt['pretrained_parsing_gen'])

        self.attr_embedder.load_state_dict(checkpoint['embedder'], strict=True)
        self.attr_embedder.eval()

        self.parsing_encoder.load_state_dict(
            checkpoint['encoder'], strict=True)
        self.parsing_encoder.eval()

        self.parsing_decoder.load_state_dict(
            checkpoint['decoder'], strict=True)
        self.parsing_decoder.eval()

    def palette_result(self, result):
        seg = result[0]
        palette = np.array(self.palette)
        assert palette.shape[1] == 3
        assert len(palette.shape) == 2
        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
        for label, color in enumerate(palette):
            color_seg[seg == label, :] = color
        # convert to BGR
        color_seg = color_seg[..., ::-1]
        return color_seg