coreml-community/ControlNet-v1-1-Annotators-cpu
15
1import torch2from collections import OrderedDict3 4import torch5import torch.nn as nn6 7def make_layers(block, no_relu_layers):8 layers = []9 for layer_name, v in block.items():10 if 'pool' in layer_name:11 layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1],12 padding=v[2])13 layers.append((layer_name, layer))14 else:15 conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1],16 kernel_size=v[2], stride=v[3],17 padding=v[4])18 layers.append((layer_name, conv2d))19 if layer_name not in no_relu_layers:20 layers.append(('relu_'+layer_name, nn.ReLU(inplace=True)))21 22 return nn.Sequential(OrderedDict(layers))23 24class bodypose_model(nn.Module):25 def __init__(self):26 super(bodypose_model, self).__init__()27 28 # these layers have no relu layer29 no_relu_layers = ['conv5_5_CPM_L1', 'conv5_5_CPM_L2', 'Mconv7_stage2_L1',\30 'Mconv7_stage2_L2', 'Mconv7_stage3_L1', 'Mconv7_stage3_L2',\31 'Mconv7_stage4_L1', 'Mconv7_stage4_L2', 'Mconv7_stage5_L1',\32 'Mconv7_stage5_L2', 'Mconv7_stage6_L1', 'Mconv7_stage6_L1']33 blocks = {}34 block0 = OrderedDict([35 ('conv1_1', [3, 64, 3, 1, 1]),36 ('conv1_2', [64, 64, 3, 1, 1]),37 ('pool1_stage1', [2, 2, 0]),38 ('conv2_1', [64, 128, 3, 1, 1]),39 ('conv2_2', [128, 128, 3, 1, 1]),40 ('pool2_stage1', [2, 2, 0]),41 ('conv3_1', [128, 256, 3, 1, 1]),42 ('conv3_2', [256, 256, 3, 1, 1]),43 ('conv3_3', [256, 256, 3, 1, 1]),44 ('conv3_4', [256, 256, 3, 1, 1]),45 ('pool3_stage1', [2, 2, 0]),46 ('conv4_1', [256, 512, 3, 1, 1]),47 ('conv4_2', [512, 512, 3, 1, 1]),48 ('conv4_3_CPM', [512, 256, 3, 1, 1]),49 ('conv4_4_CPM', [256, 128, 3, 1, 1])50 ])51 52 53 # Stage 154 block1_1 = OrderedDict([55 ('conv5_1_CPM_L1', [128, 128, 3, 1, 1]),56 ('conv5_2_CPM_L1', [128, 128, 3, 1, 1]),57 ('conv5_3_CPM_L1', [128, 128, 3, 1, 1]),58 ('conv5_4_CPM_L1', [128, 512, 1, 1, 0]),59 ('conv5_5_CPM_L1', [512, 38, 1, 1, 0])60 ])61 62 block1_2 = OrderedDict([63 ('conv5_1_CPM_L2', [128, 128, 3, 1, 1]),64 ('conv5_2_CPM_L2', [128, 128, 3, 1, 1]),65 ('conv5_3_CPM_L2', [128, 128, 3, 1, 1]),66 ('conv5_4_CPM_L2', [128, 512, 1, 1, 0]),67 ('conv5_5_CPM_L2', [512, 19, 1, 1, 0])68 ])69 blocks['block1_1'] = block1_170 blocks['block1_2'] = block1_271 72 self.model0 = make_layers(block0, no_relu_layers)73 74 # Stages 2 - 675 for i in range(2, 7):76 blocks['block%d_1' % i] = OrderedDict([77 ('Mconv1_stage%d_L1' % i, [185, 128, 7, 1, 3]),78 ('Mconv2_stage%d_L1' % i, [128, 128, 7, 1, 3]),79 ('Mconv3_stage%d_L1' % i, [128, 128, 7, 1, 3]),80 ('Mconv4_stage%d_L1' % i, [128, 128, 7, 1, 3]),81 ('Mconv5_stage%d_L1' % i, [128, 128, 7, 1, 3]),82 ('Mconv6_stage%d_L1' % i, [128, 128, 1, 1, 0]),83 ('Mconv7_stage%d_L1' % i, [128, 38, 1, 1, 0])84 ])85 86 blocks['block%d_2' % i] = OrderedDict([87 ('Mconv1_stage%d_L2' % i, [185, 128, 7, 1, 3]),88 ('Mconv2_stage%d_L2' % i, [128, 128, 7, 1, 3]),89 ('Mconv3_stage%d_L2' % i, [128, 128, 7, 1, 3]),90 ('Mconv4_stage%d_L2' % i, [128, 128, 7, 1, 3]),91 ('Mconv5_stage%d_L2' % i, [128, 128, 7, 1, 3]),92 ('Mconv6_stage%d_L2' % i, [128, 128, 1, 1, 0]),93 ('Mconv7_stage%d_L2' % i, [128, 19, 1, 1, 0])94 ])95 96 for k in blocks.keys():97 blocks[k] = make_layers(blocks[k], no_relu_layers)98 99 self.model1_1 = blocks['block1_1']100 self.model2_1 = blocks['block2_1']101 self.model3_1 = blocks['block3_1']102 self.model4_1 = blocks['block4_1']103 self.model5_1 = blocks['block5_1']104 self.model6_1 = blocks['block6_1']105 106 self.model1_2 = blocks['block1_2']107 self.model2_2 = blocks['block2_2']108 self.model3_2 = blocks['block3_2']109 self.model4_2 = blocks['block4_2']110 self.model5_2 = blocks['block5_2']111 self.model6_2 = blocks['block6_2']112 113 114 def forward(self, x):115 116 out1 = self.model0(x)117 118 out1_1 = self.model1_1(out1)119 out1_2 = self.model1_2(out1)120 out2 = torch.cat([out1_1, out1_2, out1], 1)121 122 out2_1 = self.model2_1(out2)123 out2_2 = self.model2_2(out2)124 out3 = torch.cat([out2_1, out2_2, out1], 1)125 126 out3_1 = self.model3_1(out3)127 out3_2 = self.model3_2(out3)128 out4 = torch.cat([out3_1, out3_2, out1], 1)129 130 out4_1 = self.model4_1(out4)131 out4_2 = self.model4_2(out4)132 out5 = torch.cat([out4_1, out4_2, out1], 1)133 134 out5_1 = self.model5_1(out5)135 out5_2 = self.model5_2(out5)136 out6 = torch.cat([out5_1, out5_2, out1], 1)137 138 out6_1 = self.model6_1(out6)139 out6_2 = self.model6_2(out6)140 141 return out6_1, out6_2142 143class handpose_model(nn.Module):144 def __init__(self):145 super(handpose_model, self).__init__()146 147 # these layers have no relu layer148 no_relu_layers = ['conv6_2_CPM', 'Mconv7_stage2', 'Mconv7_stage3',\149 'Mconv7_stage4', 'Mconv7_stage5', 'Mconv7_stage6']150 # stage 1151 block1_0 = OrderedDict([152 ('conv1_1', [3, 64, 3, 1, 1]),153 ('conv1_2', [64, 64, 3, 1, 1]),154 ('pool1_stage1', [2, 2, 0]),155 ('conv2_1', [64, 128, 3, 1, 1]),156 ('conv2_2', [128, 128, 3, 1, 1]),157 ('pool2_stage1', [2, 2, 0]),158 ('conv3_1', [128, 256, 3, 1, 1]),159 ('conv3_2', [256, 256, 3, 1, 1]),160 ('conv3_3', [256, 256, 3, 1, 1]),161 ('conv3_4', [256, 256, 3, 1, 1]),162 ('pool3_stage1', [2, 2, 0]),163 ('conv4_1', [256, 512, 3, 1, 1]),164 ('conv4_2', [512, 512, 3, 1, 1]),165 ('conv4_3', [512, 512, 3, 1, 1]),166 ('conv4_4', [512, 512, 3, 1, 1]),167 ('conv5_1', [512, 512, 3, 1, 1]),168 ('conv5_2', [512, 512, 3, 1, 1]),169 ('conv5_3_CPM', [512, 128, 3, 1, 1])170 ])171 172 block1_1 = OrderedDict([173 ('conv6_1_CPM', [128, 512, 1, 1, 0]),174 ('conv6_2_CPM', [512, 22, 1, 1, 0])175 ])176 177 blocks = {}178 blocks['block1_0'] = block1_0179 blocks['block1_1'] = block1_1180 181 # stage 2-6182 for i in range(2, 7):183 blocks['block%d' % i] = OrderedDict([184 ('Mconv1_stage%d' % i, [150, 128, 7, 1, 3]),185 ('Mconv2_stage%d' % i, [128, 128, 7, 1, 3]),186 ('Mconv3_stage%d' % i, [128, 128, 7, 1, 3]),187 ('Mconv4_stage%d' % i, [128, 128, 7, 1, 3]),188 ('Mconv5_stage%d' % i, [128, 128, 7, 1, 3]),189 ('Mconv6_stage%d' % i, [128, 128, 1, 1, 0]),190 ('Mconv7_stage%d' % i, [128, 22, 1, 1, 0])191 ])192 193 for k in blocks.keys():194 blocks[k] = make_layers(blocks[k], no_relu_layers)195 196 self.model1_0 = blocks['block1_0']197 self.model1_1 = blocks['block1_1']198 self.model2 = blocks['block2']199 self.model3 = blocks['block3']200 self.model4 = blocks['block4']201 self.model5 = blocks['block5']202 self.model6 = blocks['block6']203 204 def forward(self, x):205 out1_0 = self.model1_0(x)206 out1_1 = self.model1_1(out1_0)207 concat_stage2 = torch.cat([out1_1, out1_0], 1)208 out_stage2 = self.model2(concat_stage2)209 concat_stage3 = torch.cat([out_stage2, out1_0], 1)210 out_stage3 = self.model3(concat_stage3)211 concat_stage4 = torch.cat([out_stage3, out1_0], 1)212 out_stage4 = self.model4(concat_stage4)213 concat_stage5 = torch.cat([out_stage4, out1_0], 1)214 out_stage5 = self.model5(concat_stage5)215 concat_stage6 = torch.cat([out_stage5, out1_0], 1)216 out_stage6 = self.model6(concat_stage6)217 return out_stage6218 219 220