o  i4@sdZddlmZddlZddlmZmZddlmZm Z ddl m Z ddl Z ddl m Z ddlmZmZmZmZdd lmZed ZGd d d eZGd ddeZe dedZdS)a The functions in this script are adapted from nnDetection, https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/matcher.py which is adapted from torchvision. These are the changes compared with nndetection: 1) comments and docstrings; 2) reformat; 3) add a debug option to ATSSMatcher to help the users to tune parameters; 4) add a corner case return in ATSSMatcher.compute_matches; 5) add support for float16 cpu ) annotationsN)ABCabstractmethod)CallableSequence)TypeVar)Tensor) COMPUTE_DTYPEbox_iouboxes_center_distancecenters_in_boxes)convert_to_tensorinfc@sPeZdZUdZdZded<dZded<efdd d ZdddZ e dddZ dS)Matcherz Base class of Matcher, which matches boxes and anchors to each other Args: similarity_fn: function for similarity computation between boxes and anchors intBELOW_LOW_THRESHOLDBETWEEN_THRESHOLDS similarity_fn"Callable[[Tensor, Tensor], Tensor]cCs ||_dS)Nr)selfrri/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/detection/utils/ATSS_matcher.py__init__hs zMatcher.__init__boxes torch.Tensoranchorsnum_anchors_per_level Sequence[int]num_anchors_per_locreturn!tuple[torch.Tensor, torch.Tensor]cCsX|dkr#|jd}tg|}tj|tjd|j}||fS|j ||||dS)a Compute matches for a single image Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` anchors: anchors to match Mx4 or Mx6, also assumed to be ``StandardMode``. num_anchors_per_level: number of anchors per feature pyramid level num_anchors_per_loc: number of anchors per position Returns: - matrix which contains the similarity from each boxes to each anchor [N, M] - vector which contains the matched box index for all anchors (if background `BELOW_LOW_THRESHOLD` is used and if it should be ignored `BETWEEN_THRESHOLDS` is used) [M] Note: ``StandardMode`` = :class:`~monai.data.box_utils.CornerCornerModeTypeA`, also represented as "xyxy" ([xmin, ymin, xmax, ymax]) for 2D and "xyzxyz" ([xmin, ymin, zmin, xmax, ymax, zmax]) for 3D. r)dtype)rrrr!) numelshapetorchtensortoemptyint64fill_rcompute_matches)rrrrr! num_anchorsmatch_quality_matrixmatchesrrr__call__ks  zMatcher.__call__cCst)a Compute matches Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` anchors: anchors to match Mx4 or Mx6, also assumed to be ``StandardMode``. num_anchors_per_level: number of anchors per feature pyramid level num_anchors_per_loc: number of anchors per position Returns: - matrix which contains the similarity from each boxes to each anchor [N, M] - vector which contains the matched box index for all anchors (if background `BELOW_LOW_THRESHOLD` is used and if it should be ignored `BETWEEN_THRESHOLDS` is used) [M] )NotImplementedError)rrrrr!rrrr-szMatcher.compute_matchesN)rr rrrrrr r!rr"r#) __name__ __module__ __qualname____doc__r__annotations__rr rr1rr-rrrrr\s    %rcs2eZdZdeddfdfd d ZdddZZS) ATSSMatcherTFnum_candidatesrrr center_in_gtbooldebugcsFtj|d||_d|_||_||_td|jd|jddS)an Compute matching based on ATSS https://arxiv.org/abs/1912.02424 `Bridging the Gap Between Anchor-based and Anchor-free Detection via Adaptive Training Sample Selection` Args: num_candidates: number of positions to select candidates from. Smaller value will result in a higher matcher threshold and less matched candidates. similarity_fn: function for similarity computation between boxes and anchors center_in_gt: If False (default), matched anchor center points do not need to lie withing the ground truth box. Recommend False for small objects. If True, will result in a strict matcher and less matched candidates. debug: if True, will print the matcher threshold in order to tune ``num_candidates`` and ``center_in_gt``. rg{Gz?z*Running ATSS Matching with num_candidates=z and center_in_gt .N)superrr;min_distr<r>logginginfo)rr;rr<r> __class__rrrszATSSMatcher.__init__rrrrr r!r"r#c CsJ|jd}|jd}t||\}}} t|} g} d} t|D]0\}} | | |}t|j|| }| dd| |ftj|ddd\}}| || |} qt j | dd}| ||}| d|}|jddkr~dt j|ft j|jd}d||<||fS|jdd}|jdd}||}||dddfk}|jrtd ||jrt j||jt jd dddf|}t| |d||d|jd }t|}|||@}t|D]}||ddf||7<qt |t d}|d|d}|d|||<||}|tj dd\}}|j!||t k<||fS) aa Compute matches according to ATTS for a single image Adapted from (https://github.com/sfzhang15/ATSS/blob/79dfb28bd1/atss_core/modeling/rpn/atss/loss.py#L180-L184) Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` anchors: anchors to match Mx4 or Mx6, also assumed to be ``StandardMode``. num_anchors_per_level: number of anchors per feature pyramid level num_anchors_per_loc: number of anchors per position Returns: - matrix which contains the similarity from each boxes to each anchor [N, M] - vector which contains the matched box index for all anchors (if background `BELOW_LOW_THRESHOLD` is used and if it should be ignored `BETWEEN_THRESHOLDS` is used) [M] Note: ``StandardMode`` = :class:`~monai.data.box_utils.CornerCornerModeTypeA`, also represented as "xyxy" ([xmin, ymin, xmax, ymax]) for 2D and "xyzxyz" ([xmin, ymin, zmin, xmax, ymax, zmax]) for 3D. rNF)dimlargest)rGr)r$devicezAnchor matcher threshold: )rIr$)eps)"r&r r enumerateminr;r)r topkappendr'catrgatheroneslongrImeanstdr>printr<arange expand_as contiguousr viewrAview_asrange full_likeINFmaxr) rrrrr!num_gtr. distances__Zanchors_center distancesZcandidate_idx_list start_idxZaplend_idxrMidxZ candidate_idxr/Zcandidate_iousr0Ziou_mean_per_gtZiou_std_per_gtZiou_thresh_per_gtis_posZ boxes_idxZ is_in_gt_Zis_in_gtngZious_infindex matched_valsrrrr-sX   *      zATSSMatcher.compute_matches)r;rrrr<r=r>r=r3)r4r5r6r rr- __classcell__rrrDrr9sr9 MatcherType)bound)r7 __future__rrBabcrrcollections.abcrrtypingrr'rmonai.data.box_utilsr r r r monai.utils.type_conversionr floatr]rr9rkrrrrs=    J}