(self, x, mask, background=None)
| 325 | return out |
| 326 | |
| 327 | def processImage(self, x, mask, background=None): |
| 328 | if background is not None: |
| 329 | x = x*mask + background * (1 - mask) |
| 330 | if self.input_nc == 4: |
| 331 | x = torch.cat([x, mask], dim=1) # (bs, 4, 256, 256) |
| 332 | pred = self.forward(x, mask) |
| 333 | |
| 334 | return pred * mask + x[:,:3,:,:] * (1 - mask) |
| 335 | |
| 336 | class UnetBlockCodec(nn.Module): |
| 337 | """Defines the Unet submodule with skip connection. |