o  iI@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_optioncspeZdZdZdeejiZ   d*d+fdd Zej dfd,ddZ d-ddZ ddZ d.d#d$Z d/d(d)ZZS)0AnchorGeneratora 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)())?)rrijsizesSequence[Sequence[int]] aspect_ratiosrindexingstrreturnNonecstt|dstdd|D_nt|_t|ds)|ftj}tjt|kr6tdtt|ddd}t|ddg}|_ t|dd g_ |_ fd d t j|D_ dS) Nrcss|]}|fVqdSN.0srri/home/dell461/cl/sdc2/last_ska_mid/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)rrrr0 __class__r(r r,os$       zAnchorGenerator.__init__Nscalesdtype torch.dtypedevicetorch.device | None torch.TensorcCstj|||d}tj|||d}|jdkr1t|jdkr1td|jdt|jdd|jd|jdkrS|jd|jdkrStd|jd |jdd |jd|jdkrdt|}d|}|} n,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} | S| 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. r5r7r#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_tensorr0r.shaper/sqrtpowviewstackround)r)r4rr5r7Zscales_tZaspect_ratios_tZ area_scalew_ratiosh_ratiosZd_ratioswshs base_anchorsdsrrr r%sD     (&& &$z AnchorGenerator.generate_anchors torch.devicecsfdd|jD|_dS)z` Convert each element in self.cell_anchors to ``dtype`` and send to ``device``. csg|] }|jdqSr:)to)r cell_anchorr7r5rr r*z4AnchorGenerator.set_cell_anchors..Nr )r)r5r7rrRr set_cell_anchorssz AnchorGenerator.set_cell_anchorscCsdd|jDS)zF Return number of anchor shapes for each feature map. cSsg|]}|jdqS)r)rB)rcrrr r*sz.rTr(rrr num_anchors_per_locationsz(AnchorGenerator.num_anchors_per_location grid_sizeslist[list[int]]strideslist[list[Tensor]] list[Tensor]c s4g}|j}|dur tt|t|krt|ks#tdtdt|||D]n\}|jfddt|jD}tt j |d|jdd}t|jD] }|| d||<qR|j dkrp|d |d |d <|d <t j |d d d }||dd |jd |d d|jd  d|jd q)|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)rr:)r@arangeint32raxisr7r&striderr r*sz0AnchorGenerator.grid_anchors..r)rr<r$rrr"r=)r AssertionErrorr.r/r1r7ranger0listr@meshgridreshaperrFappendrE) r)rXrZanchorsr rLZshifts_centersr`shiftsrrar grid_anchorss: ( zAnchorGenerator.grid_anchorsimagesr feature_mapsc 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)rerBr0)r feature_mapr(rr r*1s"z+AnchorGenerator.forward..Nrcs(g|]fddtjDqS)cs*g|]}tj||tjdqSrO)r@tensorint64r_)r7g image_sizerr r*6sz6AnchorGenerator.forward...)rdr0)rr7rrr))rqr r*5s ) rBr0r5r7rUrkr@catre) r)rlrmrX batchsizer5rZanchors_over_all_feature_mapsanchors_per_imagerrsr forwards    zAnchorGenerator.forward)r rr)rrrrrrrr) r4rrrr5r6r7r8rr9)r5r6r7rNrr)rXrYrZr[rr\)rlrrmr\rr\)__name__ __module__ __qualname____doc__rr@r__annotations__r,float32r%rUrWrkrx __classcell__rrr2r r 4s8) < @r c@sFeZdZdZdeejiZ   ddddZe ej dfdddZ dS)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 rrfeature_map_scalesSequence[int] | Sequence[float]base_anchor_shapes3Sequence[Sequence[int]] | Sequence[Sequence[float]]rrrrcs`tjt|d}t|ddg}|_t|ddg_t|fdd|D_ dS)Nrr"r#rr$csg|] }|qSr)generate_anchors_using_shaperZbase_anchor_shapes_tr)rr r*rSz;AnchorGeneratorWithAnchorShape.__init__..) rModuler,r.r r0rr@rr )r)rrrr0rrr r,qs  z'AnchorGeneratorWithAnchorShape.__init__N anchor_shapesr9r5r6r7r8cCs.|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=r:)r@rtrGrP)rr5r7Zhalf_anchor_shapesrLrrr rsz;AnchorGeneratorWithAnchorShape.generate_anchors_using_shape)rrr)rrrrrrrr)rr9r5r6r7r8rr9) ryrzr{r|rr@rr}r, staticmethodr~rrrrr rDs* r)r| __future__rtypingrrr@rr monai.utilsrmonai.utils.miscrmonai.utils.moduler rr rrrrr s"