<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function

from collections import namedtuple

import tensorflow as tf

from niftynet.layer import layer_util
from niftynet.layer.base_layer import TrainableLayer
from niftynet.layer.convolution import ConvolutionalLayer
from niftynet.layer.channel_sparse_convolution import ChannelSparseConvolutionalLayer
from niftynet.layer.bn import BNLayer
from niftynet.layer.spatial_transformer import ResamplerLayer
from niftynet.layer.grid_warper import AffineGridWarperLayer

from niftynet.network.base_net import BaseNet
from niftynet.io.misc_io import image3_axial


# Create a structure with all the fields of a DenseVNet network
DenseVNetDesc = namedtuple(
    'DenseVNetParts',
    ['initial_bn', 'initial_conv', 'dense_vblocks', 'seg_layer']
)


class DenseVNet(BaseNet):
    """
    implementation of Dense-V-Net:
       <NAME> al.
       Automatic multi-organ segmentation on abdominal CT with dense V-networks

    ### Diagram

    DFS = Dense Feature Stack Block

    - Initial image is first downsampled to a given size.
    - Each DFS+SD outputs a skip link + a downsampled output.
    - All outputs are upscaled to the initial downsampled size.
    - If initial prior is given add it to the output prediction.

    Input
      |
      --[ DFS ]-----------------------[ Conv ]------------[ Conv ]------[+]-->
           |                                       |  |              |
           -----[ DFS ]---------------[ Conv ]------  |              |
                   |                                  |              |
                   -----[ DFS ]-------[ Conv ]---------              |
                                                          [ Prior ]---

    The layer DenseFeatureStackBlockWithSkipAndDownsample layer implements
    [DFS + Conv + Downsampling] in a single module, and outputs 2 elements:
        - Skip layer:          [ DFS + Conv]
        - Downsampled output:  [ DFS + Down]

    """

    __hyper_params__ = dict(
        prior_size=12,
        p_channels_selected=0.5,
        n_dense_channels=(4, 8, 16),
        n_seg_channels=(12, 24, 24),
        n_input_channels=(24, 24, 24),
        dilation_rates=([1] * 5, [1] * 10, [1] * 10),
        final_kernel=3,
        augmentation_scale=0.1
    )

    __net_params__ = dict(
        use_bdo=False,
        use_prior=False,
        use_dense_connections=True,
        use_coords=False
    )

    def __init__(self,
                 num_classes,
                 hyperparameters={},
                 architecture_parameters={},
                 w_initializer=None,
                 w_regularizer=None,
                 b_initializer=None,
                 b_regularizer=None,
                 acti_func='relu',
                 name='DenseVNet'):

        super(DenseVNet, self).__init__(
            num_classes=num_classes,
            w_initializer=w_initializer,
            w_regularizer=w_regularizer,
            b_initializer=b_initializer,
            b_regularizer=b_regularizer,
            acti_func=acti_func,
            name=name)

        # Override default Hyperparameters
        self.hyperparameters = dict(self.__hyper_params__)
        self.hyperparameters.update(hyperparameters)

        # Check for dilation rates
        if any([d != 1 for ds in self.hyperparameters['dilation_rates']
                for d in ds]):
            raise NotImplementedError(
                'Dilated convolutions are not yet implemented')

        # Override default architectural parameters
        self.architecture_parameters = dict(self.__net_params__)
        self.architecture_parameters.update(architecture_parameters)

        # Check available modes
        if self.architecture_parameters['use_dense_connections'] is False:
            raise NotImplementedError(
                'Non-dense connections are not yet implemented')
        if self.architecture_parameters['use_coords'] is True:
            raise NotImplementedError(
                'Image coordinate augmentation is not yet implemented')

    def create_network(self):
        hp = self.hyperparameters

        # Initial Convolution
        net_initial_conv = ConvolutionalLayer(
            hp['n_input_channels'][0],
            kernel_size=5, stride=2
        )

        # Dense Block Params
        downsample_channels = list(hp['n_input_channels'][1:]) + [None]
        num_blocks = len(hp["n_dense_channels"])
        use_bdo = self.architecture_parameters['use_bdo']

        # Create DenseBlocks
        net_dense_vblocks = []

        for idx in range(num_blocks):
            dense_ch = hp["n_dense_channels"][idx]  # Number or dense channels
            seg_ch = hp["n_seg_channels"][idx]      # Number of segmentation ch
            down_ch = downsample_channels[idx]      # Number of downsampling ch
            dil_rate = hp["dilation_rates"][idx]    # Dilation rate

            # Dense feature block
            dblock = DenseFeatureStackBlockWithSkipAndDownsample(
                dense_ch, 3, dil_rate, seg_ch, down_ch, use_bdo,
                acti_func='relu'
            )

            net_dense_vblocks.append(dblock)

        # Segmentation
        net_seg_layer = ConvolutionalLayer(
            self.num_classes, kernel_size=hp['final_kernel'],
            with_bn=False, with_bias=True
        )

        return DenseVNetDesc(initial_bn=BNLayer(),
                             initial_conv=net_initial_conv,
                             dense_vblocks=net_dense_vblocks,
                             seg_layer=net_seg_layer)

    def downsample_input(self, input_tensor, n_spatial_dims):
        # Initial downsampling params
        d_size1 = (1,) + (3,) * n_spatial_dims + (1,)
        d_size2 = (1,) + (2,) * n_spatial_dims + (1,)

        # Downsample input
        return tf.nn.avg_pool3d(input_tensor, d_size1, d_size2, 'SAME')

    def layer_op(self, input_tensor, is_training, layer_id=-1):
        hp = self.hyperparameters

        # Initialize DenseVNet network layers
        net = self.create_network()

        #
        # Parameter handling
        #

        # Shape and dimension variable shortcuts
        channel_dim = len(input_tensor.shape) - 1
        input_size = input_tensor.shape.as_list()
        spatial_size = input_size[1:-1]
        n_spatial_dims = input_tensor.shape.ndims - 2

        # Quick access to hyperparams
        pkeep = hp['p_channels_selected']

        # Validate input dimension with dilation rates
        modulo = 2 ** (len(hp['dilation_rates']))
        assert layer_util.check_spatial_dims(input_tensor,
                                             lambda x: x % modulo == 0)

        #
        # Augmentation + Downsampling + Initial Layers
        #

        # On the fly data augmentation
        if is_training and hp['augmentation_scale'] > 0:
            augment_layer = Affine3DAugmentationLayer(hp['augmentation_scale'],
                                                      'LINEAR', 'ZERO')
            input_tensor = augment_layer(input_tensor)

        # Variable storing all intermediate results -- VLinks
        all_segmentation_features = []

        # Downsample input to the network
        down_tensor = self.downsample_input(input_tensor, n_spatial_dims)
        downsampled_img = net.initial_bn(down_tensor, is_training=is_training)

        # Add initial downsampled image VLink
        all_segmentation_features.append(downsampled_img)

        # All results should match the downsampled input's shape
        output_shape = downsampled_img.shape.as_list()[1:-1]

        init_features = net.initial_conv(input_tensor, is_training=is_training)

        #
        # Dense VNet Main Block
        #

        # `down` will handle the input of each Dense VNet block<commit_msg>enable 2d slices for dense_vnet<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function

from collections import namedtuple

import tensorflow as tf

from niftynet.layer import layer_util
from niftynet.layer.base_layer import TrainableLayer
from niftynet.layer.convolution import ConvolutionalLayer
from niftynet.layer.channel_sparse_convolution import ChannelSparseConvolutionalLayer
from niftynet.layer.bn import BNLayer
from niftynet.layer.spatial_transformer import ResamplerLayer
from niftynet.layer.grid_warper import AffineGridWarperLayer

from niftynet.network.base_net import BaseNet
from niftynet.io.misc_io import image3_axial


# Create a structure with all the fields of a DenseVNet network
DenseVNetDesc = namedtuple(
    'DenseVNetParts',
    ['initial_bn', 'initial_conv', 'dense_vblocks', 'seg_layer']
)


class DenseVNet(BaseNet):
    """
    implementation of Dense-V-Net:
       Gibson et al.
       Automatic multi-organ segmentation on abdominal CT with dense V-networks

    ### Diagram

    DFS = Dense Feature Stack Block

    - Initial image is first downsampled to a given size.
    - Each DFS+SD outputs a skip link + a downsampled output.
    - All outputs are upscaled to the initial downsampled size.
    - If initial prior is given add it to the output prediction.

    Input
      |
      --[ DFS ]-----------------------[ Conv ]------------[ Conv ]------[+]-->
           |                                       |  |              |
           -----[ DFS ]---------------[ Conv ]------  |              |
                   |                                  |              |
                   -----[ DFS ]-------[ Conv ]---------              |
                                                          [ Prior ]---

    The layer DenseFeatureStackBlockWithSkipAndDownsample layer implements
    [DFS + Conv + Downsampling] in a single module, and outputs 2 elements:
        - Skip layer:          [ DFS + Conv]
        - Downsampled output:  [ DFS + Down]

    """

    __hyper_params__ = dict(
        prior_size=12,
        p_channels_selected=0.5,
        n_dense_channels=(4, 8, 16),
        n_seg_channels=(12, 24, 24),
        n_input_channels=(24, 24, 24),
        dilation_rates=([1] * 5, [1] * 10, [1] * 10),
        final_kernel=3,
        augmentation_scale=0.1
    )

    __net_params__ = dict(
        use_bdo=False,
        use_prior=False,
        use_dense_connections=True,
        use_coords=False
    )

    def __init__(self,
                 num_classes,
                 hyperparameters={},
                 architecture_parameters={},
                 w_initializer=None,
                 w_regularizer=None,
                 b_initializer=None,
                 b_regularizer=None,
                 acti_func='relu',
                 name='DenseVNet'):

        super(DenseVNet, self).__init__(
            num_classes=num_classes,
            w_initializer=w_initializer,
            w_regularizer=w_regularizer,
            b_initializer=b_initializer,
            b_regularizer=b_regularizer,
            acti_func=acti_func,
            name=name)

        # Override default Hyperparameters
        self.hyperparameters = dict(self.__hyper_params__)
        self.hyperparameters.update(hyperparameters)

        # Check for dilation rates
        if any([d != 1 for ds in self.hyperparameters['dilation_rates']
                for d in ds]):
            raise NotImplementedError(
                'Dilated convolutions are not yet implemented')

        # Override default architectural parameters
        self.architecture_parameters = dict(self.__net_params__)
        self.architecture_parameters.update(architecture_parameters)

        # Check available modes
        if self.architecture_parameters['use_dense_connections'] is False:
            raise NotImplementedError(
                'Non-dense connections are not yet implemented')
        if self.architecture_parameters['use_coords'] is True:
            raise NotImplementedError(
                'Image coordinate augmentation is not yet implemented')

    def create_network(self):
        hp = self.hyperparameters

        # Initial Convolution
        net_initial_conv = ConvolutionalLayer(
            hp['n_input_channels'][0],
            kernel_size=5, stride=2
        )

        # Dense Block Params
        downsample_channels = list(hp['n_input_channels'][1:]) + [None]
        num_blocks = len(hp["n_dense_channels"])
        use_bdo = self.architecture_parameters['use_bdo']

        # Create DenseBlocks
        net_dense_vblocks = []

        for idx in range(num_blocks):
            dense_ch = hp["n_dense_channels"][idx]  # Number or dense channels
            seg_ch = hp["n_seg_channels"][idx]      # Number of segmentation ch
            down_ch = downsample_channels[idx]      # Number of downsampling ch
            dil_rate = hp["dilation_rates"][idx]    # Dilation rate

            # Dense feature block
            dblock = DenseFeatureStackBlockWithSkipAndDownsample(
                dense_ch, 3, dil_rate, seg_ch, down_ch, use_bdo,
                acti_func='relu'
            )

            net_dense_vblocks.append(dblock)

        # Segmentation
        net_seg_layer = ConvolutionalLayer(
            self.num_classes, kernel_size=hp['final_kernel'],
            with_bn=False, with_bias=True
        )

        return DenseVNetDesc(initial_bn=BNLayer(),
                             initial_conv=net_initial_conv,
                             dense_vblocks=net_dense_vblocks,
                             seg_layer=net_seg_layer)

    def downsample_input(self, input_tensor, n_spatial_dims):
        # Initial downsampling params
        d_size1 = (1,) + (3,) * n_spatial_dims + (1,)
        d_size2 = (1,) + (2,) * n_spatial_dims + (1,)

        # Downsample input
        return tf.nn.avg_pool(input_tensor, d_size1, d_size2, 'SAME')

    def layer_op(self, input_tensor, is_training, layer_id=-1):
        hp = self.hyperparameters

        # Initialize DenseVNet network layers
        net = self.create_network()

        #
        # Parameter handling
        #

        # Shape and dimension variable shortcuts
        channel_dim = len(input_tensor.shape) - 1
        input_size = input_tensor.shape.as_list()
        spatial_size = input_size[1:-1]
        n_spatial_dims = input_tensor.shape.ndims - 2

        # Quick access to hyperparams
        pkeep = hp['p_channels_selected']

        # Validate input dimension with dilation rates
        modulo = 2 ** (len(hp['dilation_rates']))
        assert layer_util.check_spatial_dims(input_tensor,
                                             lambda x: x % modulo == 0)

        #
        # Augmentation + Downsampling + Initial Layers
        #

        # On the fly data augmentation
        if n_spatial_dims == 2:
            # augmentation for 2d slice not implemented
            hp['augmentation_scale'] = 0

        if is_training and hp['augmentation_scale'] > 0:
            augment_layer = Affine3DAugmentationLayer(hp['augmentation_scale'],
                                                      'LINEAR', 'ZERO')
            input_tensor = augment_layer(input_tensor)

        # Variable storing all intermediate results -- VLinks
        all_segmentation_features = []

        # Downsample input to the network
        down_tensor = self.downsample_input(input_tensor, n_spatial_dims)
        downsampled_img = net.initial_bn(down_tensor, is_training=is_training)

        # Add initial downsampled image VLink
        all_segmentation_features.append(downsampled_img)

        # All results should match the downsampled input's shape
        output_shape = downsampled_img.shape.as_list()[1:-1]

        init_features = net.initial_conv(input_tensor, is_training=is_training)

        #
        # Dense VNet Main Block
        #

        # `down` will handle the input of each Dense VNet block
{'fraction_non_alphanumeric': 0.08731716583296398, 'fraction_numerical': 0.008041537389982903, 'mean_word_length': 1.9975327386600874, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3688438', 'n_tokens_mistral': 4626, 'n_tokens_neox': 4322, 'n_words': 1007}