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 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