# Path: net_modules/keypoints_2d.py
# def yx_grid_map(h, w, dtype, aspect_ratio=1.):
# def keypoint_map_to_gaussian_coordinate(keypoint_map, diag_cov=None, use_hard_max_as_anchors=None):
#     def cyclic_clip(a):
# def gaussian_coordinate_to_keypoint_map(yx_mean_stddev_corr, km_h, km_w, dtype=None):
# def gaussian2dparam_to_real(p):
# def real_to_gaussian2dparam(r):
# def gaussian2dparam_to_recon_code(p):
# def real_to_recon_code(r):
# def gaussian2d_det(p, no_mean_input=False):
# def gaussian2d_entropy(p, no_mean_input=False):
# def gaussian2d_exp_entropy(p, no_mean_input=False, stddev_scaling=1.):
# def gaussian2d_axis_balancing(p, no_mean_input=False):
# def keypoint_map_boundary_mask(h, w, padding_h, padding_w, dtype=np.float32):
# def keypoint_map_boundary_mask_with_bg(h, w, padding_h, padding_w, keypoint_num, dtype=np.float32):
# def keypoint_map_depth_normalization_with_fake_bg(keypoint_map):
# def keypoint_map_depth_entropy_with_fake_bg(keypoint_map):
# def keypoint_map_depth_normalization_with_real_bg(keypoint_map):
# def keypoint_map_depth_entropy_with_real_bg(keypoint_map):
# def scale_keypoint_param(keypoint_param, scaling_factor, src_aspect_ratio):
# def get_map_aspect_ratio(a):
# def parse_landmark_condition_tensor_single(condition_tensor, img_size=None, full_img_size=None):
# def parse_landmark_condition_struct(keypoint_struct, img_size=None, full_img_size=None):
# 
# Path: net_modules/hourglass.py
# class HourglassOptions:
#     def layer(p):
# def _hourglass_builder(x, layer_def, net_type=None, extra_highlevel_feature=None):
# def hourglass_builder(x, layer_def, net_type=None, scope=None, extra_highlevel_feature=None):
# 
# Path: net_modules/auto_struct/keypoint_encoder.py
# class BasicFactory(BaseFactory):
#     def __init__(self, output_channels, options):
#     def keypoint_init(self):
#     def use_random_transform(self):
#     def image_size(self, img_h, img_w):
#     def augment_images(self, image_tensor):
#     def heatmap_postprocess(self, heatmap):
#     def heatmap_postpostprocess(self, heatmap, image_tensor=None, heatmap_extra=None):
#     def heatmap2structure_internal(self, heatmap_tensor):
#     def heatmap2structure_basic(self, heatmap_tensor):
#     def heatmap2structure(self, heatmap_tensor):
#     def heatmap2structure_poststep(self, structure_pack):
#     def cleanup_augmentation_patchfeatures(self, patch_features, aug_cache):
#     def cleanup_augmentation_structure(self, structure_param, aug_cache, condition_tensor=None):
#     def structure_param2euclidean(self, structure_param):
#     def coordinate_to_stn(self, keypoint_param, aspect_ratio):


import numpy as np
import tensorflow as tf
import prettytensor as pt
import math
import zutils.tf_math_funcs as tmf
from net_modules.common import *
from net_modules import keypoints_2d
from net_modules.hourglass import hourglass
from net_modules.auto_struct.keypoint_encoder import Factory as BaseFactory
    def __init__(self, output_channels, options):
        """
        :param output_channels: output_channels for the encoding net
        """
        super().__init__(output_channels, options)

        self.target_input_size = [96, 96]
        self.input_size = [80, 80]

    def image2heatmap(self, image_tensor):
        hgd = [
            {"type": "conv2d", "depth": 32, "decoder_depth": self.options["keypoint_num"] + 1,
             "decoder_activation_fn": None},
            # plus one for bg
            {"type": "conv2d", "depth": 32},
            {"type": "skip", "layer_num": 3, },
            {"type": "pool", "pool": "max"},
            {"type": "conv2d", "depth": 64},
            {"type": "conv2d", "depth": 64},
            {"type": "skip", "layer_num": 3, },
            {"type": "pool", "pool": "max"},
            {"type": "conv2d", "depth": 64},
            {"type": "conv2d", "depth": 64},
            {"type": "skip", "layer_num": 3, },
            {"type": "pool", "pool": "max"},
            {"type": "conv2d", "depth": 64},
            {"type": "conv2d", "depth": 64},
        ]

        with pt.defaults_scope(**self.pt_defaults_scope_value()):

# Path: net_modules/keypoints_2d.py
# def yx_grid_map(h, w, dtype, aspect_ratio=1.):
# def keypoint_map_to_gaussian_coordinate(keypoint_map, diag_cov=None, use_hard_max_as_anchors=None):
#     def cyclic_clip(a):
# def gaussian_coordinate_to_keypoint_map(yx_mean_stddev_corr, km_h, km_w, dtype=None):
# def gaussian2dparam_to_real(p):
# def real_to_gaussian2dparam(r):
# def gaussian2dparam_to_recon_code(p):
# def real_to_recon_code(r):
# def gaussian2d_det(p, no_mean_input=False):
# def gaussian2d_entropy(p, no_mean_input=False):
# def gaussian2d_exp_entropy(p, no_mean_input=False, stddev_scaling=1.):
# def gaussian2d_axis_balancing(p, no_mean_input=False):
# def keypoint_map_boundary_mask(h, w, padding_h, padding_w, dtype=np.float32):
# def keypoint_map_boundary_mask_with_bg(h, w, padding_h, padding_w, keypoint_num, dtype=np.float32):
# def keypoint_map_depth_normalization_with_fake_bg(keypoint_map):
# def keypoint_map_depth_entropy_with_fake_bg(keypoint_map):
# def keypoint_map_depth_normalization_with_real_bg(keypoint_map):
# def keypoint_map_depth_entropy_with_real_bg(keypoint_map):
# def scale_keypoint_param(keypoint_param, scaling_factor, src_aspect_ratio):
# def get_map_aspect_ratio(a):
# def parse_landmark_condition_tensor_single(condition_tensor, img_size=None, full_img_size=None):
# def parse_landmark_condition_struct(keypoint_struct, img_size=None, full_img_size=None):
# 
# Path: net_modules/hourglass.py
# class HourglassOptions:
#     def layer(p):
# def _hourglass_builder(x, layer_def, net_type=None, extra_highlevel_feature=None):
# def hourglass_builder(x, layer_def, net_type=None, scope=None, extra_highlevel_feature=None):
# 
# Path: net_modules/auto_struct/keypoint_encoder.py
# class BasicFactory(BaseFactory):
#     def __init__(self, output_channels, options):
#     def keypoint_init(self):
#     def use_random_transform(self):
#     def image_size(self, img_h, img_w):
#     def augment_images(self, image_tensor):
#     def heatmap_postprocess(self, heatmap):
#     def heatmap_postpostprocess(self, heatmap, image_tensor=None, heatmap_extra=None):
#     def heatmap2structure_internal(self, heatmap_tensor):
#     def heatmap2structure_basic(self, heatmap_tensor):
#     def heatmap2structure(self, heatmap_tensor):
#     def heatmap2structure_poststep(self, structure_pack):
#     def cleanup_augmentation_patchfeatures(self, patch_features, aug_cache):
#     def cleanup_augmentation_structure(self, structure_param, aug_cache, condition_tensor=None):
#     def structure_param2euclidean(self, structure_param):
#     def coordinate_to_stn(self, keypoint_param, aspect_ratio):


# Path: nets/encoder/general_80x80.py
import numpy as np
import tensorflow as tf
import prettytensor as pt
import math
import zutils.tf_math_funcs as tmf
from net_modules.common import *
from net_modules import keypoints_2d
from net_modules.hourglass import hourglass
from net_modules.auto_struct.keypoint_encoder import Factory as BaseFactory

    def __init__(self, output_channels, options):
        """
        :param output_channels: output_channels for the encoding net
        """
        super().__init__(output_channels, options)

        self.target_input_size = [96, 96]
        self.input_size = [80, 80]

    def image2heatmap(self, image_tensor):
        hgd = [
            {"type": "conv2d", "depth": 32, "decoder_depth": self.options["keypoint_num"] + 1,
             "decoder_activation_fn": None},
            # plus one for bg
            {"type": "conv2d", "depth": 32},
            {"type": "skip", "layer_num": 3, },
            {"type": "pool", "pool": "max"},
            {"type": "conv2d", "depth": 64},
            {"type": "conv2d", "depth": 64},
            {"type": "skip", "layer_num": 3, },
            {"type": "pool", "pool": "max"},
            {"type": "conv2d", "depth": 64},
            {"type": "conv2d", "depth": 64},
            {"type": "skip", "layer_num": 3, },
            {"type": "pool", "pool": "max"},
            {"type": "conv2d", "depth": 64},
            {"type": "conv2d", "depth": 64},
        ]

        with pt.defaults_scope(**self.pt_defaults_scope_value()):

            raw_heatmap = hourglass(