datasets/t2m_dataset.py

import torch
from torch.utils import data
import numpy as np
from os.path import join as pjoin
import random
import codecs as cs
from tqdm.auto import tqdm
from utils.word_vectorizer import WordVectorizer, POS_enumerator
from utils.motion_process import recover_from_ric


class Text2MotionDataset(data.Dataset):
    """
    Dataset for Text2Motion generation task.
    """

    data_root = ""
    min_motion_len = 40
    joints_num = None
    dim_pose = None
    max_motion_length = 196

    def __init__(self, opt, split, mode="train", accelerator=None):
        self.max_text_len = getattr(opt, "max_text_len", 20)
        self.unit_length = getattr(opt, "unit_length", 4)
        self.mode = mode
        motion_dir = pjoin(self.data_root, "new_joint_vecs")
        text_dir = pjoin(self.data_root, "texts")

        if mode not in ["train", "eval", "gt_eval", "xyz_gt", "hml_gt"]:
            raise ValueError(
                f"Mode '{mode}' is not supported. Please use one of: 'train', 'eval', 'gt_eval', 'xyz_gt','hml_gt'."
            )

        mean, std = None, None
        if mode == "gt_eval":
            print(pjoin(opt.eval_meta_dir, f"{opt.dataset_name}_std.npy"))
            # used by T2M models (including evaluators)
            mean = np.load(pjoin(opt.eval_meta_dir, f"{opt.dataset_name}_mean.npy"))
            std = np.load(pjoin(opt.eval_meta_dir, f"{opt.dataset_name}_std.npy"))
        elif mode in ["eval"]:
            print(pjoin(opt.meta_dir, "std.npy"))
            # used by our models during inference
            mean = np.load(pjoin(opt.meta_dir, "mean.npy"))
            std = np.load(pjoin(opt.meta_dir, "std.npy"))
        else:
            # used by our models during train
            mean = np.load(pjoin(self.data_root, "Mean.npy"))
            std = np.load(pjoin(self.data_root, "Std.npy"))

        if mode == "eval":
            # used by T2M models (including evaluators)
            # this is to translate ours norms to theirs
            self.mean_for_eval = np.load(
                pjoin(opt.eval_meta_dir, f"{opt.dataset_name}_mean.npy")
            )
            self.std_for_eval = np.load(
                pjoin(opt.eval_meta_dir, f"{opt.dataset_name}_std.npy")
            )
        if mode in ["gt_eval", "eval"]:
            self.w_vectorizer = WordVectorizer(opt.glove_dir, "our_vab")

        data_dict = {}
        id_list = []
        split_file = pjoin(self.data_root, f"{split}.txt")
        with cs.open(split_file, "r") as f:
            for line in f.readlines():
                id_list.append(line.strip())

        if opt.debug == True:
            id_list = id_list[:1000]

        new_name_list = []
        length_list = []
        for name in tqdm(
            id_list,
            disable=(
                not accelerator.is_local_main_process
                if accelerator is not None
                else False
            ),
        ):
            motion = np.load(pjoin(motion_dir, name + ".npy"))
            if (len(motion)) < self.min_motion_len or (len(motion) >= 200):
                continue
            text_data = []
            flag = False
            with cs.open(pjoin(text_dir, name + ".txt")) as f:
                for line in f.readlines():
                    text_dict = {}
                    line_split = line.strip().split("#")
                    caption = line_split[0]
                    try:
                        tokens = line_split[1].split(" ")
                        f_tag = float(line_split[2])
                        to_tag = float(line_split[3])
                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
                        to_tag = 0.0 if np.isnan(to_tag) else to_tag
                    except:
                        tokens = ["a/NUM", "a/NUM"]
                        f_tag = 0.0
                        to_tag = 8.0
                    text_dict["caption"] = caption
                    text_dict["tokens"] = tokens
                    if f_tag == 0.0 and to_tag == 0.0:
                        flag = True
                        text_data.append(text_dict)
                    else:
                        n_motion = motion[int(f_tag * 20) : int(to_tag * 20)]
                        if (len(n_motion)) < self.min_motion_len or (
                            len(n_motion) >= 200
                        ):
                            continue
                        new_name = random.choice("ABCDEFGHIJKLMNOPQRSTUVW") + "_" + name
                        while new_name in data_dict:
                            new_name = (
                                random.choice("ABCDEFGHIJKLMNOPQRSTUVW") + "_" + name
                            )
                        data_dict[new_name] = {
                            "motion": n_motion,
                            "length": len(n_motion),
                            "text": [text_dict],
                        }
                        new_name_list.append(new_name)
                        length_list.append(len(n_motion))
            if flag:
                data_dict[name] = {
                    "motion": motion,
                    "length": len(motion),
                    "text": text_data,
                }
                new_name_list.append(name)
                length_list.append(len(motion))

        name_list, length_list = zip(
            *sorted(zip(new_name_list, length_list), key=lambda x: x[1])
        )

        if mode == "train":
            if opt.dataset_name != "amass":
                joints_num = self.joints_num
                # root_rot_velocity (B, seq_len, 1)
                std[0:1] = std[0:1] / opt.feat_bias
                # root_linear_velocity (B, seq_len, 2)
                std[1:3] = std[1:3] / opt.feat_bias
                # root_y (B, seq_len, 1)
                std[3:4] = std[3:4] / opt.feat_bias
                # ric_data (B, seq_len, (joint_num - 1)*3)
                std[4 : 4 + (joints_num - 1) * 3] = (
                    std[4 : 4 + (joints_num - 1) * 3] / 1.0
                )
                # rot_data (B, seq_len, (joint_num - 1)*6)
                std[4 + (joints_num - 1) * 3 : 4 + (joints_num - 1) * 9] = (
                    std[4 + (joints_num - 1) * 3 : 4 + (joints_num - 1) * 9] / 1.0
                )
                # local_velocity (B, seq_len, joint_num*3)
                std[
                    4 + (joints_num - 1) * 9 : 4 + (joints_num - 1) * 9 + joints_num * 3
                ] = (
                    std[
                        4
                        + (joints_num - 1) * 9 : 4
                        + (joints_num - 1) * 9
                        + joints_num * 3
                    ]
                    / 1.0
                )
                # foot contact (B, seq_len, 4)
                std[4 + (joints_num - 1) * 9 + joints_num * 3 :] = (
                    std[4 + (joints_num - 1) * 9 + joints_num * 3 :] / opt.feat_bias
                )

                assert 4 + (joints_num - 1) * 9 + joints_num * 3 + 4 == mean.shape[-1]

            if accelerator is not None and accelerator.is_main_process:
                np.save(pjoin(opt.meta_dir, "mean.npy"), mean)
                np.save(pjoin(opt.meta_dir, "std.npy"), std)

        self.mean = mean
        self.std = std
        self.data_dict = data_dict
        self.name_list = name_list

    def inv_transform(self, data):
        return data * self.std + self.mean

    def __len__(self):
        return len(self.data_dict)

    def __getitem__(self, idx):
        data = self.data_dict[self.name_list[idx]]
        motion, m_length, text_list = data["motion"], data["length"], data["text"]

        # Randomly select a caption
        text_data = random.choice(text_list)
        caption = text_data["caption"]

        "Z Normalization"
        if self.mode not in ["xyz_gt", "hml_gt"]:
            motion = (motion - self.mean) / self.std

        "crop motion"
        if self.mode in ["eval", "gt_eval"]:
            # Crop the motions in to times of 4, and introduce small variations
            if self.unit_length < 10:
                coin2 = np.random.choice(["single", "single", "double"])
            else:
                coin2 = "single"
            if coin2 == "double":
                m_length = (m_length // self.unit_length - 1) * self.unit_length
            elif coin2 == "single":
                m_length = (m_length // self.unit_length) * self.unit_length
            idx = random.randint(0, len(motion) - m_length)
            motion = motion[idx : idx + m_length]
        elif m_length >= self.max_motion_length:
            idx = random.randint(0, len(motion) - self.max_motion_length)
            motion = motion[idx : idx + self.max_motion_length]
            m_length = self.max_motion_length

        "pad motion"
        if m_length < self.max_motion_length:
            motion = np.concatenate(
                [
                    motion,
                    np.zeros((self.max_motion_length - m_length, motion.shape[1])),
                ],
                axis=0,
            )
        assert len(motion) == self.max_motion_length

        if self.mode in ["gt_eval", "eval"]:
            "word embedding for text-to-motion evaluation"
            tokens = text_data["tokens"]
            if len(tokens) < self.max_text_len:
                # pad with "unk"
                tokens = ["sos/OTHER"] + tokens + ["eos/OTHER"]
                sent_len = len(tokens)
                tokens = tokens + ["unk/OTHER"] * (self.max_text_len + 2 - sent_len)
            else:
                # crop
                tokens = tokens[: self.max_text_len]
                tokens = ["sos/OTHER"] + tokens + ["eos/OTHER"]
                sent_len = len(tokens)
            pos_one_hots = []
            word_embeddings = []
            for token in tokens:
                word_emb, pos_oh = self.w_vectorizer[token]
                pos_one_hots.append(pos_oh[None, :])
                word_embeddings.append(word_emb[None, :])
            pos_one_hots = np.concatenate(pos_one_hots, axis=0)
            word_embeddings = np.concatenate(word_embeddings, axis=0)
            return (
                word_embeddings,
                pos_one_hots,
                caption,
                sent_len,
                motion,
                m_length,
                "_".join(tokens),
            )
        elif self.mode in ["xyz_gt"]:
            "Convert motion hml representation to skeleton points xyz"
            # 1. Use kn to get the keypoints position (the padding position after kn is all zero)
            motion = torch.from_numpy(motion).float()
            pred_joints = recover_from_ric(
                motion, self.joints_num
            )  # (nframe, njoints, 3)

            # 2. Put on Floor (Y axis)
            floor_height = pred_joints.min(dim=0)[0].min(dim=0)[0][1]
            pred_joints[:, :, 1] -= floor_height
            return pred_joints

        return caption, motion, m_length


class HumanML3D(Text2MotionDataset):
    def __init__(self, opt, split="train", mode="train", accelerator=None):
        self.data_root = "./data/HumanML3D"
        self.min_motion_len = 40
        self.joints_num = 22
        self.dim_pose = 263
        self.max_motion_length = 196
        if accelerator:
            accelerator.print(
                "\n Loading %s mode HumanML3D %s dataset ..." % (mode, split)
            )
        else:
            print("\n Loading %s mode HumanML3D dataset ..." % mode)
        super(HumanML3D, self).__init__(opt, split, mode, accelerator)


class KIT(Text2MotionDataset):
    def __init__(self, opt, split="train", mode="train", accelerator=None):
        self.data_root = "./data/KIT-ML"
        self.min_motion_len = 24
        self.joints_num = 21
        self.dim_pose = 251
        self.max_motion_length = 196
        if accelerator:
            accelerator.print("\n Loading %s mode KIT %s dataset ..." % (mode, split))
        else:
            print("\n Loading %s mode KIT dataset ..." % mode)
        super(KIT, self).__init__(opt, split, mode, accelerator)