Per-frame pkl export saves entire batch

There's a minor bug in the main function of the model transfer code. When the per-frame fitting result is exported to pkl files, it saves the entire batch instead of just a single frame. I simply fixed this as follows (file transfer_model/__main__.py):

var_dict = run_fitting(
    exp_cfg, batch, body_model, def_matrix, mask_ids)
paths = batch['paths']

for ii, path in enumerate(paths):
    _, fname = osp.split(path)

    var_dict_ii = {}
    for k in var_dict.keys():
        if k != 'faces':
            var_dict_ii[k] = var_dict[k][ii] if var_dict[k] is not None else None
    var_dict_ii['faces'] = var_dict['faces']

    output_path = osp.join(
        output_folder, f'{osp.splitext(fname)[0]}.pkl')
    with open(output_path, 'wb') as f:
        pickle.dump(var_dict_ii, f)

    output_path = osp.join(
        output_folder, f'{osp.splitext(fname)[0]}.obj')
    mesh = np_mesh_to_o3d(
        var_dict_ii['vertices'], var_dict_ii['faces'])
    o3d.io.write_triangle_mesh(output_path, mesh)

vchoutas / smplx

Per-frame pkl export saves entire batch #105