videoavatars/step2_consensus.py

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import h5py
import argparse
import numpy as np
import chumpy as ch
import cPickle as pkl

from opendr.camera import ProjectPoints
from opendr.renderer import BoundaryRenderer, ColoredRenderer
from tqdm import tqdm

from util import im, mesh
from util.logger import log
from lib.frame import setup_frame_rays
from lib.rays import ray_objective
from lib.geometry import laplacian
from lib.ch import sp_dot
from models.smpl import Smpl
from models.bodyparts import faces_no_hands, regularize_laplace, regularize_model, regularize_symmetry


def get_cb(frame, base_smpl, camera, frustum):
    viz_mask = frame.mask / 255.
    base_smpl.pose[:] = frame.pose
    camera.t[:] = frame.trans
    camera.rt[:] = 0

    rn = ColoredRenderer(camera=camera, v=base_smpl, f=base_smpl.f, vc=np.ones_like(base_smpl),
                         frustum=frustum, bgcolor=0, num_channels=1)

    def cb(_):
        silh_diff = (rn.r - viz_mask + 1) / 2.
        im.show(silh_diff, waittime=1)

    return cb


def fit_consensus(frames, base_smpl, camera, frustum, model_data, nohands, icp_count, naked, display):
    if nohands:
        faces = faces_no_hands(base_smpl.f)
    else:
        faces = base_smpl.f

    vis_rn_b = BoundaryRenderer(camera=camera, frustum=frustum, f=faces, num_channels=1)
    vis_rn_m = ColoredRenderer(camera=camera, frustum=frustum, f=faces, vc=np.zeros_like(base_smpl), bgcolor=1,
                               num_channels=1)

    model_template = Smpl(model_data)
    model_template.betas[:] = base_smpl.betas.r

    g_laplace = regularize_laplace()
    g_model = regularize_model()
    g_symmetry = regularize_symmetry()

    for step, (w_laplace, w_model, w_symmetry, sigma) in enumerate(zip(
            np.linspace(6.5, 4.0, icp_count) if naked else np.linspace(4.0, 2.0, icp_count),
            np.linspace(0.9, 0.6, icp_count) if naked else np.linspace(0.6, 0.3, icp_count),
            np.linspace(3.6, 1.8, icp_count),
            np.linspace(0.06, 0.003, icp_count),
    )):
        log.info('# Step {}'.format(step))

        L = laplacian(model_template.r, base_smpl.f)
        delta = L.dot(model_template.r)

        w_laplace *= g_laplace.reshape(-1, 1)
        w_model *= g_model.reshape(-1, 1)
        w_symmetry *= g_symmetry.reshape(-1, 1)

        E = {
            'laplace': (sp_dot(L, base_smpl.v_shaped_personal) - delta) * w_laplace,
            'model': (base_smpl.v_shaped_personal - model_template) * w_model,
            'symmetry': (base_smpl.v_personal + np.array([1, -1, -1])
                         * base_smpl.v_personal[model_data['vert_sym_idxs']]) * w_symmetry,
        }

        log.info('## Matching rays with contours')
        for current, f in enumerate(tqdm(frames)):
            E['silh_{}'.format(current)] = ray_objective(f, sigma, base_smpl, camera, vis_rn_b, vis_rn_m)

        log.info('## Run optimization')
        ch.minimize(
            E,
            [base_smpl.v_personal, model_template.betas],
            method='dogleg',
            options={'maxiter': 15, 'e_3': 0.001},
            callback=get_cb(frames[0], base_smpl, camera, frustum) if display else None
        )


def main(pose_file, masks_file, camera_file, out, obj_out, num, icp_count, model_file, first_frame, last_frame,
         nohands, naked, display):

    # load data
    with open(model_file, 'rb') as fp:
        model_data = pkl.load(fp)

    with open(camera_file, 'rb') as fp:
        camera_data = pkl.load(fp)

    pose_data = h5py.File(pose_file, 'r')
    poses = pose_data['pose'][first_frame:last_frame]
    trans = pose_data['trans'][first_frame:last_frame]
    masks = h5py.File(masks_file, 'r')['masks'][first_frame:last_frame]
    num_frames = masks.shape[0]

    indices_consensus = np.ceil(np.arange(num) * num_frames * 1. / num).astype(np.int)

    # init
    base_smpl = Smpl(model_data)
    base_smpl.betas[:] = np.array(pose_data['betas'], dtype=np.float32)

    camera = ProjectPoints(t=np.zeros(3), rt=np.zeros(3), c=camera_data['camera_c'],
                           f=camera_data['camera_f'], k=camera_data['camera_k'], v=base_smpl)
    camera_t = camera_data['camera_t']
    camera_rt = camera_data['camera_rt']
    frustum = {'near': 0.1, 'far': 1000., 'width': int(camera_data['width']), 'height': int(camera_data['height'])}
    frames = []

    for i in indices_consensus:
        log.info('Set up frame {}...'.format(i))

        mask = np.array(masks[i] * 255, dtype=np.uint8)
        pose_i = np.array(poses[i], dtype=np.float32)
        trans_i = np.array(trans[i], dtype=np.float32)

        frames.append(setup_frame_rays(base_smpl, camera, camera_t, camera_rt, pose_i, trans_i, mask))

    log.info('Set up complete.')
    log.info('Begin consensus fit...')
    fit_consensus(frames, base_smpl, camera, frustum, model_data, nohands, icp_count, naked, display)

    with open(out, 'wb') as fp:
        pkl.dump({
            'v_personal': base_smpl.v_personal.r,
            'betas': base_smpl.betas.r,
        }, fp, protocol=2)

    if obj_out is not None:
        base_smpl.pose[:] = 0
        vt = np.load('assets/basicModel_vt.npy')
        ft = np.load('assets/basicModel_ft.npy')
        mesh.write(obj_out, base_smpl.r, base_smpl.f, vt=vt, ft=ft)

    log.info('Done.')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'pose_file',
        type=str,
        help="File that contains poses")
    parser.add_argument(
        'masks_file',
        type=str,
        help="File that contains segmentations")
    parser.add_argument(
        'camera',
        type=str,
        help="pkl file that contains camera settings")
    parser.add_argument(
        'out',
        type=str,
        help="Out file path")
    parser.add_argument(
        '--obj_out', '-oo',
        default=None,
        help='obj out file name (optional)')
    parser.add_argument(
        '--num', '-n', default=120, type=int,
        help="Number of used frames")
    parser.add_argument(
        '--icp', '-i', default=3, type=int,
        help="ICP Iterations")
    parser.add_argument(
        '--model', '-m',
        default='vendor/smpl/models/basicmodel_m_lbs_10_207_0_v1.0.0.pkl',
        help='Path to SMPL model')
    parser.add_argument(
        '--first_frame', '-f', default=0, type=int,
        help="First frame to use")
    parser.add_argument(
        '--last_frame', '-l', default=2000, type=int,
        help="Last frame to use")
    parser.add_argument(
        '--nohands', '-nh',
        action='store_true',
        help="Exclude hands from optimization")
    parser.add_argument(
        '--naked', '-nk',
        action='store_true',
        help="Person wears (almost) no clothing")
    parser.add_argument(
        '--display', '-d',
        action='store_true',
        help="Enable visualization")

    args = parser.parse_args()

    main(args.pose_file, args.masks_file, args.camera, args.out, args.obj_out, args.num, args.icp, args.model,
         args.first_frame, args.last_frame, args.nohands, args.naked, args.display)