| 6 | from network_swinir import SwinIR |
| 7 | |
| 8 | class reconstructor(nn.Module): |
| 9 | # the version of reconstructor is set as FBP-DuDo-SwinIR to test the effectiveness of dual domain SwinIR in Recon |
| 10 | def __init__(self, dset, img_h=512, img_w=512): |
| 11 | super(reconstructor, self).__init__() |
| 12 | self.Image_recon_module=SwinIR(upscale=1, img_size=(img_h, img_w), in_chans=2, |
| 13 | window_size=8, img_range=1., depths=[2, 2, 2, 2], |
| 14 | embed_dim=60, num_heads=[2, 2, 2, 2], mlp_ratio=2, upsampler='') |
| 15 | self.Sinogram_recon_module=SwinIR(upscale=1, img_size=(96, 800), |
| 16 | window_size=8, img_range=1., depths=[1, 1, 1], |
| 17 | embed_dim=60, num_heads=[2, 2, 2], mlp_ratio=2, upsampler='') |
| 18 | self.dset=dset |
| 19 | self.fp_senet_gt, _, self.ril_odl, __ = self.dset.ril() |
| 20 | def forward(self, img, img_gt, sinos): |
| 21 | sinos_gt=self.radon_senet_gt(img_gt) |
| 22 | sinos_enhanced=self.Sinogram_recon_module(sinos) |
| 23 | img_ril=self.ril(sinos_enhanced) |
| 24 | img_input=torch.cat((img, img_ril), 1) |
| 25 | reconstructed_img=self.Image_recon_module(img_input) |
| 26 | return sinos_gt, sinos_enhanced, img_ril, reconstructed_img |
| 27 | |
| 28 | def ril(self, img): |
| 29 | # return ril results of enhanced sinograms |
| 30 | if len(img.shape) == 4: |
| 31 | img = img.squeeze(1) |
| 32 | return self.ril_odl(img).unsqueeze(1) |
| 33 | |
| 34 | def radon_senet_gt(self, img): |
| 35 | # return supervision of Sinogram_recon_module |
| 36 | if len(img.shape)==4: |
| 37 | img=img.squeeze(1) |
| 38 | return self.fp_senet_gt(img).unsqueeze(1) |
| 39 | |
| 40 | class reconstructor_loss(nn.Module): |
| 41 | def __init__(self): |
no outgoing calls
no test coverage detected