| 448 | |
| 449 | |
| 450 | class MscaleBasic(MscaleBase): |
| 451 | """ |
| 452 | """ |
| 453 | def __init__(self, num_classes, trunk='hrnetv2', criterion=None): |
| 454 | super(MscaleBasic, self).__init__() |
| 455 | self.criterion = criterion |
| 456 | self.backbone, _, _, high_level_ch = get_trunk( |
| 457 | trunk_name=trunk, output_stride=8) |
| 458 | |
| 459 | self.cls_head = make_seg_head(in_ch=high_level_ch, |
| 460 | out_ch=num_classes) |
| 461 | self.scale_attn = make_attn_head(in_ch=high_level_ch, |
| 462 | out_ch=1) |
| 463 | |
| 464 | def _fwd(self, x, aspp_lo=None, aspp_attn=None, scale_float=None): |
| 465 | _, _, final_features = self.backbone(x) |
| 466 | attn = self.scale_attn(final_features) |
| 467 | pred = self.cls_head(final_features) |
| 468 | attn = scale_as(attn, x) |
| 469 | pred = scale_as(pred, x) |
| 470 | |
| 471 | return pred, attn, None, None |
| 472 | |
| 473 | |
| 474 | def HRNet(num_classes, criterion, s2s4=None): |