python-wml 3.0.0__py3-none-any.whl

wml/wtorch/nets/simple_fpn.py

@@ -0,0 +1,101 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import math
+import fvcore.nn.weight_init as weight_init
+import torch
+import torch.nn.functional as F
+from torch import nn
+from wml.wtorch.nn import Conv2d
+from wml.wtorch.nn import get_norm
+
+
+class SimpleFPN(torch.nn.Module):
+    """
+    This module implements :paper:`FPN`.
+    It creates pyramid features built on top of some input feature maps.
+    """
+
+    _fuse_type: torch.jit.Final[str]
+
+    def __init__( self, in_channels, out_channels, norm="",  fuse_type="sum",interpolate_mode="nearest"
+    ):
+        """
+        Args:
+            out_channels (int): number of channels in the output feature maps.
+            norm (str): the normalization to use.
+            fuse_type (str): types for fusing the top down features and the lateral
+                ones. It can be "sum" (default), which sums up element-wise; or "avg",
+                which takes the element-wise mean of the two.
+        """
+        super().__init__()
+
+        # Feature map strides and channels from the bottom up network (e.g. ResNet)
+        in_channels_per_feature = in_channels
+        self.interpolate_mode = interpolate_mode
+
+        lateral_convs = []
+        output_convs = []
+
+        use_bias = norm == ""
+        for idx, in_channels in enumerate(in_channels_per_feature):
+            lateral_norm = get_norm(norm, out_channels)
+            output_norm = get_norm(norm, out_channels)
+
+            lateral_conv = Conv2d(
+                in_channels, out_channels, kernel_size=1, bias=use_bias, norm=lateral_norm
+            )
+            output_conv = Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=use_bias,
+                norm=output_norm,
+            )
+            weight_init.c2_xavier_fill(lateral_conv)
+            weight_init.c2_xavier_fill(output_conv)
+            self.add_module("fpn_lateral{}".format(idx), lateral_conv)
+            self.add_module("fpn_output{}".format(idx), output_conv)
+
+            lateral_convs.append(lateral_conv)
+            output_convs.append(output_conv)
+        # Place convs into top-down order (from low to high resolution)
+        # to make the top-down computation in forward clearer.
+        self.lateral_convs = lateral_convs[::-1]
+        self.output_convs = output_convs[::-1]
+        # Return feature names are "p<stage>", like ["p2", "p3", ..., "p6"]
+
+        assert fuse_type in {"avg", "sum"}
+        self._fuse_type = fuse_type
+
+    def forward(self, xs):
+        """
+        Args:
+            xs: list of tensors, shape like [H,W],[H//2,W//2],[H//4,W//4],...
+
+        Returns:
+            list of tensors
+        """
+        results = []
+        prev_features = self.lateral_convs[0](xs[-1])
+        results.append(self.output_convs[0](prev_features))
+
+        # Reverse feature maps into top-down order (from low to high resolution)
+        for idx, (lateral_conv, output_conv) in enumerate(
+            zip(self.lateral_convs, self.output_convs)
+        ):
+            # Slicing of ModuleList is not supported https://github.com/pytorch/pytorch/issues/47336
+            # Therefore we loop over all modules but skip the first one
+            if idx > 0:
+                features = xs[-idx - 1]
+                if self.interpolate_mode == "nearest":
+                    top_down_features = F.interpolate(prev_features, scale_factor=2.0, mode="nearest")
+                else:
+                    top_down_features = F.interpolate(prev_features, size=features.shape[-2:],mode="bilinear")
+                lateral_features = lateral_conv(features)
+                prev_features = lateral_features + top_down_features
+                if self._fuse_type == "avg":
+                    prev_features /= 2
+                results.insert(0, output_conv(prev_features))
+
+        return results

wml/wtorch/nms.py

@@ -0,0 +1,109 @@
+import time
+import numpy as np
+import torch
+import torchvision
+from typing import Tuple
+from torch import Tensor
+
+
+def __soft_nms(dets, box_scores, thresh=0.001, sigma=0.5):
+    """
+    Build a pytorch implement of Soft NMS algorithm.
+    # Augments
+        dets:        boxes coordinate tensor (format:[y1, x1, y2, x2])
+        box_scores:  box score tensors
+        sigma:       variance of Gaussian function, decay=exp(-(iou*iou)/sigma), smaller sigma, decay faster.
+        thresh:      score thresh
+        cuda:        CUDA flag
+    # Return
+        the index of the selected boxes
+    """
+
+    # Indexes concatenate boxes with the last column
+    N = dets.shape[0]
+    device = dets.device
+    indexes = torch.arange(0, N, dtype=torch.float).view(N, 1)
+    indexes = indexes.to(device)
+    dets = torch.cat((dets, indexes), dim=1)
+
+    # The order of boxes coordinate is [y1,x1,y2,x2]
+    y1 = dets[:, 0]
+    x1 = dets[:, 1]
+    y2 = dets[:, 2]
+    x2 = dets[:, 3]
+    scores = box_scores
+    areas = (x2 - x1) * (y2 - y1)
+
+    for i in range(N):
+        # intermediate parameters for later parameters exchange
+        tscore = scores[i].clone()
+        pos = i + 1
+
+        if i != N - 1:
+            maxscore, maxpos = torch.max(scores[pos:], dim=0)
+            if tscore < maxscore:
+                dets[i], dets[maxpos + i + 1] = dets[maxpos + i + 1].clone(), dets[i].clone()
+                scores[i], scores[maxpos + i + 1] = scores[maxpos + i + 1].clone(), scores[i].clone()
+                areas[i], areas[maxpos + i + 1] = areas[maxpos + i + 1].clone(), areas[i].clone()
+
+        # IoU calculate
+        yy1 = torch.maximum(dets[i, 0][None], dets[pos:, 0])
+        xx1 = torch.maximum(dets[i, 1][None], dets[pos:, 1])
+        yy2 = torch.minimum(dets[i, 2][None], dets[pos:, 2])
+        xx2 = torch.minimum(dets[i, 3][None], dets[pos:, 3])
+
+        zeros = torch.zeros_like(yy1)
+        w = torch.maximum(zeros, xx2 - xx1)
+        h = torch.maximum(zeros, yy2 - yy1)
+        inter = w * h
+        ious = torch.div(inter, (areas[i] + areas[pos:] - inter))
+
+        # Gaussian decay
+        weight = torch.exp(-(ious* ious) / sigma)
+        scores[pos:] = weight * scores[pos:]
+
+    # select the boxes and keep the corresponding indexes
+    valid_mask = scores>thresh
+    keep = dets[:, 4][valid_mask].long()
+    valid_scores = scores[valid_mask]
+
+    return keep,valid_scores
+
+def soft_nms(dets, box_scores, thresh=0.001, sigma=0.5):
+    keep,_ = __soft_nms(dets,box_scores,thresh=thresh,sigma=sigma)
+    return keep
+
+def soft_nmsv2(dets, box_scores, thresh=0.001, sigma=0.5, cuda=0):
+    keep,scores  = __soft_nms(dets,box_scores,thresh=thresh,sigma=sigma)
+
+    return keep,scores
+
+def group_nms(bboxes,scores,ids,nms_threshold:float=0.7,max_value:float=20000.0):
+    '''
+        boxes (Tensor[N, 4])): boxes to perform NMS on. They
+            are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and
+            ``0 <= y1 < y2``.
+    '''
+    if max_value is None:
+        max_value = torch.max(bboxes)+1.0
+    bboxes = bboxes.float()
+    tmp_bboxes = bboxes+ids[:,None].to(bboxes.dtype)*max_value
+    idxs = torchvision.ops.nms(tmp_bboxes,scores,nms_threshold)
+    return idxs
+
+def nms(bboxes:Tensor,scores:Tensor,labels:Tensor,nms_threshold:float=0.5,max_num:int=1000)->Tuple[Tensor,Tensor,Tensor]:
+    """
+    boxes (Tensor[N, 4])): boxes to perform NMS on. They
+    are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and
+    ``0 <= y1 < y2``.
+    """
+    #if bboxes.numel() == 0:
+        #return bboxes.new_zeros([0,5]), labels,labels.new_zeros([0]).to(torch.int64)
+
+    keep = torchvision.ops.nms(bboxes,scores,iou_threshold=nms_threshold)
+    if max_num>0:
+        keep = keep[:max_num]
+
+    dets = torch.cat([bboxes[keep],scores[keep].unsqueeze(-1)],dim=-1)
+    
+    return dets,labels[keep],keep.to(torch.int64)