o  i@s<dZddlmZddlZddlmZmZddlmZddl Z ddl m Z m Z ddl m Z mZddlmZdd lmZdd lmZdd lmZdd lmZmZdd lmZddlmZmZddlm Z ddl!m"Z"ddl#m$Z$ddl%m&Z&m'Z'm(Z(m)Z)e)ddd\Z*Z+e)ddd\Z,Z+Gddde j-Z.   d*d+d(d)Z/dS),z{ Part of this script is adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py ) annotationsN)CallableSequence)Any)Tensornn) RetinaNetresnet_fpn_feature_extractor)AnchorGenerator) ATSSMatcher)BoxCoder) BoxSelector)check_training_targetspreprocess_images)HardNegativeSampler)ensure_dict_value_to_list_predict_with_inferer)box_iou)SlidingWindowInferer)resnet) BlendModePytorchPadModeensure_tuple_repoptional_importz#torchvision.models.detection._utilsBalancedPositiveNegativeSampler)nameMatcherc sFeZdZdZeddddddfdfdd ZdddZddd Zdd#d$Zdd&d'Z dd+d,Z -ddd2d3Z ddd8d9Z :ddd?d@Z ddAdBZddCejdDejdEddddf ddUdVZ W X C Y -ddd_d`Z  dddgdhZdidjZddodpZddsdtZ -ddd~dZdddZdddZdddZdddZdddZdddZZS)RetinaNetDetectora Retinanet detector, expandable to other one stage anchor based box detectors in the future. An example of construction can found in the source code of :func:`~monai.apps.detection.networks.retinanet_detector.retinanet_resnet50_fpn_detector` . The input to the model is expected to be a list of tensors, each of shape (C, H, W) or (C, H, W, D), one for each image, and should be in 0-1 range. Different images can have different sizes. Or it can also be a Tensor sized (B, C, H, W) or (B, C, H, W, D). In this case, all images have same size. The behavior of the model changes depending if it is in training or evaluation mode. During training, the model expects both the input tensors, as well as a targets (list of dictionary), containing: - boxes (``FloatTensor[N, 4]`` or ``FloatTensor[N, 6]``): the ground-truth boxes in ``StandardMode``, i.e., ``[xmin, ymin, xmax, ymax]`` or ``[xmin, ymin, zmin, xmax, ymax, zmax]`` format, with ``0 <= xmin < xmax <= H``, ``0 <= ymin < ymax <= W``, ``0 <= zmin < zmax <= D``. - labels: the class label for each ground-truth box The model returns a Dict[str, Tensor] during training, containing the classification and regression losses. When saving the model, only self.network contains trainable parameters and needs to be saved. During inference, the model requires only the input tensors, and returns the post-processed predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as follows: - boxes (``FloatTensor[N, 4]`` or ``FloatTensor[N, 6]``): the predicted boxes in ``StandardMode``, i.e., ``[xmin, ymin, xmax, ymax]`` or ``[xmin, ymin, zmin, xmax, ymax, zmax]`` format, with ``0 <= xmin < xmax <= H``, ``0 <= ymin < ymax <= W``, ``0 <= zmin < zmax <= D``. - labels (Int64Tensor[N]): the predicted labels for each image - labels_scores (Tensor[N]): the scores for each prediction Args: network: a network that takes an image Tensor sized (B, C, H, W) or (B, C, H, W, D) as input and outputs a dictionary Dict[str, List[Tensor]] or Dict[str, Tensor]. anchor_generator: anchor generator. box_overlap_metric: func that compute overlap between two sets of boxes, default is Intersection over Union (IoU). debug: whether to print out internal parameters, used for debugging and parameter tuning. Notes: Input argument ``network`` can be a monai.apps.detection.networks.retinanet_network.RetinaNet(*) object, but any network that meets the following rules is a valid input ``network``. 1. It should have attributes including spatial_dims, num_classes, cls_key, box_reg_key, num_anchors, size_divisible. - spatial_dims (int) is the spatial dimension of the network, we support both 2D and 3D. - num_classes (int) is the number of classes, excluding the background. - size_divisible (int or Sequence[int]) is the expectation on the input image shape. The network needs the input spatial_size to be divisible by size_divisible, length should be 2 or 3. - cls_key (str) is the key to represent classification in the output dict. - box_reg_key (str) is the key to represent box regression in the output dict. - num_anchors (int) is the number of anchor shapes at each location. it should equal to ``self.anchor_generator.num_anchors_per_location()[0]``. If network does not have these attributes, user needs to provide them for the detector. 2. Its input should be an image Tensor sized (B, C, H, W) or (B, C, H, W, D). 3. About its output ``head_outputs``, it should be either a list of tensors or a dictionary of str: List[Tensor]: - If it is a dictionary, it needs to have at least two keys: ``network.cls_key`` and ``network.box_reg_key``, representing predicted classification maps and box regression maps. ``head_outputs[network.cls_key]`` should be List[Tensor] or Tensor. Each Tensor represents classification logits map at one resolution level, sized (B, num_classes*num_anchors, H_i, W_i) or (B, num_classes*num_anchors, H_i, W_i, D_i). ``head_outputs[network.box_reg_key]`` should be List[Tensor] or Tensor. Each Tensor represents box regression map at one resolution level, sized (B, 2*spatial_dims*num_anchors, H_i, W_i)or (B, 2*spatial_dims*num_anchors, H_i, W_i, D_i). ``len(head_outputs[network.cls_key]) == len(head_outputs[network.box_reg_key])``. - If it is a list of 2N tensors, the first N tensors should be the predicted classification maps, and the second N tensors should be the predicted box regression maps. Example: .. code-block:: python # define a naive network import torch class NaiveNet(torch.nn.Module): def __init__(self, spatial_dims: int, num_classes: int): super().__init__() self.spatial_dims = spatial_dims self.num_classes = num_classes self.size_divisible = 2 self.cls_key = "cls" self.box_reg_key = "box_reg" self.num_anchors = 1 def forward(self, images: torch.Tensor): spatial_size = images.shape[-self.spatial_dims:] out_spatial_size = tuple(s//self.size_divisible for s in spatial_size) # half size of input out_cls_shape = (images.shape[0],self.num_classes*self.num_anchors) + out_spatial_size out_box_reg_shape = (images.shape[0],2*self.spatial_dims*self.num_anchors) + out_spatial_size return {self.cls_key: [torch.randn(out_cls_shape)], self.box_reg_key: [torch.randn(out_box_reg_shape)]} # create a RetinaNetDetector detector spatial_dims = 3 num_classes = 5 anchor_generator = monai.apps.detection.utils.anchor_utils.AnchorGeneratorWithAnchorShape( feature_map_scales=(1, ), base_anchor_shapes=((8,) * spatial_dims) ) net = NaiveNet(spatial_dims, num_classes) detector = RetinaNetDetector(net, anchor_generator) # only detector.network may contain trainable parameters. optimizer = torch.optim.SGD( detector.network.parameters(), 1e-3, momentum=0.9, weight_decay=3e-5, nesterov=True, ) torch.save(detector.network.state_dict(), 'model.pt') # save model detector.network.load_state_dict(torch.load('model.pt', weights_only=True)) # load model Nclassificationbox_regressionFnetwork nn.Moduleanchor_generatorr box_overlap_metricr spatial_dims int | None num_classessize_divisibleSequence[int] | intcls_keystr box_reg_keydebugboolc sRt||_|jd|d|_|jd|d|_|jd|d|_t|j|j|_|jd|d|_|jd|d|_ ||_ |j d|_ |jd|j d} |j | kr\t d | d |j d d|_d|_||_| |_d|_|tjjd d |jtjjdd ddddtd|jd|_d|_d|_|jd|_d|_t|jdddddd|_ dS)Nr%) default_valuer'r(r*r,r num_anchorsz Number of feature map channels (z8) should match with number of anchors at each location (z).mean) reductiongqq?)betar2TF) encode_gt decode_pred)?r6weightsboxeslabels_scores皙??,)r$ score_threshtopk_candidates_per_level nms_threshdetections_per_img apply_sigmoid)!super__init__r!get_attribute_from_networkr%r'r(rr*r,r#num_anchors_per_locationnum_anchors_per_loc ValueErroranchorsprevious_image_shaper$r- fg_bg_sampler set_cls_losstorchrBCEWithLogitsLossset_box_regression_loss SmoothL1Lossr box_codertarget_box_keytarget_label_keypred_score_keyinfererr box_selector) selfr!r#r$r%r'r(r*r,r-Znetwork_num_anchors __class__r/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/detection/networks/retinanet_detector.pyrFsN    zRetinaNetDetector.__init__cCs4t|j|r t|j|S|dur|Std|d)Nz network does not have attribute z$, please provide it in the detector.)hasattrr!getattrrJ)rY attr_namer/r\r\r]rGs  z,RetinaNetDetector.get_attribute_from_networkr8 tuple[float]returnNonecCs>t|d|jkrtdd|jd|dt|d|_dS)z Set the weights for box coder. Args: weights: a list/tuple with length of 2*self.spatial_dims zlen(weights) should be z, got weights=.r7N)lenr%rJr rS)rYr8r\r\r]set_box_coder_weights sz'RetinaNetDetector.set_box_coder_weightsbox_key label_keycCs||_||_|d|_dS)aB Set keys for the training targets and inference outputs. During training, both box_key and label_key should be keys in the targets when performing ``self.forward(input_images, targets)``. During inference, they will be the keys in the output dict of `self.forward(input_images)``. r;N)rTrUrV)rYrhrir\r\r]set_target_keyssz!RetinaNetDetector.set_target_keyscls_losscCs ||_dS)a Using for training. Set loss for classification that takes logits as inputs, make sure sigmoid/softmax is built in. Args: cls_loss: loss module for classification Example: .. code-block:: python detector.set_cls_loss(torch.nn.BCEWithLogitsLoss(reduction="mean")) detector.set_cls_loss(FocalLoss(reduction="mean", gamma=2.0)) N) cls_loss_func)rYrkr\r\r]rN!s zRetinaNetDetector.set_cls_lossbox_lossr4r5cCs||_||_||_dS)a Using for training. Set loss for box regression. Args: box_loss: loss module for box regression encode_gt: if True, will encode ground truth boxes to target box regression before computing the losses. Should be True for L1 loss and False for GIoU loss. decode_pred: if True, will decode predicted box regression into predicted boxes before computing losses. Should be False for L1 loss and True for GIoU loss. Example: .. code-block:: python detector.set_box_regression_loss( torch.nn.SmoothL1Loss(beta=1.0 / 9, reduction="mean"), encode_gt = True, decode_pred = False ) detector.set_box_regression_loss( monai.losses.giou_loss.BoxGIoULoss(reduction="mean"), encode_gt = False, decode_pred = True ) N) box_loss_funcr4r5)rYrmr4r5r\r\r]rQ0s z)RetinaNetDetector.set_box_regression_lossT fg_iou_threshfloat bg_iou_threshallow_low_quality_matchescCs2||krtd|d|dt|||d|_dS)a Using for training. Set torchvision matcher that matches anchors with ground truth boxes. Args: fg_iou_thresh: foreground IoU threshold for Matcher, considered as matched if IoU > fg_iou_thresh bg_iou_thresh: background IoU threshold for Matcher, considered as not matched if IoU < bg_iou_thresh allow_low_quality_matches: if True, produce additional matches for predictions that have only low-quality match candidates. z:Require fg_iou_thresh >= bg_iou_thresh. Got fg_iou_thresh=z, bg_iou_thresh=re)rrN)rJrproposal_matcher)rYrorqrrr\r\r]set_regular_matcherKs  z%RetinaNetDetector.set_regular_matchernum_candidatesint center_in_gtcCst||j||jd|_dS)a& Using for training. Set ATSS matcher that matches anchors with ground truth boxes Args: num_candidates: number of positions to select candidates from. Smaller value will result in a higher matcher threshold and less matched candidates. center_in_gt: If False (default), matched anchor center points do not need to lie within the ground truth box. Recommend False for small objects. If True, will result in a strict matcher and less matched candidates. )r-N)r r$r-rs)rYrvrxr\r\r]set_atss_matcher`s z"RetinaNetDetector.set_atss_matcher batch_size_per_imagepositive_fractionmin_neg pool_sizecCst||||d|_dS)a Using for training. Set hard negative sampler that samples part of the anchors for training. HardNegativeSampler is used to suppress false positive rate in classification tasks. During training, it select negative samples with high prediction scores. Args: batch_size_per_image: number of elements to be selected per image positive_fraction: percentage of positive elements in the selected samples min_neg: minimum number of negative samples to select if possible. pool_size: when we need ``num_neg`` hard negative samples, they will be randomly selected from ``num_neg * pool_size`` negative samples with the highest prediction scores. Larger ``pool_size`` gives more randomness, yet selects negative samples that are less 'hard', i.e., negative samples with lower prediction scores. )r{r|r}r~N)rrM)rYr{r|r}r~r\r\r]set_hard_negative_samplerms  z+RetinaNetDetector.set_hard_negative_samplercCst||d|_dS)a Using for training. Set torchvision balanced sampler that samples part of the anchors for training. Args: batch_size_per_image: number of elements to be selected per image positive_fraction: percentage of positive elements per batch )r{r|N)rrM)rYr{r|r\r\r]set_balanced_samplers  z&RetinaNetDetector.set_balanced_samplerr>g?roi_size sw_batch_sizeoverlapmodeBlendMode | str sigma_scaleSequence[float] | float padding_modePytorchPadMode | strcval sw_devicetorch.device | str | Nonedeviceprogresscache_roi_weight_mapc Cs"t||||||||| | | |_dS)zM Define sliding window inferer and store it to self.inferer. N)rrW) rYrrrrrrrrrrrr\r\r]set_sliding_window_inferers z,RetinaNetDetector.set_sliding_window_infererr<r=r?r@rArBrCrDcCst|j|||||d|_dS)aW Using for inference. Set the parameters that are used for box selection during inference. The box selection is performed with the following steps: #. For each level, discard boxes with scores less than self.score_thresh. #. For each level, keep boxes with top self.topk_candidates_per_level scores. #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overlapping threshold nms_thresh. #. For the whole image, keep boxes with top self.detections_per_img scores. Args: score_thresh: no box with scores less than score_thresh will be kept topk_candidates_per_level: max number of boxes to keep for each level nms_thresh: box overlapping threshold for NMS detections_per_img: max number of boxes to keep for each image )r$rDr@rArBrCN)r r$rX)rYr@rArBrCrDr\r\r]set_box_selector_parameterss z-RetinaNetDetector.set_box_selector_parameters input_imageslist[Tensor] | Tensortargetslist[dict[str, Tensor]] | None use_inferer+dict[str, Tensor] | list[dict[str, Tensor]]c CsF|jrt|||j|j|j}|t||j|j\}}|js!|sQ||}t |t t frLi}|dt |d||j <|t |dd||j<|}nt|n|jdurZtdt||j|j |jg|jd}|||dd||j D}|j |jfD] } ||| || <q~|jr||||j|} | S|||j||} | S)a  Returns a dict of losses during training, or a list predicted dict of boxes and labels during inference. Args: input_images: The input to the model is expected to be a list of tensors, each of shape (C, H, W) or (C, H, W, D), one for each image, and should be in 0-1 range. Different images can have different sizes. Or it can also be a Tensor sized (B, C, H, W) or (B, C, H, W, D). In this case, all images have same size. targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image (optional). use_inferer: whether to use self.inferer, a sliding window inferer, to do the inference. If False, will simply forward the network. If True, will use self.inferer, and requires ``self.set_sliding_window_inferer(*args)`` to have been called before. Return: If training mode, will return a dict with at least two keys, including self.cls_key and self.box_reg_key, representing classification loss and box regression loss. If evaluation mode, will return a list of detection results. Each element corresponds to an images in ``input_images``, is a dict with at least three keys, including self.target_box_key, self.target_label_key, self.pred_score_key, representing predicted boxes, classification labels, and classification scores. NrdzZ`self.inferer` is not defined.Please refer to function self.set_sliding_window_inferer(*).)keysrWcSsg|] }|jddqS)rdN)shapenumel).0xr\r\r] sz-RetinaNetDetector.forward..)trainingrr%rUrT#_check_detector_training_componentsrr(r! isinstancetuplelistrfr*r,rrWrJrgenerate_anchors _reshape_maps compute_lossrKpostprocess_detections) rYrrrimages image_sizes head_outputsZtmp_dictnum_anchor_locs_per_levelkeylosses detectionsr\r\r]forwards@      zRetinaNetDetector.forwardcCs8t|ds td|jdur|jrtddSdSdS)zc Check if self.proposal_matcher and self.fg_bg_sampler have been set for training. rsz\Matcher is not set. Please refer to self.set_regular_matcher(*) or self.set_atss_matcher(*).NaNo balanced sampler is used. Negative samples are likely to be much more than positive samples. Please set balanced samplers with self.set_balanced_sampler(*) or self.set_hard_negative_sampler(*), or set classification loss function as Focal loss with self.set_cls_loss(*))r^AttributeErrorrMr-warningswarnrYr\r\r]r(s z5RetinaNetDetector._check_detector_training_componentsrrrdict[str, list[Tensor]]cCs:|jdus |j|jkr||||j|_|j|_dSdS)aA Generate anchors and store it in self.anchors: List[Tensor]. We generate anchors only when there is no stored anchors, or the new coming images has different shape with self.previous_image_shape Args: images: input images, a (B, C, H, W) or (B, C, H, W, D) Tensor. head_outputs: head_outputs. ``head_output_reshape[self.cls_key]`` is a Tensor sized (B, sum(HW(D)A), self.num_classes). ``head_output_reshape[self.box_reg_key]`` is a Tensor sized (B, sum(HW(D)A), 2*self.spatial_dims) N)rKrLrr#r*)rYrrr\r\r]r8s  z"RetinaNetDetector.generate_anchors result_maps list[Tensor]c Csg}|D]n}|jd}|jd|j}|j|j d}|d|f|}||}|jdkr7|ddddd}n|jdkrG|dddddd}ntd ||d|}t| s`t | rmt rhtd t d ||qtj|dd S) a Concat network output map list to a single Tensor. This function is used in both training and inference. Args: result_maps: a list of Tensor, each Tensor is a (B, num_channel*A, H, W) or (B, num_channel*A, H, W, D) map. A = self.num_anchors_per_loc Return: reshaped and concatenated result, sized (B, sum(HWA), num_channel) or (B, sum(HWDA), num_channel) rrNrdruzImages can only be 2D or 3D.z"Concatenated result is NaN or Inf.dim)rrIr%viewpermuterJreshaperOisnananyisinfis_grad_enabledrrappendcat) rYrZall_reshaped_result_map result_map batch_sizeZ num_channel spatial_size view_shapeZreshaped_result_mapr\r\r]rHs&       zRetinaNetDetector._reshape_mapshead_outputs_reshapedict[str, Tensor]rKrlist[list[int]]r Sequence[int] need_sigmoidlist[dict[str, Tensor]]c s fdd|Di}|D]}t||jdd||<q fdd|D}|j} |j} | djt|} g} t| D]Cfdd| D} fdd| D}||}}fd dt| |D}j |||\}}}| j |j |j |iq?| S) a Postprocessing to generate detection result from classification logits and box regression. Use self.box_selector to select the final output boxes for each image. Args: head_outputs_reshape: reshaped head_outputs. ``head_output_reshape[self.cls_key]`` is a Tensor sized (B, sum(HW(D)A), self.num_classes). ``head_output_reshape[self.box_reg_key]`` is a Tensor sized (B, sum(HW(D)A), 2*self.spatial_dims) targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image. anchors: a list of Tensor. Each Tensor represents anchors for each image, sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). A = self.num_anchors_per_loc. Return: a list of dict, each dict corresponds to detection result on image. csg|]}|jqSr\)rI)rZnum_anchor_locsrr\r]rs z.rrcsg|] }t|qSr\)rsplit)ra)num_anchors_per_levelr\r]rsrcg|]}|qSr\r\)rbrindexr\r]rscrr\r\)rclrr\r]rscs,g|]\}}j|tj|qSr\)rS decode_singletorOfloat32)rbr) compute_dtyperYr\r]rs)rrr*r,dtyperfrangeziprXZselect_boxes_per_imagerrTrVrU)rYrrKrrrsplit_head_outputsk split_anchors class_logitsr num_imagesrbox_regression_per_imagelogits_per_imageanchors_per_imageZimg_spatial_sizeboxes_per_imageZselected_boxesZselected_scoresZselected_labelsr\)rrrrYr]rus>        z(RetinaNetDetector.postprocess_detectionscCsH||||}|||j||}|||j|||}|j||j|iS)a Compute losses. Args: head_outputs_reshape: reshaped head_outputs. ``head_output_reshape[self.cls_key]`` is a Tensor sized (B, sum(HW(D)A), self.num_classes). ``head_output_reshape[self.box_reg_key]`` is a Tensor sized (B, sum(HW(D)A), 2*self.spatial_dims) targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image. anchors: a list of Tensor. Each Tensor represents anchors for each image, sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). A = self.num_anchors_per_loc. Return: a dict of several kinds of losses. )compute_anchor_matched_idxscompute_cls_lossr*compute_box_lossr,)rYrrrKr matched_idxsZ losses_clsZlosses_box_regressionr\r\r]rs zRetinaNetDetector.compute_lossc Csg}t||D]\}}||jdkr'|tj|dfdtj|jdqt |j t r@| ||j |j|}| |}nt |j trY| ||j |j|||j\}}ntd|jrotdtj|ddddt|dkrtd ||jd||q|S) a  Compute the matched indices between anchors and ground truth (gt) boxes in targets. output[k][i] represents the matched gt index for anchor[i] in image k. Suppose there are M gt boxes for image k. The range of it output[k][i] value is [-2, -1, 0, ..., M-1]. [0, M - 1] indicates this anchor is matched with a gt box, while a negative value indicating that it is not matched. Args: anchors: a list of Tensor. Each Tensor represents anchors for each image, sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). A = self.num_anchors_per_loc. targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image. num_anchor_locs_per_level: each element represents HW or HWD at this level. Return: a list of matched index `matched_idxs_per_image` (Tensor[int64]), Tensor sized (sum(HWA),) or (sum(HWDA),). Suppose there are M gt boxes. `matched_idxs_per_image[i]` is a matched gt index in [0, M - 1] or a negative value indicating that anchor i could not be matched. BELOW_LOW_THRESHOLD = -1, BETWEEN_THRESHOLDS = -2 rr)rrzCurrently support torchvision Matcher and monai ATSS matcher. Other types of matcher not supported. Please override self.compute_anchor_matched_idxs(*) for your own matcher.z.Max box overlap between anchors and gt boxes: rrrezNo anchor is matched with GT boxes. Please adjust matcher setting, anchor setting, or the network setting to change zoom scale between network output and input images.GT boxes are )rrTrrrOfullsizeint64rrrsrr$rr rINotImplementedErrorr-printmaxrr) rYrKrrrrtargets_per_imagematch_quality_matrixmatched_idxs_per_imager\r\r]rs@     z-RetinaNetDetector.compute_anchor_matched_idxs cls_logitsrcCszg}g}t|||D]\}}}||||\} } || || q tj|dd} tj|dd} || | | j} | S)a Compute classification losses. Args: cls_logits: classification logits, sized (B, sum(HW(D)A), self.num_classes) targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image. matched_idxs: a list of matched index. each element is sized (sum(HWA),) or (sum(HWDA),) Return: classification losses. rr)rget_cls_train_sample_per_imagerrOrrlrr)rYrrrZtotal_cls_logits_listZtotal_gt_classes_target_listrcls_logits_per_imagerZsampled_cls_logits_per_imageZsampled_gt_classes_targetZtotal_cls_logitsZtotal_gt_classes_targetrr\r\r]r"s  z"RetinaNetDetector.compute_cls_losscCsg}g}t||||D]\}}} } |||| | \} } || || q tj|dd} tj|dd}| jddkrBtd}|S|| || j }|S)a Compute box regression losses. Args: box_regression: box regression results, sized (B, sum(HWA), 2*self.spatial_dims) targets: a list of dict. Each dict with two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image. anchors: a list of Tensor. Each Tensor represents anchors for each image, sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). A = self.num_anchors_per_loc. matched_idxs: a list of matched index. each element is sized (sum(HWA),) or (sum(HWDA),) Return: box regression losses. rrr) rget_box_train_sample_per_imagerrOrrtensorrnrr)rYr rrKrZtotal_box_regression_listZtotal_target_regression_listrrrrZdecode_box_regression_per_imagematched_gt_boxes_per_imageZtotal_box_regressionZtotal_target_regressionrr\r\r]r@s"   z"RetinaNetDetector.compute_box_lossrrrtuple[Tensor, Tensor]cCst|st|rtrtdtd|dk}t| }||j j d}|j rPt d|d|d|dkrP|d|krPt d|d|d t|}d ||||j||f<|jd urn||jjk}nUt|jtrtj|tjd d d} ||d g| \} } nt|jtr||d g\} } ntdttj| dd d} ttj| dd d} tj| | gdd }||d d f||d d ffS)a; Get samples from one image for classification losses computation. Args: cls_logits_per_image: classification logits for one image, (sum(HWA), self.num_classes) targets_per_image: a dict with at least two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image. matched_idxs_per_image: matched index, Tensor sized (sum(HWA),) or (sum(HWDA),) Suppose there are M gt boxes. matched_idxs_per_image[i] is a matched gt index in [0, M - 1] or a negative value indicating that anchor i could not be matched. BELOW_LOW_THRESHOLD = -1, BETWEEN_THRESHOLDS = -2 Return: paired predicted and GT samples from one image for classification losses computation z.NaN or Inf in predicted classification logits.rz&Number of positive (matched) anchors: z; Number of GT box: rerdzOnly z anchors are matched with z GT boxes. Please consider adjusting matcher setting, anchor setting, or the network setting to change zoom scale between network output and input images.r6NrrzCurrently support torchvision BalancedPositiveNegativeSampler and monai HardNegativeSampler matcher. Other types of sampler not supported. Please override self.get_cls_train_sample_per_image(*) for your own sampler.)rOrrrrrJrrrwsumrTrr-r zeros_likerUrMrsBETWEEN_THRESHOLDSrrrrrrrwherer)rYrrrforeground_idxs_per_imagenum_foreground num_gt_boxgt_classes_targetvalid_idxs_per_imageZmax_cls_logits_per_imageZsampled_pos_inds_listZsampled_neg_inds_listsampled_pos_indssampled_neg_indsr\r\r]rosL        z0RetinaNetDetector.get_cls_train_sample_per_imagerrc Cst|st|rtrtdtdt|dkd}||j j d}|dkrD|ddddf|ddddffS||j || |j }||ddf}||ddf}|}|} |j ro|j||}|jry|j| |} | |fS)a Get samples from one image for box regression losses computation. Args: box_regression_per_image: box regression result for one image, (sum(HWA), 2*self.spatial_dims) targets_per_image: a dict with at least two keys: self.target_box_key and self.target_label_key, ground-truth boxes present in the image. anchors_per_image: anchors of one image, sized (sum(HWA), 2*spatial_dims) or (sum(HWDA), 2*spatial_dims). A = self.num_anchors_per_loc. matched_idxs_per_image: matched index, sized (sum(HWA),) or (sum(HWDA),) Return: paired predicted and GT samples from one image for box regression losses computation z'NaN or Inf in predicted box regression.rN)rOrrrrrJrrrrTrrrr4rS encode_singler5r) rYrrrrrrrZmatched_gt_boxes_per_image_Zbox_regression_per_image_r\r\r]rs, (z0RetinaNetDetector.get_box_train_sample_per_image)r!r"r#r r$rr%r&r'r&r(r)r*r+r,r+r-r.)N)r8rarbrc)rhr+rir+rbrc)rkr"rbrc)rmr"r4r.r5r.rbrc)T)rorprqrprrr.rbrc)ruF)rvrwrxr.rbrc)rrz) r{rwr|rpr}rwr~rprbrc)r{rwr|rprbrc)rr)rrwrrprrrrrrrrprrrrrr.rr.rbrc)r<r=r>r?T) r@rprArwrBrprCrwrDr.rbrc)NF)rrrrrr.rbr)rrrrrbrc)rrrbr) rrrKrrrrrrr.rbr) rrrrrKrrrrbr)rKrrrrrrbr)rrrrrrrbr) r rrrrKrrrrbr)rrrrrrrbr) rrrrrrrrrbr) __name__ __module__ __qualname____doc__rrFrGrgrjrNrQrtryrrrCONSTANTrrrrrrrrrrrrrr __classcell__r\r\rZr]rDshy H     ! $ U  3 G  H  /OrrrdrFTr'rwr#r returned_layersr pretrainedr.rkwargsrrbc  srtj||fi|}t|jj}t|||dd}|d} fdd|jjjD} t||| || d} t | |S)aX Returns a RetinaNet detector using a ResNet-50 as backbone, which can be pretrained from `Med3D: Transfer Learning for 3D Medical Image Analysis ` _. Args: num_classes: number of output classes of the model (excluding the background). anchor_generator: AnchorGenerator, returned_layers: returned layers to extract feature maps. Each returned layer should be in the range [1,4]. len(returned_layers)+1 will be the number of extracted feature maps. There is an extra maxpooling layer LastLevelMaxPool() appended. pretrained: If True, returns a backbone pre-trained on 23 medical datasets progress: If True, displays a progress bar of the download to stderr Return: A RetinaNetDetector object with resnet50 as backbone Example: .. code-block:: python # define a naive network resnet_param = { "pretrained": False, "spatial_dims": 3, "n_input_channels": 2, "num_classes": 3, "conv1_t_size": 7, "conv1_t_stride": (2, 2, 2) } returned_layers = [1] anchor_generator = monai.apps.detection.utils.anchor_utils.AnchorGeneratorWithAnchorShape( feature_map_scales=(1, 2), base_anchor_shapes=((8,) * resnet_param["spatial_dims"]) ) detector = retinanet_resnet50_fpn_detector( **resnet_param, anchor_generator=anchor_generator, returned_layers=returned_layers ) N)backboner%pretrained_backbonetrainable_backbone_layersrrcs g|] }|ddtqS)rd)r)rsrr\r]r1s z3retinanet_resnet50_fpn_detector..)r%r'r0feature_extractorr() rresnet50rfconv1strider rHbodyrr) r'r#rrrrrr%rr0r(r!r\rr]retinanet_resnet50_fpn_detectors&/   r)rFT)r'rwr#r rrrr.rr.rrrbr)0r  __future__rrcollections.abcrrtypingrrOrrZ/monai.apps.detection.networks.retinanet_networkrr Z'monai.apps.detection.utils.anchor_utilsr Z'monai.apps.detection.utils.ATSS_matcherr Z$monai.apps.detection.utils.box_coderr Z'monai.apps.detection.utils.box_selectorr Z)monai.apps.detection.utils.detector_utilsrrZ0monai.apps.detection.utils.hard_negative_samplerrZ(monai.apps.detection.utils.predict_utilsrrmonai.data.box_utilsrmonai.inferersrmonai.networks.netsr monai.utilsrrrrr_rModulerrr\r\r\r]sF"           =