import torch import numpy as np import torch.nn as nn import torch.nn.functional as F from utils import clones, is_list_or_tuple from torchvision.ops import RoIAlign class HorizontalPoolingPyramid(): """ Horizontal Pyramid Matching for Person Re-identification Arxiv: https://arxiv.org/abs/1804.05275 Github: https://github.com/SHI-Labs/Horizontal-Pyramid-Matching """ def __init__(self, bin_num=None): if bin_num is None: bin_num = [16, 8, 4, 2, 1] self.bin_num = bin_num def __call__(self, x): """ x : [n, c, h, w] ret: [n, c, p] """ n, c = x.size()[:2] features = [] for b in self.bin_num: z = x.view(n, c, b, -1) z = z.mean(-1) + z.max(-1)[0] features.append(z) return torch.cat(features, -1) class SetBlockWrapper(nn.Module): def __init__(self, forward_block): super(SetBlockWrapper, self).__init__() self.forward_block = forward_block def forward(self, x, *args, **kwargs): """ In x: [n, c_in, s, h_in, w_in] Out x: [n, c_out, s, h_out, w_out] """ n, c, s, h, w = x.size() x = self.forward_block(x.transpose( 1, 2).reshape(-1, c, h, w), *args, **kwargs) output_size = x.size() return x.reshape(n, s, *output_size[1:]).transpose(1, 2).contiguous() class PackSequenceWrapper(nn.Module): def __init__(self, pooling_func): super(PackSequenceWrapper, self).__init__() self.pooling_func = pooling_func def forward(self, seqs, seqL, dim=2, options={}): """ In seqs: [n, c, s, ...] Out rets: [n, ...] """ if seqL is None: return self.pooling_func(seqs, **options) seqL = seqL[0].data.cpu().numpy().tolist() start = [0] + np.cumsum(seqL).tolist()[:-1] rets = [] for curr_start, curr_seqL in zip(start, seqL): narrowed_seq = seqs.narrow(dim, curr_start, curr_seqL) rets.append(self.pooling_func(narrowed_seq, **options)) if len(rets) > 0 and is_list_or_tuple(rets[0]): return [torch.cat([ret[j] for ret in rets]) for j in range(len(rets[0]))] return torch.cat(rets) class BasicConv2d(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride, padding, **kwargs): super(BasicConv2d, self).__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=False, **kwargs) def forward(self, x): x = self.conv(x) return x class SeparateFCs(nn.Module): def __init__(self, parts_num, in_channels, out_channels, norm=False): super(SeparateFCs, self).__init__() self.p = parts_num self.fc_bin = nn.Parameter( nn.init.xavier_uniform_( torch.zeros(parts_num, in_channels, out_channels))) self.norm = norm def forward(self, x): """ x: [n, c_in, p] out: [n, c_out, p] """ x = x.permute(2, 0, 1).contiguous() if self.norm: out = x.matmul(F.normalize(self.fc_bin, dim=1)) else: out = x.matmul(self.fc_bin) return out.permute(1, 2, 0).contiguous() class SeparateBNNecks(nn.Module): """ Bag of Tricks and a Strong Baseline for Deep Person Re-Identification CVPR Workshop: https://openaccess.thecvf.com/content_CVPRW_2019/papers/TRMTMCT/Luo_Bag_of_Tricks_and_a_Strong_Baseline_for_Deep_Person_CVPRW_2019_paper.pdf Github: https://github.com/michuanhaohao/reid-strong-baseline """ def __init__(self, parts_num, in_channels, class_num, norm=True, parallel_BN1d=True): super(SeparateBNNecks, self).__init__() self.p = parts_num self.class_num = class_num self.norm = norm self.fc_bin = nn.Parameter( nn.init.xavier_uniform_( torch.zeros(parts_num, in_channels, class_num))) if parallel_BN1d: self.bn1d = nn.BatchNorm1d(in_channels * parts_num) else: self.bn1d = clones(nn.BatchNorm1d(in_channels), parts_num) self.parallel_BN1d = parallel_BN1d def forward(self, x): """ x: [n, c, p] """ if self.parallel_BN1d: n, c, p = x.size() x = x.view(n, -1) # [n, c*p] x = self.bn1d(x) x = x.view(n, c, p) else: x = torch.cat([bn(_x) for _x, bn in zip( x.split(1, 2), self.bn1d)], 2) # [p, n, c] feature = x.permute(2, 0, 1).contiguous() if self.norm: feature = F.normalize(feature, dim=-1) # [p, n, c] logits = feature.matmul(F.normalize( self.fc_bin, dim=1)) # [p, n, c] else: logits = feature.matmul(self.fc_bin) return feature.permute(1, 2, 0).contiguous(), logits.permute(1, 2, 0).contiguous() class FocalConv2d(nn.Module): """ GaitPart: Temporal Part-based Model for Gait Recognition CVPR2020: https://openaccess.thecvf.com/content_CVPR_2020/papers/Fan_GaitPart_Temporal_Part-Based_Model_for_Gait_Recognition_CVPR_2020_paper.pdf Github: https://github.com/ChaoFan96/GaitPart """ def __init__(self, in_channels, out_channels, kernel_size, halving, **kwargs): super(FocalConv2d, self).__init__() self.halving = halving self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False, **kwargs) def forward(self, x): if self.halving == 0: z = self.conv(x) else: h = x.size(2) split_size = int(h // 2**self.halving) z = x.split(split_size, 2) z = torch.cat([self.conv(_) for _ in z], 2) return z class BasicConv3d(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False, **kwargs): super(BasicConv3d, self).__init__() self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias, **kwargs) def forward(self, ipts): ''' ipts: [n, c, s, h, w] outs: [n, c, s, h, w] ''' outs = self.conv3d(ipts) return outs class GaitAlign(nn.Module): """ GaitEdge: Beyond Plain End-to-end Gait Recognition for Better Practicality ECCV2022: https://arxiv.org/pdf/2203.03972v2.pdf Github: https://github.com/ShiqiYu/OpenGait/tree/master/configs/gaitedge """ def __init__(self, H=64, W=44, eps=1, **kwargs): super(GaitAlign, self).__init__() self.H, self.W, self.eps = H, W, eps self.Pad = nn.ZeroPad2d((int(self.W / 2), int(self.W / 2), 0, 0)) self.RoiPool = RoIAlign((self.H, self.W), 1, sampling_ratio=-1) def forward(self, feature_map, binary_mask, w_h_ratio): """ In sils: [n, c, h, w] w_h_ratio: [n, 1] Out aligned_sils: [n, c, H, W] """ n, c, h, w = feature_map.size() # w_h_ratio = w_h_ratio.repeat(1, 1) # [n, 1] w_h_ratio = w_h_ratio.view(-1, 1) # [n, 1] h_sum = binary_mask.sum(-1) # [n, c, h] _ = (h_sum >= self.eps).float().cumsum(axis=-1) # [n, c, h] h_top = (_ == 0).float().sum(-1) # [n, c] h_bot = (_ != torch.max(_, dim=-1, keepdim=True) [0]).float().sum(-1) + 1. # [n, c] w_sum = binary_mask.sum(-2) # [n, c, w] w_cumsum = w_sum.cumsum(axis=-1) # [n, c, w] w_h_sum = w_sum.sum(-1).unsqueeze(-1) # [n, c, 1] w_center = (w_cumsum < w_h_sum / 2.).float().sum(-1) # [n, c] p1 = self.W - self.H * w_h_ratio p1 = p1 / 2. p1 = torch.clamp(p1, min=0) # [n, c] t_w = w_h_ratio * self.H / w p2 = p1 / t_w # [n, c] height = h_bot - h_top # [n, c] width = height * w / h # [n, c] width_p = int(self.W / 2) feature_map = self.Pad(feature_map) w_center = w_center + width_p # [n, c] w_left = w_center - width / 2 - p2 # [n, c] w_right = w_center + width / 2 + p2 # [n, c] w_left = torch.clamp(w_left, min=0., max=w+2*width_p) w_right = torch.clamp(w_right, min=0., max=w+2*width_p) boxes = torch.cat([w_left, h_top, w_right, h_bot], dim=-1) # index of bbox in batch box_index = torch.arange(n, device=feature_map.device) rois = torch.cat([box_index.view(-1, 1), boxes], -1) crops = self.RoiPool(feature_map, rois) # [n, c, H, W] return crops def RmBN2dAffine(model): for m in model.modules(): if isinstance(m, nn.BatchNorm2d): m.weight.requires_grad = False m.bias.requires_grad = False ''' Modifed from https://github.com/BNU-IVC/FastPoseGait/blob/main/fastposegait/modeling/components/units ''' class Graph(): """ # Thanks to YAN Sijie for the released code on Github (https://github.com/yysijie/st-gcn) """ def __init__(self, joint_format='coco', max_hop=2, dilation=1): self.joint_format = joint_format self.max_hop = max_hop self.dilation = dilation # get edges self.num_node, self.edge, self.connect_joint, self.parts = self._get_edge() # get adjacency matrix self.A = self._get_adjacency() def __str__(self): return self.A def _get_edge(self): if self.joint_format == 'coco': # keypoints = { # 0: "nose", # 1: "left_eye", # 2: "right_eye", # 3: "left_ear", # 4: "right_ear", # 5: "left_shoulder", # 6: "right_shoulder", # 7: "left_elbow", # 8: "right_elbow", # 9: "left_wrist", # 10: "right_wrist", # 11: "left_hip", # 12: "right_hip", # 13: "left_knee", # 14: "right_knee", # 15: "left_ankle", # 16: "right_ankle" # } num_node = 17 self_link = [(i, i) for i in range(num_node)] neighbor_link = [(0, 1), (0, 2), (1, 3), (2, 4), (3, 5), (4, 6), (5, 6), (5, 7), (7, 9), (6, 8), (8, 10), (5, 11), (6, 12), (11, 12), (11, 13), (13, 15), (12, 14), (14, 16)] self.edge = self_link + neighbor_link self.center = 0 self.flip_idx = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] connect_joint = np.array([5,0,0,1,2,0,0,5,6,7,8,5,6,11,12,13,14]) parts = [ np.array([5, 7, 9]), # left_arm np.array([6, 8, 10]), # right_arm np.array([11, 13, 15]), # left_leg np.array([12, 14, 16]), # right_leg np.array([0, 1, 2, 3, 4]), # head ] elif self.joint_format == 'coco-no-head': num_node = 12 self_link = [(i, i) for i in range(num_node)] neighbor_link = [(0, 1), (0, 2), (2, 4), (1, 3), (3, 5), (0, 6), (1, 7), (6, 7), (6, 8), (8, 10), (7, 9), (9, 11)] self.edge = self_link + neighbor_link self.center = 0 connect_joint = np.array([3,1,0,2,4,0,6,8,10,7,9,11]) parts =[ np.array([0, 2, 4]), # left_arm np.array([1, 3, 5]), # right_arm np.array([6, 8, 10]), # left_leg np.array([7, 9, 11]) # right_leg ] elif self.joint_format =='alphapose' or self.joint_format =='openpose': num_node = 18 self_link = [(i, i) for i in range(num_node)] neighbor_link = [(0, 1), (0, 14), (0, 15), (14, 16), (15, 17), (1, 2), (2, 3), (3, 4), (1, 5), (5, 6), (6, 7), (1, 8), (8, 9), (9, 10), (1, 11), (11, 12), (12, 13)] self.edge = self_link + neighbor_link self.center = 1 self.flip_idx = [0, 1, 5, 6, 7, 2, 3, 4, 11, 12, 13, 8, 9, 10, 15, 14, 17, 16] connect_joint = np.array([1,1,1,2,3,1,5,6,2,8,9,5,11,12,0,0,14,15]) parts = [ np.array([5, 6, 7]), # left_arm np.array([2, 3, 4]), # right_arm np.array([11, 12, 13]), # left_leg np.array([8, 9, 10]), # right_leg np.array([0, 1, 14, 15, 16, 17]), # head ] else: num_node, neighbor_link, connect_joint, parts = 0, [], [], [] logging.info('') logging.error('Error: Do NOT exist this dataset: {}!'.format(self.dataset)) raise ValueError() self_link = [(i, i) for i in range(num_node)] edge = self_link + neighbor_link return num_node, edge, connect_joint, parts def _get_hop_distance(self): A = np.zeros((self.num_node, self.num_node)) for i, j in self.edge: A[j, i] = 1 A[i, j] = 1 hop_dis = np.zeros((self.num_node, self.num_node)) + np.inf transfer_mat = [np.linalg.matrix_power(A, d) for d in range(self.max_hop + 1)] arrive_mat = (np.stack(transfer_mat) > 0) for d in range(self.max_hop, -1, -1): hop_dis[arrive_mat[d]] = d return hop_dis def _get_adjacency(self): hop_dis = self._get_hop_distance() valid_hop = range(0, self.max_hop + 1, self.dilation) adjacency = np.zeros((self.num_node, self.num_node)) for hop in valid_hop: adjacency[hop_dis == hop] = 1 normalize_adjacency = self._normalize_digraph(adjacency) A = np.zeros((len(valid_hop), self.num_node, self.num_node)) for i, hop in enumerate(valid_hop): A[i][hop_dis == hop] = normalize_adjacency[hop_dis == hop] return A def _normalize_digraph(self, A): Dl = np.sum(A, 0) num_node = A.shape[0] Dn = np.zeros((num_node, num_node)) for i in range(num_node): if Dl[i] > 0: Dn[i, i] = Dl[i]**(-1) AD = np.dot(A, Dn) return AD class TemporalBasicBlock(nn.Module): """ TemporalConv_Res_Block Arxiv: https://arxiv.org/abs/2010.09978 Github: https://github.com/Thomas-yx/ResGCNv1 """ def __init__(self, channels, temporal_window_size, stride=1, residual=False,reduction=0,get_res=False,tcn_stride=False): super(TemporalBasicBlock, self).__init__() padding = ((temporal_window_size - 1) // 2, 0) if not residual: self.residual = lambda x: 0 elif stride == 1: self.residual = lambda x: x else: self.residual = nn.Sequential( nn.Conv2d(channels, channels, 1, (stride,1)), nn.BatchNorm2d(channels), ) self.conv = nn.Conv2d(channels, channels, (temporal_window_size,1), (stride,1), padding) self.bn = nn.BatchNorm2d(channels) self.relu = nn.ReLU(inplace=True) def forward(self, x, res_module): res_block = self.residual(x) x = self.conv(x) x = self.bn(x) x = self.relu(x + res_block + res_module) return x class TemporalBottleneckBlock(nn.Module): """ TemporalConv_Res_Bottleneck Arxiv: https://arxiv.org/abs/2010.09978 Github: https://github.com/Thomas-yx/ResGCNv1 """ def __init__(self, channels, temporal_window_size, stride=1, residual=False, reduction=4,get_res=False, tcn_stride=False): super(TemporalBottleneckBlock, self).__init__() tcn_stride =False padding = ((temporal_window_size - 1) // 2, 0) inter_channels = channels // reduction if get_res: if tcn_stride: stride =2 self.residual = nn.Sequential( nn.Conv2d(channels, channels, 1, (2,1)), nn.BatchNorm2d(channels), ) tcn_stride= True else: if not residual: self.residual = lambda x: 0 elif stride == 1: self.residual = lambda x: x else: self.residual = nn.Sequential( nn.Conv2d(channels, channels, 1, (2,1)), nn.BatchNorm2d(channels), ) tcn_stride= True self.conv_down = nn.Conv2d(channels, inter_channels, 1) self.bn_down = nn.BatchNorm2d(inter_channels) if tcn_stride: stride=2 self.conv = nn.Conv2d(inter_channels, inter_channels, (temporal_window_size,1), (stride,1), padding) self.bn = nn.BatchNorm2d(inter_channels) self.conv_up = nn.Conv2d(inter_channels, channels, 1) self.bn_up = nn.BatchNorm2d(channels) self.relu = nn.ReLU(inplace=True) def forward(self, x, res_module): res_block = self.residual(x) x = self.conv_down(x) x = self.bn_down(x) x = self.relu(x) x = self.conv(x) x = self.bn(x) x = self.relu(x) x = self.conv_up(x) x = self.bn_up(x) x = self.relu(x + res_block + res_module) return x class SpatialGraphConv(nn.Module): """ SpatialGraphConv_Basic_Block Arxiv: https://arxiv.org/abs/1801.07455 Github: https://github.com/yysijie/st-gcn """ def __init__(self, in_channels, out_channels, max_graph_distance): super(SpatialGraphConv, self).__init__() # spatial class number (distance = 0 for class 0, distance = 1 for class 1, ...) self.s_kernel_size = max_graph_distance + 1 # weights of different spatial classes self.gcn = nn.Conv2d(in_channels, out_channels*self.s_kernel_size, 1) def forward(self, x, A): # numbers in same class have same weight x = self.gcn(x) # divide nodes into different classes n, kc, t, v = x.size() x = x.view(n, self.s_kernel_size, kc//self.s_kernel_size, t, v).contiguous() # spatial graph convolution x = torch.einsum('nkctv,kvw->nctw', (x, A[:self.s_kernel_size])).contiguous() return x class SpatialBasicBlock(nn.Module): """ SpatialGraphConv_Res_Block Arxiv: https://arxiv.org/abs/2010.09978 Github: https://github.com/Thomas-yx/ResGCNv1 """ def __init__(self, in_channels, out_channels, max_graph_distance, residual=False,reduction=0): super(SpatialBasicBlock, self).__init__() if not residual: self.residual = lambda x: 0 elif in_channels == out_channels: self.residual = lambda x: x else: self.residual = nn.Sequential( nn.Conv2d(in_channels, out_channels, 1), nn.BatchNorm2d(out_channels), ) self.conv = SpatialGraphConv(in_channels, out_channels, max_graph_distance) self.bn = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) def forward(self, x, A): res_block = self.residual(x) x = self.conv(x, A) x = self.bn(x) x = self.relu(x + res_block) return x class SpatialBottleneckBlock(nn.Module): """ SpatialGraphConv_Res_Bottleneck Arxiv: https://arxiv.org/abs/2010.09978 Github: https://github.com/Thomas-yx/ResGCNv1 """ def __init__(self, in_channels, out_channels, max_graph_distance, residual=False, reduction=4): super(SpatialBottleneckBlock, self).__init__() inter_channels = out_channels // reduction if not residual: self.residual = lambda x: 0 elif in_channels == out_channels: self.residual = lambda x: x else: self.residual = nn.Sequential( nn.Conv2d(in_channels, out_channels, 1), nn.BatchNorm2d(out_channels), ) self.conv_down = nn.Conv2d(in_channels, inter_channels, 1) self.bn_down = nn.BatchNorm2d(inter_channels) self.conv = SpatialGraphConv(inter_channels, inter_channels, max_graph_distance) self.bn = nn.BatchNorm2d(inter_channels) self.conv_up = nn.Conv2d(inter_channels, out_channels, 1) self.bn_up = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) def forward(self, x, A): res_block = self.residual(x) x = self.conv_down(x) x = self.bn_down(x) x = self.relu(x) x = self.conv(x, A) x = self.bn(x) x = self.relu(x) x = self.conv_up(x) x = self.bn_up(x) x = self.relu(x + res_block) return x class SpatialAttention(nn.Module): """ This class implements Spatial Transformer. Function adapted from: https://github.com/leaderj1001/Attention-Augmented-Conv2d """ def __init__(self, in_channels, out_channel, A, num_point, dk_factor=0.25, kernel_size=1, Nh=8, num=4, stride=1): super(SpatialAttention, self).__init__() self.in_channels = in_channels self.kernel_size = kernel_size self.dk = int(dk_factor * out_channel) self.dv = int(out_channel) self.num = num self.Nh = Nh self.num_point=num_point self.A = A[0] + A[1] + A[2] self.stride = stride self.padding = (self.kernel_size - 1) // 2 assert self.Nh != 0, "integer division or modulo by zero, Nh >= 1" assert self.dk % self.Nh == 0, "dk should be divided by Nh. (example: out_channels: 20, dk: 40, Nh: 4)" assert self.dv % self.Nh == 0, "dv should be divided by Nh. (example: out_channels: 20, dv: 4, Nh: 4)" assert stride in [1, 2], str(stride) + " Up to 2 strides are allowed." self.qkv_conv = nn.Conv2d(self.in_channels, 2 * self.dk + self.dv, kernel_size=self.kernel_size, stride=stride, padding=self.padding) self.attn_out = nn.Conv2d(self.dv, self.dv, kernel_size=1, stride=1) def forward(self, x): # Input x # (batch_size, channels, 1, joints) B, _, T, V = x.size() # flat_q, flat_k, flat_v # (batch_size, Nh, dvh or dkh, joints) # dvh = dv / Nh, dkh = dk / Nh # q, k, v obtained by doing 2D convolution on the input (q=XWq, k=XWk, v=XWv) flat_q, flat_k, flat_v, q, k, v = self.compute_flat_qkv(x, self.dk, self.dv, self.Nh) # Calculate the scores, obtained by doing q*k # (batch_size, Nh, joints, dkh)*(batch_size, Nh, dkh, joints) = (batch_size, Nh, joints,joints) # The multiplication can also be divided (multi_matmul) in case of space problems logits = torch.matmul(flat_q.transpose(2, 3), flat_k) weights = F.softmax(logits, dim=-1) # attn_out # (batch, Nh, joints, dvh) # weights*V # (batch, Nh, joints, joints)*(batch, Nh, joints, dvh)=(batch, Nh, joints, dvh) attn_out = torch.matmul(weights, flat_v.transpose(2, 3)) attn_out = torch.reshape(attn_out, (B, self.Nh, T, V, self.dv // self.Nh)) attn_out = attn_out.permute(0, 1, 4, 2, 3) # combine_heads_2d, combine heads only after having calculated each Z separately # (batch, Nh*dv, 1, joints) attn_out = self.combine_heads_2d(attn_out) # Multiply for W0 (batch, out_channels, 1, joints) with out_channels=dv attn_out = self.attn_out(attn_out) return attn_out def compute_flat_qkv(self, x, dk, dv, Nh): qkv = self.qkv_conv(x) # T=1 in this case, because we are considering each frame separately N, _, T, V = qkv.size() q, k, v = torch.split(qkv, [dk, dk, dv], dim=1) q = self.split_heads_2d(q, Nh) k = self.split_heads_2d(k, Nh) v = self.split_heads_2d(v, Nh) dkh = dk // Nh q = q*(dkh ** -0.5) flat_q = torch.reshape(q, (N, Nh, dkh, T * V)) flat_k = torch.reshape(k, (N, Nh, dkh, T * V)) flat_v = torch.reshape(v, (N, Nh, dv // self.Nh, T * V)) return flat_q, flat_k, flat_v, q, k, v def split_heads_2d(self, x, Nh): B, channels, T, V = x.size() ret_shape = (B, Nh, channels // Nh, T, V) split = torch.reshape(x, ret_shape) return split def combine_heads_2d(self, x): batch, Nh, dv, T, V = x.size() ret_shape = (batch, Nh * dv, T, V) return torch.reshape(x, ret_shape)