U Ph,I@sdZddlmZddlmZmZddlZddlmZmZddl m Z ddl m Z ddl mZGd d d ejZGd d d eZdS) z~ This script is adapted from https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py ) annotations)ListSequenceN)Tensornn) ensure_tuple)issequenceiterable)look_up_optioncseZdZdZdeejiZd$dddd d fd d Zej d fddddddddZ ddd dddZ ddZ ddddddZ d ddd!d"d#ZZS)%AnchorGeneratora This module is modified from torchvision to support both 2D and 3D images. Module that generates anchors for a set of feature maps and image sizes. The module support computing anchors at multiple sizes and aspect ratios per feature map. sizes and aspect_ratios should have the same number of elements, and it should correspond to the number of feature maps. sizes[i] and aspect_ratios[i] can have an arbitrary number of elements. For 2D images, anchor width and height w:h = 1:aspect_ratios[i,j] For 3D images, anchor width, height, and depth w:h:d = 1:aspect_ratios[i,j,0]:aspect_ratios[i,j,1] AnchorGenerator will output a set of sizes[i] * aspect_ratios[i] anchors per spatial location for feature map i. Args: sizes: base size of each anchor. len(sizes) is the number of feature maps, i.e., the number of output levels for the feature pyramid network (FPN). Each element of ``sizes`` is a Sequence which represents several anchor sizes for each feature map. aspect_ratios: the aspect ratios of anchors. ``len(aspect_ratios) = len(sizes)``. For 2D images, each element of ``aspect_ratios[i]`` is a Sequence of float. For 3D images, each element of ``aspect_ratios[i]`` is a Sequence of 2 value Sequence. indexing: choose from {``'ij'``, ``'xy'``}, optional, Matrix (``'ij'``, default and recommended) or Cartesian (``'xy'``) indexing of output. - Matrix (``'ij'``, default and recommended) indexing keeps the original axis not changed. - To use other monai detection components, please set ``indexing = 'ij'``. - Cartesian (``'xy'``) indexing swaps axis 0 and 1. - For 2D cases, monai ``AnchorGenerator(sizes, aspect_ratios, indexing='xy')`` and ``torchvision.models.detection.anchor_utils.AnchorGenerator(sizes, aspect_ratios)`` are equivalent. Reference:. https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py Example: .. code-block:: python # 2D example inputs for a 2-level feature maps sizes = ((10,12,14,16), (20,24,28,32)) base_aspect_ratios = (1., 0.5, 2.) aspect_ratios = (base_aspect_ratios, base_aspect_ratios) anchor_generator = AnchorGenerator(sizes, aspect_ratios) # 3D example inputs for a 2-level feature maps sizes = ((10,12,14,16), (20,24,28,32)) base_aspect_ratios = ((1., 1.), (1., 0.5), (0.5, 1.), (2., 2.)) aspect_ratios = (base_aspect_ratios, base_aspect_ratios) anchor_generator = AnchorGenerator(sizes, aspect_ratios) cell_anchors)())?)rrijzSequence[Sequence[int]]rstrNone)sizes aspect_ratiosindexingreturncstt|ds,tdd|D_n t|_t|dsR|ftj}tjt|krltdtt|ddd}t|ddg}|_ t|dd g_ |_ fd d t j|D_ dS) Nrcss|] }|fVqdSN.0srr\/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/detection/utils/anchor_utils.py xsz+AnchorGenerator.__init__..zllen(sizes) and len(aspect_ratios) should be equal. It represents the number of feature maps.rrxycsg|]\}}||qSr)generate_anchors)rsize aspect_ratioselfrr sz,AnchorGenerator.__init__..)super__init__rtuplerrlen ValueErrorr spatial_dimsrrzipr )r(rrrr/ __class__r'rr+os$      zAnchorGenerator.__init__N torch.dtypetorch.device | None torch.Tensor)scalesrdtypedevicercCstj|||d}tj|||d}|jdkrbt|jdkrbtd|jdt|jdd|jd|jdkr|jd|jdkrtd|jd |jdd |jd|jdkrt|}d|}|} nd|jdkr,t|d d d f|d d dfd }d|}|d d d f|} |d d df|} |d d d f|d d d fd} | d d d f|d d d fd} |jdkrtj | | | | gddd} nV|jdkr| d d d f|d d d fd}tj | | | | | |gddd} | S)a Compute cell anchor shapes at multiple sizes and aspect ratios for the current feature map. Args: scales: a sequence which represents several anchor sizes for the current feature map. aspect_ratios: a sequence which represents several aspect_ratios for the current feature map. For 2D images, it is a Sequence of float aspect_ratios[j], anchor width and height w:h = 1:aspect_ratios[j]. For 3D images, it is a Sequence of 2 value Sequence aspect_ratios[j,0] and aspect_ratios[j,1], anchor width, height, and depth w:h:d = 1:aspect_ratios[j,0]:aspect_ratios[j,1] dtype: target data type of the output Tensor. device: target device to put the output Tensor data. Returns: For each s in scales, returns [s, s*aspect_ratios[j]] for 2D images, and [s, s*aspect_ratios[j,0],s*aspect_ratios[j,1]] for 3D images. r7r8r"r!zIn zA-D image, aspect_ratios for each level should be rz%-D. But got aspect_ratios with shape .zK-D image, aspect_ratios for each level should has shape (_,z$). But got aspect_ratios with shape NrgUUUUUU?dim@) torch as_tensorr/r-shaper.sqrtpowviewstackround)r(r6rr7r8Zscales_tZaspect_ratios_tZ area_scalew_ratiosh_ratiosZd_ratioswshs base_anchorsdsrrrr$sF      (&&   &$z AnchorGenerator.generate_anchorsz torch.device)r7r8rcsfdd|jD|_dS)z` Convert each element in self.cell_anchors to ``dtype`` and send to ``device``. csg|]}|jdqSr9)to)r cell_anchorr8r7rrr)sz4AnchorGenerator.set_cell_anchors..Nr )r(r7r8rrPrset_cell_anchorssz AnchorGenerator.set_cell_anchorscCsdd|jDS)zF Return number of anchor shapes for each feature map. cSsg|]}|jdqS)r)rA)rcrrrr)sz.rQr'rrrnum_anchors_per_locationsz(AnchorGenerator.num_anchors_per_locationzlist[list[int]]zlist[list[Tensor]]z list[Tensor]) grid_sizesstridesrc s.g}|j}|dkrtt|t|kr6t|ks@ntdt|||D]\}|jfddt|jD}tt j |d|jdd}t|jD]}|| d||<q|j dkr|d |d |d <|d <t j |d d d }||dd |jd |d d|jd  d|jd qL|S) ai Every combination of (a, (g, s), i) in (self.cell_anchors, zip(grid_sizes, strides), 0:spatial_dims) corresponds to a feature map. It outputs g[i] anchors that are s[i] distance apart in direction i, with the same dimensions as a. Args: grid_sizes: spatial size of the feature maps strides: strides of the feature maps regarding to the original image Example: .. code-block:: python grid_sizes = [[100,100],[50,50]] strides = [[torch.tensor(2),torch.tensor(2)], [torch.tensor(4),torch.tensor(4)]] NzAnchors should be Tuple[Tuple[int]] because each feature map could potentially have different sizes and aspect ratios. There needs to be a match between the number of feature maps passed and the number of sizes / aspect ratios specified.cs,g|]$}tjd|tjd|qS)rr9)r?arangeint32raxisr8r%striderrr)sz0AnchorGenerator.grid_anchors..r)rr;r#rrr!r<)r AssertionErrorr-r.r0r8ranger/listr?meshgridreshaperrEappendrD) r(rUrVanchorsr rKZshifts_centersrZshiftsrr[r grid_anchorss4" (zAnchorGenerator.grid_anchorsr)images feature_mapsrc sfdd|D}|jj d|jd}|dj|dj}fdd|D}|||}tt|}|g|S)aF Generate anchor boxes for each image. Args: images: sized (B, C, W, H) or (B, C, W, H, D) feature_maps: for FPN level i, feature_maps[i] is sized (B, C_i, W_i, H_i) or (B, C_i, W_i, H_i, D_i). This input argument does not have to be the actual feature maps. Any list variable with the same (C_i, W_i, H_i) or (C_i, W_i, H_i, D_i) as feature maps works. Return: A list with length of B. Each element represents the anchors for this image. The B elements are identical. Example: .. code-block:: python images = torch.zeros((3,1,128,128,128)) feature_maps = [torch.zeros((3,6,64,64,32)), torch.zeros((3,6,32,32,16))] anchor_generator(images, feature_maps) cs"g|]}t|jj dqSr)r_rAr/)r feature_mapr'rrr)1sz+AnchorGenerator.forward..Nrcs(g|] fddtjDqS)cs*g|]"}tj||tjdqSrM)r?tensorint64rY)r8g image_sizerrr)6sz6AnchorGenerator.forward...)r^r/)rr8rlr()rkrr)5s) rAr/r7r8rRrer?catr_) r(rfrgrU batchsizer7rVanchors_over_all_feature_mapsanchors_per_imagerrmrforwards   zAnchorGenerator.forward)r rr)__name__ __module__ __qualname____doc__rr?r__annotations__r+float32r$rRrTrerr __classcell__rrr1rr 4s8)<@r c@sTeZdZdZdeejiZddddd d d d Ze ej d fdddddddZ d S)AnchorGeneratorWithAnchorShapea Module that generates anchors for a set of feature maps and image sizes, inherited from :py:class:`~monai.apps.detection.networks.utils.anchor_utils.AnchorGenerator` The module support computing anchors at multiple base anchor shapes per feature map. ``feature_map_scales`` should have the same number of elements with the number of feature maps. base_anchor_shapes can have an arbitrary number of elements. For 2D images, each element represents anchor width and height [w,h]. For 2D images, each element represents anchor width, height, and depth [w,h,d]. AnchorGenerator will output a set of ``len(base_anchor_shapes)`` anchors per spatial location for feature map ``i``. Args: feature_map_scales: scale of anchors for each feature map, i.e., each output level of the feature pyramid network (FPN). ``len(feature_map_scales)`` is the number of feature maps. ``scale[i]*base_anchor_shapes`` represents the anchor shapes for feature map ``i``. base_anchor_shapes: a sequence which represents several anchor shapes for one feature map. For N-D images, it is a Sequence of N value Sequence. indexing: choose from {'xy', 'ij'}, optional Cartesian ('xy') or matrix ('ij', default) indexing of output. Cartesian ('xy') indexing swaps axis 0 and 1, which is the setting inside torchvision. matrix ('ij', default) indexing keeps the original axis not changed. See also indexing in https://pytorch.org/docs/stable/generated/torch.meshgrid.html Example: .. code-block:: python # 2D example inputs for a 2-level feature maps feature_map_scales = (1, 2) base_anchor_shapes = ((10, 10), (6, 12), (12, 6)) anchor_generator = AnchorGeneratorWithAnchorShape(feature_map_scales, base_anchor_shapes) # 3D example inputs for a 2-level feature maps feature_map_scales = (1, 2) base_anchor_shapes = ((10, 10, 10), (12, 12, 8), (10, 10, 6), (16, 16, 10)) anchor_generator = AnchorGeneratorWithAnchorShape(feature_map_scales, base_anchor_shapes) r rr!) rr)0r r )r rr )r r rrzSequence[int] | Sequence[float]z3Sequence[Sequence[int]] | Sequence[Sequence[float]]rr)feature_map_scalesbase_anchor_shapesrrcs`tjt|d}t|ddg}|_t|ddg_t|fdd|D_ dS)Nrr!r"rr#csg|]}|qSr)generate_anchors_using_shaperZbase_anchor_shapes_tr(rrr)sz;AnchorGeneratorWithAnchorShape.__init__..) rModuler+r-r r/rr?rr )r(rrrr/rrrr+qs   z'AnchorGeneratorWithAnchorShape.__init__Nr5r3r4) anchor_shapesr7r8rcCs.|d}tj| |gdd}|j||dS)a Compute cell anchor shapes at multiple sizes and aspect ratios for the current feature map. Args: anchor_shapes: [w, h] or [w, h, d], sized (N, spatial_dims), represents N anchor shapes for the current feature map. dtype: target data type of the output Tensor. device: target device to put the output Tensor data. Returns: For 2D images, returns [-w/2, -h/2, w/2, h/2]; For 3D images, returns [-w/2, -h/2, -d/2, w/2, h/2, d/2] r>rr<r9)r?rnrFrN)rr7r8Zhalf_anchor_shapesrKrrrrsz;AnchorGeneratorWithAnchorShape.generate_anchors_using_shape)r{r~r) rsrtrurvrr?rrwr+ staticmethodrxrrrrrrzDs*rz)rv __future__rtypingrrr?rr monai.utilsrmonai.utils.miscrmonai.utils.moduler rr rzrrrr#s