import copy import torch from torch import nn from .loss import Loss from ..backbone import ResUNet from ..roi_heads import BoxHead, BoxModule from ..rpn import RpnHead, RpnModule from ..util import convert_to_roi_format, pad2factor, norm, get_anchors, add_rcnn_probability from ...layers.roi_align_3d import ROIAlign3D @torch.jit.script def make_rpn_windows(fs, stride, anchors): """ Generating anchor boxes at each voxel on the feature map, the center of the anchor box on each voxel corresponds to center on the original input image. return windows: list of anchor boxes, [z, y, x, d, h, w] """ offset = (stride - 1) / 2 device = fs.device # anchors = torch.tensor(anchors, dtype=torch.float32, device=device) # D, H , W =(i/cfg['stride'] for i in cfg['crop_size']) _, _, D, H, W = fs.shape oz = torch.arange(offset, offset + stride * (D - 1) + 1, step=stride, dtype=torch.float32, device=device) oh = torch.arange(offset, offset + stride * (H - 1) + 1, step=stride, dtype=torch.float32, device=device) ow = torch.arange(offset, offset + stride * (W - 1) + 1, step=stride, dtype=torch.float32, device=device) oanchor = torch.arange(anchors.size()[0], dtype=torch.float32, device=device) shift_z, shift_y, shift_x, anchor_idx = torch.meshgrid(oz, oh, ow, oanchor) shift_z = shift_z.reshape(-1) shift_y = shift_y.reshape(-1) shift_x = shift_x.reshape(-1) anchor_idx = anchor_idx.reshape(-1).long() shift_anchors = torch.index_select(anchors, 0, anchor_idx) shift_anchor_z = shift_anchors[:, 0].squeeze() shift_anchor_y = shift_anchors[:, 1].squeeze() shift_anchor_x = shift_anchors[:, 2].squeeze() windows = torch.stack((shift_z, shift_y, shift_x, shift_anchor_z, shift_anchor_y, shift_anchor_x), dim=1) return windows class NoduleNet(nn.Module): def __init__(self, cfg, mode='train'): super(NoduleNet, self).__init__() self.cfg = cfg self.mode = mode if cfg.MODEL.BACKBONE.CONV_BODY == 'ResUNet': self.feature_net = ResUNet(cfg, 1, 128) else: raise ValueError('Detector backbone %s is not implemented.' % self.cfg.MODEL.BACKBONE.CONV_BODY) self.stride = torch.tensor(cfg.DATA.DATA_PROCESS.STRIDE, dtype=torch.float32) self.anchors = torch.tensor(get_anchors(cfg.MODEL.ANCHOR.BASES, cfg.MODEL.ANCHOR.ASPECT_RATIOS), dtype=torch.float32).cuda() self.rpn_head = RpnHead(len(cfg.MODEL.ANCHOR.BASES), in_channels=128) self.box_head = BoxHead(cfg.MODEL.ROI_BOX_HEAD.NUM_CLASS, cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION, in_channels=128) self.rpn = RpnModule(cfg, self.mode) self.roi_box = BoxModule(cfg) self.roi_crop = ROIAlign3D(cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION, spatial_scale=1 / cfg.DATA.DATA_PROCESS.STRIDE, sampling_ratio=cfg.MODEL.ROI_BOX_HEAD.SAMPLING_RATIO) self.use_rcnn = False self.loss = Loss(self.rpn.loss, self.roi_box.loss, self.use_rcnn) def forward(self, inputs, truth_bboxes=None): if truth_bboxes is None: return self.forward_deploy(inputs) features, feat_4 = self.feature_net(inputs) self.rpn_window = make_rpn_windows(features[-1], self.stride, self.anchors) self.rpn_logits, self.rpn_deltas = self.rpn_head(features[-1]) if self.mode in ['train', 'valid']: self.rpn_labels, self.rpn_label_assigns, self.rpn_label_weights, \ self.rpn_targets, self.rpn_target_weights = self.rpn.make_rpn_target(self.rpn_window, truth_bboxes) if self.use_rcnn: self.rpn_proposals = self.rpn.rpn_nms(inputs, self.rpn_window, self.rpn_logits, self.rpn_deltas) self.rpn_proposals, self.rcnn_labels, self.rcnn_targets = self.roi_box.make_box_target( self.rpn_proposals, truth_bboxes) if len(self.rpn_proposals) > 0: self.rpn_proposal_rois = convert_to_roi_format(self.rpn_proposals) self.rcnn_logits, self.rcnn_deltas = self.box_head(self.roi_crop(feat_4, self.rpn_proposal_rois)) else: self.rcnn_logits, self.rcnn_deltas, self.rcnn_labels, self.rcnn_targets = None, None, None, None return self.rpn_logits, self.rpn_deltas, \ self.rpn_labels, self.rpn_label_weights, self.rpn_targets, self.rpn_target_weights, \ self.rcnn_logits, self.rcnn_deltas, self.rcnn_labels, self.rcnn_targets elif self.mode in ['test']: self.rpn_proposals = self.rpn.rpn_nms(inputs, self.rpn_window, self.rpn_logits, self.rpn_deltas) if self.use_rcnn: if len(self.rpn_proposals) > 0: self.ensemble_proposals = copy.deepcopy(self.rpn_proposals) self.rpn_proposal_rois = convert_to_roi_format(self.rpn_proposals) self.rcnn_logits, self.rcnn_deltas = self.box_head(self.roi_crop(feat_4, self.rpn_proposal_rois)) self.fpr_res = self.roi_box.get_probability(inputs, self.rpn_proposals, self.rcnn_logits, self.rcnn_deltas) self.ensemble_proposals[:, 1] = (self.ensemble_proposals[:, 1] + self.fpr_res[:, 0]) / 2 return self.rpn_proposals, self.fpr_res, self.ensemble_proposals else: raise RuntimeError("The nms pre score threshold {} is too high, change for a lower one.".format( self.pre_nms_score_threshold)) else: return self.rpn_proposals else: raise ValueError('rpn_nms(): invalid mode = %s?' % self.mode) def forward_deploy(self, inputs): inputs = pad2factor(inputs) inputs = norm(inputs) features, feat_4 = self.feature_net(inputs) self.rpn_window = make_rpn_windows(features[-1], self.stride, self.anchors) self.rpn_logits, self.rpn_deltas = self.rpn_head(features[-1]) self.rpn_proposals = self.rpn.rpn_nms_deploy(inputs, self.rpn_window, self.rpn_logits, self.rpn_deltas) # self.rpn_proposal_rois = convert_to_roi_format(self.rpn_proposals) # self.rcnn_logits, self.rcnn_deltas = self.box_head( # self.roi_crop(feat_4.float(), self.rpn_proposal_rois.float()).half()) # self.rpn_proposals = add_rcnn_probability(self.rpn_proposals, self.rcnn_logits) return self.rpn_proposals def set_mode(self, mode): assert mode in ['train', 'valid', 'test'] self.mode = mode self.rpn.update_mode(self.mode) # model.eval() 在模型测试阶段使用，让model自动把 BN 和 dropout 固定住，不会取平均，而是用训练好的值。 # model.train() 在模型训练阶段使用，此时 dropout 和 BN 的操作在训练起到防止网络过拟合的问题。 if mode in ['train']: self.train() else: self.eval()