georgefen/Face-Landmark-ControlNet

1import os2import sys3import torch4import torch.nn as nn5import torch.utils.model_zoo as model_zoo6from  torch.nn import  functional as F7 8 9class BlockTypeA(nn.Module):10    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):11        super(BlockTypeA, self).__init__()12        self.conv1 = nn.Sequential(13            nn.Conv2d(in_c2, out_c2, kernel_size=1),14            nn.BatchNorm2d(out_c2),15            nn.ReLU(inplace=True)16        )17        self.conv2 = nn.Sequential(18            nn.Conv2d(in_c1, out_c1, kernel_size=1),19            nn.BatchNorm2d(out_c1),20            nn.ReLU(inplace=True)21        )22        self.upscale = upscale23 24    def forward(self, a, b):25        b = self.conv1(b)26        a = self.conv2(a)27        b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)28        return torch.cat((a, b), dim=1)29 30 31class BlockTypeB(nn.Module):32    def __init__(self, in_c, out_c):33        super(BlockTypeB, self).__init__()34        self.conv1 = nn.Sequential(35            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),36            nn.BatchNorm2d(in_c),37            nn.ReLU()38        )39        self.conv2 = nn.Sequential(40            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),41            nn.BatchNorm2d(out_c),42            nn.ReLU()43        )44 45    def forward(self, x):46        x = self.conv1(x) + x47        x = self.conv2(x)48        return x49 50class BlockTypeC(nn.Module):51    def __init__(self, in_c, out_c):52        super(BlockTypeC, self).__init__()53        self.conv1 = nn.Sequential(54            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),55            nn.BatchNorm2d(in_c),56            nn.ReLU()57        )58        self.conv2 = nn.Sequential(59            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),60            nn.BatchNorm2d(in_c),61            nn.ReLU()62        )63        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)64 65    def forward(self, x):66        x = self.conv1(x)67        x = self.conv2(x)68        x = self.conv3(x)69        return x70 71def _make_divisible(v, divisor, min_value=None):72    """73    This function is taken from the original tf repo.74    It ensures that all layers have a channel number that is divisible by 875    It can be seen here:76    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py77    :param v:78    :param divisor:79    :param min_value:80    :return:81    """82    if min_value is None:83        min_value = divisor84    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)85    # Make sure that round down does not go down by more than 10%.86    if new_v < 0.9 * v:87        new_v += divisor88    return new_v89 90 91class ConvBNReLU(nn.Sequential):92    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):93        self.channel_pad = out_planes - in_planes94        self.stride = stride95        #padding = (kernel_size - 1) // 296 97        # TFLite uses slightly different padding than PyTorch98        if stride == 2:99            padding = 0100        else:101            padding = (kernel_size - 1) // 2102 103        super(ConvBNReLU, self).__init__(104            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),105            nn.BatchNorm2d(out_planes),106            nn.ReLU6(inplace=True)107        )108        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)109 110 111    def forward(self, x):112        # TFLite uses  different padding113        if self.stride == 2:114            x = F.pad(x, (0, 1, 0, 1), "constant", 0)115            #print(x.shape)116 117        for module in self:118            if not isinstance(module, nn.MaxPool2d):119                x = module(x)120        return x121 122 123class InvertedResidual(nn.Module):124    def __init__(self, inp, oup, stride, expand_ratio):125        super(InvertedResidual, self).__init__()126        self.stride = stride127        assert stride in [1, 2]128 129        hidden_dim = int(round(inp * expand_ratio))130        self.use_res_connect = self.stride == 1 and inp == oup131 132        layers = []133        if expand_ratio != 1:134            # pw135            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))136        layers.extend([137            # dw138            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),139            # pw-linear140            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),141            nn.BatchNorm2d(oup),142        ])143        self.conv = nn.Sequential(*layers)144 145    def forward(self, x):146        if self.use_res_connect:147            return x + self.conv(x)148        else:149            return self.conv(x)150 151 152class MobileNetV2(nn.Module):153    def __init__(self, pretrained=True):154        """155        MobileNet V2 main class156        Args:157            num_classes (int): Number of classes158            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount159            inverted_residual_setting: Network structure160            round_nearest (int): Round the number of channels in each layer to be a multiple of this number161            Set to 1 to turn off rounding162            block: Module specifying inverted residual building block for mobilenet163        """164        super(MobileNetV2, self).__init__()165 166        block = InvertedResidual167        input_channel = 32168        last_channel = 1280169        width_mult = 1.0170        round_nearest = 8171 172        inverted_residual_setting = [173            # t, c, n, s174            [1, 16, 1, 1],175            [6, 24, 2, 2],176            [6, 32, 3, 2],177            [6, 64, 4, 2],178            #[6, 96, 3, 1],179            #[6, 160, 3, 2],180            #[6, 320, 1, 1],181        ]182 183        # only check the first element, assuming user knows t,c,n,s are required184        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:185            raise ValueError("inverted_residual_setting should be non-empty "186                             "or a 4-element list, got {}".format(inverted_residual_setting))187 188        # building first layer189        input_channel = _make_divisible(input_channel * width_mult, round_nearest)190        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)191        features = [ConvBNReLU(4, input_channel, stride=2)]192        # building inverted residual blocks193        for t, c, n, s in inverted_residual_setting:194            output_channel = _make_divisible(c * width_mult, round_nearest)195            for i in range(n):196                stride = s if i == 0 else 1197                features.append(block(input_channel, output_channel, stride, expand_ratio=t))198                input_channel = output_channel199        self.features = nn.Sequential(*features)200 201        self.fpn_selected = [3, 6, 10]202        # weight initialization203        for m in self.modules():204            if isinstance(m, nn.Conv2d):205                nn.init.kaiming_normal_(m.weight, mode='fan_out')206                if m.bias is not None:207                    nn.init.zeros_(m.bias)208            elif isinstance(m, nn.BatchNorm2d):209                nn.init.ones_(m.weight)210                nn.init.zeros_(m.bias)211            elif isinstance(m, nn.Linear):212                nn.init.normal_(m.weight, 0, 0.01)213                nn.init.zeros_(m.bias)214 215        #if pretrained:216        #    self._load_pretrained_model()217 218    def _forward_impl(self, x):219        # This exists since TorchScript doesn't support inheritance, so the superclass method220        # (this one) needs to have a name other than `forward` that can be accessed in a subclass221        fpn_features = []222        for i, f in enumerate(self.features):223            if i > self.fpn_selected[-1]:224                break225            x = f(x)226            if i in self.fpn_selected:227                fpn_features.append(x)228 229        c2, c3, c4 = fpn_features230        return c2, c3, c4231 232 233    def forward(self, x):234        return self._forward_impl(x)235 236    def _load_pretrained_model(self):237        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')238        model_dict = {}239        state_dict = self.state_dict()240        for k, v in pretrain_dict.items():241            if k in state_dict:242                model_dict[k] = v243        state_dict.update(model_dict)244        self.load_state_dict(state_dict)245 246 247class MobileV2_MLSD_Tiny(nn.Module):248    def __init__(self):249        super(MobileV2_MLSD_Tiny, self).__init__()250 251        self.backbone = MobileNetV2(pretrained=True)252 253        self.block12 = BlockTypeA(in_c1= 32, in_c2= 64,254                                  out_c1= 64, out_c2=64)255        self.block13 = BlockTypeB(128, 64)256 257        self.block14 = BlockTypeA(in_c1 = 24,  in_c2 = 64,258                                  out_c1= 32,  out_c2= 32)259        self.block15 = BlockTypeB(64, 64)260 261        self.block16 = BlockTypeC(64, 16)262 263    def forward(self, x):264        c2, c3, c4 = self.backbone(x)265 266        x = self.block12(c3, c4)267        x = self.block13(x)268        x = self.block14(c2, x)269        x = self.block15(x)270        x = self.block16(x)271        x = x[:, 7:, :, :]272        #print(x.shape)273        x = F.interpolate(x, scale_factor=2.0, mode='bilinear', align_corners=True)274 275        return x