# -*- coding: utf-8 -*-
"""
Created on Mon Sep 16 14:24:03 2019

@author: shirhe-lyh


Implementation of DeepLabV3+:
    Encoder-Decoder with atrous separable convolution for semantic image
    segmentation, Liang-Chieh Chen, YuKun Zhu, George Papandreou, Florian
    Schroff, Hartwig Adam, arxiv:1802.02611 (https://arxiv.org/abs/1802.02611).
   
Official implementation:
    https://github.com/tensorflow/models/tree/master/research/deeplab
"""

import os
import torch

import common
import core.feature_extractor as extractor


_BATCH_NORM_PARAMS = {
    'momentum': 0.9997,
    'eps': 1e-5,
    'affine': True,
}


class DeepLab(torch.nn.Module):
    """Implementation of DeepLab V3+."""
   
    def __init__(self, feature_extractor, model_options):
        """Constructor.
       
        Args:
            feature_extractor: The backbone of DeepLab model.
            model_options: A ModelOptions instance to configure models.
        """
        super(DeepLab, self).__init__()
        self._model_options = model_options
       
        # Feature extractor
        self._feature_extractor = feature_extractor

        # Atrous spatial pyramid pooling
        self._aspp = AtrousSpatialPyramidPooling(
            in_channels=feature_extractor.out_channels,
            crop_size=model_options.crop_size,
            output_stride=model_options.output_stride,
            atrous_rates=model_options.atrous_rates,
            use_bounded_activation=model_options.use_bounded_activation,
            add_image_level_feature=model_options.add_image_level_feature,
            image_pooling_stride=model_options.image_pooling_stride,
            image_pooling_crop_size=model_options.image_pooling_crop_size,
            aspp_with_separable_conv=model_options.aspp_with_separable_conv)
       
        # Refine by decoder
        self._refine_decoder = None
        if model_options.decoder_output_stride:
            self._refine_decoder = RefineDecoder(
                feature_extractor=self._feature_extractor,
                crop_size=model_options.crop_size,
                decoder_output_stride=model_options.decoder_output_stride[0],
                decoder_use_separable_conv=model_options.decoder_use_separable_conv,
                model_variant=model_options.model_variant,
                use_bounded_activation=model_options.use_bounded_activation)
           
        # Branch logits
        num_classes = model_options.outputs_to_num_classes[common.OUTPUT_TYPE]
        self._logits_layer = torch.nn.Conv2d(
            in_channels=256, out_channels=num_classes,
            kernel_size=model_options.logits_kernel_size)
       
    def forward(self, x):
        features = self._feature_extractor(x)
        features = self._aspp(features)
        if self._refine_decoder is not None:
            features = self._refine_decoder(features)
        logits = self._logits_layer(features)
        return logits
       
       
class AtrousSpatialPyramidPooling(torch.nn.Module):
    """Atrous Spatial Pyramid Pooling."""
   
    def __init__(self,
                 in_channels,
                 out_channels=256,
                 output_stride=16,
                 crop_size=[513, 513],
                 atrous_rates=[12, 24, 36],
                 use_bounded_activation=False,
                 add_image_level_feature=True,
                 image_pooling_stride=[1, 1],
                 image_pooling_crop_size=None,
                 aspp_with_separable_conv=True):
        """Constructor.
       
        Args:
            in_channels: Number of input filters.
            out_channels: Number of filters in the 1x1 pointwise convolution.
            atrous_rates: A list of atrous convolution rates for ASPP.
            use_bounded_activation: Whether or not to use bounded activations.
            crop_size: A tuple [crop_height, crop_width].
            image_pooling_crop_size: Image pooling crop size [height, width]
                used in the ASPP module.
        """
        super(AtrousSpatialPyramidPooling, self).__init__()
        activation_fn = (
            torch.nn.ReLU6(inplace=False) if use_bounded_activation else
            torch.nn.ReLU(inplace=False))
       
        depth = out_channels
        branches = []
        if add_image_level_feature:
            layers = []
            if crop_size is not None:
                # If image_pooling_crop_size is not specified, use crop_size
                if image_pooling_crop_size is None:
                    image_pooling_crop_size = crop_size
                pool_height = scale_dimension(image_pooling_crop_size[0],
                                              1. / output_stride)
                pool_width = scale_dimension(image_pooling_crop_size[1],
                                             1. / output_stride)
                layers += [torch.nn.AvgPool2d(
                    (pool_height, pool_width), image_pooling_stride)]
                resize_height = scale_dimension(
                    crop_size[0], 1. / output_stride)
                resize_width = scale_dimension(
                    crop_size[1], 1. / output_stride)
            else:
                # If crop_size is None, we simply do global pooling
                layers += [torch.nn.AdaptiveAvgPool2d(output_size=(1, 1))]
                resize_height, resize_width = None, None
            self._resize_height = resize_height
            self._resize_width = resize_width
            layers += [torch.nn.Conv2d(in_channels, depth, 1, bias=False),
                       torch.nn.BatchNorm2d(depth, **_BATCH_NORM_PARAMS),
                       activation_fn]
            branches.append(torch.nn.Sequential(*layers))
       
        # Employ a 1x1 convolution.
        branches.append(torch.nn.Sequential(
            torch.nn.Conv2d(in_channels, depth, kernel_size=1, bias=False),
            torch.nn.BatchNorm2d(depth, **_BATCH_NORM_PARAMS),
            activation_fn))
       
        if atrous_rates:
            # Employ 3x3 convolutions with different atrous rates.
            for i, rate in enumerate(atrous_rates):
                layers = []
                if aspp_with_separable_conv:
                    layers += [
                        SplitSeparableConv2d(in_channels, depth, rate=rate,
                                             activation_fn=activation_fn)]
                else:
                    layers += [
                        torch.nn.Conv2d(in_channels, depth, kernel_size=3,
                                        rate=rate, padding=1, bias=False),
                        torch.nn.BatchNorm2d(depth, **_BATCH_NORM_PARAMS),
                        activation_fn]
                branches.append(torch.nn.Sequential(*layers))
        self._branches = torch.nn.Sequential(*branches)
       
        # Merge branch logits
        self._conv_concat = torch.nn.Sequential(
            torch.nn.Conv2d(len(self._branches) * depth,
                            depth, kernel_size=1, bias=False),
            torch.nn.BatchNorm2d(depth, **_BATCH_NORM_PARAMS),
            activation_fn)
        self._dropout = torch.nn.Dropout2d(p=0.9, inplace=True)
   
    def forward(self, x):
        branch_logits = []
        conv_branches = self._branches
        if len(self._branches) > 4:
            conv_branches = self._branches[1:]
            image_feature = self._branches[0](x)
            if self._resize_height is None:
                _, _, height, width = x.size()
                self._resize_height, self._resize_width = height, width
            image_feature = resize_bilinear(
                image_feature, (self._resize_height, self._resize_width))
            branch_logits.append(image_feature)
        branch_logits += [branch(x) for branch in conv_branches]
        cancat_logits = torch.cat(branch_logits, dim=1)
        cancat_logits = self._conv_concat(cancat_logits)
        cancat_logits = self._dropout(cancat_logits)
        return cancat_logits
   
   
class SplitSeparableConv2d(torch.nn.Module):
    """Splits a seperable conv2d into depthwise and pointwise conv2d."""
   
    def __init__(self, in_channels, out_channels, kernel_size=3, rate=1,
                 use_batch_norm=True, activation_fn=None):
        """Constructor.
       
        Args:
            in_channels: Number of input filters.
            out_channels: Number of filters in the 1x1 pointwise convolution.
            kernel_size: A list of length 2: [kernel_height, kernel_width]
                of the filters. Can be an int if both values are the same.
            rate: Atrous convolution rate for the depthwise convolution.
            with_batch_norm: Whether or not to use batch normalization.
            activation_fn: The activation function to be applied.
        """
        super(SplitSeparableConv2d, self).__init__()
        # For the shape of output of Conv2d, see details at:
        # https://pytorch.org/docs/stable/nn.html#convolution-layers
        # Here, we assume that floor(padding) = padding
        padding = (kernel_size - 1) * rate // 2
        self._conv_depthwise = torch.nn.Conv2d(in_channels,
                                               in_channels,
                                               kernel_size=kernel_size,
                                               dilation=rate,
                                               padding=padding,
                                               groups=in_channels,
                                               bias=False)
        self._conv_pointwise = torch.nn.Conv2d(in_channels,
                                               out_channels,
                                               kernel_size=1,
                                               bias=False)
        self._batch_norm_depthwise = None
        self._batch_norm_pointwise = None
        if use_batch_norm:
            self._batch_norm_depthwise = torch.nn.BatchNorm2d(
                num_features=in_channels, **_BATCH_NORM_PARAMS)
            self._batch_norm_pointwise = torch.nn.BatchNorm2d(
                num_features=out_channels, **_BATCH_NORM_PARAMS)
        self._activation_fn = activation_fn
       
    def forward(self, x):
        x = self._conv_depthwise(x)
        if self._batch_norm_depthwise is not None:
            x = self._batch_norm_depthwise(x)
        if self._activation_fn is not None:
            x = self._activation_fn(x)
        x = self._conv_pointwise(x)
        if self._batch_norm_pointwise is not None:
            x = self._batch_norm_pointwise(x)
        if self._activation_fn is not None:
            x = self._activation_fn(x)
        return x
   
   
class RefineDecoder(torch.nn.Module):
    """Adds the decoder to obtain sharper segmentation results."""
   
    def __init__(self, feature_extractor, aspp_channels=256, crop_size=None,
                 decoder_output_stride=None, decoder_use_separable_conv=False,
                 model_variant=None, use_bounded_activation=False):
        """Constructor.
       
        Args:
            feature_extractor: The backbone of the DeepLab model.
            aspp_channels: The out channels of ASPP.
            crop_size: A tuple [crop_height, crop_width] specifying whole
                patch crop size.
            decoder_output_stride: An integer specifying the output stride of
                low-level features used in the decoder module.
            decoder_use_separable_conv: Employ separable convolution for
                decoder or not.
            model_variant: Model variant for feature extractor.
            use_bounded_activation: Whether or not to use bounded activations.
                Bounded activations better lend themselves to quantized
                inference.
           
        Raises:
            ValueError: If crop_size is None.
        """
        super(RefineDecoder, self).__init__()
       
        if crop_size is None:
            raise ValueError('crop_size must be provided when using decoder.')
           
        self._crop_size = crop_size
        self._output_stride = decoder_output_stride
        activation_fn = (
            torch.nn.ReLU6(inplace=False) if use_bounded_activation else
            torch.nn.ReLU(inplace=False))
        self._extractor = feature_extractor
        feature_names = extractor.networks_to_feature_maps[
            model_variant][extractor.DECODER_END_POINTS][decoder_output_stride]
        self._feature_name = '{}/{}'.format(model_variant, feature_names[0])
        self._decoder = torch.nn.Sequential(
            torch.nn.Conv2d(extractor.feature_out_channels_map[model_variant],
                            out_channels=48, kernel_size=1, bias=False),
            torch.nn.BatchNorm2d(num_features=48, **_BATCH_NORM_PARAMS),
            activation_fn)
        concat_layers = []
        decoder_depth = 256
        if decoder_use_separable_conv:
            concat_layers += [
                SplitSeparableConv2d(in_channels=aspp_channels + 48,
                                     out_channels=decoder_depth, rate=1,
                                     activation_fn=activation_fn),
                SplitSeparableConv2d(in_channels=decoder_depth,
                                     out_channels=decoder_depth,
                                     rate=1, activation_fn=activation_fn)]
        else:
            concat_layers += [
                torch.nn.Conv2d(in_channels=aspp_channels + 48,
                                out_channels=decoder_depth, kernel_size=3,
                                padding=1, bias=False),
                torch.nn.BatchNorm2d(num_features=decoder_depth,
                                     **_BATCH_NORM_PARAMS),
                activation_fn,
                torch.nn.Conv2d(in_channels=decoder_depth,
                                out_channels=decoder_depth,
                                kernel_size=3, padding=1, bias=False),
                torch.nn.BatchNorm2d(num_features=decoder_depth,
                                     **_BATCH_NORM_PARAMS),
                activation_fn]
        self._concat_layers = torch.nn.Sequential(*concat_layers)
       
    def forward(self, x):
        decoder_features_list = [x]
        inter_feature = self._extractor.end_points()[self._feature_name]
        decoder_features_list.append(self._decoder(inter_feature))
        # Determine the output size
        decoder_height = scale_dimension(self._crop_size[0],
                                         1.0 / self._output_stride)
        decoder_width = scale_dimension(self._crop_size[1],
                                        1.0 / self._output_stride)
        # Resize to decoder_height/decoder_width
        for j, feature in enumerate(decoder_features_list):
            decoder_features_list[j] = resize_bilinear(
                feature, (decoder_height, decoder_width))
        x = self._concat_layers(torch.cat(decoder_features_list, dim=1))
        return x
   

def scale_dimension(dim, scale):
    """Scales the input dimension.
   
    Args:
        dim: Input dimension (a scalar).
        scale: The amount of scaling applied to the input.
       
    Returns:
        scaled dimension.
    """
    return int((float(dim) - 1.0) * scale + 1.0)


def resize_bilinear(images, size):
    """Returns resized images.
   
    Args:
        images: A tensor of size [batch, height_in, width_in, channels].
        size: A tuple (height, width).
       
    Returns:
        A tensor of shape [batch, height_out, height_width, channels].
    """
    return torch.nn.functional.interpolate(
        images, size, mode='bilinear', align_corners=True)
   
   
def deeplab(num_classes, crop_size=[513, 513], atrous_rates=[12, 24, 36],
            output_stride=16, pretrained=True, pretained_num_classes=21,
            checkpoint_path='./pretrained_models/deeplabv3_pascal_trainval.pth'):
    """DeepLab v3+ for semantic segmentation."""
    outputs_to_num_classes = {'semantic': num_classes}
    model_options = common.ModelOptions(outputs_to_num_classes,
                                        crop_size=crop_size,
                                        atrous_rates=atrous_rates,
                                        output_stride=output_stride)
    feature_extractor = extractor.feature_extractor(
        model_options.model_variant, pretrained=False,
        output_stride=model_options.output_stride)
    model = DeepLab(feature_extractor, model_options)
   
    if pretrained:
        _load_state_dict(model, num_classes, pretained_num_classes,
                         checkpoint_path)
    return model


def _load_state_dict(model, num_classes, pretained_num_classes,
                     checkpoint_path):
    """Load pretrained weights."""
    if os.path.exists(checkpoint_path):
        state_dict = torch.load(checkpoint_path)
        if num_classes is None or num_classes != pretained_num_classes:
            state_dict.pop('_logits_layer.weight')
            state_dict.pop('_logits_layer.bias')
        model.load_state_dict(state_dict, strict=False)
        print('Load pretrained weights successfully.')
    else:
        raise ValueError('`checkpoint_path` does not exist.')

# -*- coding: utf-8 -*-
"""
Created on Wed Oct  9 17:46:13 2019

@author: shirhe-lyh


Convert tensorflow weights to pytorch weights for DeepLab V3+ models.

Reference:
    https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/
        tf_to_pytorch/convert_tf_to_pt/load_tf_weights.py
"""

import numpy as np
import tensorflow as tf
import torch


_BLOCK_UNIT_COUNT_MAP = {
    'xception_41': [[3, 1], [1, 8], [2, 1]],
    'xception_65': [[3, 1], [1, 16], [2, 1]],
    'xception_71': [[5, 1], [1, 16], [2, 1]],
}


def load_param(checkpoint_path, conversion_map, model_name):
    """Load parameters according to conversion_map.
   
    Args:
        checkpoint_path: Path to tensorflow's checkpoint file.
        conversion_map: A dictionary with format
            {pytorch tensor in a model: checkpoint variable name}
        model_name: The name of Xception model, only supports 'xception_41',
            'xception_65', or 'xception_71'.
    """
    for pth_param, tf_param_name in conversion_map.items():
        param_name_strs =  tf_param_name.split('_')
        if len(param_name_strs) > 1 and param_name_strs[1].startswith('flow'):
            tf_param_name = str(model_name) + '/' + tf_param_name
        tf_param = tf.train.load_variable(checkpoint_path, tf_param_name)
        if 'conv' in tf_param_name and 'weights' in tf_param_name:
            tf_param = np.transpose(tf_param, (3, 2, 0, 1))
            if 'depthwise' in tf_param_name:
                tf_param = np.transpose(tf_param, (1, 0, 2, 3))
        elif 'depthwise_weights' in tf_param_name:
            tf_param = np.transpose(tf_param, (3, 2, 0, 1))
            tf_param = np.transpose(tf_param, (1, 0, 2, 3))
        elif tf_param_name.endswith('weights'):
            tf_param = np.transpose(tf_param)
        assert pth_param.size() == tf_param.shape, ('Dimension mismatch: ' +
            '{} vs {}; {}'.format(pth_param.size(), tf_param.shape,
                 tf_param_name))
        pth_param.data = torch.from_numpy(tf_param)


def convert(model, checkpoint_path):
    """Load Pytorch Xception from TensorFlow checkpoint file.
   
    Args:
        model: The pytorch Xception model, only supports 'xception_41',
            'xception_65', or 'xception_71'.
        checkpoint_path: Path to tensorflow's checkpoint file.
       
    Returns:
        The pytorch Xception model with pretrained parameters.
    """
    block_unit_counts = _BLOCK_UNIT_COUNT_MAP.get(
        model._feature_extractor.scope, None)
    if block_unit_counts is None:
        raise ValueError('Unsupported Xception model name.')
    flow_names = []
    block_indices = []
    unit_indices = []
    flow_names_unique = ['entry_flow', 'middle_flow', 'exit_flow']
    for i, [block_count, unit_count] in enumerate(block_unit_counts):
        flow_names += [flow_names_unique[i]] * (block_count * unit_count)
        for i in range(block_count):
            block_indices += [i + 1] * unit_count
            unit_indices += [j + 1 for j in range(unit_count)]
   
    conversion_map = {}
    # Feature extractor: Root block
    conversion_map_for_root_block = {
        model._feature_extractor._layers[0]._conv.weight:
            'entry_flow/conv1_1/weights',
        model._feature_extractor._layers[0]._batch_norm.bias:
            'entry_flow/conv1_1/BatchNorm/beta',
        model._feature_extractor._layers[0]._batch_norm.weight:
            'entry_flow/conv1_1/BatchNorm/gamma',
        model._feature_extractor._layers[0]._batch_norm.running_mean:
            'entry_flow/conv1_1/BatchNorm/moving_mean',
        model._feature_extractor._layers[0]._batch_norm.running_var:
            'entry_flow/conv1_1/BatchNorm/moving_variance',
        model._feature_extractor._layers[1]._conv.weight:
            'entry_flow/conv1_2/weights',
        model._feature_extractor._layers[1]._batch_norm.bias:
            'entry_flow/conv1_2/BatchNorm/beta',
        model._feature_extractor._layers[1]._batch_norm.weight:
            'entry_flow/conv1_2/BatchNorm/gamma',
        model._feature_extractor._layers[1]._batch_norm.running_mean:
            'entry_flow/conv1_2/BatchNorm/moving_mean',
        model._feature_extractor._layers[1]._batch_norm.running_var:
            'entry_flow/conv1_2/BatchNorm/moving_variance',
    }
    conversion_map.update(conversion_map_for_root_block)
   
    # Feature extractor: Xception block
    for i in range(len(model._feature_extractor._layers[2]._blocks)):
        block = model._feature_extractor._layers[2]._blocks[i]
        ind = [1, 3, 5]
        if len(block._separable_conv_block) < 6:
            ind = [0, 1, 2]
        for j in range(3):
            conversion_map_for_separable_block = {
                block._separable_conv_block[ind[j]]._conv_depthwise.weight:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_depthwise/depthwise_weights').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._conv_pointwise.weight:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_pointwise/weights').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_depthwise.bias:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_depthwise/BatchNorm/beta').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_depthwise.weight:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_depthwise/BatchNorm/gamma').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_depthwise.running_mean:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_depthwise/BatchNorm/moving_mean').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_depthwise.running_var:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_depthwise/BatchNorm/moving_variance').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_pointwise.bias:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_pointwise/BatchNorm/beta').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_pointwise.weight:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_pointwise/BatchNorm/gamma').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_pointwise.running_mean:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_pointwise/BatchNorm/moving_mean').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
                block._separable_conv_block[ind[j]]._batch_norm_pointwise.running_var:
                    ('{}/block{}/unit_{}/xception_module/' +
                     'separable_conv{}_pointwise/BatchNorm/moving_variance').format(
                        flow_names[i], block_indices[i], unit_indices[i], j+1),
            }
            conversion_map.update(conversion_map_for_separable_block)
           
            if getattr(block, '_conv_skip_connection', None) is not None:
                conversion_map_for_shortcut = {
                    block._conv_skip_connection.weight:
                       ('{}/block{}/unit_{}/xception_module/shortcut/' +
                        'weights').format(
                            flow_names[i], block_indices[i], unit_indices[i]),
                    block._batch_norm_shortcut.bias:
                        ('{}/block{}/unit_{}/xception_module/shortcut/' +
                         'BatchNorm/beta').format(
                            flow_names[i], block_indices[i], unit_indices[i]),
                    block._batch_norm_shortcut.weight:
                        ('{}/block{}/unit_{}/xception_module/shortcut/' +
                         'BatchNorm/gamma').format(
                            flow_names[i], block_indices[i], unit_indices[i]),
                    block._batch_norm_shortcut.running_mean:
                        ('{}/block{}/unit_{}/xception_module/shortcut/' +
                         'BatchNorm/moving_mean').format(
                            flow_names[i], block_indices[i], unit_indices[i]),
                    block._batch_norm_shortcut.running_var:
                        ('{}/block{}/unit_{}/xception_module/shortcut/' +
                         'BatchNorm/moving_variance').format(
                            flow_names[i], block_indices[i], unit_indices[i]),
                }
                conversion_map.update(conversion_map_for_shortcut)
       
    # Atrous Spatial Pyramid Pooling: Image feature
    branches = model._aspp._branches
    conversion_map_for_aspp_image_feature = {
        branches[0][1].weight: 'image_pooling/weights',
        branches[0][2].bias: 'image_pooling/BatchNorm/beta',
        branches[0][2].weight: 'image_pooling/BatchNorm/gamma',
        branches[0][2].running_mean: 'image_pooling/BatchNorm/moving_mean',
        branches[0][2].running_var: 'image_pooling/BatchNorm/moving_variance',
        branches[1][0].weight: 'aspp0/weights',
        branches[1][1].bias: 'aspp0/BatchNorm/beta',
        branches[1][1].weight: 'aspp0/BatchNorm/gamma',
        branches[1][1].running_mean: 'aspp0/BatchNorm/moving_mean',
        branches[1][1].running_var: 'aspp0/BatchNorm/moving_variance',
    }
    conversion_map.update(conversion_map_for_aspp_image_feature)
   
    # Atrous Spatial Pyramid Pooling: Atrous convolution
    for i in range(3):
        branch = branches[i+2][0]
        conversion_map_for_atrous_conv = {
            branch._conv_depthwise.weight:
                'aspp{}_depthwise/depthwise_weights'.format(i+1),
            branch._conv_pointwise.weight:
                'aspp{}_pointwise/weights'.format(i+1),
            branch._batch_norm_depthwise.bias:
                'aspp{}_depthwise/BatchNorm/beta'.format(i+1),
            branch._batch_norm_depthwise.weight:
                'aspp{}_depthwise/BatchNorm/gamma'.format(i+1),
            branch._batch_norm_depthwise.running_mean:
                'aspp{}_depthwise/BatchNorm/moving_mean'.format(i+1),
            branch._batch_norm_depthwise.running_var:
                'aspp{}_depthwise/BatchNorm/moving_variance'.format(i+1),
            branch._batch_norm_pointwise.bias:
                'aspp{}_pointwise/BatchNorm/beta'.format(i+1),
            branch._batch_norm_pointwise.weight:
                'aspp{}_pointwise/BatchNorm/gamma'.format(i+1),
            branch._batch_norm_pointwise.running_mean:
                'aspp{}_pointwise/BatchNorm/moving_mean'.format(i+1),
            branch._batch_norm_pointwise.running_var:
                'aspp{}_pointwise/BatchNorm/moving_variance'.format(i+1),
        }
        conversion_map.update(conversion_map_for_atrous_conv)
       
    # Atrous Spatial Pyramid Pooling: Concat projection
    conversion_map_for_concat_projection = {
        model._aspp._conv_concat[0].weight:
            'concat_projection/weights',
        model._aspp._conv_concat[1].bias:
            'concat_projection/BatchNorm/beta',
        model._aspp._conv_concat[1].weight:
            'concat_projection/BatchNorm/gamma',
        model._aspp._conv_concat[1].running_mean:
            'concat_projection/BatchNorm/moving_mean',
        model._aspp._conv_concat[1].running_var:
            'concat_projection/BatchNorm/moving_variance',
    }
    conversion_map.update(conversion_map_for_concat_projection)
   
    # Refine decoder: Feature projection
    conversion_map_for_decoder = {
        model._refine_decoder._decoder[0].weight:
            'decoder/feature_projection0/weights',
        model._refine_decoder._decoder[1].bias:
            'decoder/feature_projection0/BatchNorm/beta',
        model._refine_decoder._decoder[1].weight:
            'decoder/feature_projection0/BatchNorm/gamma',
        model._refine_decoder._decoder[1].running_mean:
            'decoder/feature_projection0/BatchNorm/moving_mean',
        model._refine_decoder._decoder[1].running_var:
            'decoder/feature_projection0/BatchNorm/moving_variance',
    }
    conversion_map.update(conversion_map_for_decoder)
   
    # Refine decoder: Concat
    layers = model._refine_decoder._concat_layers
    for i in range(2):
        layer = layers[i]
        conversion_map_decoder = {
            layer._conv_depthwise.weight:
                'decoder/decoder_conv{}_depthwise/depthwise_weights'.format(i),
            layer._conv_pointwise.weight:
                'decoder/decoder_conv{}_pointwise/weights'.format(i),
            layer._batch_norm_depthwise.bias:
                'decoder/decoder_conv{}_depthwise/BatchNorm/beta'.format(i),
            layer._batch_norm_depthwise.weight:
                'decoder/decoder_conv{}_depthwise/BatchNorm/gamma'.format(i),
            layer._batch_norm_depthwise.running_mean:
                'decoder/decoder_conv{}_depthwise/BatchNorm/moving_mean'.format(i),
            layer._batch_norm_depthwise.running_var:
                'decoder/decoder_conv{}_depthwise/BatchNorm/moving_variance'.format(i),
            layer._batch_norm_pointwise.bias:
                'decoder/decoder_conv{}_pointwise/BatchNorm/beta'.format(i),
            layer._batch_norm_pointwise.weight:
                'decoder/decoder_conv{}_pointwise/BatchNorm/gamma'.format(i),
            layer._batch_norm_pointwise.running_mean:
                'decoder/decoder_conv{}_pointwise/BatchNorm/moving_mean'.format(i),
            layer._batch_norm_pointwise.running_var:
                'decoder/decoder_conv{}_pointwise/BatchNorm/moving_variance'.format(i),
        }
        conversion_map.update(conversion_map_decoder)
       
    # Prediction logits
    conversion_map_for_logits = {
        model._logits_layer.weight: 'logits/semantic/weights',
        model._logits_layer.bias: 'logits/semantic/biases',
    }
    conversion_map.update(conversion_map_for_logits)
       
    # Load TensorFlow parameters into PyTorch model
    load_param(checkpoint_path, conversion_map, model._feature_extractor.scope)

python3 tf_weights_to_pth.py --tf_checkpoint_path "xxxx/deeplabv3_pascal_trainval/model.ckpt"

python3 model_test.py

python3 model_test.py --image_path "./test/tf_vis2.jpg"
python3 model_test.py --image_path "./test/tf_vis2.jpg" --val_output_stride 16 \
    --val_atrous_rates 6 12 18

python3 tf_weights_to_pth.py --tf_checkpoint_path "xxxx/model.ckpt" \
    --output_dir "xxx" --output_name "xxx.pth"  --pretained_num_classes xxx

--model_variant "xception_71"

import model

deeplab = model.deeplab(num_classes,
                        crop_size=[513, 513],
                        atrous_rates=[12, 24, 36],
                        output_stride=8,
                        pretrained=True,
                        pretained_num_classes=21,
                        checkpoint_path='./pretrained_models/deeplabv3_pascal_trainval.pth')

python3 train.py --root_dir "xxx/Matting_Human_Half" [--gpu_indices 0 1 ...]

--gpu_indices 0

--gpu_indices 1 3

tensorboard --logdir ./models/logs

python3 predict.py

[[image_path, mask_path],
 [image_path, mask_path],
...
 [image_path, mask_path]]