Upload 8 files

AESKConv_240_100.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cd9566b24264f34d44003b3de62cdfd50aa85b7cdde2d369214599023c40f55
+size 17558653

decoders.py

@@ -0,0 +1,56 @@
+# This script is modified from https://github.com/EricGuo5513/TM2T
+# Licensed under: https://github.com/EricGuo5513/TM2T/blob/main/LICENSE
+
+import torch.nn as nn
+
+
+class VQDecoderV3(nn.Module):
+    def __init__(self, args):
+        super(VQDecoderV3, self).__init__()
+        n_up = args.vae_layer
+        channels = []
+        for i in range(n_up - 1):
+            channels.append(args.vae_length)
+        channels.append(args.vae_length)
+        channels.append(args.vae_test_dim)
+        input_size = args.vae_length
+        n_resblk = 2
+        assert len(channels) == n_up + 1
+        if input_size == channels[0]:
+            layers = []
+        else:
+            layers = [nn.Conv1d(input_size, channels[0], kernel_size=3, stride=1, padding=1)]
+
+        for i in range(n_resblk):
+            layers += [ResBlock(channels[0])]
+        # channels = channels
+        for i in range(n_up):
+            layers += [
+                nn.Upsample(scale_factor=2, mode="nearest"),
+                nn.Conv1d(channels[i], channels[i + 1], kernel_size=3, stride=1, padding=1),
+                nn.LeakyReLU(0.2, inplace=True),
+            ]
+        layers += [nn.Conv1d(channels[-1], channels[-1], kernel_size=3, stride=1, padding=1)]
+        self.main = nn.Sequential(*layers)
+        # self.main.apply(init_weight)
+
+    def forward(self, inputs):
+        inputs = inputs.permute(0, 2, 1)
+        outputs = self.main(inputs).permute(0, 2, 1)
+        return outputs
+
+
+class ResBlock(nn.Module):
+    def __init__(self, channel):
+        super(ResBlock, self).__init__()
+        self.model = nn.Sequential(
+            nn.Conv1d(channel, channel, kernel_size=3, stride=1, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv1d(channel, channel, kernel_size=3, stride=1, padding=1),
+        )
+
+    def forward(self, x):
+        residual = x
+        out = self.model(x)
+        out += residual
+        return out

mean_vel_smplxflame_30.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53b5e48f2a7bf78c41a6de6395d6bb4f29018465ca5d0ee2820a2be3eebb7137
+size 348

mertic.py

@@ -0,0 +1,357 @@
+import os
+import wget
+import math
+import numpy as np
+import librosa
+import librosa.display
+import matplotlib.pyplot as plt
+from scipy.signal import argrelextrema
+from scipy import linalg
+import torch
+
+from .motion_encoder import VAESKConv
+
+class LVDFace(object):
+    def __init__(self):
+        self.counter = 0
+        self.sum = 0
+
+    def compute(self, pred_vertices, target_vertices):
+        t, c = pred_vertices.shape
+        diff_pred = pred_vertices[1:, :] - pred_vertices[:-1, :]  
+        diff_target = target_vertices[1:, :] - target_vertices[:-1, :] 
+        loss = np.abs(diff_pred - diff_target) 
+        loss = np.sum(loss)
+        self.counter += t * c
+        self.sum += loss
+
+    def avg(self):
+        return self.sum / self.counter
+
+    def reset(self):
+        self.counter = 0
+        self.sum = 0
+
+
+class MSEFace(object):
+    def __init__(self):
+        self.counter = 0  
+        self.sum = 0      
+
+    def compute(self, pred_vertices, target_vertices):
+        t, c = pred_vertices.shape
+        loss = np.square(pred_vertices - target_vertices)  
+        self.sum += np.sum(loss) 
+        self.counter += t * c   
+
+    def avg(self):
+        if self.counter == 0:
+            return 0  
+        return self.sum / self.counter
+
+    def reset(self):
+        self.counter = 0
+        self.sum = 0
+
+
+class L1div(object):
+    def __init__(self):
+        self.counter = 0  
+        self.sum = 0     
+
+    def compute(self, results):
+        self.counter += results.shape[0] 
+        mean = np.mean(results, axis=0)  
+        sum_l1 = np.sum(np.abs(results - mean), axis=None)  
+        self.sum += sum_l1
+
+    def avg(self):
+        if self.counter == 0:
+            return 0 
+        return self.sum / self.counter
+
+    def reset(self):
+        self.counter = 0
+        self.sum = 0
+
+
+class SRGR(object):
+    def __init__(self, threshold=0.1, joints=47, joint_dim=3):
+        self.threshold = threshold
+        self.pose_dimes = joints
+        self.joint_dim = joint_dim
+        self.counter = 0
+        self.sum = 0
+
+    def run(self, results, targets, semantic=None, verbose=False):
+        if semantic is None:
+            semantic = np.ones(results.shape[0])
+            avg_weight = 1.0
+        else:
+            # srgr == 0.165 when all success, scale range to [0, 1]
+            avg_weight = 0.165
+        results = results.reshape(-1, self.pose_dimes, self.joint_dim)
+        targets = targets.reshape(-1, self.pose_dimes, self.joint_dim)
+        semantic = semantic.reshape(-1)
+        diff = np.linalg.norm(results - targets, axis=2)  # T, J
+        if verbose:
+            print(diff)
+        success = np.where(diff < self.threshold, 1.0, 0.0)
+        for i in range(success.shape[0]):
+            success[i, :] *= semantic[i] * (1 / avg_weight)
+        rate = np.sum(success) / (success.shape[0] * success.shape[1])
+        self.counter += success.shape[0]
+        self.sum += rate * success.shape[0]
+        return rate
+
+    def avg(self):
+        return self.sum / self.counter
+
+    def reset(self):
+        self.counter = 0
+        self.sum = 0
+
+
+class BC(object):
+    def __init__(self, download_path=None, sigma=0.3, order=7, upper_body=[3, 6, 9, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]):
+        self.sigma = sigma
+        self.order = order
+        self.upper_body = upper_body
+        self.pose_data = []
+        if download_path is not None:
+            os.makedirs(download_path, exist_ok=True)
+            model_file_path = os.path.join(download_path, "mean_vel_smplxflame_30.npy")
+            if not os.path.exists(model_file_path):
+                print(f"Downloading {model_file_path}")
+                wget.download(
+                    "https://huggingface.co/spaces/H-Liu1997/EMAGE/resolve/main/EMAGE/test_sequences/weights/mean_vel_smplxflame_30.npy",
+                    model_file_path,
+                )
+        self.mmae = np.load(os.path.join(download_path, "mean_vel_smplxflame_30.npy")) if download_path is not None else None
+        self.threshold = 0.10
+        self.counter = 0
+        self.sum = 0
+
+    def load_audio(self, wave, t_start=None, t_end=None, without_file=False, sr_audio=16000):
+        hop_length = 512
+        if without_file:
+            y = wave
+        else:
+            y, sr = librosa.load(wave, sr=sr_audio)
+    
+        short_y = y[t_start:t_end] if t_start is not None else y
+        short_y = short_y.astype(np.float32)
+        onset_t = librosa.onset.onset_detect(y=short_y, sr=sr_audio, hop_length=hop_length, units="time")
+        return onset_t
+
+    def load_motion(self, pose, t_start, t_end, pose_fps, without_file=False):
+        data_each_file = []
+        if without_file:
+            data_each_file = pose
+        else:
+            with open(pose, "r") as f:
+                for i, line_data in enumerate(f.readlines()):
+                    if i < 432:
+                        continue
+                    line_data_np = np.fromstring(line_data, sep=" ")
+                    if pose_fps == 15 and i % 2 == 0:
+                        continue
+                    data_each_file.append(np.concatenate([line_data_np[30:39], line_data_np[112:121]], 0))
+            data_each_file = np.array(data_each_file)  # T*165
+        # print(data_each_file.shape)
+        joints = data_each_file.transpose(1, 0)
+        dt = 1 / pose_fps
+        init_vel = (joints[:, 1:2] - joints[:, :1]) / dt
+        middle_vel = (joints[:, 2:] - joints[:, 0:-2]) / (2 * dt)
+        final_vel = (joints[:, -1:] - joints[:, -2:-1]) / dt
+        vel = np.concatenate([init_vel, middle_vel, final_vel], 1).transpose(1, 0).reshape(data_each_file.shape[0], -1, 3)
+        # print(vel.shape)
+
+        if self.mmae is not None:
+            vel = np.linalg.norm(vel, axis=2) / self.mmae
+        else:
+            print("Warning: mmae is not provided, using max value of vel as mmae")
+            self.mmae = np.linalg.norm(vel, axis=2).max()
+            vel = np.linalg.norm(vel, axis=2) / self.mmae
+        # print(vel.shape) # T*J
+
+        beat_vel_all = []
+        for i in range(vel.shape[1]):
+            vel_mask = np.where(vel[:, i] > self.threshold)
+            beat_vel = argrelextrema(vel[t_start:t_end, i], np.less, order=self.order)
+            beat_vel_list = [j for j in beat_vel[0] if j in vel_mask[0]]
+            beat_vel_all.append(np.array(beat_vel_list))
+        return beat_vel_all
+
+    def eval_random_pose(self, wave, pose, t_start, t_end, pose_fps, num_random=60):
+        onset_raw = self.load_audio(wave, t_start, t_end)
+        dur = t_end - t_start
+        for i in range(num_random):
+            beat_vel_all = self.load_motion(pose, i, i + dur, pose_fps)
+            dis_all_b2a = self.compute(onset_raw, beat_vel_all)
+            print(f"{i}s: ", dis_all_b2a)
+
+    @staticmethod
+    def plot_onsets(audio, sr, onset_times_1, onset_times_2):
+        fig, axarr = plt.subplots(2, 1, figsize=(10, 10), sharex=True)
+        librosa.display.waveshow(audio, sr=sr, alpha=0.7, ax=axarr[0])
+        librosa.display.waveshow(audio, sr=sr, alpha=0.7, ax=axarr[1])
+
+        for onset in onset_times_1:
+            axarr[0].axvline(onset, color="r", linestyle="--", alpha=0.9, label="Onset Method 1")
+        axarr[0].legend()
+        axarr[0].set(title="Onset Method 1", xlabel="", ylabel="Amplitude")
+
+        for onset in onset_times_2:
+            axarr[1].axvline(onset, color="b", linestyle="-", alpha=0.7, label="Onset Method 2")
+        axarr[1].legend()
+        axarr[1].set(title="Onset Method 2", xlabel="Time (s)", ylabel="Amplitude")
+
+        handles, labels = plt.gca().get_legend_handles_labels()
+        by_label = dict(zip(labels, handles))
+        plt.legend(by_label.values(), by_label.keys())
+        plt.title("Audio waveform with Onsets")
+        plt.savefig("./onset.png", dpi=500)
+
+    def audio_beat_vis(self, onset_raw, onset_bt, onset_bt_rms):
+        fig, ax = plt.subplots(nrows=4, sharex=True)
+        librosa.display.specshow(librosa.amplitude_to_db(self.S, ref=np.max), y_axis="log", x_axis="time", ax=ax[0])
+        ax[1].plot(self.times, self.oenv, label="Onset strength")
+        ax[1].vlines(librosa.frames_to_time(onset_raw), 0, self.oenv.max(), label="Raw onsets", color="r")
+        ax[1].legend()
+        ax[2].vlines(librosa.frames_to_time(onset_bt), 0, self.oenv.max(), label="Backtracked", color="r")
+        ax[2].legend()
+        ax[3].vlines(librosa.frames_to_time(onset_bt_rms), 0, self.oenv.max(), label="Backtracked (RMS)", color="r")
+        ax[3].legend()
+        fig.savefig("./onset.png", dpi=500)
+
+    @staticmethod
+    def motion_frames2time(vel, offset, pose_fps):
+        return vel / pose_fps + offset
+
+    @staticmethod
+    def GAHR(a, b, sigma):
+        dis_all_b2a = 0
+        for b_each in b:
+            l2_min = min(abs(a_each - b_each) for a_each in a)
+            dis_all_b2a += math.exp(-(l2_min**2) / (2 * sigma**2))
+        return dis_all_b2a / len(b)
+
+    @staticmethod
+    def fix_directed_GAHR(a, b, sigma):
+        a = BC.motion_frames2time(a, 0, 30)
+        b = BC.motion_frames2time(b, 0, 30)
+        a = [0] + a + [len(a) / 30]
+        b = [0] + b + [len(b) / 30]
+        return BC.GAHR(a, b, sigma)
+
+    def compute(self, onset_bt_rms, beat_vel, length=1, pose_fps=30):
+        avg_dis_all_b2a_list = []
+        for its, beat_vel_each in enumerate(beat_vel):
+            if its not in self.upper_body:
+                continue
+            if beat_vel_each.size == 0:
+                avg_dis_all_b2a_list.append(0)
+                continue
+            pose_bt = self.motion_frames2time(beat_vel_each, 0, pose_fps)
+            avg_dis_all_b2a_list.append(self.GAHR(pose_bt, onset_bt_rms, self.sigma))
+        self.sum += (sum(avg_dis_all_b2a_list) / len(self.upper_body)) * length
+        self.counter += length
+
+    def avg(self):
+        return self.sum / self.counter
+
+    def reset(self):
+        self.counter = 0
+        self.sum = 0
+
+
+class Arg(object):
+    def __init__(self):
+        self.vae_length = 240
+        self.vae_test_dim = 330
+        self.vae_test_len = 32
+        self.vae_layer = 4
+        self.vae_test_stride = 20
+        self.vae_grow = [1, 1, 2, 1]
+        self.variational = False
+
+
+class FGD(object):
+    def __init__(self, download_path="./emage/", device="cuda"):
+        if download_path is not None:
+            os.makedirs(download_path, exist_ok=True)
+            model_file_path = os.path.join(download_path, "AESKConv_240_100.bin")
+            smplx_model_dir = os.path.join(download_path, "smplx_models", "smplx")
+            smplx_model_file_path = os.path.join(smplx_model_dir, "SMPLX_NEUTRAL_2020.npz")
+            if not os.path.exists(model_file_path):
+                print(f"Downloading {model_file_path}")
+                wget.download(
+                    "https://huggingface.co/spaces/H-Liu1997/EMAGE/resolve/main/EMAGE/test_sequences/weights/AESKConv_240_100.bin",
+                    model_file_path,
+                )
+
+            os.makedirs(smplx_model_dir, exist_ok=True)
+            if not os.path.exists(smplx_model_file_path):
+                print(f"Downloading {smplx_model_file_path}")
+                wget.download(
+                    "https://huggingface.co/spaces/H-Liu1997/EMAGE/resolve/main/EMAGE/smplx_models/smplx/SMPLX_NEUTRAL_2020.npz",
+                    smplx_model_file_path,
+                )
+        args = Arg()
+        self.eval_model = VAESKConv(args, model_save_path=download_path)  # Assumes LocalEncoder is defined elsewhere
+        old_stat = torch.load(download_path + "AESKConv_240_100.bin")["model_state"]
+        new_stat = {}
+        for k, v in old_stat.items():
+            # If 'module.' is in the key, remove it
+            new_key = k.replace("module.", "") if "module." in k else k
+            new_stat[new_key] = v
+        self.eval_model.load_state_dict(new_stat)
+
+        self.eval_model.eval()
+        if torch.cuda.is_available():
+            self.eval_model.to(device)
+
+        self.pred_features = []
+        self.target_features = []
+        self.device = device
+
+    def reset(self):
+        self.pred_features = []
+        self.target_features = []
+
+    def get_feature(self, data):
+        assert len(data.shape) == 3
+        if data.shape[1] % 32 != 0:
+            drop_len = data.shape[1] % 32 
+            data = data[:, :-drop_len]
+            # print(data.shape)
+        with torch.no_grad():
+            if torch.cuda.is_available():
+                data = data.to(self.device)
+            feature = self.eval_model.map2latent(data).cpu().numpy()
+        # print(feature.shape)
+        return feature
+
+    def update(self, pred, target):
+        self.pred_features.append(self.get_feature(pred))
+        self.target_features.append(self.get_feature(target))
+
+    def compute(self):
+        pred_features = np.concatenate([x.reshape(-1, x.shape[-1]) for x in self.pred_features], axis=0)
+        target_features = np.concatenate([x.reshape(-1, x.shape[-1]) for x in self.target_features], axis=0)
+        # print(pred_features.shape, target_features.shape)
+        return self.frechet_distance(pred_features, target_features)
+
+    @staticmethod
+    def frechet_distance(samples_A, samples_B, eps=1e-6):
+        mu1 = np.mean(samples_A, axis=0)
+        sigma1 = np.cov(samples_A, rowvar=False)
+        mu2 = np.mean(samples_B, axis=0)
+        sigma2 = np.cov(samples_B, rowvar=False)
+        diff = mu1 - mu2
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+        if np.iscomplexobj(covmean):
+            covmean = covmean.real
+        return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * np.trace(covmean)

motion_encoder.py

@@ -0,0 +1,193 @@
+import torch.nn as nn
+import torch
+import numpy as np
+from .skeleton_DME import SkeletonConv, SkeletonPool, find_neighbor, build_edge_topology
+from .skeleton import SkeletonResidual
+from .decoders import VQDecoderV3
+
+
+class LocalEncoder(nn.Module):
+    def __init__(self, args, topology):
+        super(LocalEncoder, self).__init__()
+        args.channel_base = 6
+        args.activation = "tanh"
+        args.use_residual_blocks = True
+        args.z_dim = 1024
+        args.temporal_scale = 8
+        args.kernel_size = 4
+        args.num_layers = args.vae_layer
+        args.skeleton_dist = 2
+        args.extra_conv = 0
+        # check how to reflect in 1d
+        args.padding_mode = "constant"
+        args.skeleton_pool = "mean"
+        args.upsampling = "linear"
+
+        self.topologies = [topology]
+        self.channel_base = [args.channel_base]
+
+        self.channel_list = []
+        self.edge_num = [len(topology)]
+        self.pooling_list = []
+        self.layers = nn.ModuleList()
+        self.args = args
+        # self.convs = []
+
+        kernel_size = args.kernel_size
+        kernel_even = False if kernel_size % 2 else True
+        padding = (kernel_size - 1) // 2
+        bias = True
+        self.grow = args.vae_grow
+        for i in range(args.num_layers):
+            self.channel_base.append(self.channel_base[-1] * self.grow[i])
+
+        for i in range(args.num_layers):
+            seq = []
+            neighbour_list = find_neighbor(self.topologies[i], args.skeleton_dist)
+            in_channels = self.channel_base[i] * self.edge_num[i]
+            out_channels = self.channel_base[i + 1] * self.edge_num[i]
+            if i == 0:
+                self.channel_list.append(in_channels)
+            self.channel_list.append(out_channels)
+            last_pool = True if i == args.num_layers - 1 else False
+
+            # (T, J, D) => (T, J', D)
+            pool = SkeletonPool(
+                edges=self.topologies[i],
+                pooling_mode=args.skeleton_pool,
+                channels_per_edge=out_channels // len(neighbour_list),
+                last_pool=last_pool,
+            )
+
+            if args.use_residual_blocks:
+                # (T, J, D) => (T/2, J', 2D)
+                seq.append(
+                    SkeletonResidual(
+                        self.topologies[i],
+                        neighbour_list,
+                        joint_num=self.edge_num[i],
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        kernel_size=kernel_size,
+                        stride=2,
+                        padding=padding,
+                        padding_mode=args.padding_mode,
+                        bias=bias,
+                        extra_conv=args.extra_conv,
+                        pooling_mode=args.skeleton_pool,
+                        activation=args.activation,
+                        last_pool=last_pool,
+                    )
+                )
+            else:
+                for _ in range(args.extra_conv):
+                    # (T, J, D) => (T, J, D)
+                    seq.append(
+                        SkeletonConv(
+                            neighbour_list,
+                            in_channels=in_channels,
+                            out_channels=in_channels,
+                            joint_num=self.edge_num[i],
+                            kernel_size=kernel_size - 1 if kernel_even else kernel_size,
+                            stride=1,
+                            padding=padding,
+                            padding_mode=args.padding_mode,
+                            bias=bias,
+                        )
+                    )
+                    seq.append(nn.PReLU() if args.activation == "relu" else nn.Tanh())
+                # (T, J, D) => (T/2, J, 2D)
+                seq.append(
+                    SkeletonConv(
+                        neighbour_list,
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        joint_num=self.edge_num[i],
+                        kernel_size=kernel_size,
+                        stride=2,
+                        padding=padding,
+                        padding_mode=args.padding_mode,
+                        bias=bias,
+                        add_offset=False,
+                        in_offset_channel=3 * self.channel_base[i] // self.channel_base[0],
+                    )
+                )
+                # self.convs.append(seq[-1])
+
+                seq.append(pool)
+                seq.append(nn.PReLU() if args.activation == "relu" else nn.Tanh())
+            self.layers.append(nn.Sequential(*seq))
+
+            self.topologies.append(pool.new_edges)
+            self.pooling_list.append(pool.pooling_list)
+            self.edge_num.append(len(self.topologies[-1]))
+
+        # in_features = self.channel_base[-1] * len(self.pooling_list[-1])
+        # in_features *= int(args.temporal_scale / 2)
+        # self.reduce = nn.Linear(in_features, args.z_dim)
+        # self.mu = nn.Linear(in_features, args.z_dim)
+        # self.logvar = nn.Linear(in_features, args.z_dim)
+
+    def forward(self, input):
+        # bs, n, c = input.shape[0], input.shape[1], input.shape[2]
+        output = input.permute(0, 2, 1)  # input.reshape(bs, n, -1, 6)
+        for layer in self.layers:
+            output = layer(output)
+        # output = output.view(output.shape[0], -1)
+        output = output.permute(0, 2, 1)
+        return output
+
+
+def reparameterize(mu, logvar):
+    std = torch.exp(0.5 * logvar)
+    eps = torch.randn_like(std)
+    return mu + eps * std
+
+
+class VAEConv(nn.Module):
+    def __init__(self, args):
+        super(VAEConv, self).__init__()
+        # self.encoder = VQEncoderV3(args)
+        # self.decoder = VQDecoderV3(args)
+        self.fc_mu = nn.Linear(args.vae_length, args.vae_length)
+        self.fc_logvar = nn.Linear(args.vae_length, args.vae_length)
+        self.variational = args.variational
+
+    def forward(self, inputs):
+        pre_latent = self.encoder(inputs)
+        mu, logvar = None, None
+        if self.variational:
+            mu = self.fc_mu(pre_latent)
+            logvar = self.fc_logvar(pre_latent)
+            pre_latent = reparameterize(mu, logvar)
+        rec_pose = self.decoder(pre_latent)
+        return {
+            "poses_feat": pre_latent,
+            "rec_pose": rec_pose,
+            "pose_mu": mu,
+            "pose_logvar": logvar,
+        }
+
+    def map2latent(self, inputs):
+        pre_latent = self.encoder(inputs)
+        if self.variational:
+            mu = self.fc_mu(pre_latent)
+            logvar = self.fc_logvar(pre_latent)
+            pre_latent = reparameterize(mu, logvar)
+        return pre_latent
+
+    def decode(self, pre_latent):
+        rec_pose = self.decoder(pre_latent)
+        return rec_pose
+
+
+class VAESKConv(VAEConv):
+    def __init__(self, args, model_save_path="./emage/"):
+        # args = args()
+        super(VAESKConv, self).__init__(args)
+        smpl_fname = model_save_path + "smplx_models/smplx/SMPLX_NEUTRAL_2020.npz"
+        smpl_data = np.load(smpl_fname, encoding="latin1")
+        parents = smpl_data["kintree_table"][0].astype(np.int32)
+        edges = build_edge_topology(parents)
+        self.encoder = LocalEncoder(args, edges)
+        self.decoder = VQDecoderV3(args)

skeleton.py

@@ -0,0 +1,298 @@
+import torch
+import torch.nn as nn
+
+from .skeleton_DME import SkeletonConv, SkeletonPool, SkeletonUnpool
+
+
+def calc_node_depth(topology):
+    def dfs(node, topology):
+        if topology[node] < 0:
+            return 0
+        return 1 + dfs(topology[node], topology)
+
+    depth = []
+    for i in range(len(topology)):
+        depth.append(dfs(i, topology))
+
+    return depth
+
+
+def residual_ratio(k):
+    return 1 / (k + 1)
+
+
+class Affine(nn.Module):
+    def __init__(self, num_parameters, scale=True, bias=True, scale_init=1.0):
+        super(Affine, self).__init__()
+        if scale:
+            self.scale = nn.Parameter(torch.ones(num_parameters) * scale_init)
+        else:
+            self.register_parameter("scale", None)
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(num_parameters))
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, input):
+        output = input
+        if self.scale is not None:
+            scale = self.scale.unsqueeze(0)
+            while scale.dim() < input.dim():
+                scale = scale.unsqueeze(2)
+        output = output.mul(scale)
+
+        if self.bias is not None:
+            bias = self.bias.unsqueeze(0)
+            while bias.dim() < input.dim():
+                bias = bias.unsqueeze(2)
+        output += bias
+
+        return output
+
+
+class BatchStatistics(nn.Module):
+    def __init__(self, affine=-1):
+        super(BatchStatistics, self).__init__()
+        self.affine = nn.Sequential() if affine == -1 else Affine(affine)
+        self.loss = 0
+
+    def clear_loss(self):
+        self.loss = 0
+
+    def compute_loss(self, input):
+        input_flat = input.view(input.size(1), input.numel() // input.size(1))
+        mu = input_flat.mean(1)
+        logvar = (input_flat.pow(2).mean(1) - mu.pow(2)).sqrt().log()
+
+        self.loss = mu.pow(2).mean() + logvar.pow(2).mean()
+
+    def forward(self, input):
+        self.compute_loss(input)
+        return self.affine(input)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride, padding, residual_ratio, activation, batch_statistics=False, last_layer=False
+    ):
+        super(ResidualBlock, self).__init__()
+
+        self.residual_ratio = residual_ratio
+        self.shortcut_ratio = 1 - residual_ratio
+
+        residual = []
+        residual.append(nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding))
+        if batch_statistics:
+            residual.append(BatchStatistics(out_channels))
+        if not last_layer:
+            residual.append(nn.PReLU() if activation == "relu" else nn.Tanh())
+        self.residual = nn.Sequential(*residual)
+
+        self.shortcut = nn.Sequential(
+            nn.AvgPool1d(kernel_size=2) if stride == 2 else nn.Sequential(),
+            nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=1, padding=0),
+            BatchStatistics(out_channels) if (in_channels != out_channels and batch_statistics is True) else nn.Sequential(),
+        )
+
+    def forward(self, input):
+        return self.residual(input).mul(self.residual_ratio) + self.shortcut(input).mul(self.shortcut_ratio)
+
+
+class ResidualBlockTranspose(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, residual_ratio, activation):
+        super(ResidualBlockTranspose, self).__init__()
+
+        self.residual_ratio = residual_ratio
+        self.shortcut_ratio = 1 - residual_ratio
+
+        self.residual = nn.Sequential(
+            nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride, padding), nn.PReLU() if activation == "relu" else nn.Tanh()
+        )
+
+        self.shortcut = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode="linear", align_corners=False) if stride == 2 else nn.Sequential(),
+            nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=1, padding=0),
+        )
+
+    def forward(self, input):
+        return self.residual(input).mul(self.residual_ratio) + self.shortcut(input).mul(self.shortcut_ratio)
+
+
+class SkeletonResidual(nn.Module):
+    def __init__(
+        self,
+        topology,
+        neighbour_list,
+        joint_num,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        padding_mode,
+        bias,
+        extra_conv,
+        pooling_mode,
+        activation,
+        last_pool,
+    ):
+        super(SkeletonResidual, self).__init__()
+
+        kernel_even = False if kernel_size % 2 else True
+
+        seq = []
+        for _ in range(extra_conv):
+            # (T, J, D) => (T, J, D)
+            seq.append(
+                SkeletonConv(
+                    neighbour_list,
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    joint_num=joint_num,
+                    kernel_size=kernel_size - 1 if kernel_even else kernel_size,
+                    stride=1,
+                    padding=padding,
+                    padding_mode=padding_mode,
+                    bias=bias,
+                )
+            )
+            seq.append(nn.PReLU() if activation == "relu" else nn.Tanh())
+        # (T, J, D) => (T/2, J, 2D)
+        seq.append(
+            SkeletonConv(
+                neighbour_list,
+                in_channels=in_channels,
+                out_channels=out_channels,
+                joint_num=joint_num,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                padding_mode=padding_mode,
+                bias=bias,
+                add_offset=False,
+            )
+        )
+        seq.append(nn.GroupNorm(10, out_channels))  # FIXME: REMEMBER TO CHANGE BACK !!!
+        self.residual = nn.Sequential(*seq)
+
+        # (T, J, D) => (T/2, J, 2D)
+        self.shortcut = SkeletonConv(
+            neighbour_list,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            joint_num=joint_num,
+            kernel_size=1,
+            stride=stride,
+            padding=0,
+            bias=True,
+            add_offset=False,
+        )
+
+        seq = []
+        # (T/2, J, 2D) => (T/2, J', 2D)
+        pool = SkeletonPool(
+            edges=topology, pooling_mode=pooling_mode, channels_per_edge=out_channels // len(neighbour_list), last_pool=last_pool
+        )
+        if len(pool.pooling_list) != pool.edge_num:
+            seq.append(pool)
+        seq.append(nn.PReLU() if activation == "relu" else nn.Tanh())
+        self.common = nn.Sequential(*seq)
+
+    def forward(self, input):
+        output = self.residual(input) + self.shortcut(input)
+
+        return self.common(output)
+
+
+class SkeletonResidualTranspose(nn.Module):
+    def __init__(
+        self,
+        neighbour_list,
+        joint_num,
+        in_channels,
+        out_channels,
+        kernel_size,
+        padding,
+        padding_mode,
+        bias,
+        extra_conv,
+        pooling_list,
+        upsampling,
+        activation,
+        last_layer,
+    ):
+        super(SkeletonResidualTranspose, self).__init__()
+
+        kernel_even = False if kernel_size % 2 else True
+
+        seq = []
+        # (T, J, D) => (2T, J, D)
+        if upsampling is not None:
+            seq.append(nn.Upsample(scale_factor=2, mode=upsampling, align_corners=False))
+        # (2T, J, D) => (2T, J', D)
+        unpool = SkeletonUnpool(pooling_list, in_channels // len(neighbour_list))
+        if unpool.input_edge_num != unpool.output_edge_num:
+            seq.append(unpool)
+        self.common = nn.Sequential(*seq)
+
+        seq = []
+        for _ in range(extra_conv):
+            # (2T, J', D) => (2T, J', D)
+            seq.append(
+                SkeletonConv(
+                    neighbour_list,
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    joint_num=joint_num,
+                    kernel_size=kernel_size - 1 if kernel_even else kernel_size,
+                    stride=1,
+                    padding=padding,
+                    padding_mode=padding_mode,
+                    bias=bias,
+                )
+            )
+            seq.append(nn.PReLU() if activation == "relu" else nn.Tanh())
+        # (2T, J', D) => (2T, J', D/2)
+        seq.append(
+            SkeletonConv(
+                neighbour_list,
+                in_channels=in_channels,
+                out_channels=out_channels,
+                joint_num=joint_num,
+                kernel_size=kernel_size - 1 if kernel_even else kernel_size,
+                stride=1,
+                padding=padding,
+                padding_mode=padding_mode,
+                bias=bias,
+                add_offset=False,
+            )
+        )
+        self.residual = nn.Sequential(*seq)
+
+        # (2T, J', D) => (2T, J', D/2)
+        self.shortcut = SkeletonConv(
+            neighbour_list,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            joint_num=joint_num,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+            add_offset=False,
+        )
+
+        if activation == "relu":
+            self.activation = nn.PReLU() if not last_layer else None
+        else:
+            self.activation = nn.Tanh() if not last_layer else None
+
+    def forward(self, input):
+        output = self.common(input)
+        output = self.residual(output) + self.shortcut(output)
+
+        if self.activation is not None:
+            return self.activation(output)
+        else:
+            return output

skeleton_DME.py

@@ -0,0 +1,473 @@
+# This script is modified from https://github.com/DeepMotionEditing/deep-motion-editing
+# Licensed under:
+"""
+Copyright (c) 2020, Kfir Aberman, Peizhuo Li, Yijia Weng, Dani Lischinski, Olga Sorkine-Hornung, Daniel Cohen-Or and Baoquan Chen.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SkeletonConv(nn.Module):
+    def __init__(
+        self,
+        neighbour_list,
+        in_channels,
+        out_channels,
+        kernel_size,
+        joint_num,
+        stride=1,
+        padding=0,
+        bias=True,
+        padding_mode="zeros",
+        add_offset=False,
+        in_offset_channel=0,
+    ):
+        self.in_channels_per_joint = in_channels // joint_num
+        self.out_channels_per_joint = out_channels // joint_num
+        if in_channels % joint_num != 0 or out_channels % joint_num != 0:
+            raise Exception("BAD")
+        super(SkeletonConv, self).__init__()
+
+        if padding_mode == "zeros":
+            padding_mode = "constant"
+        if padding_mode == "reflection":
+            padding_mode = "reflect"
+
+        self.expanded_neighbour_list = []
+        self.expanded_neighbour_list_offset = []
+        self.neighbour_list = neighbour_list
+        self.add_offset = add_offset
+        self.joint_num = joint_num
+
+        self.stride = stride
+        self.dilation = 1
+        self.groups = 1
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self._padding_repeated_twice = (padding, padding)
+
+        for neighbour in neighbour_list:
+            expanded = []
+            for k in neighbour:
+                for i in range(self.in_channels_per_joint):
+                    expanded.append(k * self.in_channels_per_joint + i)
+            self.expanded_neighbour_list.append(expanded)
+
+        if self.add_offset:
+            self.offset_enc = SkeletonLinear(neighbour_list, in_offset_channel * len(neighbour_list), out_channels)
+
+            for neighbour in neighbour_list:
+                expanded = []
+                for k in neighbour:
+                    for i in range(add_offset):
+                        expanded.append(k * in_offset_channel + i)
+                self.expanded_neighbour_list_offset.append(expanded)
+
+        self.weight = torch.zeros(out_channels, in_channels, kernel_size)
+        if bias:
+            self.bias = torch.zeros(out_channels)
+        else:
+            self.register_parameter("bias", None)
+
+        self.mask = torch.zeros_like(self.weight)
+        for i, neighbour in enumerate(self.expanded_neighbour_list):
+            self.mask[self.out_channels_per_joint * i : self.out_channels_per_joint * (i + 1), neighbour, ...] = 1
+        self.mask = nn.Parameter(self.mask, requires_grad=False)
+
+        self.description = (
+            "SkeletonConv(in_channels_per_armature={}, out_channels_per_armature={}, kernel_size={}, "
+            "joint_num={}, stride={}, padding={}, bias={})".format(
+                in_channels // joint_num, out_channels // joint_num, kernel_size, joint_num, stride, padding, bias
+            )
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for i, neighbour in enumerate(self.expanded_neighbour_list):
+            """ Use temporary variable to avoid assign to copy of slice, which might lead to unexpected result """
+            tmp = torch.zeros_like(self.weight[self.out_channels_per_joint * i : self.out_channels_per_joint * (i + 1), neighbour, ...])
+            nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
+            self.weight[self.out_channels_per_joint * i : self.out_channels_per_joint * (i + 1), neighbour, ...] = tmp
+            if self.bias is not None:
+                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(
+                    self.weight[self.out_channels_per_joint * i : self.out_channels_per_joint * (i + 1), neighbour, ...]
+                )
+                bound = 1 / math.sqrt(fan_in)
+                tmp = torch.zeros_like(self.bias[self.out_channels_per_joint * i : self.out_channels_per_joint * (i + 1)])
+                nn.init.uniform_(tmp, -bound, bound)
+                self.bias[self.out_channels_per_joint * i : self.out_channels_per_joint * (i + 1)] = tmp
+
+        self.weight = nn.Parameter(self.weight)
+        if self.bias is not None:
+            self.bias = nn.Parameter(self.bias)
+
+    def set_offset(self, offset):
+        if not self.add_offset:
+            raise Exception("Wrong Combination of Parameters")
+        self.offset = offset.reshape(offset.shape[0], -1)
+
+    def forward(self, input):
+        # print('SkeletonConv')
+        weight_masked = self.weight * self.mask
+        # print(f'input: {input.size()}')
+        res = F.conv1d(
+            F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
+            weight_masked,
+            self.bias,
+            self.stride,
+            0,
+            self.dilation,
+            self.groups,
+        )
+
+        if self.add_offset:
+            offset_res = self.offset_enc(self.offset)
+            offset_res = offset_res.reshape(offset_res.shape + (1,))
+            res += offset_res / 100
+        # print(f'res: {res.size()}')
+        return res
+
+
+class SkeletonLinear(nn.Module):
+    def __init__(self, neighbour_list, in_channels, out_channels, extra_dim1=False):
+        super(SkeletonLinear, self).__init__()
+        self.neighbour_list = neighbour_list
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.in_channels_per_joint = in_channels // len(neighbour_list)
+        self.out_channels_per_joint = out_channels // len(neighbour_list)
+        self.extra_dim1 = extra_dim1
+        self.expanded_neighbour_list = []
+
+        for neighbour in neighbour_list:
+            expanded = []
+            for k in neighbour:
+                for i in range(self.in_channels_per_joint):
+                    expanded.append(k * self.in_channels_per_joint + i)
+            self.expanded_neighbour_list.append(expanded)
+
+        self.weight = torch.zeros(out_channels, in_channels)
+        self.mask = torch.zeros(out_channels, in_channels)
+        self.bias = nn.Parameter(torch.Tensor(out_channels))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for i, neighbour in enumerate(self.expanded_neighbour_list):
+            tmp = torch.zeros_like(self.weight[i * self.out_channels_per_joint : (i + 1) * self.out_channels_per_joint, neighbour])
+            self.mask[i * self.out_channels_per_joint : (i + 1) * self.out_channels_per_joint, neighbour] = 1
+            nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
+            self.weight[i * self.out_channels_per_joint : (i + 1) * self.out_channels_per_joint, neighbour] = tmp
+
+        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+        bound = 1 / math.sqrt(fan_in)
+        nn.init.uniform_(self.bias, -bound, bound)
+
+        self.weight = nn.Parameter(self.weight)
+        self.mask = nn.Parameter(self.mask, requires_grad=False)
+
+    def forward(self, input):
+        input = input.reshape(input.shape[0], -1)
+        weight_masked = self.weight * self.mask
+        res = F.linear(input, weight_masked, self.bias)
+        if self.extra_dim1:
+            res = res.reshape(res.shape + (1,))
+        return res
+
+
+class SkeletonPool(nn.Module):
+    def __init__(self, edges, pooling_mode, channels_per_edge, last_pool=False):
+        super(SkeletonPool, self).__init__()
+
+        if pooling_mode != "mean":
+            raise Exception("Unimplemented pooling mode in matrix_implementation")
+
+        self.channels_per_edge = channels_per_edge
+        self.pooling_mode = pooling_mode
+        self.edge_num = len(edges)
+        # self.edge_num = len(edges) + 1
+        self.seq_list = []
+        self.pooling_list = []
+        self.new_edges = []
+        degree = [0] * 100  # each element represents the degree of the corresponding joint
+
+        for edge in edges:
+            degree[edge[0]] += 1
+            degree[edge[1]] += 1
+
+        # seq_list contains multiple sub-lists where each sub-list is an edge chain from the joint whose degree > 2 to the end effectors or joints whose degree > 2.
+        def find_seq(j, seq):
+            nonlocal self, degree, edges
+
+            if degree[j] > 2 and j != 0:
+                self.seq_list.append(seq)
+                seq = []
+
+            if degree[j] == 1:
+                self.seq_list.append(seq)
+                return
+
+            for idx, edge in enumerate(edges):
+                if edge[0] == j:
+                    find_seq(edge[1], seq + [idx])
+
+        find_seq(0, [])
+        # print(f'self.seq_list: {self.seq_list}')
+
+        for seq in self.seq_list:
+            if last_pool:
+                self.pooling_list.append(seq)
+                continue
+            if len(seq) % 2 == 1:
+                self.pooling_list.append([seq[0]])
+                self.new_edges.append(edges[seq[0]])
+                seq = seq[1:]
+            for i in range(0, len(seq), 2):
+                self.pooling_list.append([seq[i], seq[i + 1]])
+                self.new_edges.append([edges[seq[i]][0], edges[seq[i + 1]][1]])
+        # print(f'self.pooling_list: {self.pooling_list}')
+        # print(f'self.new_egdes: {self.new_edges}')
+
+        # add global position
+        # self.pooling_list.append([self.edge_num - 1])
+
+        self.description = "SkeletonPool(in_edge_num={}, out_edge_num={})".format(len(edges), len(self.pooling_list))
+
+        self.weight = torch.zeros(len(self.pooling_list) * channels_per_edge, self.edge_num * channels_per_edge)
+
+        for i, pair in enumerate(self.pooling_list):
+            for j in pair:
+                for c in range(channels_per_edge):
+                    self.weight[i * channels_per_edge + c, j * channels_per_edge + c] = 1.0 / len(pair)
+
+        self.weight = nn.Parameter(self.weight, requires_grad=False)
+
+    def forward(self, input: torch.Tensor):
+        # print('SkeletonPool')
+        # print(f'input: {input.size()}')
+        # print(f'self.weight: {self.weight.size()}')
+        return torch.matmul(self.weight, input)
+
+
+class SkeletonUnpool(nn.Module):
+    def __init__(self, pooling_list, channels_per_edge):
+        super(SkeletonUnpool, self).__init__()
+        self.pooling_list = pooling_list
+        self.input_edge_num = len(pooling_list)
+        self.output_edge_num = 0
+        self.channels_per_edge = channels_per_edge
+        for t in self.pooling_list:
+            self.output_edge_num += len(t)
+
+        self.description = "SkeletonUnpool(in_edge_num={}, out_edge_num={})".format(
+            self.input_edge_num,
+            self.output_edge_num,
+        )
+
+        self.weight = torch.zeros(self.output_edge_num * channels_per_edge, self.input_edge_num * channels_per_edge)
+
+        for i, pair in enumerate(self.pooling_list):
+            for j in pair:
+                for c in range(channels_per_edge):
+                    self.weight[j * channels_per_edge + c, i * channels_per_edge + c] = 1
+
+        self.weight = nn.Parameter(self.weight)
+        self.weight.requires_grad_(False)
+
+    def forward(self, input: torch.Tensor):
+        # print('SkeletonUnpool')
+        # print(f'input: {input.size()}')
+        # print(f'self.weight: {self.weight.size()}')
+        return torch.matmul(self.weight, input)
+
+
+"""
+Helper functions for skeleton operation
+"""
+
+
+def dfs(x, fa, vis, dist):
+    vis[x] = 1
+    for y in range(len(fa)):
+        if (fa[y] == x or fa[x] == y) and vis[y] == 0:
+            dist[y] = dist[x] + 1
+            dfs(y, fa, vis, dist)
+
+
+"""
+def find_neighbor_joint(fa, threshold):
+    neighbor_list = [[]]
+    for x in range(1, len(fa)):
+        vis = [0 for _ in range(len(fa))]
+        dist = [0 for _ in range(len(fa))]
+        dist[0] = 10000
+        dfs(x, fa, vis, dist)
+        neighbor = []
+        for j in range(1, len(fa)):
+            if dist[j] <= threshold:
+                neighbor.append(j)
+        neighbor_list.append(neighbor)
+
+    neighbor = [0]
+    for i, x in enumerate(neighbor_list):
+        if i == 0: continue
+        if 1 in x:
+            neighbor.append(i)
+            neighbor_list[i] = [0] + neighbor_list[i]
+    neighbor_list[0] = neighbor
+    return neighbor_list
+
+
+def build_edge_topology(topology, offset):
+    # get all edges (pa, child, offset)
+    edges = []
+    joint_num = len(topology)
+    for i in range(1, joint_num):
+        edges.append((topology[i], i, offset[i]))
+    return edges
+"""
+
+
+def build_edge_topology(topology):
+    # get all edges (pa, child)
+    edges = []
+    joint_num = len(topology)
+    edges.append((0, joint_num))  # add an edge between the root joint and a virtual joint
+    for i in range(1, joint_num):
+        edges.append((topology[i], i))
+    return edges
+
+
+def build_joint_topology(edges, origin_names):
+    parent = []
+    offset = []
+    names = []
+    edge2joint = []
+    joint_from_edge = []  # -1 means virtual joint
+    joint_cnt = 0
+    out_degree = [0] * (len(edges) + 10)
+    for edge in edges:
+        out_degree[edge[0]] += 1
+
+    # add root joint
+    joint_from_edge.append(-1)
+    parent.append(0)
+    offset.append(np.array([0, 0, 0]))
+    names.append(origin_names[0])
+    joint_cnt += 1
+
+    def make_topology(edge_idx, pa):
+        nonlocal edges, parent, offset, names, edge2joint, joint_from_edge, joint_cnt
+        edge = edges[edge_idx]
+        if out_degree[edge[0]] > 1:
+            parent.append(pa)
+            offset.append(np.array([0, 0, 0]))
+            names.append(origin_names[edge[1]] + "_virtual")
+            edge2joint.append(-1)
+            pa = joint_cnt
+            joint_cnt += 1
+
+        parent.append(pa)
+        offset.append(edge[2])
+        names.append(origin_names[edge[1]])
+        edge2joint.append(edge_idx)
+        pa = joint_cnt
+        joint_cnt += 1
+
+        for idx, e in enumerate(edges):
+            if e[0] == edge[1]:
+                make_topology(idx, pa)
+
+    for idx, e in enumerate(edges):
+        if e[0] == 0:
+            make_topology(idx, 0)
+
+    return parent, offset, names, edge2joint
+
+
+def calc_edge_mat(edges):
+    edge_num = len(edges)
+    # edge_mat[i][j] = distance between edge(i) and edge(j)
+    edge_mat = [[100000] * edge_num for _ in range(edge_num)]
+    for i in range(edge_num):
+        edge_mat[i][i] = 0
+
+    # initialize edge_mat with direct neighbor
+    for i, a in enumerate(edges):
+        for j, b in enumerate(edges):
+            link = 0
+            for x in range(2):
+                for y in range(2):
+                    if a[x] == b[y]:
+                        link = 1
+            if link:
+                edge_mat[i][j] = 1
+
+    # calculate all the pairs distance
+    for k in range(edge_num):
+        for i in range(edge_num):
+            for j in range(edge_num):
+                edge_mat[i][j] = min(edge_mat[i][j], edge_mat[i][k] + edge_mat[k][j])
+    return edge_mat
+
+
+def find_neighbor(edges, d):
+    """
+    Args:
+        edges: The list contains N elements, each element represents (parent, child).
+        d: Distance between edges (the distance of the same edge is 0 and the distance of adjacent edges is 1).
+
+    Returns:
+        The list contains N elements, each element is a list of edge indices whose distance <= d.
+    """
+    edge_mat = calc_edge_mat(edges)
+    neighbor_list = []
+    edge_num = len(edge_mat)
+    for i in range(edge_num):
+        neighbor = []
+        for j in range(edge_num):
+            if edge_mat[i][j] <= d:
+                neighbor.append(j)
+        neighbor_list.append(neighbor)
+
+    # # add neighbor for global part
+    # global_part_neighbor = neighbor_list[0].copy()
+    # """
+    # Line #373 is buggy. Thanks @crissallan!!
+    # See issue #30 (https://github.com/DeepMotionEditing/deep-motion-editing/issues/30)
+    # However, fixing this bug will make it unable to load the pretrained model and
+    # affect the reproducibility of quantitative error reported in the paper.
+    # It is not a fatal bug so we didn't touch it and we are looking for possible solutions.
+    # """
+    # for i in global_part_neighbor:
+    #     neighbor_list[i].append(edge_num)
+    # neighbor_list.append(global_part_neighbor)
+
+    return neighbor_list