From 6bfc471480bd52c3dbd70792e238520d1302b42e Mon Sep 17 00:00:00 2001 From: "Kin-Yiu, Wong" <102582011@cc.ncu.edu.tw> Date: Mon, 18 Jul 2022 13:31:25 +0800 Subject: [PATCH] main code update pose demo --- models/yolo.py | 121 ++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 120 insertions(+), 1 deletion(-) diff --git a/models/yolo.py b/models/yolo.py index 7e1b3da1..fb7f1eca 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -108,6 +108,107 @@ class IDetect(nn.Module): return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() +class IKeypoint(nn.Module): + stride = None # strides computed during build + export = False # onnx export + + def __init__(self, nc=80, anchors=(), nkpt=17, ch=(), inplace=True, dw_conv_kpt=False): # detection layer + super(IKeypoint, self).__init__() + self.nc = nc # number of classes + self.nkpt = nkpt + self.dw_conv_kpt = dw_conv_kpt + self.no_det=(nc + 5) # number of outputs per anchor for box and class + self.no_kpt = 3*self.nkpt ## number of outputs per anchor for keypoints + self.no = self.no_det+self.no_kpt + self.nl = len(anchors) # number of detection layers + self.na = len(anchors[0]) // 2 # number of anchors + self.grid = [torch.zeros(1)] * self.nl # init grid + self.flip_test = False + a = torch.tensor(anchors).float().view(self.nl, -1, 2) + self.register_buffer('anchors', a) # shape(nl,na,2) + self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) + self.m = nn.ModuleList(nn.Conv2d(x, self.no_det * self.na, 1) for x in ch) # output conv + + self.ia = nn.ModuleList(ImplicitA(x) for x in ch) + self.im = nn.ModuleList(ImplicitM(self.no_det * self.na) for _ in ch) + + if self.nkpt is not None: + if self.dw_conv_kpt: #keypoint head is slightly more complex + self.m_kpt = nn.ModuleList( + nn.Sequential(DWConv(x, x, k=3), Conv(x,x), + DWConv(x, x, k=3), Conv(x, x), + DWConv(x, x, k=3), Conv(x,x), + DWConv(x, x, k=3), Conv(x, x), + DWConv(x, x, k=3), Conv(x, x), + DWConv(x, x, k=3), nn.Conv2d(x, self.no_kpt * self.na, 1)) for x in ch) + else: #keypoint head is a single convolution + self.m_kpt = nn.ModuleList(nn.Conv2d(x, self.no_kpt * self.na, 1) for x in ch) + + self.inplace = inplace # use in-place ops (e.g. slice assignment) + + def forward(self, x): + # x = x.copy() # for profiling + z = [] # inference output + self.training |= self.export + for i in range(self.nl): + if self.nkpt is None or self.nkpt==0: + x[i] = self.im[i](self.m[i](self.ia[i](x[i]))) # conv + else : + x[i] = torch.cat((self.im[i](self.m[i](self.ia[i](x[i]))), self.m_kpt[i](x[i])), axis=1) + + bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) + x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() + x_det = x[i][..., :6] + x_kpt = x[i][..., 6:] + + if not self.training: # inference + if self.grid[i].shape[2:4] != x[i].shape[2:4]: + self.grid[i] = self._make_grid(nx, ny).to(x[i].device) + kpt_grid_x = self.grid[i][..., 0:1] + kpt_grid_y = self.grid[i][..., 1:2] + + if self.nkpt == 0: + y = x[i].sigmoid() + else: + y = x_det.sigmoid() + + if self.inplace: + xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy + wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh + if self.nkpt != 0: + x_kpt[..., 0::3] = (x_kpt[..., ::3] * 2. - 0.5 + kpt_grid_x.repeat(1,1,1,1,17)) * self.stride[i] # xy + x_kpt[..., 1::3] = (x_kpt[..., 1::3] * 2. - 0.5 + kpt_grid_y.repeat(1,1,1,1,17)) * self.stride[i] # xy + #x_kpt[..., 0::3] = (x_kpt[..., ::3] + kpt_grid_x.repeat(1,1,1,1,17)) * self.stride[i] # xy + #x_kpt[..., 1::3] = (x_kpt[..., 1::3] + kpt_grid_y.repeat(1,1,1,1,17)) * self.stride[i] # xy + #print('=============') + #print(self.anchor_grid[i].shape) + #print(self.anchor_grid[i][...,0].unsqueeze(4).shape) + #print(x_kpt[..., 0::3].shape) + #x_kpt[..., 0::3] = ((x_kpt[..., 0::3].tanh() * 2.) ** 3 * self.anchor_grid[i][...,0].unsqueeze(4).repeat(1,1,1,1,self.nkpt)) + kpt_grid_x.repeat(1,1,1,1,17) * self.stride[i] # xy + #x_kpt[..., 1::3] = ((x_kpt[..., 1::3].tanh() * 2.) ** 3 * self.anchor_grid[i][...,1].unsqueeze(4).repeat(1,1,1,1,self.nkpt)) + kpt_grid_y.repeat(1,1,1,1,17) * self.stride[i] # xy + #x_kpt[..., 0::3] = (((x_kpt[..., 0::3].sigmoid() * 4.) ** 2 - 8.) * self.anchor_grid[i][...,0].unsqueeze(4).repeat(1,1,1,1,self.nkpt)) + kpt_grid_x.repeat(1,1,1,1,17) * self.stride[i] # xy + #x_kpt[..., 1::3] = (((x_kpt[..., 1::3].sigmoid() * 4.) ** 2 - 8.) * self.anchor_grid[i][...,1].unsqueeze(4).repeat(1,1,1,1,self.nkpt)) + kpt_grid_y.repeat(1,1,1,1,17) * self.stride[i] # xy + x_kpt[..., 2::3] = x_kpt[..., 2::3].sigmoid() + + y = torch.cat((xy, wh, y[..., 4:], x_kpt), dim = -1) + + else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953 + xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy + wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh + if self.nkpt != 0: + y[..., 6:] = (y[..., 6:] * 2. - 0.5 + self.grid[i].repeat((1,1,1,1,self.nkpt))) * self.stride[i] # xy + y = torch.cat((xy, wh, y[..., 4:]), -1) + + z.append(y.view(bs, -1, self.no)) + + return x if self.training else (torch.cat(z, 1), x) + + @staticmethod + def _make_grid(nx=20, ny=20): + yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) + return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() + + class IAuxDetect(nn.Module): stride = None # strides computed during build export = False # onnx export @@ -292,6 +393,14 @@ class Model(nn.Module): self.stride = m.stride self._initialize_biases_bin() # only run once # print('Strides: %s' % m.stride.tolist()) + if isinstance(m, IKeypoint): + s = 256 # 2x min stride + m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward + m.anchors /= m.stride.view(-1, 1, 1) + check_anchor_order(m) + self.stride = m.stride + self._initialize_biases_kpt() # only run once + # print('Strides: %s' % m.stride.tolist()) # Init weights, biases initialize_weights(self) @@ -389,6 +498,16 @@ class Model(nn.Module): b[:, (0,1,2,bc+3)].data = old mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) + def _initialize_biases_kpt(self, cf=None): # initialize biases into Detect(), cf is class frequency + # https://arxiv.org/abs/1708.02002 section 3.3 + # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. + m = self.model[-1] # Detect() module + for mi, s in zip(m.m, m.stride): # from + b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) + b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) + b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls + mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) + def _print_biases(self): m = self.model[-1] # Detect() module for mi in m.m: # from @@ -494,7 +613,7 @@ def parse_model(d, ch): # model_dict, input_channels(3) c2 = ch[f[0]] elif m is Foldcut: c2 = ch[f] // 2 - elif m in [Detect, IDetect, IAuxDetect, IBin]: + elif m in [Detect, IDetect, IAuxDetect, IBin, IKeypoint]: args.append([ch[x] for x in f]) if isinstance(args[1], int): # number of anchors args[1] = [list(range(args[1] * 2))] * len(f)