OCR net
| 123 | |
| 124 | |
| 125 | class OCRNetASPP(nn.Module): |
| 126 | """ |
| 127 | OCR net |
| 128 | """ |
| 129 | def __init__(self, num_classes, trunk='hrnetv2', criterion=None): |
| 130 | super(OCRNetASPP, self).__init__() |
| 131 | self.criterion = criterion |
| 132 | self.backbone, _, _, high_level_ch = get_trunk(trunk) |
| 133 | self.aspp, aspp_out_ch = get_aspp(high_level_ch, |
| 134 | bottleneck_ch=256, |
| 135 | output_stride=8) |
| 136 | self.ocr = OCR_block(aspp_out_ch) |
| 137 | |
| 138 | def forward(self, inputs): |
| 139 | assert 'images' in inputs |
| 140 | x = inputs['images'] |
| 141 | |
| 142 | _, _, high_level_features = self.backbone(x) |
| 143 | aspp = self.aspp(high_level_features) |
| 144 | cls_out, aux_out, _ = self.ocr(aspp) |
| 145 | aux_out = scale_as(aux_out, x) |
| 146 | cls_out = scale_as(cls_out, x) |
| 147 | |
| 148 | if self.training: |
| 149 | gts = inputs['gts'] |
| 150 | loss = cfg.LOSS.OCR_ALPHA * self.criterion(aux_out, gts) + \ |
| 151 | self.criterion(cls_out, gts) |
| 152 | return loss |
| 153 | else: |
| 154 | output_dict = {'pred': cls_out} |
| 155 | return output_dict |
| 156 | |
| 157 | |
| 158 | class MscaleOCR(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected