o  isF@s$ddlmZddlmZddlmZddlZddlZddl m Z m Z m Z ddl mZmZmZddlmZddlmZdd lmZdd lmZmZdd lmZmZdGdHddZdIddZdJddZdKd!d"Z dLdMd&d'Z! ( )dNdOd.d/Z"d(ej#ej$fdPd5d6Z%dQd:d;Z&dRd>d?Z' AdSdTdEdFZ(dS)U) annotations)Sequence)deepcopyN) DtypeLikeNdarrayOrTensor NdarrayTensor) COMPUTE_DTYPE TO_REMOVEget_spatial_dims)Resize) create_scale)look_up_option) ensure_tupleensure_tuple_rep)convert_data_typeconvert_to_dst_typeTpoints torch.Tensoraffine include_shiftboolreturncCst|d}|r6tj|tj|jdd|j|jdgdddd}t||}|d|ddfdd}|S|dd}t|d|d|f|}|dd}|S)a This internal function applies affine matrices to the point coordinate Args: points: point coordinates, Nx2 or Nx3 torch tensor or ndarray, representing [x, y] or [x, y, z] affine: affine matrix to be applied to the point coordinates, sized (spatial_dims+1,spatial_dims+1) include_shift: default True, whether the function apply translation (shift) in the affine transform Returns: transformed point coordinates, with same data type as ``points``, does not share memory with ``points`` )rr)devicedtypedimN) r torchcatonesshaperr transposematmul)rrr spatial_dimsZ points_affiner$i/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/detection/transforms/box_ops.py_apply_affine_to_pointss     r&boxesrrc Cst|tj^}}|jtd}t||d^}}t|d}t|ddd|f|dd}t|dd|df|dd}tjtj ||gdddd\}}tj tj ||gdddd\} }tj || gd d} t| |d^} }| S) a This function applies affine matrices to the boxes Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be StandardMode affine: affine matrix to be applied to the box coordinates, sized (spatial_dims+1,spatial_dims+1) Returns: returned affine transformed boxes, with same data type as ``boxes``, does not share memory with ``boxes`` rsrcdstr'NT)rrr) rrTensortorrr r&minstackmaxr) r'rboxes_t_Zaffine_tr#ltrbZlt_newZrb_newZboxes_t_affineZ boxes_affiner$r$r%apply_affine_to_boxes>s    r7zoomSequence[float] | floatcCs"t|d}t||d}t||dS)al Zoom boxes Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be StandardMode zoom: The zoom factor along the spatial axes. If a float, zoom is the same for each spatial axis. If a sequence, zoom should contain one value for each spatial axis. Returns: zoomed boxes, with same data type as ``boxes``, does not share memory with ``boxes`` Example: .. code-block:: python boxes = torch.ones(1,4) zoom_boxes(boxes, zoom=[0.5,2.2]) # will return tensor([[0.5, 2.2, 0.5, 2.2]]) r,)r#scaling_factor)r'r)r r r7)r'r8r#rr$r$r% zoom_boxesfs   r;src_spatial_sizeSequence[int] | intdst_spatial_sizecsBt|d}t|t|fddt|D}t||dS)a Resize boxes when the corresponding image is resized Args: boxes: source bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` src_spatial_size: source image spatial size. dst_spatial_size: target image spatial size. Returns: resized boxes, with same data type as ``boxes``, does not share memory with ``boxes`` Example: .. code-block:: python boxes = torch.ones(1,4) src_spatial_size = [100, 100] dst_spatial_size = [128, 256] resize_boxes(boxes, src_spatial_size, dst_spatial_size) # will return tensor([[1.28, 2.56, 1.28, 2.56]]) r,cs g|] }|t|qSr$)float.0axisr>r<r$r% s z resize_boxes..)r'r8)r rranger;)r'r<r>r#r8r$rCr% resize_boxess    rF spatial_size flip_axesSequence[int] | int | NonecCst|d}t||}|durttd|}t|}t|tjr#|nt |}|D].}|||dd|ft |dd||f<|||dd||ft |dd|f<q)|S)a Flip boxes when the corresponding image is flipped Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` spatial_size: image spatial size. flip_axes: spatial axes along which to flip over. Default is None. The default `axis=None` will flip over all of the axes of the input array. If axis is negative it counts from the last to the first axis. If axis is a tuple of ints, flipping is performed on all of the axes specified in the tuple. Returns: flipped boxes, with same data type as ``boxes``, does not share memory with ``boxes`` r,Nr) r rtuplerEr isinstancerr.clonerr )r'rGrHr#Z _flip_boxesrBr$r$r% flip_boxess  ,.rMFlabelsbg_labelint ellipse_maskcsRt|dt|}|jddkr,tjd|tjdt|}t||tjd^}}|S|t|kr>t dt|d||jd|jdkrLt dtj|jdf|tjdt|}t |tj tj ddt d d d ft|krt d t|d ^}}tjdD]fd d tD} |rt| d} ddtjgtjdt|} tfddtD} tfddtj| D} t|| | | dk<t| ddd}|| d d} ntj| tjdt|} g}|fddtD| |t|<qt||tjddS)au Convert box to int16 mask image, which has the same size with the input image. Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``. labels: classification foreground(fg) labels corresponding to `boxes`, dtype should be int, sized (N,). spatial_size: image spatial size. bg_label: background labels for the output mask image, make sure it is smaller than any fg labels. ellipse_mask: bool. - If True, it assumes the object shape is close to ellipse or ellipsoid. - If False, it assumes the object shape is close to rectangle or cube and well occupies the bounding box. - If the users are going to apply random rotation as data augmentation, we suggest setting ellipse_mask=True See also Kalra et al. "Towards Rotation Invariance in Object Detection", ICCV 2021. Return: - int16 array, sized (num_box, H, W). Each channel represents a box. The foreground region in channel c has intensity of labels[c]. The background intensity is bg_label. r,r)rr(r*r+rzGbg_label should be smaller than any foreground box labels. min(labels)=z, while bg_label=z1Number of labels should equal to number of boxes.Nz%Some boxes are larger than the image.r)cs(g|]}|f|fqSr$r$r@bboxes_npr#r$r%rDs(z'convert_box_to_mask..g@rc3s|]}tdVqdS)rNslice)rAr4) max_box_sizer$r% sz&convert_box_to_mask..c3s|] }|dVqdS)r-Nr$)rAgrid)centerr$r%rZsr-nearestF)rGmode anti_aliasingc3s.|]}t|f|fVqdSNrW)rAdrTr$r%rZs,)r rr nprint16rrr0 ValueErrorrndarrayint32anyarrayrEr2rJsumogridr extend)r'rOrGrPrR boxes_mask_np boxes_maskr4 labels_npbox_sizeradiusZboxes_only_maskrangesdist_from_centerresizerZslicingr$)rUrVr\rYr#r%convert_box_to_masksL   &$   rtrm box_dtypeDtypeLike | torch.dtype label_dtype'tuple[NdarrayOrTensor, NdarrayOrTensor]c Cstt|jddgt|jdd}t|d}t|tj^}}g}g} t|jdD]k} t || df|} | djddkrBq+g} | D] } | t | qF| D] } | t | dt qR| | |dkr{| || | dd| ddf|dkr| || | dd| dd| ddfq+t|dkrtdd|gtdg}}n t|t| }}t|||d ^}}t|||d ^}}||fS) a Convert int16 mask image to box, which has the same size with the input image Args: boxes_mask: int16 array, sized (num_box, H, W). Each channel represents a box. The foreground region in channel c has intensity of labels[c]. The background intensity is bg_label. bg_label: background labels for the boxes_mask box_dtype: output dtype for boxes label_dtype: output dtype for labels Return: - bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``. - classification foreground(fg) labels, dtype should be int, sized (N,). rN)rGr.r-rS)r lenr listr rrbrerEnonzeroappendr0r2r zerosasarrayr)rmrPrurwrGr#rlr4 boxes_listZ labels_listrU fg_indicesZboxes_bZfd_irVrnr'rOr$r$r%convert_mask_to_boxs6  $. "r+Sequence[NdarrayOrTensor] | NdarrayOrTensorkeeptuple | NdarrayOrTensorcCs|t|d}g}t|tjd}|D]}t|tjd}||df}|t||ddqt|tjtjfr:|dSt |S)au For element in labels, select indices keep from it. Args: labels: Sequence of array. Each element represents classification labels or scores corresponding to ``boxes``, sized (N,). keep: the indices to keep, same length with each element in labels. Return: selected labels, does not share memory with original labels. Tr.r)) rrrr.r~rrKrbrerJ)rOrZ labels_tupleZlabels_select_listkeep_titemlabels_tr$r$r% select_labelsHs  raxis1axis2cCst|d}t|tjr|}nt|}|dd||gf|dd||gf<|dd||||gf|dd||||gf<|S)a Interchange two axes of boxes. Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` axis1: First axis. axis2: Second axis. Returns: boxes with two axes interchanged. r,N)r rKrr.rLr)r'rrr#Z boxes_swapr$r$r%swapaxes_boxeses  $rrrrkaxestuple[int, int]cCs`t|d}tt||}t|}t|dkrtd|d|dks.t|d|d|kr2td|d|ksL|d| ksL|d|ksL|d| krWtd|d|d |d ;}|dkra|S|dkrstt|||d||dS|dkrt|||d}t||d|dSt||d|d}||d||d||d<||d<t|||dS) a@ Rotate boxes by 90 degrees in the plane specified by axes. Rotation direction is from the first towards the second axis. Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` spatial_size: image spatial size. k : number of times the array is rotated by 90 degrees. axes: (2,) array_like The array is rotated in the plane defined by the axes. Axes must be different. Returns: A rotated view of `boxes`. Notes: ``rot90_boxes(boxes, spatial_size, k=1, axes=(1,0))`` is the reverse of ``rot90_boxes(boxes, spatial_size, k=1, axes=(0,1))`` ``rot90_boxes(boxes, spatial_size, k=1, axes=(1,0))`` is equivalent to ``rot90_boxes(boxes, spatial_size, k=-1, axes=(0,1))`` r,r-zlen(axes) must be 2.rrzAxes must be different.zAxes=z out of range for array of ndim=.rz) r r|rrr{rdabsrMr)r'rGrrr# spatial_size_Zboxes_r$r$r% rot90_boxess(  (4*r)T)rrrrrrrr)r'rrrrr)r'rr8r9rr)r'rr<r=r>r=rrr`)r'rrGr=rHrIrr)rNF) r'rrOrrGr=rPrQrRrrr) rmrrPrQrurvrwrvrrx)rOrrrrr)r'rrrQrrQrr)rr) r'rrGr=rrQrrrr)) __future__rcollections.abcrcopyrnumpyrbrmonai.config.type_definitionsrrrmonai.data.box_utilsrr r monai.transformsr monai.transforms.utilsr monai.utilsr monai.utils.miscrrmonai.utils.type_conversionrrr&r7r;rFrMrtfloat32longrrrrr$r$r$r%s:       ! ( ! & R 5