GPS.py

import os
import os.path as osp
import torch
import torch.nn.functional as F
import numpy as np
import itertools
from tensorboardX import SummaryWriter

from NN.losses import make_criteria
from utils.base import logger

class GPS:
    def __init__(self,
                 init_mode: str = 'random_synthesis',
                 noise_sigma: float = 1.0,
                 coarse_ratio: float = 0.2,
                 coarse_ratio_factor: float = 6,
                 pyr_factor: float = 0.75,
                 num_stages_limit: int = -1,
                 device: str = 'cuda:0',
                 silent: bool = False
                 ):
        '''
        Args:
            init_mode:
                - 'random_synthesis': init with random seed
                - 'random': init with random seed
            noise_sigma: float = 1.0, random noise.
            coarse_ratio: float = 0.2, ratio at the coarse level.
            pyr_factor: float = 0.75, pyramid factor.
            num_stages_limit: int = -1, no limit.
            device: str = 'cuda:0', default device.
            silent: bool = False, mute the output.
        '''
        self.init_mode = init_mode
        self.noise_sigma = noise_sigma
        self.coarse_ratio = coarse_ratio
        self.coarse_ratio_factor = coarse_ratio_factor
        self.pyr_factor = pyr_factor
        self.num_stages_limit = num_stages_limit
        self.device = torch.device(device)
        self.silent = silent

    def _get_pyramid_lengths(self, dest, ext=None):
        """Get a list of pyramid lengths"""
        if self.coarse_ratio == -1:
            self.coarse_ratio = np.around(ext['criteria']['patch_size'] * self.coarse_ratio_factor / dest, 2)

        lengths = [int(np.round(dest * self.coarse_ratio))]
        while lengths[-1] < dest:
            lengths.append(int(np.round(lengths[-1] / self.pyr_factor)))
            if lengths[-1] == lengths[-2]:
                lengths[-1] += 1
        lengths[-1] = dest

        return lengths

    def _get_target_pyramid(self, target, ext=None):
        """Reads a target motion(s) and create a pyraimd out of it. Ordered in increatorch.sing size"""
        self._num_target = len(target)
        lengths = []
        min_len = 10000
        for i in range(len(target)):
            new_length = self._get_pyramid_lengths(len(target[i]), ext)
            min_len = min(min_len, len(new_length))
            if self.num_stages_limit != -1:
                new_length = new_length[:self.num_stages_limit]
            lengths.append(new_length)
        for i in range(len(target)):
            lengths[i] = lengths[i][-min_len:]
        self.pyraimd_lengths = lengths

        target_pyramid = [[] for _ in range(len(lengths[0]))]
        for step in range(len(lengths[0])):
            for i in range(len(target)):
                length = lengths[i][step]
                motion = target[i]
                target_pyramid[step].append(motion.sample(size=length).to(self.device))
                # target_pyramid[step].append(motion.pos2velo(motion.sample(size=length)))
                # motion.motion_data = motion.pos2velo(motion.motion_data)
                # target_pyramid[step].append(motion.sample(size=length))
                # motion.motion_data = motion.velo2pos(motion.motion_data)    

        if not self.silent:
            print('Levels:', lengths)
            for i in range(len(target_pyramid)):
                print(f'Number of clips in target pyramid {i} is {len(target_pyramid[i])}: {[[tgt.min(), tgt.max()] for tgt in target_pyramid[i]]}')

        return target_pyramid

    def _get_initial_motion(self):
        """Prepare the initial motion for optimization"""
        if 'random_synthesis' in str(self.init_mode):
            m = self.init_mode.split('/')[-1]
            if m =='random_synthesis':
                final_length = sum([i[-1] for i in self.pyraimd_lengths])
            elif 'x' in m:
                final_length = int(m.replace('x', '')) * sum([i[-1] for i in self.pyraimd_lengths])
            elif (self.init_mode.split('/')[-1]).isdigit():
                final_length = int(self.init_mode.split('/')[-1])
            else:
                raise ValueError(f'incorrect init_mode: {self.init_mode}')

            self.synthesized_lengths = self._get_pyramid_lengths(final_length)

        else:
            raise ValueError(f'Unsupported init_mode {self.init_mode}')
            
        initial_motion = F.interpolate(torch.cat([self.target_pyramid[0][i] for i in range(self._num_target)], dim=-1),
                                       size=self.synthesized_lengths[0], mode='linear', align_corners=True)
        if self.noise_sigma > 0:
            initial_motion_w_noise = initial_motion + torch.randn_like(initial_motion) * self.noise_sigma
            initial_motion_w_noise = torch.fmod(initial_motion_w_noise, 1.0)
        else:
            initial_motion_w_noise = initial_motion

        if not self.silent:
            print('Synthesized lengths:', self.synthesized_lengths)
            print('Initial motion:', initial_motion.min(), initial_motion.max())
            print('Initial motion with noise:', initial_motion_w_noise.min(), initial_motion_w_noise.max())

        return initial_motion_w_noise

    def run(self, target, mode="backpropagate", ext=None, debug_dir=None):
        '''
        Run the patch-based motion synthesis.

        Args:
            target (torch.Tensor): Target data.
            mode (str): Optimization mode. Support ['backpropagate', 'match_and_blend']
            ext (dict): extra data or constrain.
            debug_dir (str): Debug directory.
        '''
        # preprare data
        self.target_pyramid = self._get_target_pyramid(target, ext)
        self.synthesized = self._get_initial_motion()
        if debug_dir is not None:
            writer = SummaryWriter(log_dir=debug_dir)

        # prepare configuration
        if mode == "backpropagate":
            self.synthesized.requires_grad_(True)
            assert 'criteria' in ext.keys(), 'Please specify a criteria for synthsis.'
            criteria = make_criteria(ext['criteria']).to(self.device)
        elif mode == "match_and_blend":
            self.synthesized.requires_grad_(False)
            assert 'criteria' in ext.keys(), 'Please specify a criteria for synthsis.'
            criteria = make_criteria(ext['criteria']).to(self.device)
        else:
            raise ValueError(f'Unsupported mode: {mode}')

        # perform synthsis
        self.pbar = logger(ext['num_itrs'], len(self.target_pyramid))
        ext['pbar'] = self.pbar
        for lvl, lvl_target in enumerate(self.target_pyramid):
            self.pbar.new_lvl()
            if lvl > 0:
                with torch.no_grad():
                    self.synthesized = F.interpolate(self.synthesized.detach(), size=self.synthesized_lengths[lvl], mode='linear')
                if mode == "backpropagate":
                    self.synthesized.requires_grad_(True)

            if mode == "backpropagate": # direct optimize the synthesized motion
                self.synthesized, losses = GPS.backpropagate(self.synthesized, lvl_target, criteria, ext=ext)
            elif mode == "match_and_blend":
                self.synthesized, losses = GPS.match_and_blend(self.synthesized, lvl_target, criteria, ext=ext)

            criteria.clean_cache()
            if debug_dir:
                for itr in range(len(losses)):
                    writer.add_scalar(f'optimize/losses_lvl{lvl}', losses[itr], itr)
        self.pbar.pbar.close()


        return self.synthesized.detach()

    @staticmethod
    def backpropagate(synthesized, targets, criteria=None, ext=None):
        """
        Minimizes criteria(synthesized, target) for num_steps SGD steps
        Args:
            targets (torch.Tensor): Target data.
            ext (dict): extra configurations.
        """
        if criteria is None:
            assert 'criteria' in ext.keys(), 'Criteria is not set'
            criteria = make_criteria(ext['criteria']).to(synthesized.device)

        optim = None
        if 'optimizer' in ext.keys():
            if ext['optimizer'] == 'Adam':
                optim = torch.optim.Adam([synthesized], lr=ext['lr'])
            elif ext['optimizer'] == 'SGD':
                optim = torch.optim.SGD([synthesized], lr=ext['lr'])
            elif ext['optimizer'] == 'RMSprop':
                optim = torch.optim.RMSprop([synthesized], lr=ext['lr'])
            else:
                print(f'use default RMSprop optimizer')
        optim = torch.optim.RMSprop([synthesized], lr=ext['lr']) if optim is None else optim
        # optim = torch.optim.Adam([synthesized], lr=ext['lr']) if optim is None else optim
        lr_decay = np.exp(np.log(0.333) / ext['num_itrs'])

        # other constraints
        trajectory = ext['trajectory'] if 'trajectory' in ext.keys() else None

        losses = []
        for _i in range(ext['num_itrs']):
            optim.zero_grad()
            
            loss = criteria(synthesized, targets)

            if trajectory is not None: ## velo constrain
                target_traj = F.interpolate(trajectory, size=synthesized.shape[-1], mode='linear')
                # target_traj = F.interpolate(trajectory, size=synthesized.shape[-1], mode='linear', align_corners=False)
                target_velo = ext['pos2velo'](target_traj)
                
                velo_mask = [-3, -1]
                loss += 1 * F.l1_loss(synthesized[:, velo_mask, :], target_velo[:, velo_mask, :])

            loss.backward()
            optim.step()

            # Update staus
            losses.append(loss.item())
            if 'pbar' in ext.keys():
                ext['pbar'].step()
                ext['pbar'].print()

        return synthesized, losses

    @staticmethod
    @torch.no_grad()
    def match_and_blend(synthesized, targets, criteria, ext):
        """
        Minimizes criteria(synthesized, target)
        Args:
            targets (torch.Tensor): Target data.
            ext (dict): extra configurations.
        """
        losses = []
        for _i in range(ext['num_itrs']):
            if 'parts_list' in ext.keys():
                def extract_part_motions(motion, parts_list):
                    part_motions = []
                    n_frames = motion.shape[-1]
                    rot, pos = motion[:, :-3, :].reshape(-1, 6, n_frames), motion[:, -3:, :]

                    for part in parts_list:
                        # part -= 1
                        part = [i -1 for i in part]

                        # print(part)
                        if 0 in part:
                            part_motions += [torch.cat([rot[part].view(1, -1, n_frames), pos.view(1, -1, n_frames)], dim=1)]
                        else:
                            part_motions += [rot[part].view(1, -1, n_frames)]

                    return part_motions
                def combine_part_motions(part_motions, parts_list):
                    assert len(part_motions) == len(parts_list)
                    n_frames = part_motions[0].shape[-1]
                    l = max(list(itertools.chain(*parts_list)))
                    # print(l, n_frames)
                    # motion = torch.zeros((1, (l+1)*6 + 3, n_frames), device=part_motions[0].device)
                    rot = torch.zeros(((l+1), 6, n_frames), device=part_motions[0].device)
                    pos = torch.zeros((1, 3, n_frames), device=part_motions[0].device)
                    div_rot = torch.zeros((l+1), device=part_motions[0].device)
                    div_pos = torch.zeros(1, device=part_motions[0].device)

                    for part_motion, part in zip(part_motions, parts_list):
                        part = [i -1 for i in part]

                        if 0 in part:
                            # print(part_motion.shape)
                            pos += part_motion[:, -3:, :]
                            div_pos += 1
                            rot[part] += part_motion[:, :-3, :].view(-1, 6, n_frames)
                            div_rot[part] += 1
                        else:
                            rot[part] += part_motion.view(-1, 6, n_frames)
                            div_rot[part] += 1
                            
                    # print(div_rot, div_pos)
                    # print(rot.shape)
                    rot = (rot.permute(1, 2, 0) / div_rot).permute(2, 0, 1)
                    pos = pos / div_pos

                    return torch.cat([rot.view(1, -1, n_frames), pos.view(1, 3, n_frames)], dim=1)

                # raw_synthesized = synthesized
                # print(synthesized, synthesized.shape)
                synthesized_part_motions = extract_part_motions(synthesized, ext['parts_list'])
                targets_part_motions = [extract_part_motions(target, ext['parts_list']) for target in targets]

                synthesized = []
                for _j in range(len(synthesized_part_motions)):
                    synthesized_part_motion = synthesized_part_motions[_j]
                    # synthesized += [synthesized_part_motion]
                    targets_part_motion = [target[_j] for target in targets_part_motions]
                    # # print(synthesized_part_motion.shape, targets_part_motion[0].shape)
                    synthesized += [criteria(synthesized_part_motion, targets_part_motion, ext=ext, return_blended_results=True)[0]]

                # print(len(synthesized))
                
                synthesized = combine_part_motions(synthesized, ext['parts_list'])
                # print(synthesized, synthesized.shape)
                # print((raw_synthesized-synthesized > 0.00001).sum())
                # exit()
                # print(synthesized.shape)
                losses = 0

                # exit()
       
            else:
                synthesized, loss = criteria(synthesized, targets, ext=ext, return_blended_results=True)

                # Update staus
                losses.append(loss.item())
                if 'pbar' in ext.keys():
                    ext['pbar'].step()
                    ext['pbar'].print()

        return synthesized, losses