coreml-community/ControlNet-v1-1-Annotators-cpu

1import logging2import numpy as np3from torchvision.transforms import ToTensor, ToPILImage4import torch5import torch.nn.functional as F6import cv27 8from . import util9from torch.nn import Conv2d, Module, ReLU, MaxPool2d, init10 11 12class FaceNet(Module):13    """Model the cascading heatmaps. """14    def __init__(self):15        super(FaceNet, self).__init__()16        # cnn to make feature map17        self.relu = ReLU()18        self.max_pooling_2d = MaxPool2d(kernel_size=2, stride=2)19        self.conv1_1 = Conv2d(in_channels=3, out_channels=64,20                              kernel_size=3, stride=1, padding=1)21        self.conv1_2 = Conv2d(22            in_channels=64, out_channels=64, kernel_size=3, stride=1,23            padding=1)24        self.conv2_1 = Conv2d(25            in_channels=64, out_channels=128, kernel_size=3, stride=1,26            padding=1)27        self.conv2_2 = Conv2d(28            in_channels=128, out_channels=128, kernel_size=3, stride=1,29            padding=1)30        self.conv3_1 = Conv2d(31            in_channels=128, out_channels=256, kernel_size=3, stride=1,32            padding=1)33        self.conv3_2 = Conv2d(34            in_channels=256, out_channels=256, kernel_size=3, stride=1,35            padding=1)36        self.conv3_3 = Conv2d(37            in_channels=256, out_channels=256, kernel_size=3, stride=1,38            padding=1)39        self.conv3_4 = Conv2d(40            in_channels=256, out_channels=256, kernel_size=3, stride=1,41            padding=1)42        self.conv4_1 = Conv2d(43            in_channels=256, out_channels=512, kernel_size=3, stride=1,44            padding=1)45        self.conv4_2 = Conv2d(46            in_channels=512, out_channels=512, kernel_size=3, stride=1,47            padding=1)48        self.conv4_3 = Conv2d(49            in_channels=512, out_channels=512, kernel_size=3, stride=1,50            padding=1)51        self.conv4_4 = Conv2d(52            in_channels=512, out_channels=512, kernel_size=3, stride=1,53            padding=1)54        self.conv5_1 = Conv2d(55            in_channels=512, out_channels=512, kernel_size=3, stride=1,56            padding=1)57        self.conv5_2 = Conv2d(58            in_channels=512, out_channels=512, kernel_size=3, stride=1,59            padding=1)60        self.conv5_3_CPM = Conv2d(61            in_channels=512, out_channels=128, kernel_size=3, stride=1,62            padding=1)63 64        # stage165        self.conv6_1_CPM = Conv2d(66            in_channels=128, out_channels=512, kernel_size=1, stride=1,67            padding=0)68        self.conv6_2_CPM = Conv2d(69            in_channels=512, out_channels=71, kernel_size=1, stride=1,70            padding=0)71 72        # stage273        self.Mconv1_stage2 = Conv2d(74            in_channels=199, out_channels=128, kernel_size=7, stride=1,75            padding=3)76        self.Mconv2_stage2 = Conv2d(77            in_channels=128, out_channels=128, kernel_size=7, stride=1,78            padding=3)79        self.Mconv3_stage2 = Conv2d(80            in_channels=128, out_channels=128, kernel_size=7, stride=1,81            padding=3)82        self.Mconv4_stage2 = Conv2d(83            in_channels=128, out_channels=128, kernel_size=7, stride=1,84            padding=3)85        self.Mconv5_stage2 = Conv2d(86            in_channels=128, out_channels=128, kernel_size=7, stride=1,87            padding=3)88        self.Mconv6_stage2 = Conv2d(89            in_channels=128, out_channels=128, kernel_size=1, stride=1,90            padding=0)91        self.Mconv7_stage2 = Conv2d(92            in_channels=128, out_channels=71, kernel_size=1, stride=1,93            padding=0)94 95        # stage396        self.Mconv1_stage3 = Conv2d(97            in_channels=199, out_channels=128, kernel_size=7, stride=1,98            padding=3)99        self.Mconv2_stage3 = Conv2d(100            in_channels=128, out_channels=128, kernel_size=7, stride=1,101            padding=3)102        self.Mconv3_stage3 = Conv2d(103            in_channels=128, out_channels=128, kernel_size=7, stride=1,104            padding=3)105        self.Mconv4_stage3 = Conv2d(106            in_channels=128, out_channels=128, kernel_size=7, stride=1,107            padding=3)108        self.Mconv5_stage3 = Conv2d(109            in_channels=128, out_channels=128, kernel_size=7, stride=1,110            padding=3)111        self.Mconv6_stage3 = Conv2d(112            in_channels=128, out_channels=128, kernel_size=1, stride=1,113            padding=0)114        self.Mconv7_stage3 = Conv2d(115            in_channels=128, out_channels=71, kernel_size=1, stride=1,116            padding=0)117 118        # stage4119        self.Mconv1_stage4 = Conv2d(120            in_channels=199, out_channels=128, kernel_size=7, stride=1,121            padding=3)122        self.Mconv2_stage4 = Conv2d(123            in_channels=128, out_channels=128, kernel_size=7, stride=1,124            padding=3)125        self.Mconv3_stage4 = Conv2d(126            in_channels=128, out_channels=128, kernel_size=7, stride=1,127            padding=3)128        self.Mconv4_stage4 = Conv2d(129            in_channels=128, out_channels=128, kernel_size=7, stride=1,130            padding=3)131        self.Mconv5_stage4 = Conv2d(132            in_channels=128, out_channels=128, kernel_size=7, stride=1,133            padding=3)134        self.Mconv6_stage4 = Conv2d(135            in_channels=128, out_channels=128, kernel_size=1, stride=1,136            padding=0)137        self.Mconv7_stage4 = Conv2d(138            in_channels=128, out_channels=71, kernel_size=1, stride=1,139            padding=0)140 141        # stage5142        self.Mconv1_stage5 = Conv2d(143            in_channels=199, out_channels=128, kernel_size=7, stride=1,144            padding=3)145        self.Mconv2_stage5 = Conv2d(146            in_channels=128, out_channels=128, kernel_size=7, stride=1,147            padding=3)148        self.Mconv3_stage5 = Conv2d(149            in_channels=128, out_channels=128, kernel_size=7, stride=1,150            padding=3)151        self.Mconv4_stage5 = Conv2d(152            in_channels=128, out_channels=128, kernel_size=7, stride=1,153            padding=3)154        self.Mconv5_stage5 = Conv2d(155            in_channels=128, out_channels=128, kernel_size=7, stride=1,156            padding=3)157        self.Mconv6_stage5 = Conv2d(158            in_channels=128, out_channels=128, kernel_size=1, stride=1,159            padding=0)160        self.Mconv7_stage5 = Conv2d(161            in_channels=128, out_channels=71, kernel_size=1, stride=1,162            padding=0)163 164        # stage6165        self.Mconv1_stage6 = Conv2d(166            in_channels=199, out_channels=128, kernel_size=7, stride=1,167            padding=3)168        self.Mconv2_stage6 = Conv2d(169            in_channels=128, out_channels=128, kernel_size=7, stride=1,170            padding=3)171        self.Mconv3_stage6 = Conv2d(172            in_channels=128, out_channels=128, kernel_size=7, stride=1,173            padding=3)174        self.Mconv4_stage6 = Conv2d(175            in_channels=128, out_channels=128, kernel_size=7, stride=1,176            padding=3)177        self.Mconv5_stage6 = Conv2d(178            in_channels=128, out_channels=128, kernel_size=7, stride=1,179            padding=3)180        self.Mconv6_stage6 = Conv2d(181            in_channels=128, out_channels=128, kernel_size=1, stride=1,182            padding=0)183        self.Mconv7_stage6 = Conv2d(184            in_channels=128, out_channels=71, kernel_size=1, stride=1,185            padding=0)186 187        for m in self.modules():188            if isinstance(m, Conv2d):189                init.constant_(m.bias, 0)190 191    def forward(self, x):192        """Return a list of heatmaps."""193        heatmaps = []194 195        h = self.relu(self.conv1_1(x))196        h = self.relu(self.conv1_2(h))197        h = self.max_pooling_2d(h)198        h = self.relu(self.conv2_1(h))199        h = self.relu(self.conv2_2(h))200        h = self.max_pooling_2d(h)201        h = self.relu(self.conv3_1(h))202        h = self.relu(self.conv3_2(h))203        h = self.relu(self.conv3_3(h))204        h = self.relu(self.conv3_4(h))205        h = self.max_pooling_2d(h)206        h = self.relu(self.conv4_1(h))207        h = self.relu(self.conv4_2(h))208        h = self.relu(self.conv4_3(h))209        h = self.relu(self.conv4_4(h))210        h = self.relu(self.conv5_1(h))211        h = self.relu(self.conv5_2(h))212        h = self.relu(self.conv5_3_CPM(h))213        feature_map = h214 215        # stage1216        h = self.relu(self.conv6_1_CPM(h))217        h = self.conv6_2_CPM(h)218        heatmaps.append(h)219 220        # stage2221        h = torch.cat([h, feature_map], dim=1)  # channel concat222        h = self.relu(self.Mconv1_stage2(h))223        h = self.relu(self.Mconv2_stage2(h))224        h = self.relu(self.Mconv3_stage2(h))225        h = self.relu(self.Mconv4_stage2(h))226        h = self.relu(self.Mconv5_stage2(h))227        h = self.relu(self.Mconv6_stage2(h))228        h = self.Mconv7_stage2(h)229        heatmaps.append(h)230 231        # stage3232        h = torch.cat([h, feature_map], dim=1)  # channel concat233        h = self.relu(self.Mconv1_stage3(h))234        h = self.relu(self.Mconv2_stage3(h))235        h = self.relu(self.Mconv3_stage3(h))236        h = self.relu(self.Mconv4_stage3(h))237        h = self.relu(self.Mconv5_stage3(h))238        h = self.relu(self.Mconv6_stage3(h))239        h = self.Mconv7_stage3(h)240        heatmaps.append(h)241 242        # stage4243        h = torch.cat([h, feature_map], dim=1)  # channel concat244        h = self.relu(self.Mconv1_stage4(h))245        h = self.relu(self.Mconv2_stage4(h))246        h = self.relu(self.Mconv3_stage4(h))247        h = self.relu(self.Mconv4_stage4(h))248        h = self.relu(self.Mconv5_stage4(h))249        h = self.relu(self.Mconv6_stage4(h))250        h = self.Mconv7_stage4(h)251        heatmaps.append(h)252 253        # stage5254        h = torch.cat([h, feature_map], dim=1)  # channel concat255        h = self.relu(self.Mconv1_stage5(h))256        h = self.relu(self.Mconv2_stage5(h))257        h = self.relu(self.Mconv3_stage5(h))258        h = self.relu(self.Mconv4_stage5(h))259        h = self.relu(self.Mconv5_stage5(h))260        h = self.relu(self.Mconv6_stage5(h))261        h = self.Mconv7_stage5(h)262        heatmaps.append(h)263 264        # stage6265        h = torch.cat([h, feature_map], dim=1)  # channel concat266        h = self.relu(self.Mconv1_stage6(h))267        h = self.relu(self.Mconv2_stage6(h))268        h = self.relu(self.Mconv3_stage6(h))269        h = self.relu(self.Mconv4_stage6(h))270        h = self.relu(self.Mconv5_stage6(h))271        h = self.relu(self.Mconv6_stage6(h))272        h = self.Mconv7_stage6(h)273        heatmaps.append(h)274 275        return heatmaps276 277 278LOG = logging.getLogger(__name__)279TOTEN = ToTensor()280TOPIL = ToPILImage()281 282 283params = {284    'gaussian_sigma': 2.5,285    'inference_img_size': 736,  # 368, 736, 1312286    'heatmap_peak_thresh': 0.1,287    'crop_scale': 1.5,288    'line_indices': [289        [0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6],290        [6, 7], [7, 8], [8, 9], [9, 10], [10, 11], [11, 12], [12, 13],291        [13, 14], [14, 15], [15, 16],292        [17, 18], [18, 19], [19, 20], [20, 21],293        [22, 23], [23, 24], [24, 25], [25, 26],294        [27, 28], [28, 29], [29, 30],295        [31, 32], [32, 33], [33, 34], [34, 35],296        [36, 37], [37, 38], [38, 39], [39, 40], [40, 41], [41, 36],297        [42, 43], [43, 44], [44, 45], [45, 46], [46, 47], [47, 42],298        [48, 49], [49, 50], [50, 51], [51, 52], [52, 53], [53, 54],299        [54, 55], [55, 56], [56, 57], [57, 58], [58, 59], [59, 48],300        [60, 61], [61, 62], [62, 63], [63, 64], [64, 65], [65, 66],301        [66, 67], [67, 60]302    ],303}304 305 306class Face(object):307    """308    The OpenPose face landmark detector model.309 310    Args:311        inference_size: set the size of the inference image size, suggested:312            368, 736, 1312, default 736313        gaussian_sigma: blur the heatmaps, default 2.5314        heatmap_peak_thresh: return landmark if over threshold, default 0.1315 316    """317    def __init__(self, face_model_path,318                 inference_size=None,319                 gaussian_sigma=None,320                 heatmap_peak_thresh=None):321        self.inference_size = inference_size or params["inference_img_size"]322        self.sigma = gaussian_sigma or params['gaussian_sigma']323        self.threshold = heatmap_peak_thresh or params["heatmap_peak_thresh"]324        self.model = FaceNet()325        self.model.load_state_dict(torch.load(face_model_path))326        if torch.cuda.is_available():327            self.model = self.model.cuda()328            print('cuda')329        self.model.eval()330 331    def __call__(self, face_img):332        H, W, C = face_img.shape333 334        w_size = 384335        x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5336 337        if torch.cuda.is_available():338            x_data = x_data.cuda()339 340        with torch.no_grad():341            hs = self.model(x_data[None, ...])342            heatmaps = F.interpolate(343                hs[-1],344                (H, W),345                mode='bilinear', align_corners=True).cpu().numpy()[0]346        return heatmaps347 348    def compute_peaks_from_heatmaps(self, heatmaps):349        all_peaks = []350        for part in range(heatmaps.shape[0]):351            map_ori = heatmaps[part].copy()352            binary = np.ascontiguousarray(map_ori > 0.05, dtype=np.uint8)353 354            if np.sum(binary) == 0:355                continue356 357            positions = np.where(binary > 0.5)358            intensities = map_ori[positions]359            mi = np.argmax(intensities)360            y, x = positions[0][mi], positions[1][mi]361            all_peaks.append([x, y])362 363        return np.array(all_peaks)364