OCR net
| 156 | |
| 157 | |
| 158 | class MscaleOCR(nn.Module): |
| 159 | """ |
| 160 | OCR net |
| 161 | """ |
| 162 | def __init__(self, num_classes, trunk='hrnetv2', criterion=None): |
| 163 | super(MscaleOCR, self).__init__() |
| 164 | self.criterion = criterion |
| 165 | self.backbone, _, _, high_level_ch = get_trunk(trunk) |
| 166 | self.ocr = OCR_block(high_level_ch) |
| 167 | self.scale_attn = make_attn_head( |
| 168 | in_ch=cfg.MODEL.OCR.MID_CHANNELS, out_ch=1) |
| 169 | |
| 170 | def _fwd(self, x): |
| 171 | x_size = x.size()[2:] |
| 172 | |
| 173 | _, _, high_level_features = self.backbone(x) |
| 174 | cls_out, aux_out, ocr_mid_feats = self.ocr(high_level_features) |
| 175 | attn = self.scale_attn(ocr_mid_feats) |
| 176 | |
| 177 | aux_out = Upsample(aux_out, x_size) |
| 178 | cls_out = Upsample(cls_out, x_size) |
| 179 | attn = Upsample(attn, x_size) |
| 180 | |
| 181 | return {'cls_out': cls_out, |
| 182 | 'aux_out': aux_out, |
| 183 | 'logit_attn': attn} |
| 184 | |
| 185 | def nscale_forward(self, inputs, scales): |
| 186 | """ |
| 187 | Hierarchical attention, primarily used for getting best inference |
| 188 | results. |
| 189 | |
| 190 | We use attention at multiple scales, giving priority to the lower |
| 191 | resolutions. For example, if we have 4 scales {0.5, 1.0, 1.5, 2.0}, |
| 192 | then evaluation is done as follows: |
| 193 | |
| 194 | p_joint = attn_1.5 * p_1.5 + (1 - attn_1.5) * down(p_2.0) |
| 195 | p_joint = attn_1.0 * p_1.0 + (1 - attn_1.0) * down(p_joint) |
| 196 | p_joint = up(attn_0.5 * p_0.5) * (1 - up(attn_0.5)) * p_joint |
| 197 | |
| 198 | The target scale is always 1.0, and 1.0 is expected to be part of the |
| 199 | list of scales. When predictions are done at greater than 1.0 scale, |
| 200 | the predictions are downsampled before combining with the next lower |
| 201 | scale. |
| 202 | |
| 203 | Inputs: |
| 204 | scales - a list of scales to evaluate |
| 205 | inputs - dict containing 'images', the input, and 'gts', the ground |
| 206 | truth mask |
| 207 | |
| 208 | Output: |
| 209 | If training, return loss, else return prediction + attention |
| 210 | """ |
| 211 | x_1x = inputs['images'] |
| 212 | |
| 213 | assert 1.0 in scales, 'expected 1.0 to be the target scale' |
| 214 | # Lower resolution provides attention for higher rez predictions, |
| 215 | # so we evaluate in order: high to low |