#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__author__ = "Christian Heider Nielsen"
__doc__ = r"""
Created on 19-09-2021
"""
from typing import Dict, Iterable, Sequence, Tuple, Union
import numpy
import torch
from draugr.torch_utilities import to_tensor
from torch import nn
from torch.nn import init
from neodroidvision.multitask.fission.skip_hourglass.factory import (
fcn_decoder,
fcn_encoder,
)
from neodroidvision.multitask.fission.skip_hourglass.modes import MergeMode, UpscaleMode
__all__ = ["SkipHourglassFission"]
[docs]class SkipHourglassFission(nn.Module):
"""
Multi Headed Skip Fully Convolutional Network
Based on https://arxiv.org/abs/1505.04597
Contextual spatial information (from the decoding, expansive pathway) about an input tensor is merged
with information representing the localization of details (from the encoding, compressive pathway).
Modifications to the original paper:
(1) padding is used in 3x3 convolutions to prevent loss of border pixels
(2) merging outputs does not require cropping due to (1)
(3) residual connections can be used by specifying UNet(merge_mode='add')
(4) if non-parametric upsampling is used in the decoder
pathway (specified by upmode='upsample'), then an
additional 1x1 2d convolution occurs after upsampling
to reduce channel dimensionality by a factor of 2.
This channel halving happens with the convolution in
the tranpose convolution (specified by upmode='fractional')"""
[docs] def parse_arguments(self, up_mode: UpscaleMode, merge_mode: MergeMode):
"""
Args:
up_mode:
merge_mode:
"""
if up_mode in UpscaleMode:
self.up_mode = up_mode
else:
raise ValueError(
f'"{up_mode}" is not a valid mode for upsampling.'
f' Only "fractional" and "upsample" are allowed.'
)
if merge_mode in MergeMode:
self.merge_mode = merge_mode
else:
raise ValueError(
f'"{up_mode}" is not a valid mode for merging up and down paths.'
f' Only "concat" and "add" are allowed.'
)
if self.up_mode == UpscaleMode.Upsample and self.merge_mode == MergeMode.Add:
# NOTE: up_mode 'upsample' is incompatible with merge_mode 'add'
raise ValueError(
'up_mode "upsample" is incompatible with merge_mode "add"'
"can not use nearest neighbour to reduce depth channels (by half)."
)
[docs] def __init__(
self,
*,
input_channels: int,
output_heads: Union[Dict, Iterable],
encoding_depth: int = 5,
start_channels: int = 32,
up_mode: UpscaleMode = UpscaleMode.FractionalTranspose,
merge_mode: MergeMode = MergeMode.Add,
):
"""
:type input_channels: int
:type output_heads: Union[Dict, Iterable]
:type encoding_depth: int
:type start_channels: int
:type up_mode: str
:type merge_mode: str
:param input_channels: number of channels in the input tensor. E.g. 3 for RGB images.
:param output_heads:
:param encoding_depth: number of MaxPools in the U-Net.
:param start_channels: number of convolutional filters for the first convolution.
:param up_mode: string, type of upconvolution. Choices: 'fractional' for transpose convolution or
'upsample' for nearest neighbour upsampling.
:param merge_mode:"""
super().__init__()
self.parse_arguments(up_mode, merge_mode)
self.start_channels = start_channels
self.network_depth = encoding_depth
down_convolutions, encoding_channels = fcn_encoder(
input_channels, self.network_depth, self.start_channels
)
self.down_convolutions = nn.ModuleList(down_convolutions)
self.default_prefix = "fork"
self._dict_output = False
if isinstance(output_heads, Dict):
self.forks = output_heads
self._dict_output = True
else:
self.forks = {
f"{self.default_prefix}{i}": s for i, s in enumerate(output_heads)
}
for iden, channel_size in self.forks.items():
up_convolutions_ae, ae_prev_layer_channels = fcn_decoder(
encoding_channels, self.network_depth, self.up_mode, self.merge_mode
)
up_convolutions_ae.append(
nn.Conv2d(
ae_prev_layer_channels,
channel_size,
kernel_size=1,
groups=1,
stride=1,
)
)
setattr(self, iden, nn.ModuleList(up_convolutions_ae))
self.reset_params()
[docs] @staticmethod
def weight_init(m: nn.Module) -> None:
"""
Args:
m:
"""
if isinstance(m, nn.Conv2d):
init.xavier_normal_(m.weight)
init.constant_(m.bias, 0)
[docs] def reset_params(self) -> None:
"""description"""
for i, m in enumerate(self.modules()):
self.weight_init(m)
[docs] def forward(
self, x_enc: torch.Tensor
) -> Union[Tuple[torch.Tensor], Dict[str, torch.Tensor]]:
"""
Args:
x_enc:
Returns:
"""
encoder_skips = []
for i, module in enumerate(
self.down_convolutions
): # encoder pathway, keep outputs for merging
x_enc, before_pool = module(x_enc)
encoder_skips.append(before_pool)
out = {}
for key in self.forks.keys():
x_prev = x_enc
fork_i = getattr(self, key)
for j, module in enumerate(fork_i[:-1]):
before_pool_skip = encoder_skips[-(j + 2)]
x_prev = module(before_pool_skip, x_prev)
out[key] = fork_i[-1](x_prev)
if self._dict_output:
return out
return (*out.values(),)
[docs] def trim(self, idx: Sequence[Union[str, int]]) -> None:
"""
Args:
idx:
"""
if not isinstance(idx, Sequence):
idx = list(idx)
for a in idx:
if not isinstance(a, str):
a = f"{self.default_prefix}{a}"
delattr(self, a)
if __name__ == "__main__":
def _main():
from matplotlib import pyplot
channels = 3
model = SkipHourglassFission(
input_channels=channels,
output_heads=(channels, 1),
encoding_depth=2,
merge_mode=MergeMode.Concat,
)
x = to_tensor(
numpy.random.random((1, channels, 320, 320)), device="cpu"
).float()
out, out2, *_ = model(x)
loss = torch.sum(out)
loss.backward()
im = out.detach()
print(im.shape)
pyplot.imshow((torch.tanh(im[0].transpose(2, 0)) + 1) * 0.5)
pyplot.show()
im2 = out2.detach()
print(im2.shape)
pyplot.imshow((torch.tanh(im2[0][0, :, :]) + 1) * 0.5)
pyplot.show()
_main()
