def main(args):
    from utils.animation_data import AnimationData
    from utils.animation_2d_data import AnimationData2D
    data = load_output(args.file)
    total = len(data["trans"])
    content, style, foot_contact, trans, recon = data["content"], data["style"], data["foot_contact"], data["trans"], data["recon"]
    content_meta, style_meta = data["content_meta"], data["style_meta"]
    selected = list(range(total))
    print(total)
    # for test, selected = [6, 12, 7, 11, 4]
    for i in selected:

        if style_meta[i] == 0:
            style_meta[i] = {"style": [str(i)]}
        if content_meta[i] == 0:
            content_meta[i] = {"style": [str(i)]}

        vis_dict = {}
        cur_foot_contact = foot_contact[i].transpose(1, 0)
        if style[i].shape[0] == content[i].shape[0]:  # 3d
            cur_style = AnimationData.from_network_output(to_float(style[i])).get_global_positions()
        else:  # 2d
            cur_style = AnimationData2D.from_style2d(to_float(style[i])).get_projection()
        raws = [trans[i], recon[i], content[i]]
        cur_trans, cur_recon, cur_content = [AnimationData.from_network_output(to_float(raw)).get_global_positions() for raw in raws]
        vis_dict[" ".join(("style", style_meta[i]["style"][0]))] = {"motion": cur_style, "foot_contact": None}
        vis_dict["trans"] = {"motion": cur_trans, "foot_contact": cur_foot_contact}
        vis_dict["recon"] = {"motion": cur_recon, "foot_contact": cur_foot_contact}
        vis_dict[" ".join(("content", content_meta[i]["style"][0]))] = {"motion": cur_content, "foot_contact": cur_foot_contact}

        visualize(vis_dict)
def nrot2anim(nrot):
    anim = AnimationData.from_network_output(nrot)
    bvh, names, ftime = anim.get_BVH()
    anim = AnimationData.from_rotations_and_root_positions(
        np.array(bvh.rotations), bvh.positions[:, 0, :])
    glb = anim.get_global_positions(trim=False)

    return (bvh, names, ftime), glb
def process_single_bvh(filename,
                       config,
                       norm_data_dir=None,
                       downsample=4,
                       skel=None,
                       to_batch=False,
                       panda=False):
    def to_tensor(x):
        return torch.tensor(x).float().to(config.device)

    # TODO: CHANGED trim_scale was 4 and downsample was default 4
    if panda:
        anim = AnimationData.from_BVH(filename,
                                      downsample=1,
                                      skel=skel,
                                      trim_scale=None,
                                      mode="panda")
    else:
        anim = AnimationData.from_BVH(filename,
                                      downsample=1,
                                      skel=skel,
                                      trim_scale=None)
    # foot_contact = anim.get_foot_contact(transpose=True)  # [4, T]
    content = to_tensor(anim.get_content_input())
    style3d = to_tensor(anim.get_style3d_input())

    data = {
        "meta": {
            "style": "test",
            "content": filename.split('/')[-1]
        },
        # "foot_contact": to_tensor(foot_contact),
        "contentraw": content,
        "style3draw": style3d
    }

    if norm_data_dir is None:
        norm_data_dir = config.extra_data_dir
    for key, raw in zip(["content", "style3d"], [content, style3d]):
        norm_path = os.path.join(norm_data_dir, f'train_{key}.npz')
        norm = np.load(norm_path, allow_pickle=True)
        data[key] = normalize_motion(raw,
                                     to_tensor(norm['mean']).unsqueeze(-1),
                                     to_tensor(norm['std']).unsqueeze(-1),
                                     panda=panda)
        if key == "style3d":
            data[key] = to_tensor(np.nan_to_num(data[key]))

    if to_batch:
        data = single_to_batch(data)

    return data
    def __init__(self, config, subset_name, data_path=None, extra_data_dir=None):
        super(MotionNorm, self).__init__()

        np.random.seed(2020)
        self.skel = Skel()  # TD: add config

        if data_path is None:
            data_path = config.data_path
        dataset = np.load(data_path, allow_pickle=True)[subset_name].item()
        motions, labels, metas = dataset["motion"], dataset["style"], dataset["meta"]

        self.label_i = labels
        self.len = len(self.label_i)
        self.metas = [{key: metas[key][i] for key in metas.keys()} for i in range(self.len)]
        self.motion_i, self.foot_i = [], []
        content, style3d, style2d = [], [], []

        self.labels = []
        self.data_dict = {}
        self.diff_labels_dict = {}

        for i, motion in enumerate(motions):
            label = labels[i]
            anim = AnimationData(motion, skel=self.skel)
            if label not in self.labels:
                self.labels.append(label)
                self.data_dict[label] = []
            self.data_dict[label].append(i)
            self.motion_i.append(anim)
            self.foot_i.append(anim.get_foot_contact(transpose=True))  # [4, T]
            content.append(anim.get_content_input())
            style3d.append(anim.get_style3d_input())
            view_angles, scales = [], []
            for v in range(10):
                view_angles.append(self.random_view_angle())
                scales.append(self.random_scale())
            style2d.append(anim.get_projections(view_angles, scales))

        # calc diff labels
        for x in self.labels:
            self.diff_labels_dict[x] = [y for y in self.labels if y != x]

        if extra_data_dir is None:
            extra_data_dir = config.extra_data_dir

        norm_cfg = config.dataset_norm_config
        norm_data = []
        for key, raw in zip(["content", "style3d", "style2d"], [content, style3d, style2d]):
            prefix = norm_cfg[subset_name][key]
            pre_computed = prefix is not None
            if prefix is None:
                prefix = subset_name
            norm_data.append(NormData(prefix + "_" + key, pre_computed, raw,
                                      config, extra_data_dir, keep_raw=(key != "style2d")))
        self.content, self.style3d, self.style2d = norm_data
        self.device = config.device
        self.rand = random.SystemRandom()
def bvh_to_motion_and_phase(filename, downsample, skel):
    anim = AnimationData.from_BVH(filename, downsample=downsample, skel=skel)
    full = anim.get_full()  # [T, xxx]
    phases = anim.get_phases()  # [T, 1]
    return np.concatenate((full, phases), axis=-1)
    def __init__(self,
                 config,
                 subset_name,
                 data_path=None,
                 extra_data_dir=None,
                 panda=False):
        super(MotionNorm, self).__init__()

        np.random.seed(2020)
        self.skel = Skel()  # TD: add config
        if panda:
            self.skel = PandaSkel()

        if data_path is None:
            data_path = config.data_path
        dataset = np.load(data_path, allow_pickle=True)[subset_name].item()
        '''
        motions: arrays of T x 132
        labels: array of integer values to denote the 'style' of the motion
        metas:
            - style: array of string labels e.g. 'angry', 'childlike'
            - content: array of string labels e.g. 'walk'
            - phase: array of floats
        '''
        motions, labels, metas = dataset["motion"], dataset["style"], dataset[
            "meta"]

        self.label_i = labels
        self.len = len(self.label_i)
        self.metas = [{key: metas[key][i]
                       for key in metas.keys()} for i in range(self.len)]
        self.motion_i = []
        # self.foot_i = []
        content, style3d, style2d = [], [], []

        self.labels = []
        '''data_dict contains mapping of style label to indices belonging to this label'''
        self.data_dict = {}
        self.diff_labels_dict = {}

        for i, motion in enumerate(motions):
            label = labels[i]
            anim = AnimationData(motion, skel=self.skel, panda=panda)
            if label not in self.labels:
                self.labels.append(label)
                self.data_dict[label] = []
            self.data_dict[label].append(i)
            self.motion_i.append(anim)
            # self.foot_i.append(anim.get_foot_contact(transpose=True))  # [4, T]
            content.append(anim.get_content_input())
            style3d.append(anim.get_style3d_input())
            # TODO: FIGURE OUT HOW TO CALCULATE ROOT ROTATION THEN CAN DO STYLE_2D
            # view_angles, scales = [], []
            # for v in range(10):
            #     view_angles.append(self.random_view_angle())
            #     scales.append(self.random_scale())
            # style2d.append(anim.get_projections(view_angles, scales))

        # calc diff labels
        for x in self.labels:
            self.diff_labels_dict[x] = [y for y in self.labels if y != x]

        if extra_data_dir is None:
            extra_data_dir = config.extra_data_dir

        norm_cfg = config.dataset_norm_config
        norm_data = []
        # for key, raw in zip(["content", "style3d", "style2d"], [content, style3d, style2d]):
        #     prefix = norm_cfg[subset_name][key]
        #     pre_computed = prefix is not None
        #     if prefix is None:
        #         prefix = subset_name
        #     norm_data.append(NormData(prefix + "_" + key, pre_computed, raw,
        #                               config, extra_data_dir, keep_raw=(key != "style2d")))
        # self.content, self.style3d, self.style2d = norm_data
        for key, raw in zip(["content", "style3d"], [content, style3d]):
            prefix = norm_cfg[subset_name][key]
            pre_computed = prefix is not None
            if prefix is None:
                prefix = subset_name
            norm_data.append(
                NormData(prefix + "_" + key,
                         pre_computed,
                         raw,
                         config,
                         extra_data_dir,
                         keep_raw=(key != "style2d")))
        self.content, self.style3d = norm_data

        self.device = config.device
        self.rand = random.SystemRandom()
def save_bvh_from_network_output(nrot, output_path):
    anim = AnimationData.from_network_output(nrot)
    bvh, names, ftime = anim.get_BVH()
    if not os.path.exists(os.path.dirname(output_path)):
        os.makedirs(os.path.dirname(output_path))
    BVH.save(output_path, bvh, names, ftime)