Upload 5 files

lib/pymaf/core/base_trainer.py

@@ -0,0 +1,107 @@
+# This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/utils/base_trainer.py
+from __future__ import division
+import logging
+from utils import CheckpointSaver
+from tensorboardX import SummaryWriter
+
+import torch
+from tqdm import tqdm
+
+tqdm.monitor_interval = 0
+
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTrainer(object):
+    """Base class for Trainer objects.
+    Takes care of checkpointing/logging/resuming training.
+    """
+
+    def __init__(self, options):
+        self.options = options
+        if options.multiprocessing_distributed:
+            self.device = torch.device('cuda', options.gpu)
+        else:
+            self.device = torch.device(
+                'cuda' if torch.cuda.is_available() else 'cpu')
+        # override this function to define your model, optimizers etc.
+        self.saver = CheckpointSaver(save_dir=options.checkpoint_dir,
+                                     overwrite=options.overwrite)
+        if options.rank == 0:
+            self.summary_writer = SummaryWriter(self.options.summary_dir)
+        self.init_fn()
+
+        self.checkpoint = None
+        if options.resume and self.saver.exists_checkpoint():
+            self.checkpoint = self.saver.load_checkpoint(
+                self.models_dict, self.optimizers_dict)
+
+        if self.checkpoint is None:
+            self.epoch_count = 0
+            self.step_count = 0
+        else:
+            self.epoch_count = self.checkpoint['epoch']
+            self.step_count = self.checkpoint['total_step_count']
+
+        if self.checkpoint is not None:
+            self.checkpoint_batch_idx = self.checkpoint['batch_idx']
+        else:
+            self.checkpoint_batch_idx = 0
+
+        self.best_performance = float('inf')
+
+    def load_pretrained(self, checkpoint_file=None):
+        """Load a pretrained checkpoint.
+        This is different from resuming training using --resume.
+        """
+        if checkpoint_file is not None:
+            checkpoint = torch.load(checkpoint_file)
+            for model in self.models_dict:
+                if model in checkpoint:
+                    self.models_dict[model].load_state_dict(checkpoint[model],
+                                                            strict=True)
+                    print(f'Checkpoint {model} loaded')
+
+    def move_dict_to_device(self, dict, device, tensor2float=False):
+        for k, v in dict.items():
+            if isinstance(v, torch.Tensor):
+                if tensor2float:
+                    dict[k] = v.float().to(device)
+                else:
+                    dict[k] = v.to(device)
+
+    # The following methods (with the possible exception of test) have to be implemented in the derived classes
+    def train(self, epoch):
+        raise NotImplementedError('You need to provide an train method')
+
+    def init_fn(self):
+        raise NotImplementedError('You need to provide an _init_fn method')
+
+    def train_step(self, input_batch):
+        raise NotImplementedError('You need to provide a _train_step method')
+
+    def train_summaries(self, input_batch):
+        raise NotImplementedError(
+            'You need to provide a _train_summaries method')
+
+    def visualize(self, input_batch):
+        raise NotImplementedError('You need to provide a visualize method')
+
+    def validate(self):
+        pass
+
+    def test(self):
+        pass
+
+    def evaluate(self):
+        pass
+
+    def fit(self):
+        # Run training for num_epochs epochs
+        for epoch in tqdm(range(self.epoch_count, self.options.num_epochs),
+                          total=self.options.num_epochs,
+                          initial=self.epoch_count):
+            self.epoch_count = epoch
+            self.train(epoch)
+        return

lib/pymaf/core/cfgs.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os
+import json
+from yacs.config import CfgNode as CN
+
+# Configuration variables
+cfg = CN(new_allowed=True)
+
+cfg.OUTPUT_DIR = 'results'
+cfg.DEVICE = 'cuda'
+cfg.DEBUG = False
+cfg.LOGDIR = ''
+cfg.VAL_VIS_BATCH_FREQ = 200
+cfg.TRAIN_VIS_ITER_FERQ = 1000
+cfg.SEED_VALUE = -1
+
+cfg.TRAIN = CN(new_allowed=True)
+
+cfg.LOSS = CN(new_allowed=True)
+cfg.LOSS.KP_2D_W = 300.0
+cfg.LOSS.KP_3D_W = 300.0
+cfg.LOSS.SHAPE_W = 0.06
+cfg.LOSS.POSE_W = 60.0
+cfg.LOSS.VERT_W = 0.0
+
+# Loss weights for dense correspondences
+cfg.LOSS.INDEX_WEIGHTS = 2.0
+# Loss weights for surface parts. (24 Parts)
+cfg.LOSS.PART_WEIGHTS = 0.3
+# Loss weights for UV regression.
+cfg.LOSS.POINT_REGRESSION_WEIGHTS = 0.5
+
+cfg.MODEL = CN(new_allowed=True)
+
+cfg.MODEL.PyMAF = CN(new_allowed=True)
+
+# switch
+cfg.TRAIN.VAL_LOOP = True
+
+cfg.TEST = CN(new_allowed=True)
+
+
+def get_cfg_defaults():
+    """Get a yacs CfgNode object with default values for my_project."""
+    # Return a clone so that the defaults will not be altered
+    # This is for the "local variable" use pattern
+    # return cfg.clone()
+    return cfg
+
+
+def update_cfg(cfg_file):
+    # cfg = get_cfg_defaults()
+    cfg.merge_from_file(cfg_file)
+    # return cfg.clone()
+    return cfg
+
+
+def parse_args(args):
+    cfg_file = args.cfg_file
+    if args.cfg_file is not None:
+        cfg = update_cfg(args.cfg_file)
+    else:
+        cfg = get_cfg_defaults()
+
+    # if args.misc is not None:
+    #     cfg.merge_from_list(args.misc)
+
+    return cfg
+
+
+def parse_args_extend(args):
+    if args.resume:
+        if not os.path.exists(args.log_dir):
+            raise ValueError(
+                'Experiment are set to resume mode, but log directory does not exist.'
+            )
+
+        # load log's cfg
+        cfg_file = os.path.join(args.log_dir, 'cfg.yaml')
+        cfg = update_cfg(cfg_file)
+
+        if args.misc is not None:
+            cfg.merge_from_list(args.misc)
+    else:
+        parse_args(args)

lib/pymaf/core/fits_dict.py

@@ -0,0 +1,133 @@
+'''
+This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/train/fits_dict.py
+'''
+import os
+import cv2
+import torch
+import numpy as np
+from torchgeometry import angle_axis_to_rotation_matrix, rotation_matrix_to_angle_axis
+
+from core import path_config, constants
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class FitsDict():
+    """ Dictionary keeping track of the best fit per image in the training set """
+
+    def __init__(self, options, train_dataset):
+        self.options = options
+        self.train_dataset = train_dataset
+        self.fits_dict = {}
+        self.valid_fit_state = {}
+        # array used to flip SMPL pose parameters
+        self.flipped_parts = torch.tensor(constants.SMPL_POSE_FLIP_PERM,
+                                          dtype=torch.int64)
+        # Load dictionary state
+        for ds_name, ds in train_dataset.dataset_dict.items():
+            if ds_name in ['h36m']:
+                dict_file = os.path.join(path_config.FINAL_FITS_DIR,
+                                         ds_name + '.npy')
+                self.fits_dict[ds_name] = torch.from_numpy(np.load(dict_file))
+                self.valid_fit_state[ds_name] = torch.ones(len(
+                    self.fits_dict[ds_name]),
+                    dtype=torch.uint8)
+            else:
+                dict_file = os.path.join(path_config.FINAL_FITS_DIR,
+                                         ds_name + '.npz')
+                fits_dict = np.load(dict_file)
+                opt_pose = torch.from_numpy(fits_dict['pose'])
+                opt_betas = torch.from_numpy(fits_dict['betas'])
+                opt_valid_fit = torch.from_numpy(fits_dict['valid_fit']).to(
+                    torch.uint8)
+                self.fits_dict[ds_name] = torch.cat([opt_pose, opt_betas],
+                                                    dim=1)
+                self.valid_fit_state[ds_name] = opt_valid_fit
+
+        if not options.single_dataset:
+            for ds in train_dataset.datasets:
+                if ds.dataset not in ['h36m']:
+                    ds.pose = self.fits_dict[ds.dataset][:, :72].numpy()
+                    ds.betas = self.fits_dict[ds.dataset][:, 72:].numpy()
+                    ds.has_smpl = self.valid_fit_state[ds.dataset].numpy()
+
+    def save(self):
+        """ Save dictionary state to disk """
+        for ds_name in self.train_dataset.dataset_dict.keys():
+            dict_file = os.path.join(self.options.checkpoint_dir,
+                                     ds_name + '_fits.npy')
+            np.save(dict_file, self.fits_dict[ds_name].cpu().numpy())
+
+    def __getitem__(self, x):
+        """ Retrieve dictionary entries """
+        dataset_name, ind, rot, is_flipped = x
+        batch_size = len(dataset_name)
+        pose = torch.zeros((batch_size, 72))
+        betas = torch.zeros((batch_size, 10))
+        for ds, i, n in zip(dataset_name, ind, range(batch_size)):
+            params = self.fits_dict[ds][i]
+            pose[n, :] = params[:72]
+            betas[n, :] = params[72:]
+        pose = pose.clone()
+        # Apply flipping and rotation
+        pose = self.flip_pose(self.rotate_pose(pose, rot), is_flipped)
+        betas = betas.clone()
+        return pose, betas
+
+    def get_vaild_state(self, dataset_name, ind):
+        batch_size = len(dataset_name)
+        valid_fit = torch.zeros(batch_size, dtype=torch.uint8)
+        for ds, i, n in zip(dataset_name, ind, range(batch_size)):
+            valid_fit[n] = self.valid_fit_state[ds][i]
+        valid_fit = valid_fit.clone()
+        return valid_fit
+
+    def __setitem__(self, x, val):
+        """ Update dictionary entries """
+        dataset_name, ind, rot, is_flipped, update = x
+        pose, betas = val
+        batch_size = len(dataset_name)
+        # Undo flipping and rotation
+        pose = self.rotate_pose(self.flip_pose(pose, is_flipped), -rot)
+        params = torch.cat((pose, betas), dim=-1).cpu()
+        for ds, i, n in zip(dataset_name, ind, range(batch_size)):
+            if update[n]:
+                self.fits_dict[ds][i] = params[n]
+
+    def flip_pose(self, pose, is_flipped):
+        """flip SMPL pose parameters"""
+        is_flipped = is_flipped.byte()
+        pose_f = pose.clone()
+        pose_f[is_flipped, :] = pose[is_flipped][:, self.flipped_parts]
+        # we also negate the second and the third dimension of the axis-angle representation
+        pose_f[is_flipped, 1::3] *= -1
+        pose_f[is_flipped, 2::3] *= -1
+        return pose_f
+
+    def rotate_pose(self, pose, rot):
+        """Rotate SMPL pose parameters by rot degrees"""
+        pose = pose.clone()
+        cos = torch.cos(-np.pi * rot / 180.)
+        sin = torch.sin(-np.pi * rot / 180.)
+        zeros = torch.zeros_like(cos)
+        r3 = torch.zeros(cos.shape[0], 1, 3, device=cos.device)
+        r3[:, 0, -1] = 1
+        R = torch.cat([
+            torch.stack([cos, -sin, zeros], dim=-1).unsqueeze(1),
+            torch.stack([sin, cos, zeros], dim=-1).unsqueeze(1), r3
+        ],
+            dim=1)
+        global_pose = pose[:, :3]
+        global_pose_rotmat = angle_axis_to_rotation_matrix(global_pose)
+        global_pose_rotmat_3b3 = global_pose_rotmat[:, :3, :3]
+        global_pose_rotmat_3b3 = torch.matmul(R, global_pose_rotmat_3b3)
+        global_pose_rotmat[:, :3, :3] = global_pose_rotmat_3b3
+        global_pose_rotmat = global_pose_rotmat[:, :-1, :-1].cpu().numpy()
+        global_pose_np = np.zeros((global_pose.shape[0], 3))
+        for i in range(global_pose.shape[0]):
+            aa, _ = cv2.Rodrigues(global_pose_rotmat[i])
+            global_pose_np[i, :] = aa.squeeze()
+        pose[:, :3] = torch.from_numpy(global_pose_np).to(pose.device)
+        return pose

lib/pymaf/core/path_config.py

@@ -0,0 +1,23 @@
+"""
+This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/path_config.py
+path configuration
+This file contains definitions of useful data stuctures and the paths
+for the datasets and data files necessary to run the code.
+Things you need to change: *_ROOT that indicate the path to each dataset
+"""
+import os
+from huggingface_hub import hf_hub_url, cached_download
+
+# pymaf
+pymaf_data_dir = hf_hub_url('Yuliang/PyMAF', '')
+smpl_data_dir = hf_hub_url('Yuliang/SMPL', '')
+SMPL_MODEL_DIR = os.path.join(smpl_data_dir, 'models/smpl')
+
+SMPL_MEAN_PARAMS = cached_download(os.path.join(pymaf_data_dir, 'smpl_mean_params.npz'), use_auth_token=os.environ['ICON'])
+MESH_DOWNSAMPLEING = cached_download(os.path.join(pymaf_data_dir, 'mesh_downsampling.npz'), use_auth_token=os.environ['ICON'])
+CUBE_PARTS_FILE = cached_download(os.path.join(pymaf_data_dir, 'cube_parts.npy'), use_auth_token=os.environ['ICON'])
+JOINT_REGRESSOR_TRAIN_EXTRA = cached_download(os.path.join(pymaf_data_dir, 'J_regressor_extra.npy'), use_auth_token=os.environ['ICON'])
+JOINT_REGRESSOR_H36M = cached_download(os.path.join(pymaf_data_dir, 'J_regressor_h36m.npy'), use_auth_token=os.environ['ICON'])
+VERTEX_TEXTURE_FILE = cached_download(os.path.join(pymaf_data_dir, 'vertex_texture.npy'), use_auth_token=os.environ['ICON'])
+SMPL_MEAN_PARAMS = cached_download(os.path.join(pymaf_data_dir, 'smpl_mean_params.npz'), use_auth_token=os.environ['ICON'])
+CHECKPOINT_FILE = cached_download(os.path.join(pymaf_data_dir, 'pretrained_model/PyMAF_model_checkpoint.pt'), use_auth_token=os.environ['ICON'])

lib/pymaf/core/train_options.py

@@ -0,0 +1,135 @@
+import argparse
+
+
+class TrainOptions():
+    def __init__(self):
+        self.parser = argparse.ArgumentParser()
+
+        gen = self.parser.add_argument_group('General')
+        gen.add_argument(
+            '--resume',
+            dest='resume',
+            default=False,
+            action='store_true',
+            help='Resume from checkpoint (Use latest checkpoint by default')
+
+        io = self.parser.add_argument_group('io')
+        io.add_argument('--log_dir',
+                        default='logs',
+                        help='Directory to store logs')
+        io.add_argument(
+            '--pretrained_checkpoint',
+            default=None,
+            help='Load a pretrained checkpoint at the beginning training')
+
+        train = self.parser.add_argument_group('Training Options')
+        train.add_argument('--num_epochs',
+                           type=int,
+                           default=200,
+                           help='Total number of training epochs')
+        train.add_argument('--regressor',
+                           type=str,
+                           choices=['hmr', 'pymaf_net'],
+                           default='pymaf_net',
+                           help='Name of the SMPL regressor.')
+        train.add_argument('--cfg_file',
+                           type=str,
+                           default='./configs/pymaf_config.yaml',
+                           help='config file path for PyMAF.')
+        train.add_argument(
+            '--img_res',
+            type=int,
+            default=224,
+            help='Rescale bounding boxes to size [img_res, img_res] before feeding them in the network'
+        )
+        train.add_argument(
+            '--rot_factor',
+            type=float,
+            default=30,
+            help='Random rotation in the range [-rot_factor, rot_factor]')
+        train.add_argument(
+            '--noise_factor',
+            type=float,
+            default=0.4,
+            help='Randomly multiply pixel values with factor in the range [1-noise_factor, 1+noise_factor]'
+        )
+        train.add_argument(
+            '--scale_factor',
+            type=float,
+            default=0.25,
+            help='Rescale bounding boxes by a factor of [1-scale_factor,1+scale_factor]'
+        )
+        train.add_argument(
+            '--openpose_train_weight',
+            default=0.,
+            help='Weight for OpenPose keypoints during training')
+        train.add_argument('--gt_train_weight',
+                           default=1.,
+                           help='Weight for GT keypoints during training')
+        train.add_argument('--eval_dataset',
+                           type=str,
+                           default='h36m-p2-mosh',
+                           help='Name of the evaluation dataset.')
+        train.add_argument('--single_dataset',
+                           default=False,
+                           action='store_true',
+                           help='Use a single dataset')
+        train.add_argument('--single_dataname',
+                           type=str,
+                           default='h36m',
+                           help='Name of the single dataset.')
+        train.add_argument('--eval_pve',
+                           default=False,
+                           action='store_true',
+                           help='evaluate PVE')
+        train.add_argument('--overwrite',
+                           default=False,
+                           action='store_true',
+                           help='overwrite the latest checkpoint')
+
+        train.add_argument('--distributed',
+                           action='store_true',
+                           help='Use distributed training')
+        train.add_argument('--dist_backend',
+                           default='nccl',
+                           type=str,
+                           help='distributed backend')
+        train.add_argument('--dist_url',
+                           default='tcp://127.0.0.1:10356',
+                           type=str,
+                           help='url used to set up distributed training')
+        train.add_argument('--world_size',
+                           default=1,
+                           type=int,
+                           help='number of nodes for distributed training')
+        train.add_argument("--local_rank", default=0, type=int)
+        train.add_argument('--rank',
+                           default=0,
+                           type=int,
+                           help='node rank for distributed training')
+        train.add_argument(
+            '--multiprocessing_distributed',
+            action='store_true',
+            help='Use multi-processing distributed training to launch '
+            'N processes per node, which has N GPUs. This is the '
+            'fastest way to use PyTorch for either single node or '
+            'multi node data parallel training')
+
+        misc = self.parser.add_argument_group('Misc Options')
+        misc.add_argument('--misc',
+                          help="Modify config options using the command-line",
+                          default=None,
+                          nargs=argparse.REMAINDER)
+        return
+
+    def parse_args(self):
+        """Parse input arguments."""
+        self.args = self.parser.parse_args()
+        self.save_dump()
+        return self.args
+
+    def save_dump(self):
+        """Store all argument values to a json file.
+        The default location is logs/expname/args.json.
+        """
+        pass