U PhdJ@sdZddlmZddlZddlZddlmZmZddlm Z m Z ddl Z ddl m Z m Z ddlmZmZddlmZdd lmZdd lmZmZmZed d d \ZZGddde jZGddde jZGddde jZddddddddddZdS) z{ Part of this script is adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py ) annotationsN)CallableSequence)AnyDict)Tensornn)BackboneWithFPN_resnet_fpn_extractor)Conv)resnet)ensure_tuple_replook_up_optionoptional_importz+torchvision.models.detection.backbone_utils_validate_trainable_layers)namecs@eZdZdZd ddddddfdd Zddd d d ZZS) 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?intfloat) in_channels num_anchors num_classes spatial_dimsprior_probabilityc stttj|f}g}tdD]>}||||dddd|tjd|d|tq$tj ||_ |j D]4} t | |rzt jjj| jddt jj| jd qz||||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_typer0_layer __class__d/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/detection/networks/retinanet_network.pyr)Ls"    $z$RetinaNetClassificationHead.__init__ list[Tensor]xreturncCs~g}t|tr|g}n|}|D]Z}||}||}||t|s\t|rt rnt dqt 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.) r2rr0r9r,r3isnananyisinfis_grad_enabled ValueErrorwarningswarn)r<rFZcls_logits_maps feature_mapsfeaturesr9rBrBrCforwardes      z#RetinaNetClassificationHead.forward)r__name__ __module__ __qualname____doc__r)rQ __classcell__rBrBr@rCr<srcs:eZdZdZddddfdd Zddddd 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. r)rrrc stttj|f}g}tdD]>}||||dddd|tjd|d|tq$tj ||_ |||d|dddd|_ t jj j|j jdd t jj |j j|j D]2}t||rt jj j|jdd t jj |jqdS) Nrrrrr"r#rr&)r(r)r r*r+r,rr-r.r/r0bbox_regr3r4r5r6zeros_r8r1r2)r<rrrr=r0r>r?r@rBrCr)s    z RetinaNetRegressionHead.__init__rDrEcCs~g}t|tr|g}n|}|D]Z}||}||}||t|s\t|rt rnt dqt dq|S)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.) r2rr0rZr,r3rHrIrJrKrLrMrN)r<rFZbox_regression_mapsrOrPbox_regressionrBrBrCrQs      zRetinaNetRegressionHead.forwardrRrBrBr@rCrXs rXcsBeZdZdZddddddddfd d Zd d d ddZZS) 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 rFrz nn.ModulezSequence[int] | intbool)rrrfeature_extractorsize_divisibleuse_list_outputcstt|dddgd|_||_t||j|_||_t|dsJt d||_ |j j |_ ||_ t|j |j |j|jd|_t|j |j |jd|_d|_d |_dS) NrrYr) 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)rrrr r`rahasattrrLr_rcZfeature_map_channelsrrclassification_headrXregression_headcls_key box_reg_key)r<rrrr_r`rar@rBrCr)s2   zRetinaNet.__init__rr)imagesrGcCs||}t|tr|g}n,tj|tttfr@t|}nt|}t|dts^t d|j s|j | |i}| |||j<|S| || |}|SdS)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].N)r_r2rr3jitrstrlistvaluesrLrarhrfrgri)r<rjrPrO head_outputsZhead_outputs_sequencerBrBrCrQ8s  zRetinaNet.forward)rFrRrBrBr@rCr]s H$r]FrrYrz resnet.ResNetrr^z Sequence[int]z int | Noner )backbonerpretrained_backbonereturned_layerstrainable_backbone_layersrGcCs*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)rs extra_blocks)rr rm)rqrrrrsrtZvalid_trainable_backbone_layersr_rBrBrCresnet_fpn_feature_extractorcsBry)FrpN) rV __future__rr:rMcollections.abcrrtypingrrr3rrZ(monai.networks.blocks.backbone_fpn_utilsr r monai.networks.layers.factoriesr monai.networks.netsr monai.utilsr rrrr>ModulerrXr]ryrBrBrBrC#s.    MH