o  iJ@sdZddlmZddlZddlZddlmZmZddlm Z ddl Z ddl m Z m Z ddl mZmZddlmZdd lmZdd lmZmZmZed d d \ZZGddde jZGddde jZGddde jZ   d$d%d"d#ZdS)&z{ Part of this script is adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py ) annotationsN)CallableSequence)Any)Tensornn)BackboneWithFPN_resnet_fpn_extractor)Conv)resnet)ensure_tuple_replook_up_optionoptional_importz+torchvision.models.detection.backbone_utils_validate_trainable_layers)namecs0eZdZdZ ddfd d ZdddZZS)RetinaNetClassificationHeada A classification head for use in RetinaNet. This head takes a list of feature maps as inputs, and outputs a list of classification maps. Each output map has same spatial size with the corresponding input feature map, and the number of output channel is num_anchors * num_classes. Args: in_channels: number of channels of the input feature num_anchors: number of anchors to be predicted num_classes: number of classes to be predicted spatial_dims: spatial dimension of the network, should be 2 or 3. prior_probability: prior probability to initialize classification convolutional layers. {Gz? in_channelsint num_anchors num_classes spatial_dimsprior_probabilityfloatc stttj|f}g}tdD]}||||dddd|tjd|d|tqtj ||_ |j D]} t | |rWt jjj| jddt jj| jd q=||||dddd|_t jjj|jjddt jj|jjtd|| ||_||_dS) N kernel_sizestridepadding num_groups num_channelsrstdr)super__init__r CONVrangeappendr GroupNormReLU Sequentialconvchildren isinstancetorchinitnormal_weight constant_bias cls_logitsmathlogrr) selfrrrrr conv_typer/_layer __class__q/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/detection/networks/retinanet_network.pyr(Ls$    $ z$RetinaNetClassificationHead.__init__x list[Tensor]returncC|g}t|tr |g}n|}|D],}||}||}||t|s.t|r;t r6t dt dq|S)ai It takes a list of feature maps as inputs, and outputs a list of classification maps. Each output classification map has same spatial size with the corresponding input feature map, and the number of output channel is num_anchors * num_classes. Args: x: list of feature map, x[i] is a (B, in_channels, H_i, W_i) or (B, in_channels, H_i, W_i, D_i) Tensor. Return: cls_logits_maps, list of classification map. cls_logits_maps[i] is a (B, num_anchors * num_classes, H_i, W_i) or (B, num_anchors * num_classes, H_i, W_i, D_i) Tensor. zcls_logits is NaN or Inf.) r1rr/r8r+r2isnananyisinfis_grad_enabled ValueErrorwarningswarn)r;rCZcls_logits_maps feature_mapsfeaturesr8rArArBforwarde     z#RetinaNetClassificationHead.forward)r) rrrrrrrrrrrCrDrErD__name__ __module__ __qualname____doc__r(rP __classcell__rArAr?rBr<s rcs,eZdZdZd fdd Zdd d ZZS)RetinaNetRegressionHeada A regression head for use in RetinaNet. This head takes a list of feature maps as inputs, and outputs a list of box regression maps. Each output box regression map has same spatial size with the corresponding input feature map, and the number of output channel is num_anchors * 2 * spatial_dims. Args: in_channels: number of channels of the input feature num_anchors: number of anchors to be predicted spatial_dims: spatial dimension of the network, should be 2 or 3. rrrrc stttj|f}g}tdD]}||||dddd|tjd|d|tqtj ||_ |||d|dddd|_ t jj j|j jdd t jj |j j|j D]}t||rxt jj j|jdd t jj |jq_dS) Nrrrrr!r"rr%)r'r(r r)r*r+rr,r-r.r/bbox_regr2r3r4r5zeros_r7r0r1)r;rrrr<r/r=r>r?rArBr(s"    z RetinaNetRegressionHead.__init__rCrDrEcCrF)a| It takes a list of feature maps as inputs, and outputs a list of box regression maps. Each output box regression map has same spatial size with the corresponding input feature map, and the number of output channel is num_anchors * 2 * spatial_dims. Args: x: list of feature map, x[i] is a (B, in_channels, H_i, W_i) or (B, in_channels, H_i, W_i, D_i) Tensor. Return: box_regression_maps, list of box regression map. cls_logits_maps[i] is a (B, num_anchors * 2 * spatial_dims, H_i, W_i) or (B, num_anchors * 2 * spatial_dims, H_i, W_i, D_i) Tensor. zbox_regression is NaN or Inf.) r1rr/r[r+r2rGrHrIrJrKrLrM)r;rCZbox_regression_mapsrNrObox_regressionrArArBrPrQzRetinaNetRegressionHead.forward)rrrrrrrRrSrArAr?rBrYs rYcs2eZdZdZ  ddfdd ZdddZZS) RetinaNeta The network used in RetinaNet. It takes an image tensor as inputs, and outputs either 1) a dictionary ``head_outputs``. ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor. ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor. or 2) a list of 2N tensors ``head_outputs``, with first N tensors being the predicted classification maps and second N tensors being the predicted box regression maps. Args: spatial_dims: number of spatial dimensions of the images. We support both 2D and 3D images. num_classes: number of output classes of the model (excluding the background). num_anchors: number of anchors at each location. feature_extractor: a network that outputs feature maps from the input images, each feature map corresponds to a different resolution. Its output can have a format of Tensor, Dict[Any, Tensor], or Sequence[Tensor]. It can be the output of ``resnet_fpn_feature_extractor(*args, **kwargs)``. size_divisible: the spatial size of the network input should be divisible by size_divisible, decided by the feature_extractor. use_list_output: default False. If False, the network outputs a dictionary ``head_outputs``, ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor. ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor. If True, the network outputs a list of 2N tensors ``head_outputs``, with first N tensors being the predicted classification maps and second N tensors being the predicted box regression maps. Example: .. code-block:: python from monai.networks.nets import resnet spatial_dims = 3 # 3D network conv1_t_stride = (2,2,1) # stride of first convolutional layer in backbone backbone = resnet.ResNet( spatial_dims = spatial_dims, block = resnet.ResNetBottleneck, layers = [3, 4, 6, 3], block_inplanes = resnet.get_inplanes(), n_input_channels= 1, conv1_t_stride = conv1_t_stride, conv1_t_size = (7,7,7), ) # This feature_extractor outputs 4-level feature maps. # number of output feature maps is len(returned_layers)+1 returned_layers = [1,2,3] # returned layer from feature pyramid network feature_extractor = resnet_fpn_feature_extractor( backbone = backbone, spatial_dims = spatial_dims, pretrained_backbone = False, trainable_backbone_layers = None, returned_layers = returned_layers, ) # This feature_extractor requires input image spatial size # to be divisible by (32, 32, 16). size_divisible = tuple(2*s*2**max(returned_layers) for s in conv1_t_stride) model = RetinaNet( spatial_dims = spatial_dims, num_classes = 5, num_anchors = 6, feature_extractor=feature_extractor, size_divisible = size_divisible, ).to(device) result = model(torch.rand(2, 1, 128,128,128)) cls_logits_maps = result["classification"] # a list of len(returned_layers)+1 Tensor box_regression_maps = result["box_regression"] # a list of len(returned_layers)+1 Tensor rFrrrrfeature_extractor nn.Modulesize_divisibleSequence[int] | intuse_list_outputboolcstt|gdd|_||_t||j|_||_t|ds$t d||_ |j j |_ ||_ t|j |j |j|jd|_t|j |j |jd|_d|_d|_dS)NrrZr) supported out_channelszfeature_extractor should contain an attribute out_channels specifying the number of output channels (assumed to be the same for all the levels))rclassificationr])r'r(r rrr rarchasattrrKr_rgZfeature_map_channelsrrclassification_headrYregression_headcls_key box_reg_key)r;rrrr_rarcr?rArBr(s(     zRetinaNet.__init__imagesrrErcCs||}t|tr|g}ntj|tttfr t|}nt|}t|dts/t d|j sD|j | |i}| |||j<|S| || |}|S)ag It takes an image tensor as inputs, and outputs predicted classification maps and predicted box regression maps in ``head_outputs``. Args: images: input images, sized (B, img_channels, H, W) or (B, img_channels, H, W, D). Return: 1) If self.use_list_output is False, output a dictionary ``head_outputs`` with keys including self.cls_key and self.box_reg_key. ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor. ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor. 2) if self.use_list_output is True, outputs a list of 2N tensors ``head_outputs``, with first N tensors being the predicted classification maps and second N tensors being the predicted box regression maps. rzWfeature_extractor output format must be Tensor, Dict[str, Tensor], or Sequence[Tensor].)r_r1rr2jitdictstrlistvaluesrKrcrlrjrkrm)r;rnrOrN head_outputsZhead_outputs_sequencerArArBrP8s  zRetinaNet.forward)rF) rrrrrrr_r`rarbrcrd)rnrrErrSrArAr?rBr^s H$r^Frebackbone resnet.ResNetrrpretrained_backbonerdreturned_layers Sequence[int]trainable_backbone_layers int | NonerErcCs*t||ddd}t|||t|dd}|S)ah Constructs a feature extractor network with a ResNet-FPN backbone, used as feature_extractor in RetinaNet. Reference: `"Focal Loss for Dense Object Detection" `_. The returned feature_extractor network takes an image tensor as inputs, and outputs a dictionary that maps string to the extracted feature maps (Tensor). The input to the returned feature_extractor is expected to be a list of tensors, each of shape ``[C, H, W]`` or ``[C, H, W, D]``, one for each image. Different images can have different sizes. Args: backbone: a ResNet model, used as backbone. spatial_dims: number of spatial dimensions of the images. We support both 2D and 3D images. pretrained_backbone: whether the backbone has been pre-trained. 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. trainable_backbone_layers: number of trainable (not frozen) resnet layers starting from final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable. When pretrained_backbone is False, this value is set to be 5. When pretrained_backbone is True, if ``None`` is passed (the default) this value is set to 3. Example: .. code-block:: python from monai.networks.nets import resnet spatial_dims = 3 # 3D network backbone = resnet.ResNet( spatial_dims = spatial_dims, block = resnet.ResNetBottleneck, layers = [3, 4, 6, 3], block_inplanes = resnet.get_inplanes(), n_input_channels= 1, conv1_t_stride = (2,2,1), conv1_t_size = (7,7,7), ) # This feature_extractor outputs 4-level feature maps. # number of output feature maps is len(returned_layers)+1 feature_extractor = resnet_fpn_feature_extractor( backbone = backbone, spatial_dims = spatial_dims, pretrained_backbone = False, trainable_backbone_layers = None, returned_layers = [1,2,3], ) model = RetinaNet( spatial_dims = spatial_dims, num_classes = 5, num_anchors = 6, feature_extractor=feature_extractor, size_divisible = 32, ).to(device) r) max_value default_valueN)rx extra_blocks)rr rr)rurrwrxrzZvalid_trainable_backbone_layersr_rArArBresnet_fpn_feature_extractorcsBr)FreN) rurvrrrwrdrxryrzr{rEr)rW __future__rr9rLcollections.abcrrtypingrr2rrZ(monai.networks.blocks.backbone_fpn_utilsrr monai.networks.layers.factoriesr monai.networks.netsr monai.utilsr r rrr=ModulerrYr^rrArArArBs."     MH