Delete diffusiondet/head.py

diffusiondet/head.py

@@ -1,386 +0,0 @@
-import copy
-import math
-from dataclasses import astuple
-
-import torch
-from torch import nn
-from torch.nn.modules.transformer import _get_activation_fn
-from torchvision.ops import RoIAlign
-
-_DEFAULT_SCALE_CLAMP = math.log(1000.0 / 16)
-
-def convert_boxes_to_pooler_format(bboxes):
-    bs, num_proposals = bboxes.shape[:2]
-    sizes = torch.full((bs,), num_proposals).to(bboxes.device)
-    aggregated_bboxes = bboxes.view(bs * num_proposals, -1)
-    indices = torch.repeat_interleave(
-        torch.arange(len(sizes), dtype=aggregated_bboxes.dtype, device=aggregated_bboxes.device), sizes
-    )
-    return torch.cat([indices[:, None], aggregated_bboxes], dim=1)
-
-
-def assign_boxes_to_levels(
-        bboxes,
-        min_level,
-        max_level,
-        canonical_box_size,
-        canonical_level,
-):
-    aggregated_bboxes = bboxes.view(bboxes.shape[0] * bboxes.shape[1], -1)
-    area = (aggregated_bboxes[:, 2] - aggregated_bboxes[:, 0]) * (aggregated_bboxes[:, 3] - aggregated_bboxes[:, 1])
-    box_sizes = torch.sqrt(area)
-    # Eqn.(1) in FPN paper
-    level_assignments = torch.floor(canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8))
-    # clamp level to (min, max), in case the box size is too large or too small
-    # for the available feature maps
-    level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level)
-    return level_assignments.to(torch.int64) - min_level
-
-
-class SinusoidalPositionEmbeddings(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, time):
-        device = time.device
-        half_dim = self.dim // 2
-        embeddings = math.log(10000) / (half_dim - 1)
-        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
-        embeddings = time[:, None] * embeddings[None, :]
-        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
-        return embeddings
-
-
-class HeadDynamicK(nn.Module):
-    def __init__(self, config, roi_input_shape):
-        super().__init__()
-        num_classes = config.num_labels
-
-        ddet_head = DiffusionDetHead(config, roi_input_shape, num_classes)
-        self.num_head = config.num_heads
-        self.head_series = nn.ModuleList([copy.deepcopy(ddet_head) for _ in range(self.num_head)])
-        self.return_intermediate = config.deep_supervision
-
-        # Gaussian random feature embedding layer for time
-        self.hidden_dim = config.hidden_dim
-        time_dim = self.hidden_dim * 4
-        self.time_mlp = nn.Sequential(
-            SinusoidalPositionEmbeddings(self.hidden_dim),
-            nn.Linear(self.hidden_dim, time_dim),
-            nn.GELU(),
-            nn.Linear(time_dim, time_dim),
-        )
-
-        # Init parameters.
-        self.use_focal = config.use_focal
-        self.use_fed_loss = config.use_fed_loss
-        self.num_classes = num_classes
-        if self.use_focal or self.use_fed_loss:
-            prior_prob = config.prior_prob
-            self.bias_value = -math.log((1 - prior_prob) / prior_prob)
-        self._reset_parameters()
-
-    def _reset_parameters(self):
-        # init all parameters.
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-            # initialize the bias for focal loss and fed loss.
-            if self.use_focal or self.use_fed_loss:
-                if p.shape[-1] == self.num_classes or p.shape[-1] == self.num_classes + 1:
-                    nn.init.constant_(p, self.bias_value)
-
-
-    def forward(self, features, bboxes, t):
-        # assert t shape (batch_size)
-        time = self.time_mlp(t)
-
-        inter_class_logits = []
-        inter_pred_bboxes = []
-
-        bs = len(features[0])
-
-        class_logits, pred_bboxes = None, None
-        for head_idx, ddet_head in enumerate(self.head_series):
-            class_logits, pred_bboxes, proposal_features = ddet_head(features, bboxes, time)
-            if self.return_intermediate:
-                inter_class_logits.append(class_logits)
-                inter_pred_bboxes.append(pred_bboxes)
-            bboxes = pred_bboxes.detach()
-
-        if self.return_intermediate:
-            return torch.stack(inter_class_logits), torch.stack(inter_pred_bboxes)
-
-        return class_logits[None], pred_bboxes[None]
-
-
-class DynamicConv(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-
-        self.hidden_dim = config.hidden_dim
-        self.dim_dynamic = config.dim_dynamic
-        self.num_dynamic = config.num_dynamic
-        self.num_params = self.hidden_dim * self.dim_dynamic
-        self.dynamic_layer = nn.Linear(self.hidden_dim, self.num_dynamic * self.num_params)
-
-        self.norm1 = nn.LayerNorm(self.dim_dynamic)
-        self.norm2 = nn.LayerNorm(self.hidden_dim)
-
-        self.activation = nn.ReLU(inplace=True)
-
-        pooler_resolution = config.pooler_resolution
-        num_output = self.hidden_dim * pooler_resolution ** 2
-        self.out_layer = nn.Linear(num_output, self.hidden_dim)
-        self.norm3 = nn.LayerNorm(self.hidden_dim)
-
-
-    def forward(self, pro_features, roi_features):
-        features = roi_features.permute(1, 0, 2)
-        parameters = self.dynamic_layer(pro_features).permute(1, 0, 2)
-
-        param1 = parameters[:, :, :self.num_params].view(-1, self.hidden_dim, self.dim_dynamic)
-        param2 = parameters[:, :, self.num_params:].view(-1, self.dim_dynamic, self.hidden_dim)
-
-        features = torch.bmm(features, param1)
-        features = self.norm1(features)
-        features = self.activation(features)
-
-        features = torch.bmm(features, param2)
-        features = self.norm2(features)
-        features = self.activation(features)
-
-        features = features.flatten(1)
-        features = self.out_layer(features)
-        features = self.norm3(features)
-        features = self.activation(features)
-
-        return features
-
-
-class DiffusionDetHead(nn.Module):
-    def __init__(self, config, roi_input_shape, num_classes):
-        super().__init__()
-
-        dim_feedforward = config.dim_feedforward
-        nhead = config.num_attn_heads
-        dropout = config.dropout
-        activation = config.activation
-        in_features = config.roi_head_in_features
-        pooler_resolution = config.pooler_resolution
-        pooler_scales = tuple(1.0 / roi_input_shape[k]['stride'] for k in in_features)
-        sampling_ratio = config.sampling_ratio
-
-        self.hidden_dim = config.hidden_dim
-
-        self.pooler = ROIPooler(
-            output_size=pooler_resolution,
-            scales=pooler_scales,
-            sampling_ratio=sampling_ratio,
-        )
-
-        # dynamic.
-        self.self_attn = nn.MultiheadAttention(self.hidden_dim, nhead, dropout=dropout)
-        self.inst_interact = DynamicConv(config)
-
-        self.linear1 = nn.Linear(self.hidden_dim, dim_feedforward)
-        self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(dim_feedforward, self.hidden_dim)
-
-        self.norm1 = nn.LayerNorm(self.hidden_dim)
-        self.norm2 = nn.LayerNorm(self.hidden_dim)
-        self.norm3 = nn.LayerNorm(self.hidden_dim)
-        self.dropout1 = nn.Dropout(dropout)
-        self.dropout2 = nn.Dropout(dropout)
-        self.dropout3 = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-
-        # block time mlp
-        self.block_time_mlp = nn.Sequential(nn.SiLU(), nn.Linear(self.hidden_dim * 4, self.hidden_dim * 2))
-
-        # cls.
-        num_cls = config.num_cls
-        cls_module = list()
-        for _ in range(num_cls):
-            cls_module.append(nn.Linear(self.hidden_dim, self.hidden_dim, False))
-            cls_module.append(nn.LayerNorm(self.hidden_dim))
-            cls_module.append(nn.ReLU(inplace=True))
-        self.cls_module = nn.ModuleList(cls_module)
-
-        # reg.
-        num_reg = config.num_reg
-        reg_module = list()
-        for _ in range(num_reg):
-            reg_module.append(nn.Linear(self.hidden_dim, self.hidden_dim, False))
-            reg_module.append(nn.LayerNorm(self.hidden_dim))
-            reg_module.append(nn.ReLU(inplace=True))
-        self.reg_module = nn.ModuleList(reg_module)
-
-        # pred.
-        self.use_focal = config.use_focal
-        self.use_fed_loss = config.use_fed_loss
-        if self.use_focal or self.use_fed_loss:
-            self.class_logits = nn.Linear(self.hidden_dim, num_classes)
-        else:
-            self.class_logits = nn.Linear(self.hidden_dim, num_classes + 1)
-        self.bboxes_delta = nn.Linear(self.hidden_dim, 4)
-        self.scale_clamp = _DEFAULT_SCALE_CLAMP
-        self.bbox_weights = (2.0, 2.0, 1.0, 1.0)
-
-    def forward(self, features, bboxes, time_emb):
-        bs, num_proposals = bboxes.shape[:2]
-
-        # roi_feature.
-        roi_features = self.pooler(features, bboxes)
-
-        pro_features = roi_features.view(bs, num_proposals, self.hidden_dim, -1).mean(-1)
-
-        roi_features = roi_features.view(bs * num_proposals, self.hidden_dim, -1).permute(2, 0, 1)
-
-        # self_att.
-        pro_features = pro_features.view(bs, num_proposals, self.hidden_dim).permute(1, 0, 2)
-        pro_features2 = self.self_attn(pro_features, pro_features, value=pro_features)[0]
-        pro_features = pro_features + self.dropout1(pro_features2)
-        pro_features = self.norm1(pro_features)
-
-        # inst_interact.
-        pro_features = pro_features.view(num_proposals, bs, self.hidden_dim).permute(1, 0, 2).reshape(1, bs * num_proposals,
-                                                                                      self.hidden_dim)
-        pro_features2 = self.inst_interact(pro_features, roi_features)
-        pro_features = pro_features + self.dropout2(pro_features2)
-        obj_features = self.norm2(pro_features)
-
-        # obj_feature.
-        obj_features2 = self.linear2(self.dropout(self.activation(self.linear1(obj_features))))
-        obj_features = obj_features + self.dropout3(obj_features2)
-        obj_features = self.norm3(obj_features)
-
-        fc_feature = obj_features.transpose(0, 1).reshape(bs * num_proposals, -1)
-
-        scale_shift = self.block_time_mlp(time_emb)
-        scale_shift = torch.repeat_interleave(scale_shift, num_proposals, dim=0)
-        scale, shift = scale_shift.chunk(2, dim=1)
-        fc_feature = fc_feature * (scale + 1) + shift
-
-        cls_feature = fc_feature.clone()
-        reg_feature = fc_feature.clone()
-        for cls_layer in self.cls_module:
-            cls_feature = cls_layer(cls_feature)
-        for reg_layer in self.reg_module:
-            reg_feature = reg_layer(reg_feature)
-        class_logits = self.class_logits(cls_feature)
-        bboxes_deltas = self.bboxes_delta(reg_feature)
-        pred_bboxes = self.apply_deltas(bboxes_deltas, bboxes.view(-1, 4))
-
-        return class_logits.view(bs, num_proposals, -1), pred_bboxes.view(bs, num_proposals, -1), obj_features
-
-    def apply_deltas(self, deltas, boxes):
-        """
-        Apply transformation `deltas` (dx, dy, dw, dh) to `boxes`.
-
-        Args:
-            deltas (Tensor): transformation deltas of shape (N, k*4), where k >= 1.
-                deltas[i] represents k potentially different class-specific
-                box transformations for the single box boxes[i].
-            boxes (Tensor): boxes to transform, of shape (N, 4)
-        """
-        boxes = boxes.to(deltas.dtype)
-
-        widths = boxes[:, 2] - boxes[:, 0]
-        heights = boxes[:, 3] - boxes[:, 1]
-        ctr_x = boxes[:, 0] + 0.5 * widths
-        ctr_y = boxes[:, 1] + 0.5 * heights
-
-        wx, wy, ww, wh = self.bbox_weights
-        dx = deltas[:, 0::4] / wx
-        dy = deltas[:, 1::4] / wy
-        dw = deltas[:, 2::4] / ww
-        dh = deltas[:, 3::4] / wh
-
-        # Prevent sending too large values into torch.exp()
-        dw = torch.clamp(dw, max=self.scale_clamp)
-        dh = torch.clamp(dh, max=self.scale_clamp)
-
-        pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
-        pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
-        pred_w = torch.exp(dw) * widths[:, None]
-        pred_h = torch.exp(dh) * heights[:, None]
-
-        pred_boxes = torch.zeros_like(deltas)
-        pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w  # x1
-        pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h  # y1
-        pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w  # x2
-        pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h  # y2
-
-        return pred_boxes
-
-
-class ROIPooler(nn.Module):
-    """
-    Region of interest feature map pooler that supports pooling from one or more
-    feature maps.
-    """
-
-    def __init__(
-            self,
-            output_size,
-            scales,
-            sampling_ratio,
-            canonical_box_size=224,
-            canonical_level=4,
-    ):
-        super().__init__()
-
-        min_level = -(math.log2(scales[0]))
-        max_level = -(math.log2(scales[-1]))
-
-        if isinstance(output_size, int):
-            output_size = (output_size, output_size)
-        assert len(output_size) == 2 and isinstance(output_size[0], int) and isinstance(output_size[1], int)
-        assert math.isclose(min_level, int(min_level)) and math.isclose(max_level, int(max_level))
-        assert (len(scales) == max_level - min_level + 1)
-        assert 0 <= min_level <= max_level
-        assert canonical_box_size > 0
-
-        self.output_size = output_size
-        self.min_level = int(min_level)
-        self.max_level = int(max_level)
-        self.canonical_level = canonical_level
-        self.canonical_box_size = canonical_box_size
-        self.level_poolers = nn.ModuleList(
-            RoIAlign(
-                output_size, spatial_scale=scale, sampling_ratio=sampling_ratio, aligned=True
-            )
-            for scale in scales
-        )
-
-    def forward(self, x, bboxes):
-        num_level_assignments = len(self.level_poolers)
-        assert len(x) == num_level_assignments and len(bboxes) == x[0].size(0)
-
-        pooler_fmt_boxes = convert_boxes_to_pooler_format(bboxes)
-
-        if num_level_assignments == 1:
-            return self.level_poolers[0](x[0], pooler_fmt_boxes)
-
-        level_assignments = assign_boxes_to_levels(
-            bboxes, self.min_level, self.max_level, self.canonical_box_size, self.canonical_level
-        )
-
-        batches = pooler_fmt_boxes.shape[0]
-        channels = x[0].shape[1]
-        output_size = self.output_size[0]
-        sizes = (batches, channels, output_size, output_size)
-
-        output = torch.zeros(sizes, dtype=x[0].dtype, device=x[0].device)
-
-        for level, (x_level, pooler) in enumerate(zip(x, self.level_poolers)):
-            inds = (level_assignments == level).nonzero(as_tuple=True)[0]
-            pooler_fmt_boxes_level = pooler_fmt_boxes[inds]
-            # Use index_put_ instead of advance indexing, to avoid pytorch/issues/49852
-            output.index_put_((inds,), pooler(x_level, pooler_fmt_boxes_level))
-
-        return output