U Ph4@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@sleZdZUdZdZded<dZded<efddd d Zd d d dd dddZ e d d d dd 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"Callable[[Tensor, Tensor], Tensor] similarity_fncCs ||_dS)Nr)selfrr\/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/detection/utils/ATSS_matcher.py__init__hszMatcher.__init__ torch.Tensor Sequence[int]!tuple[torch.Tensor, torch.Tensor]boxesanchorsnum_anchors_per_levelnum_anchors_per_locreturncCsX|dkrF|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)r r!r"r#) numelshapetorchtensortoemptyint64fill_rcompute_matches)rr r!r"r# num_anchorsmatch_quality_matrixmatchesrrr__call__ks  zMatcher.__call__cCstdS)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] N)NotImplementedError)rr r!r"r#rrrr.szMatcher.compute_matchesN) __name__ __module__ __qualname____doc__r__annotations__rr rr2rr.rrrrr\s   %rcsHeZdZdeddfdddddfdd Zd d d dd d ddZZS) ATSSMatcherTFrrbool)num_candidatesr center_in_gtdebugcsFtj|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__rrrrc CsP|jd}|jd}t||\}}} t|} g} d} t|D]`\}} | | |}t|j|| }| dd| |ftj|ddd\}}| || |} qprintr=arange expand_as contiguousr viewrAview_asrange full_likeINFmaxr) rr r!r"r#num_gtr/ distances__Zanchors_center distancesZcandidate_idx_list start_idxZaplend_idxrMidxZ candidate_idxr0Zcandidate_iousr1Ziou_mean_per_gtZiou_std_per_gtZiou_thresh_per_gtis_posZ boxes_idxZ is_in_gt_Zis_in_gtngZious_infindex matched_valsrrrr.s\   *         zATSSMatcher.compute_matches)r4r5r6r rr. __classcell__rrrDrr9s r9 MatcherType)bound)r7 __future__rrBabcrrcollections.abcrrtypingrr(rmonai.data.box_utilsr r r r monai.utils.type_conversionr floatr]rr9rkrrrr>s    J}