model.py

import os
from typing import Union, Tuple, List, NamedTuple

import torch
from torch import nn, Tensor
from torch.nn import functional as F
from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler

from backbone.base import Base as BackboneBase
from bbox import BBox
from nms.nms import NMS
from roi.wrapper import Wrapper as ROIWrapper
from rpn.region_proposal_network import RegionProposalNetwork


class Model(nn.Module):

    class ForwardInput(object):
        class Train(NamedTuple):
            image: Tensor
            gt_classes: Tensor
            gt_bboxes: Tensor

        class Eval(NamedTuple):
            image: Tensor

    class ForwardOutput(object):
        class Train(NamedTuple):
            anchor_objectness_loss: Tensor
            anchor_transformer_loss: Tensor
            proposal_class_loss: Tensor
            proposal_transformer_loss: Tensor

        class Eval(NamedTuple):
            detection_bboxes: Tensor
            detection_classes: Tensor
            detection_probs: Tensor

    def __init__(self, backbone: BackboneBase, num_classes: int, pooling_mode: ROIWrapper.Mode,
                 anchor_ratios: List[Tuple[int, int]], anchor_scales: List[int], rpn_pre_nms_top_n: int, rpn_post_nms_top_n: int):
        super().__init__()

        conv_layers, lateral_layers, dealiasing_layers, num_features_out = backbone.features()
        self.conv1, self.conv2, self.conv3, self.conv4, self.conv5 = conv_layers
        self.lateral_c2, self.lateral_c3, self.lateral_c4, self.lateral_c5 = lateral_layers
        self.dealiasing_p2, self.dealiasing_p3, self.dealiasing_p4 = dealiasing_layers

        self._bn_modules = [it for it in self.conv1.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.conv2.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.conv3.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.conv4.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.conv5.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.lateral_c2.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.lateral_c3.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.lateral_c4.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.lateral_c5.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.dealiasing_p2.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.dealiasing_p3.modules() if isinstance(it, nn.BatchNorm2d)] + \
                           [it for it in self.dealiasing_p4.modules() if isinstance(it, nn.BatchNorm2d)]

        self.num_classes = num_classes

        self.rpn = RegionProposalNetwork(num_features_out, anchor_ratios, anchor_scales, rpn_pre_nms_top_n, rpn_post_nms_top_n)
        self.detection = Model.Detection(pooling_mode, self.num_classes)

    def forward(self, forward_input: Union[ForwardInput.Train, ForwardInput.Eval]) -> Union[ForwardOutput.Train, ForwardOutput.Eval]:
        # freeze batch normalization modules for each forwarding process just in case model was switched to `train` at any time
        for bn_module in self._bn_modules:
            bn_module.eval()
            for parameter in bn_module.parameters():
                parameter.requires_grad = False

        image = forward_input.image.unsqueeze(dim=0)
        image_height, image_width = image.shape[2], image.shape[3]

        # Bottom-up pathway
        c1 = self.conv1(image)
        c2 = self.conv2(c1)
        c3 = self.conv3(c2)
        c4 = self.conv4(c3)
        c5 = self.conv5(c4)

        # Top-down pathway and lateral connections
        p5 = self.lateral_c5(c5)
        p4 = self.lateral_c4(c4) + F.interpolate(input=p5, size=(c4.shape[2], c4.shape[3]), mode='nearest')
        p3 = self.lateral_c3(c3) + F.interpolate(input=p4, size=(c3.shape[2], c3.shape[3]), mode='nearest')
        p2 = self.lateral_c2(c2) + F.interpolate(input=p3, size=(c2.shape[2], c2.shape[3]), mode='nearest')

        # Reduce the aliasing effect
        p4 = self.dealiasing_p4(p4)
        p3 = self.dealiasing_p3(p3)
        p2 = self.dealiasing_p2(p2)

        p6 = F.max_pool2d(input=p5, kernel_size=1, stride=2)

        # NOTE: We define the anchors to have areas of {32^2, 64^2, 128^2, 256^2, 512^2} pixels on {P2, P3, P4, P5, P6} respectively

        anchor_objectnesses = []
        anchor_transformers = []
        anchor_bboxes = []
        proposal_bboxes = []

        for p, anchor_size in zip([p2, p3, p4, p5, p6], [32, 64, 128, 256, 512]):
            p_anchor_objectnesses, p_anchor_transformers = self.rpn.forward(features=p, image_width=image_width, image_height=image_height)
            p_anchor_bboxes = self.rpn.generate_anchors(image_width, image_height,
                                                        num_x_anchors=p.shape[3], num_y_anchors=p.shape[2],
                                                        anchor_size=anchor_size).cuda()
            p_proposal_bboxes = self.rpn.generate_proposals(p_anchor_bboxes, p_anchor_objectnesses, p_anchor_transformers,
                                                            image_width, image_height)
            anchor_objectnesses.append(p_anchor_objectnesses)
            anchor_transformers.append(p_anchor_transformers)
            anchor_bboxes.append(p_anchor_bboxes)
            proposal_bboxes.append(p_proposal_bboxes)

        anchor_objectnesses = torch.cat(anchor_objectnesses, dim=0)
        anchor_transformers = torch.cat(anchor_transformers, dim=0)
        anchor_bboxes = torch.cat(anchor_bboxes, dim=0)
        proposal_bboxes = torch.cat(proposal_bboxes, dim=0)

        if self.training:
            forward_input: Model.ForwardInput.Train

            anchor_sample_fg_indices, anchor_sample_selected_indices, gt_anchor_objectnesses, gt_anchor_transformers = self.rpn.sample(anchor_bboxes, forward_input.gt_bboxes, image_width, image_height)
            anchor_objectnesses = anchor_objectnesses[anchor_sample_selected_indices]
            anchor_transformers = anchor_transformers[anchor_sample_fg_indices]
            anchor_objectness_loss, anchor_transformer_loss = self.rpn.loss(anchor_objectnesses, anchor_transformers, gt_anchor_objectnesses, gt_anchor_transformers)

            proposal_sample_fg_indices, proposal_sample_selected_indices, gt_proposal_classes, gt_proposal_transformers = self.detection.sample(proposal_bboxes, forward_input.gt_classes, forward_input.gt_bboxes)
            proposal_bboxes = proposal_bboxes[proposal_sample_selected_indices]
            proposal_classes, proposal_transformers = self.detection.forward(p2, p3, p4, p5, proposal_bboxes, image_width, image_height)
            proposal_class_loss, proposal_transformer_loss = self.detection.loss(proposal_classes, proposal_transformers, gt_proposal_classes, gt_proposal_transformers)

            forward_output = Model.ForwardOutput.Train(anchor_objectness_loss, anchor_transformer_loss, proposal_class_loss, proposal_transformer_loss)
        else:
            proposal_classes, proposal_transformers = self.detection.forward(p2, p3, p4, p5, proposal_bboxes, image_width, image_height)
            detection_bboxes, detection_classes, detection_probs = self.detection.generate_detections(proposal_bboxes, proposal_classes, proposal_transformers, image_width, image_height)
            forward_output = Model.ForwardOutput.Eval(detection_bboxes, detection_classes, detection_probs)

        return forward_output

    def save(self, path_to_checkpoints_dir: str, step: int, optimizer: Optimizer, scheduler: _LRScheduler) -> str:
        path_to_checkpoint = os.path.join(path_to_checkpoints_dir, f'model-{step}.pth')
        checkpoint = {
            'state_dict': self.state_dict(),
            'step': step,
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict()
        }
        torch.save(checkpoint, path_to_checkpoint)
        return path_to_checkpoint

    def load(self, path_to_checkpoint: str, optimizer: Optimizer = None, scheduler: _LRScheduler = None) -> 'Model':
        checkpoint = torch.load(path_to_checkpoint)
        self.load_state_dict(checkpoint['state_dict'])
        step = checkpoint['step']
        if optimizer is not None:
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        if scheduler is not None:
            scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        return step

    class Detection(nn.Module):

        def __init__(self, pooling_mode: ROIWrapper.Mode, num_classes: int):
            super().__init__()
            self._pooling_mode = pooling_mode
            self._hidden = nn.Sequential(
                nn.Linear(256 * 7 * 7, 1024),
                nn.ReLU(),
                nn.Linear(1024, 1024),
                nn.ReLU()
            )
            self.num_classes = num_classes
            self._class = nn.Linear(1024, num_classes)
            self._transformer = nn.Linear(1024, num_classes * 4)
            self._transformer_normalize_mean = torch.tensor([0., 0., 0., 0.], dtype=torch.float).cuda()
            self._transformer_normalize_std = torch.tensor([.1, .1, .2, .2], dtype=torch.float).cuda()

        def forward(self, p2: Tensor, p3: Tensor, p4: Tensor, p5: Tensor, proposal_bboxes: Tensor, image_width: int, image_height: int) -> Tuple[Tensor, Tensor]:
            w = proposal_bboxes[:, 2] - proposal_bboxes[:, 0]
            h = proposal_bboxes[:, 3] - proposal_bboxes[:, 1]
            k0 = 4
            k = torch.floor(k0 + torch.log2(torch.sqrt(w * h) / 224)).long()
            k = torch.clamp(k, min=2, max=5)

            k_to_p_dict = {2: p2, 3: p3, 4: p4, 5: p5}
            unique_k = torch.unique(k)

            # NOTE: `picked_indices` is for recording the order of selection from `proposal_bboxes`
            #       so that `pools` can be then restored to make it have a consistent correspondence
            #       with `proposal_bboxes`. For example:
            #
            #           proposal_bboxes =>  B0  B1  B2
            #            picked_indices =>   1   2   0
            #                     pools => BP1 BP2 BP0
            #            sorted_indices =>   2   0   1
            #                     pools => BP0 BP1 BP2

            pools = []
            picked_indices = []

            for uk in unique_k:
                uk = uk.item()
                p = k_to_p_dict[uk]
                uk_indices = (k == uk).nonzero().view(-1)
                uk_proposal_bboxes = proposal_bboxes[uk_indices]
                pool = ROIWrapper.apply(p, uk_proposal_bboxes, mode=self._pooling_mode, image_width=image_width, image_height=image_height)
                pools.append(pool)
                picked_indices.append(uk_indices)

            pools = torch.cat(pools, dim=0)
            picked_indices = torch.cat(picked_indices, dim=0)

            _, sorted_indices = torch.sort(picked_indices)
            pools = pools[sorted_indices]

            pools = pools.view(pools.shape[0], -1)
            hidden = self._hidden(pools)
            classes = self._class(hidden)
            transformers = self._transformer(hidden)
            return classes, transformers

        def sample(self, proposal_bboxes: Tensor, gt_classes: Tensor, gt_bboxes: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
            sample_fg_indices = torch.arange(end=len(proposal_bboxes), dtype=torch.long)
            sample_selected_indices = torch.arange(end=len(proposal_bboxes), dtype=torch.long)

            # find labels for each `proposal_bboxes`
            labels = torch.ones(len(proposal_bboxes), dtype=torch.long).cuda() * -1
            ious = BBox.iou(proposal_bboxes, gt_bboxes)
            proposal_max_ious, proposal_assignments = ious.max(dim=1)
            labels[proposal_max_ious < 0.5] = 0
            labels[proposal_max_ious >= 0.5] = gt_classes[proposal_assignments[proposal_max_ious >= 0.5]]

            # select 128 samples
            fg_indices = (labels > 0).nonzero().view(-1)
            bg_indices = (labels == 0).nonzero().view(-1)
            fg_indices = fg_indices[torch.randperm(len(fg_indices))[:min(len(fg_indices), 32)]]
            bg_indices = bg_indices[torch.randperm(len(bg_indices))[:128 - len(fg_indices)]]
            selected_indices = torch.cat([fg_indices, bg_indices])
            selected_indices = selected_indices[torch.randperm(len(selected_indices))]

            proposal_bboxes = proposal_bboxes[selected_indices]
            gt_proposal_transformers = BBox.calc_transformer(proposal_bboxes, gt_bboxes[proposal_assignments[selected_indices]])
            gt_proposal_classes = labels[selected_indices]

            gt_proposal_transformers = (gt_proposal_transformers - self._transformer_normalize_mean) / self._transformer_normalize_std

            gt_proposal_transformers = gt_proposal_transformers.cuda()
            gt_proposal_classes = gt_proposal_classes.cuda()

            sample_fg_indices = sample_fg_indices[fg_indices]
            sample_selected_indices = sample_selected_indices[selected_indices]

            return sample_fg_indices, sample_selected_indices, gt_proposal_classes, gt_proposal_transformers

        def loss(self, proposal_classes: Tensor, proposal_transformers: Tensor, gt_proposal_classes: Tensor, gt_proposal_transformers: Tensor) -> Tuple[Tensor, Tensor]:
            cross_entropy = F.cross_entropy(input=proposal_classes, target=gt_proposal_classes)

            proposal_transformers = proposal_transformers.view(-1, self.num_classes, 4)
            proposal_transformers = proposal_transformers[torch.arange(end=len(proposal_transformers), dtype=torch.long).cuda(), gt_proposal_classes]

            fg_indices = gt_proposal_classes.nonzero().view(-1)

            # NOTE: The default of `reduction` is `elementwise_mean`, which is divided by N x 4 (number of all elements), here we replaced by N for better performance
            smooth_l1_loss = F.smooth_l1_loss(input=proposal_transformers[fg_indices], target=gt_proposal_transformers[fg_indices], reduction='sum')
            smooth_l1_loss /= len(gt_proposal_transformers)

            return cross_entropy, smooth_l1_loss

        def generate_detections(self, proposal_bboxes: Tensor, proposal_classes: Tensor, proposal_transformers: Tensor, image_width: int, image_height: int) -> Tuple[Tensor, Tensor, Tensor]:
            proposal_transformers = proposal_transformers.view(-1, self.num_classes, 4)
            mean = self._transformer_normalize_mean.repeat(1, self.num_classes, 1)
            std = self._transformer_normalize_std.repeat(1, self.num_classes, 1)

            proposal_transformers = proposal_transformers * std - mean
            proposal_bboxes = proposal_bboxes.view(-1, 1, 4).repeat(1, self.num_classes, 1)
            detection_bboxes = BBox.apply_transformer(proposal_bboxes.view(-1, 4), proposal_transformers.view(-1, 4))

            detection_bboxes = detection_bboxes.view(-1, self.num_classes, 4)

            detection_bboxes[:, :, [0, 2]] = detection_bboxes[:, :, [0, 2]].clamp(min=0, max=image_width)
            detection_bboxes[:, :, [1, 3]] = detection_bboxes[:, :, [1, 3]].clamp(min=0, max=image_height)

            proposal_probs = F.softmax(proposal_classes, dim=1)

            detection_bboxes = detection_bboxes.cpu()
            proposal_probs = proposal_probs.cpu()

            generated_bboxes = []
            generated_classes = []
            generated_probs = []

            for c in range(1, self.num_classes):
                detection_class_bboxes = detection_bboxes[:, c, :]
                proposal_class_probs = proposal_probs[:, c]

                _, sorted_indices = proposal_class_probs.sort(descending=True)
                detection_class_bboxes = detection_class_bboxes[sorted_indices]
                proposal_class_probs = proposal_class_probs[sorted_indices]

                kept_indices = NMS.suppress(detection_class_bboxes.cuda(), threshold=0.3)
                detection_class_bboxes = detection_class_bboxes[kept_indices]
                proposal_class_probs = proposal_class_probs[kept_indices]

                generated_bboxes.append(detection_class_bboxes)
                generated_classes.append(torch.ones(len(kept_indices), dtype=torch.int) * c)
                generated_probs.append(proposal_class_probs)

            generated_bboxes = torch.cat(generated_bboxes, dim=0)
            generated_classes = torch.cat(generated_classes, dim=0)
            generated_probs = torch.cat(generated_probs, dim=0)
            return generated_bboxes, generated_classes, generated_probs