(self, n_class=1, img_size_s1=(256,256), img_size_s2=(224,224), model_scale='small', decoder_aggregation='additive', interpolation='bilinear')
| 423 | |
| 424 | class MERIT_Parallel(nn.Module): |
| 425 | def __init__(self, n_class=1, img_size_s1=(256,256), img_size_s2=(224,224), model_scale='small', decoder_aggregation='additive', interpolation='bilinear'): |
| 426 | super(MERIT_Parallel, self).__init__() |
| 427 | |
| 428 | self.n_class = n_class |
| 429 | self.img_size_s1 = img_size_s1 |
| 430 | self.img_size_s2 = img_size_s2 |
| 431 | self.model_scale = model_scale |
| 432 | self.decoder_aggregation = decoder_aggregation |
| 433 | self.interpolation = interpolation |
| 434 | |
| 435 | # conv block to convert single channel to 3 channels |
| 436 | self.conv = nn.Sequential( |
| 437 | nn.Conv2d(1, 3, kernel_size=1), |
| 438 | nn.BatchNorm2d(3), |
| 439 | nn.ReLU(inplace=True) |
| 440 | ) |
| 441 | |
| 442 | # backbone network initialization with pretrained weight |
| 443 | self.backbone1 = load_pretrained_weights(self.img_size_s1[0], self.model_scale) |
| 444 | self.backbone2 = load_pretrained_weights(self.img_size_s2[0], self.model_scale) |
| 445 | |
| 446 | if(self.model_scale=='tiny'): |
| 447 | self.channels = [512, 256, 128, 64] |
| 448 | elif(self.model_scale=='small'): |
| 449 | self.channels = [768, 384, 192, 96] |
| 450 | |
| 451 | # shared decoder initialization |
| 452 | if(self.decoder_aggregation=='additive'): |
| 453 | self.decoder = CASCADE_Add(channels=self.channels) |
| 454 | elif(self.decoder_aggregation=='concatenation'): |
| 455 | self.decoder = CASCADE_Cat(channels=self.channels) |
| 456 | else: |
| 457 | sys.exit("'"+self.decoder_aggregation+"' is not a valid decoder aggregation! Currently supported aggregations are 'additive' and 'concatenation'.") |
| 458 | |
| 459 | # Prediction heads initialization |
| 460 | self.out_head1 = nn.Conv2d(self.channels[0], self.n_class, 1) |
| 461 | self.out_head2 = nn.Conv2d(self.channels[1], self.n_class, 1) |
| 462 | self.out_head3 = nn.Conv2d(self.channels[2], self.n_class, 1) |
| 463 | self.out_head4 = nn.Conv2d(self.channels[3], self.n_class, 1) |
| 464 | |
| 465 | def forward(self, x): |
| 466 |
nothing calls this directly
no test coverage detected