Spaces:
Running
Running
| from torch import nn | |
| from spiga.models.cnn.layers import Conv, Residual | |
| from spiga.models.cnn.hourglass import HourglassCore | |
| from spiga.models.cnn.coord_conv import AddCoordsTh | |
| from spiga.models.cnn.transform_e2p import E2Ptransform | |
| class MultitaskCNN(nn.Module): | |
| def __init__(self, nstack=4, num_landmarks=98, num_edges=15, pose_req=True, **kwargs): | |
| super(MultitaskCNN, self).__init__() | |
| # Parameters | |
| self.img_res = 256 # WxH input resolution | |
| self.ch_dim = 256 # Default channel dimension | |
| self.out_res = 64 # WxH output resolution | |
| self.nstack = nstack # Hourglass modules stacked | |
| self.num_landmarks = num_landmarks # Number of landmarks | |
| self.num_edges = num_edges # Number of edges subsets (eyeR, eyeL, nose, etc) | |
| self.pose_required = pose_req # Multitask flag | |
| # Image preprocessing | |
| self.pre = nn.Sequential( | |
| AddCoordsTh(x_dim=self.img_res, y_dim=self.img_res, with_r=True), | |
| Conv(6, 64, 7, 2, bn=True, relu=True), | |
| Residual(64, 128), | |
| Conv(128, 128, 2, 2, bn=True, relu=True), | |
| Residual(128, 128), | |
| Residual(128, self.ch_dim) | |
| ) | |
| # Hourglass modules | |
| self.hgs = nn.ModuleList([HourglassCore(4, self.ch_dim) for i in range(self.nstack)]) | |
| self.hgs_out = nn.ModuleList([ | |
| nn.Sequential( | |
| Residual(self.ch_dim, self.ch_dim), | |
| Conv(self.ch_dim, self.ch_dim, 1, bn=True, relu=True) | |
| ) for i in range(nstack)]) | |
| if self.pose_required: | |
| self.hgs_core = nn.ModuleList([ | |
| nn.Sequential( | |
| Residual(self.ch_dim, self.ch_dim), | |
| Conv(self.ch_dim, self.ch_dim, 2, 2, bn=True, relu=True), | |
| Residual(self.ch_dim, self.ch_dim), | |
| Conv(self.ch_dim, self.ch_dim, 2, 2, bn=True, relu=True) | |
| ) for i in range(nstack)]) | |
| # Attention module (ADnet style) | |
| self.outs_points = nn.ModuleList([nn.Sequential(Conv(self.ch_dim, self.num_landmarks, 1, relu=False, bn=False), | |
| nn.Sigmoid()) for i in range(self.nstack - 1)]) | |
| self.outs_edges = nn.ModuleList([nn.Sequential(Conv(self.ch_dim, self.num_edges, 1, relu=False, bn=False), | |
| nn.Sigmoid()) for i in range(self.nstack - 1)]) | |
| self.E2Ptransform = E2Ptransform(self.num_landmarks, self.num_edges, out_dim=self.out_res) | |
| self.outs_features = nn.ModuleList([Conv(self.ch_dim, self.num_landmarks, 1, relu=False, bn=False)for i in range(self.nstack - 1)]) | |
| # Stacked Hourglass inputs (nstack > 1) | |
| self.merge_preds = nn.ModuleList([Conv(self.num_landmarks, self.ch_dim, 1, relu=False, bn=False) for i in range(self.nstack - 1)]) | |
| self.merge_features = nn.ModuleList([Conv(self.ch_dim, self.ch_dim, 1, relu=False, bn=False) for i in range(self.nstack - 1)]) | |
| def forward(self, imgs): | |
| x = self.pre(imgs) | |
| outputs = {'VisualField': [], | |
| 'HGcore': []} | |
| core_raw = [] | |
| for i in range(self.nstack): | |
| # Hourglass | |
| hg, core_raw = self.hgs[i](x, core=core_raw) | |
| if self.pose_required: | |
| core = self.hgs_core[i](core_raw[-self.hgs[i].n]) | |
| outputs['HGcore'].append(core) | |
| hg = self.hgs_out[i](hg) | |
| # Visual features | |
| outputs['VisualField'].append(hg) | |
| # Prepare next stacked input | |
| if i < self.nstack - 1: | |
| # Attentional modules | |
| points = self.outs_points[i](hg) | |
| edges = self.outs_edges[i](hg) | |
| edges_ext = self.E2Ptransform(edges) | |
| point_edges = points * edges_ext | |
| # Landmarks | |
| maps = self.outs_features[i](hg) | |
| preds = maps * point_edges | |
| # Outputs | |
| x = x + self.merge_preds[i](preds) + self.merge_features[i](hg) | |
| return outputs | |