SSD512.py

from utils.lib import *

class VGG16Base(nn.Module):
    """
    Lấy VGG16 làm base network, tuy nhiên cần có một vài thay đổi:
    - Đầu vào ảnh là 512x512 thay vì 224x224, các comment bên dưới sẽ áp dụng cho đầu vào 512x512
    - Lớp pooling thứ 3 sử dụng ceiling mode thay vì floor mode
    - Lớp pooling thứ 5 kernel size (2, 2) -> (3, 3) và stride 2 -> 1, và padding = 1
    - Ta downsample (decimate) parameter fc6 và fc7 để tạo thành conv6 và conv7, loại bỏ hoàn toàn fc8
    """

    def __init__(self):
        super().__init__()

        self.conv1_1 = nn.Conv2d(in_channels=  3, out_channels= 64, kernel_size=3, padding=1)
        self.conv1_2 = nn.Conv2d(in_channels= 64, out_channels= 64, kernel_size=3, padding=1)
        self.pool1   = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv2_1 = nn.Conv2d(in_channels= 64, out_channels=128, kernel_size=3, padding=1)
        self.conv2_2 = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1)
        self.pool2   = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv3_1 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1)
        self.conv3_2 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1)
        self.conv3_3 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1)
        self.pool3   = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)

        self.conv4_1 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1)
        self.conv4_2 = nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
        self.conv4_3 = nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
        self.pool4   = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv5_1 = nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
        self.conv5_2 = nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
        self.conv5_3 = nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
        self.pool5   = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        
        # Không còn fc layers nữa, thay vào đó là conv6 và conv7
        # atrous
        self.conv6   = nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, padding=6, dilation=6)
        self.conv7   = nn.Conv2d(in_channels=1024, out_channels=1024, kernel_size=1)

    def decimate(self, tensor, steps):
        assert(len(steps) == tensor.dim())
        
        for i in range(tensor.dim()):
            if steps[i] is not None:
                tensor = tensor.index_select(dim=i, index=torch.arange(start=0, end=tensor.shape[i], step=steps[i]))

        return tensor

    
    def load_pretrain(self):
        """
        load pretrain từ thư viện pytorch, decimate param lại để phù hợp với conv6 và conv7
        """

        state_dict  = self.state_dict() 
        param_names = list(state_dict.keys())

        # old version : torch.vision.models.vgg16(pretrain=True)
        # Load model theo API mới của pytorch, cụ thể hơn tại : https://pytorch.org/vision/stable/models.html
        pretrain_state_dict  = torchvision.models.vgg16(weights='VGG16_Weights.DEFAULT').state_dict()
        pretrain_param_names = list(pretrain_state_dict.keys())

        # Pretrain param name và custom param name không giống nhau, các param chỉ cùng thứ tự như trong architecture
        for idx, param_name in enumerate(param_names[:-4]): # 4 param cuối là weight và bias của conv6 và conv7, sẽ xử lí sau
            state_dict[param_name] = pretrain_state_dict[pretrain_param_names[idx]]

        # fc -> conv
        fc6_weight = pretrain_state_dict['classifier.0.weight'].view(4096, 512, 7, 7)
        fc6_bias   = pretrain_state_dict['classifier.0.bias'].view(4096)

        fc7_weight = pretrain_state_dict['classifier.3.weight'].view(4096, 4096, 1, 1)
        fc7_bias   = pretrain_state_dict['classifier.3.bias'].view(4096)

        # downsample parameter
        state_dict['conv6.weight'] = self.decimate(fc6_weight, steps=[4, None, 3, 3])
        state_dict['conv6.bias']   = self.decimate(fc6_bias, steps=[4])

        state_dict['conv7.weight'] = self.decimate(fc7_weight, steps=[4, 4, None, None])
        state_dict['conv7.bias']   = self.decimate(fc7_bias, steps=[4])

        self.load_state_dict(state_dict)


    def forward(self, images):
        """
        :param images, tensor [N, 3, 512, 512]

        return:
        """
        out = F.relu(self.conv1_1(images)) # [N, 64, 512, 512]
        out = F.relu(self.conv1_2(out))    # [N, 64, 512, 512]
        out = self.pool1(out)              # [N, 64, 256, 256]

        out = F.relu(self.conv2_1(out))    # [N, 128, 256, 256]
        out = F.relu(self.conv2_2(out))    # [N, 128, 256, 256]
        out = self.pool2(out)              # [N, 128, 128, 128]

        out = F.relu(self.conv3_1(out))    # [N, 256, 128, 128]
        out = F.relu(self.conv3_2(out))    # [N, 256, 128, 128]
        out = F.relu(self.conv3_3(out))    # [N, 256, 128, 128]
        out = self.pool3(out)              # [N, 256, 64, 64] 

        out = F.relu(self.conv4_1(out))    # [N, 512, 64, 64]
        out = F.relu(self.conv4_2(out))    # [N, 512, 64, 64]
        out = F.relu(self.conv4_3(out))    # [N, 512, 64, 64]
        conv4_3_feats = out                # [N, 512, 64, 64]
        out = self.pool4(out)              # [N, 512, 32, 32]

        out = F.relu(self.conv5_1(out))    # [N, 512, 32, 32]
        out = F.relu(self.conv5_2(out))    # [N, 512, 32, 32]
        out = F.relu(self.conv5_3(out))    # [N, 512, 32, 32]
        out = self.pool5(out)              # [N, 512, 32, 32], layer pooling này không làm thay đổi size features map

        out = F.relu(self.conv6(out))      # [N, 1024, 32, 32]

        conv7_feats = F.relu(self.conv7(out)) # [N, 1024, 32, 32]

        return conv4_3_feats, conv7_feats     # [N, 512, 64, 64], [N, 1024, 32, 32]
    

class AuxiliraryConvolutions(nn.Module):

    def __init__(self):
        super().__init__()
        
        self.conv8_1  = nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=1, padding=0)
        self.conv8_2  = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1)
        
        self.conv9_1  = nn.Conv2d(in_channels=512, out_channels=128, kernel_size=1, padding=0)
        self.conv9_2  = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)

        self.conv10_1 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=1, padding=0)
        self.conv10_2 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)

        self.conv11_1 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=1, padding=0)
        self.conv11_2 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)

        self.conv12_1 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=1, padding=0)
        self.conv12_2 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, padding=1)


    def init_conv2d(self):
        """
        Initialize convolution parameters.
        """
        for c in self.children():
            if isinstance(c, nn.Conv2d):
                nn.init.xavier_uniform_(c.weight)
                if c.bias is not None:
                    nn.init.constant_(c.bias, 0.)

    def forward(self, conv7_feats):
        """
        :param conv8_feats, tensor [N, 1024, 32, 32]
        """

        out = F.relu(self.conv8_1(conv7_feats))   # [N, 256, 32, 32]
        out = F.relu(self.conv8_2(out))           # [N, 512, 16, 16]
        conv8_2_feats = out                       # [N, 512, 16, 16]

        out = F.relu(self.conv9_1(out))           # [N, 128, 16, 16]
        out = F.relu(self.conv9_2(out))           # [N, 256, 8, 8]
        conv9_2_feats = out                       # [N, 256, 8, 8]

        out = F.relu(self.conv10_1(out))          # [N, 128, 8, 8]
        out = F.relu(self.conv10_2(out))          # [N, 256, 4, 4]
        conv10_2_feats = out                      # [N, 256, 4, 4]

        out = F.relu(self.conv11_1(out))          # [N, 128, 4, 4]
        out = F.relu(self.conv11_2(out))          # [N, 256, 2, 2]
        conv11_2_feats = out     

        out = F.relu(self.conv12_1(out))          # [N, 128, 2, 2]
        out = F.relu(self.conv12_2(out))          # [N, 256, 1, 1]
        conv12_2_feats = out

        return conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats, conv12_2_feats
    

class PredictionConvolutions(nn.Module):

    def __init__(self, n_classes=21):
        super().__init__()

        self.n_classes = n_classes

        n_boxes={
            'conv4_3'  : 4,
            'conv7'    : 6,
            'conv8_2'  : 6,
            'conv9_2'  : 6,
            'conv10_2' : 6,
            'conv11_2' : 4,
            'conv12_2' : 4
        }

        # kernel size = 3 và padding = 1 không làm thay đổi kích thước feature map 

        self.loc_conv4_3   = nn.Conv2d(512,  n_boxes['conv4_3']*4, kernel_size=3, padding=1)
        self.loc_conv7     = nn.Conv2d(1024, n_boxes['conv7']*4, kernel_size=3, padding=1)
        self.loc_conv8_2   = nn.Conv2d(512,  n_boxes['conv8_2']*4, kernel_size=3, padding=1)
        self.loc_conv9_2   = nn.Conv2d(256,  n_boxes['conv9_2']*4, kernel_size=3, padding=1)
        self.loc_conv10_2  = nn.Conv2d(256,  n_boxes['conv10_2']*4, kernel_size=3, padding=1)
        self.loc_conv11_2  = nn.Conv2d(256,  n_boxes['conv11_2']*4, kernel_size=3, padding=1)
        self.loc_conv12_2  = nn.Conv2d(256, n_boxes['conv12_2']*4, kernel_size=3, padding=1)


        self.conf_conv4_3  = nn.Conv2d(512,  n_boxes['conv4_3']*n_classes, kernel_size=3, padding=1)
        self.conf_conv7    = nn.Conv2d(1024, n_boxes['conv7']*n_classes, kernel_size=3, padding=1)
        self.conf_conv8_2  = nn.Conv2d(512,  n_boxes['conv8_2']*n_classes, kernel_size=3, padding=1)
        self.conf_conv9_2  = nn.Conv2d(256,  n_boxes['conv9_2']*n_classes, kernel_size=3, padding=1)
        self.conf_conv10_2 = nn.Conv2d(256,  n_boxes['conv10_2']*n_classes, kernel_size=3, padding=1)
        self.conf_conv11_2 = nn.Conv2d(256,  n_boxes['conv11_2']*n_classes, kernel_size=3, padding=1)
        self.conf_conv12_2 = nn.Conv2d(256,  n_boxes['conv12_2']*n_classes, kernel_size=3, padding=1)

    def init_conv2d(self):
        """
        Initialize convolution parameters.
        """
        for c in self.children():
            if isinstance(c, nn.Conv2d):
                nn.init.xavier_uniform_(c.weight)
                if c.bias is not None:
                    nn.init.constant_(c.bias, 0.)


    def forward(self, conv4_3_feats, conv7_feats, conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats, conv12_2_feats):

        batch_size = conv4_3_feats.shape[0]


        loc_conv4_3   = self.loc_conv4_3(conv4_3_feats)
        loc_conv4_3   = loc_conv4_3.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 4)
        
        loc_conv7     = self.loc_conv7(conv7_feats)
        loc_conv7     = loc_conv7.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 4)

        loc_conv8_2   = self.loc_conv8_2(conv8_2_feats)
        loc_conv8_2   = loc_conv8_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 4)

        loc_conv9_2   = self.loc_conv9_2(conv9_2_feats)
        loc_conv9_2   = loc_conv9_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 4)

        loc_conv10_2   = self.loc_conv10_2(conv10_2_feats)
        loc_conv10_2   = loc_conv10_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 4)

        loc_conv11_2   = self.loc_conv11_2(conv11_2_feats)
        loc_conv11_2   = loc_conv11_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 4)

        loc_conv12_2   = self.loc_conv12_2(conv12_2_feats)
        loc_conv12_2   = loc_conv12_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 4)



        conf_conv4_3   = self.conf_conv4_3(conv4_3_feats)
        conf_conv4_3   = conf_conv4_3.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.n_classes)
        
        conf_conv7     = self.conf_conv7(conv7_feats)
        conf_conv7     = conf_conv7.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.n_classes)

        conf_conv8_2   = self.conf_conv8_2(conv8_2_feats)
        conf_conv8_2   = conf_conv8_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.n_classes)

        conf_conv9_2   = self.conf_conv9_2(conv9_2_feats)
        conf_conv9_2   = conf_conv9_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.n_classes)

        conf_conv10_2   = self.conf_conv10_2(conv10_2_feats)
        conf_conv10_2   = conf_conv10_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.n_classes)

        conf_conv11_2   = self.conf_conv11_2(conv11_2_feats)
        conf_conv11_2   = conf_conv11_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.n_classes)

        conf_conv12_2   = self.conf_conv12_2(conv12_2_feats)
        conf_conv12_2   = conf_conv12_2.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.n_classes)

        loc  = torch.cat((loc_conv4_3, loc_conv7, loc_conv8_2, loc_conv9_2, loc_conv10_2, loc_conv11_2, loc_conv12_2), dim=1)
        conf = torch.cat((conf_conv4_3, conf_conv7, conf_conv8_2, conf_conv9_2, conf_conv10_2, conf_conv11_2, conf_conv12_2), dim=1)

        return loc, conf
    
class L2Norm(nn.Module):
    def __init__(self, input_channel=512, scale=20):
        super().__init__()
        self.scale_factors = nn.Parameter(torch.FloatTensor(1, input_channel, 1, 1))
        self.eps           = 1e-10
        nn.init.constant_(self.scale_factors, scale)
    
    def forward(self, tensor):
        norm   = tensor.pow(2).sum(dim=1, keepdim=True).sqrt()
        tensor = tensor/(norm + self.eps)*self.scale_factors
        return tensor
    
class SSD512(nn.Module):

    def __init__(self, pretrain_path = None, data_train_on = "VOC", n_classes = 21):
        super().__init__()

        self.n_classes   = n_classes
        self.data_train_on = data_train_on
        self.base_net    = VGG16Base()
        self.auxi_conv   = AuxiliraryConvolutions()
        self.pred_conv   = PredictionConvolutions(n_classes) 
        self.l2_norm     = L2Norm()

        if pretrain_path is not None:
            self.load_state_dict(torch.load(pretrain_path))
        else:
            self.base_net.load_pretrain()
            self.auxi_conv.init_conv2d()
            self.pred_conv.init_conv2d()

    def create_prior_boxes(self):
        """ 
        Tạo prior boxes (tensor) như trong paper
        mỗi box có dạng [cx, cy, w, h] được scale
        """
        # kích thước feature map tương ứng
        fmap_sizes    = [64, 32, 16, 8, 4, 2, 1]
        
        # scale như trong paper và được tính sẵn thay vì công thức
        # lưu ý ở conv4_3, tác giả xét như một trường hợp đặc biệt (scale 0.1):
        # Ở mục 3.1, trang 7 : 
        # "We set default box with scale 0.1 on conv4 3 .... "
        # "For SSD512 model, we add extra conv12 2 for prediction, set smin to 0.15, and 0.07 on conv4 3...""

        if self.data_train_on == "VOC":
            box_scales    = [0.07, 0.15, 0.3, 0.45, 0.6, 0.75, 0.9]
        elif self.data_train_on == "COCO":
            box_scales    = [0.04, 0.1, 0.26, 0.42, 0.58, 0.74, 0.9] 
            
        aspect_ratios = [
                [1., 2., 0.5],
                [1., 2., 3., 0.5, 0.333],
                [1., 2., 3., 0.5, 0.333],
                [1., 2., 3., 0.5, 0.333],
                [1., 2., 3., 0.5, 0.333],
                [1., 2., 0.5],
                [1., 2., 0.5]
            ]
        dboxes = []
        
        
        for idx, fmap_size in enumerate(fmap_sizes):
            for i in range(fmap_size):
                for j in range(fmap_size):

                    # lưu ý, cx trong ảnh là trục hoành, do đó j + 0.5 chứ không phải i + 0.5
                    cx = (j + 0.5) / fmap_size
                    cy = (i + 0.5) / fmap_size

                    for aspect_ratio in aspect_ratios[idx]:
                        scale = box_scales[idx]
                        dboxes.append([cx, cy, scale*sqrt(aspect_ratio), scale/sqrt(aspect_ratio)])

                        if aspect_ratio == 1.:
                            try:
                                scale = sqrt(scale*box_scales[idx + 1])
                            except IndexError:
                                scale = 1.
                            dboxes.append([cx, cy, scale*sqrt(aspect_ratio), scale/sqrt(aspect_ratio)])

        dboxes = torch.FloatTensor(dboxes)
        
        #dboxes = pascalVOC_style(dboxes)
        dboxes.clamp_(0, 1)
        #dboxes = yolo_style(dboxes)
                
        return dboxes

    def forward(self, images):
        conv4_3_feats, conv7_feats                                                   = self.base_net(images)
        conv4_3_feats                                                                = self.l2_norm(conv4_3_feats)
        conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats, conv12_2_feats = self.auxi_conv(conv7_feats)

        loc, conf = self.pred_conv(conv4_3_feats, conv7_feats, conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats, conv12_2_feats)
        return loc, conf



if __name__ == "__main__":
    T = SSD512()
    imgs = torch.Tensor(1, 3, 512, 512)
    loc, conf = T(imgs)
    print(loc.shape)
    print(conf.shape)