src/combined_model/model_shape_v7_withref_withgraphcnn.py

import pickle as pkl
import numpy as np
import torchvision.models as models
from torchvision import transforms
import torch
from torch import nn
from torch.nn.parameter import Parameter
from kornia.geometry.subpix import dsnt     # kornia 0.4.0

import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from stacked_hourglass.utils.evaluation import get_preds_soft
from stacked_hourglass import hg1, hg2, hg8
from lifting_to_3d.linear_model import LinearModelComplete, LinearModel      
from lifting_to_3d.inn_model_for_shape import INNForShape
from lifting_to_3d.utils.geometry_utils import rot6d_to_rotmat, rotmat_to_rot6d
from smal_pytorch.smal_model.smal_torch_new import SMAL
from smal_pytorch.renderer.differentiable_renderer import SilhRenderer
from bps_2d.bps_for_segmentation import SegBPS
# from configs.SMAL_configs import SMAL_MODEL_DATA_PATH as SHAPE_PRIOR
from configs.SMAL_configs import SMAL_MODEL_CONFIG
from configs.SMAL_configs import MEAN_DOG_BONE_LENGTHS_NO_RED, VERTEX_IDS_TAIL

# NEW: for graph cnn part
from smal_pytorch.smal_model.smal_torch_new import SMAL
from configs.SMAL_configs import SMAL_MODEL_CONFIG
from graph_networks.graphcmr.utils_mesh import Mesh
from graph_networks.graphcmr.graph_cnn_groundcontact_multistage import GraphCNNMS




class SmallLinear(nn.Module):
    def __init__(self, input_size=64, output_size=30, linear_size=128):
        super(SmallLinear, self).__init__()
        self.relu = nn.ReLU(inplace=True)
        self.w1 = nn.Linear(input_size, linear_size)
        self.w2 = nn.Linear(linear_size, linear_size)
        self.w3 = nn.Linear(linear_size, output_size)
    def forward(self, x):
        # pre-processing
        y = self.w1(x)
        y = self.relu(y)
        y = self.w2(y)
        y = self.relu(y)
        y = self.w3(y)
        return y


class MyConv1d(nn.Module):
    def __init__(self, input_size=37, output_size=30, start=True):
        super(MyConv1d, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.start = start
        self.weight = Parameter(torch.ones((self.output_size)))
        self.bias = Parameter(torch.zeros((self.output_size)))
    def forward(self, x):
        # pre-processing
        if self.start:
            y = x[:, :self.output_size]
        else:
            y = x[:, -self.output_size:]
        y = y * self.weight[None, :] + self.bias[None, :]
        return y


class ModelShapeAndBreed(nn.Module):
    def __init__(self, smal_model_type, n_betas=10, n_betas_limbs=13, n_breeds=121, n_z=512, structure_z_to_betas='default'):
        super(ModelShapeAndBreed, self).__init__()
        self.n_betas = n_betas
        self.n_betas_limbs = n_betas_limbs   # n_betas_logscale
        self.n_breeds = n_breeds
        self.structure_z_to_betas = structure_z_to_betas
        if self.structure_z_to_betas == '1dconv':
            if not (n_z == self.n_betas+self.n_betas_limbs):
                raise ValueError
        self.smal_model_type = smal_model_type
        # shape branch
        self.resnet = models.resnet34(pretrained=False)  
        # replace the first layer
        n_in = 3 + 1
        self.resnet.conv1 = nn.Conv2d(n_in, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        # replace the last layer
        self.resnet.fc = nn.Linear(512, n_z) 
        # softmax
        self.soft_max = torch.nn.Softmax(dim=1)
        # fc network (and other versions) to connect z with betas
        p_dropout = 0.2
        if self.structure_z_to_betas == 'default':
            self.linear_betas = LinearModel(linear_size=1024,     
                                                num_stage=1,
                                                p_dropout=p_dropout, 
                                                input_size=n_z,
                                                output_size=self.n_betas)
            self.linear_betas_limbs = LinearModel(linear_size=1024,    
                                                num_stage=1,
                                                p_dropout=p_dropout, 
                                                input_size=n_z,
                                                output_size=self.n_betas_limbs)
        elif self.structure_z_to_betas == 'lin':
            self.linear_betas = nn.Linear(n_z, self.n_betas)
            self.linear_betas_limbs = nn.Linear(n_z, self.n_betas_limbs)
        elif self.structure_z_to_betas == 'fc_0':
            self.linear_betas = SmallLinear(linear_size=128,     # 1024,
                                                input_size=n_z,
                                                output_size=self.n_betas)
            self.linear_betas_limbs = SmallLinear(linear_size=128,     # 1024,
                                                input_size=n_z,
                                                output_size=self.n_betas_limbs)
        elif structure_z_to_betas == 'fc_1':
            self.linear_betas = LinearModel(linear_size=64,     # 1024,
                                                num_stage=1,
                                                p_dropout=0, 
                                                input_size=n_z,
                                                output_size=self.n_betas)
            self.linear_betas_limbs = LinearModel(linear_size=64,     # 1024,
                                                num_stage=1,
                                                p_dropout=0, 
                                                input_size=n_z,
                                                output_size=self.n_betas_limbs)
        elif self.structure_z_to_betas == '1dconv':
            self.linear_betas = MyConv1d(n_z, self.n_betas, start=True)
            self.linear_betas_limbs = MyConv1d(n_z, self.n_betas_limbs, start=False)
        elif self.structure_z_to_betas == 'inn':
            self.linear_betas_and_betas_limbs = INNForShape(self.n_betas, self.n_betas_limbs, betas_scale=1.0, betas_limbs_scale=1.0)
        else:
            raise ValueError
        # network to connect latent shape vector z with dog breed classification
        self.linear_breeds = LinearModel(linear_size=1024,    # 1024,
                                            num_stage=1,
                                            p_dropout=p_dropout, 
                                            input_size=n_z,
                                            output_size=self.n_breeds)
        # shape multiplicator
        self.shape_multiplicator_np = np.ones(self.n_betas)
        with open(SMAL_MODEL_CONFIG[self.smal_model_type]['smal_model_data_path'], 'rb') as file:
            u = pkl._Unpickler(file)
            u.encoding = 'latin1'
            res = u.load()
        # shape predictions are centered around the mean dog of our dog model
        if 'dog_cluster_mean' in res.keys():
            self.betas_mean_np = res['dog_cluster_mean'] 
        else:
            assert res['cluster_means'].shape[0]==1
            self.betas_mean_np = res['cluster_means'][0, :]

                                        
    def forward(self, img, seg_raw=None, seg_prep=None):
        # img is the network input image 
        # seg_raw is before softmax and subtracting 0.5
        # seg_prep would be the prepared_segmentation
        if seg_prep is None:
            seg_prep = self.soft_max(seg_raw)[:, 1:2, :, :] - 0.5       
        input_img_and_seg = torch.cat((img, seg_prep), axis=1)
        res_output = self.resnet(input_img_and_seg)
        dog_breed_output = self.linear_breeds(res_output) 
        if self.structure_z_to_betas == 'inn':
            shape_output_orig, shape_limbs_output_orig = self.linear_betas_and_betas_limbs(res_output)
        else:
            shape_output_orig = self.linear_betas(res_output) * 0.1
            betas_mean = torch.tensor(self.betas_mean_np).float().to(img.device)
            shape_output = shape_output_orig + betas_mean[None, 0:self.n_betas]
            shape_limbs_output_orig = self.linear_betas_limbs(res_output)
            shape_limbs_output = shape_limbs_output_orig * 0.1
        output_dict = {'z': res_output,
                        'breeds': dog_breed_output,
                        'betas': shape_output_orig,
                        'betas_limbs': shape_limbs_output_orig}
        return output_dict



class LearnableShapedirs(nn.Module):
    def __init__(self, sym_ids_dict, shapedirs_init, n_betas, n_betas_fixed=10):
        super(LearnableShapedirs, self).__init__()
        # shapedirs_init = self.smal.shapedirs.detach()
        self.n_betas = n_betas
        self.n_betas_fixed = n_betas_fixed
        self.sym_ids_dict = sym_ids_dict
        sym_left_ids = self.sym_ids_dict['left']
        sym_right_ids = self.sym_ids_dict['right']
        sym_center_ids = self.sym_ids_dict['center']
        self.n_center = sym_center_ids.shape[0]
        self.n_left = sym_left_ids.shape[0]
        self.n_sd = self.n_betas - self.n_betas_fixed     # number of learnable shapedirs
        # get indices to go from half_shapedirs to shapedirs
        inds_back = np.zeros((3889))
        for ind in range(0, sym_center_ids.shape[0]):
            ind_in_forward = sym_center_ids[ind]
            inds_back[ind_in_forward] = ind
        for ind in range(0, sym_left_ids.shape[0]):
            ind_in_forward = sym_left_ids[ind]
            inds_back[ind_in_forward] = sym_center_ids.shape[0] + ind
        for ind in range(0, sym_right_ids.shape[0]):
            ind_in_forward = sym_right_ids[ind]
            inds_back[ind_in_forward] = sym_center_ids.shape[0] + sym_left_ids.shape[0] + ind
        self.register_buffer('inds_back_torch', torch.Tensor(inds_back).long())
        # self.smal.shapedirs: (51, 11667)
        # shapedirs: (3889, 3, n_sd)
        # shapedirs_half: (2012, 3, n_sd)
        sd = shapedirs_init[:self.n_betas, :].permute((1, 0)).reshape((-1, 3, self.n_betas))
        self.register_buffer('sd', sd)
        sd_center = sd[sym_center_ids, :, self.n_betas_fixed:]
        sd_left = sd[sym_left_ids, :, self.n_betas_fixed:]
        self.register_parameter('learnable_half_shapedirs_c0', torch.nn.Parameter(sd_center[:, 0, :].detach()))
        self.register_parameter('learnable_half_shapedirs_c2', torch.nn.Parameter(sd_center[:, 2, :].detach()))
        self.register_parameter('learnable_half_shapedirs_l0', torch.nn.Parameter(sd_left[:, 0, :].detach()))
        self.register_parameter('learnable_half_shapedirs_l1', torch.nn.Parameter(sd_left[:, 1, :].detach()))
        self.register_parameter('learnable_half_shapedirs_l2', torch.nn.Parameter(sd_left[:, 2, :].detach()))
    def forward(self):
        device = self.learnable_half_shapedirs_c0.device
        half_shapedirs_center = torch.stack((self.learnable_half_shapedirs_c0, \
                                            torch.zeros((self.n_center, self.n_sd)).to(device), \
                                            self.learnable_half_shapedirs_c2), axis=1)
        half_shapedirs_left = torch.stack((self.learnable_half_shapedirs_l0, \
                                            self.learnable_half_shapedirs_l1, \
                                            self.learnable_half_shapedirs_l2), axis=1)
        half_shapedirs_right = torch.stack((self.learnable_half_shapedirs_l0, \
                                            - self.learnable_half_shapedirs_l1, \
                                            self.learnable_half_shapedirs_l2), axis=1)
        half_shapedirs_tot = torch.cat((half_shapedirs_center, half_shapedirs_left, half_shapedirs_right))
        shapedirs = torch.index_select(half_shapedirs_tot, dim=0, index=self.inds_back_torch)
        shapedirs_complete = torch.cat((self.sd[:, :, :self.n_betas_fixed], shapedirs), axis=2)      # (3889, 3, n_sd)
        shapedirs_complete_prepared = torch.cat((self.sd[:, :, :10], shapedirs), axis=2).reshape((-1, 30)).permute((1, 0))   # (n_sd, 11667)
        return shapedirs_complete, shapedirs_complete_prepared


class ModelRefinement(nn.Module):
    def __init__(self, n_betas=10, n_betas_limbs=7, n_breeds=121, n_keyp=20, n_joints=35, ref_net_type='add', graphcnn_type='inexistent', isflat_type='inexistent', shaperef_type='inexistent'):
        super(ModelRefinement, self).__init__()
        self.n_betas = n_betas
        self.n_betas_limbs = n_betas_limbs
        self.n_breeds = n_breeds
        self.n_keyp = n_keyp
        self.n_joints = n_joints
        self.n_out_seg = 256
        self.n_out_keyp = 256
        self.n_out_enc = 256
        self.linear_size = 1024
        self.linear_size_small = 128
        self.ref_net_type = ref_net_type
        self.graphcnn_type = graphcnn_type
        self.isflat_type = isflat_type
        self.shaperef_type = shaperef_type
        p_dropout = 0.2
        # --- segmentation encoder
        if self.ref_net_type in ['multrot_res34', 'multrot01all_res34']:
            self.ref_res = models.resnet34(pretrained=False)
        else:
            self.ref_res = models.resnet18(pretrained=False)
        # replace the first layer
        self.ref_res.conv1 = nn.Conv2d(2, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        # replace the last layer
        self.ref_res.fc = nn.Linear(512, self.n_out_seg) 
        # softmax
        self.soft_max = torch.nn.Softmax(dim=1)
        # --- keypoint encoder
        self.linear_keyp = LinearModel(linear_size=self.linear_size,
                                            num_stage=1,
                                            p_dropout=p_dropout, 
                                            input_size=n_keyp*2*2,
                                            output_size=self.n_out_keyp)
        # --- decoder
        self.linear_combined = LinearModel(linear_size=self.linear_size,
                                            num_stage=1,
                                            p_dropout=p_dropout, 
                                            input_size=self.n_out_seg+self.n_out_keyp,
                                            output_size=self.n_out_enc)
        # output info
        pose = {'name': 'pose', 'n': self.n_joints*6, 'out_shape':[self.n_joints, 6]}
        trans = {'name': 'trans_notnorm', 'n': 3}
        cam = {'name': 'flength_notnorm', 'n': 1}
        betas = {'name': 'betas', 'n': self.n_betas}
        betas_limbs = {'name': 'betas_limbs', 'n': self.n_betas_limbs}
        if self.shaperef_type=='inexistent':          
            self.output_info = [pose, trans, cam]   # , betas]
        else:
            self.output_info = [pose, trans, cam, betas, betas_limbs]
        # output branches
        self.output_info_linear_models = []
        for ind_el, element in enumerate(self.output_info):
            n_in = self.n_out_enc + element['n']
            self.output_info_linear_models.append(LinearModel(linear_size=self.linear_size,
                                    num_stage=1,
                                    p_dropout=p_dropout, 
                                    input_size=n_in,
                                    output_size=element['n']))
            element['linear_model_index'] = ind_el
        self.output_info_linear_models = nn.ModuleList(self.output_info_linear_models)
        # new: predict if the ground is flat
        if not self.isflat_type=='inexistent':          
            self.linear_isflat = LinearModel(linear_size=self.linear_size_small,
                                    num_stage=1,
                                    p_dropout=p_dropout, 
                                    input_size=self.n_out_enc,
                                    output_size=2) # answer is just yes or no

    
        # new for ground contact prediction: graph cnn
        if not self.graphcnn_type=='inexistent':
            num_downsampling = 1
            smal_model_type = '39dogs_norm'
            smal = SMAL(smal_model_type=smal_model_type, template_name='neutral')     
            ROOT_smal_downsampling = os.path.join(os.path.dirname(__file__), './../../data/graphcmr_data/')
            smal_downsampling_npz_name = 'mesh_downsampling_' + os.path.basename(SMAL_MODEL_CONFIG[smal_model_type]['smal_model_path']).replace('.pkl', '_template.npz')
            smal_downsampling_npz_path = ROOT_smal_downsampling + smal_downsampling_npz_name  # 'data/mesh_downsampling.npz'
            self.my_custom_smal_dog_mesh = Mesh(filename=smal_downsampling_npz_path, num_downsampling=num_downsampling, nsize=1, body_model=smal) # , device=device)
            # create GraphCNN
            num_layers = 2  # <= len(my_custom_mesh._A)-1
            n_resnet_out = self.n_out_enc       # 256
            num_channels = 256      # 512 
            self.graph_cnn = GraphCNNMS(mesh=self.my_custom_smal_dog_mesh, 
                                num_downsample = num_downsampling, 
                                num_layers = num_layers, 
                                n_resnet_out = n_resnet_out, 
                                num_channels = num_channels)    # .to(device)
        
    
    
    def forward(self, keyp_sh, keyp_pred, in_pose_3x3, in_trans_notnorm, in_cam_notnorm, in_betas, in_betas_limbs, seg_pred_prep=None, seg_sh_raw=None, seg_sh_prep=None):
        # img is the network input image 
        # seg_raw is before softmax and subtracting 0.5
        # seg_prep would be the prepared_segmentation
        batch_size = in_pose_3x3.shape[0]
        device = in_pose_3x3.device
        dtype = in_pose_3x3.dtype
        # --- segmentation encoder
        if seg_sh_prep is None:
            seg_sh_prep = self.soft_max(seg_sh_raw)[:, 1:2, :, :] - 0.5       # class 1 is the dog
        input_seg_conc = torch.cat((seg_sh_prep, seg_pred_prep), axis=1)  
        network_output_seg = self.ref_res(input_seg_conc)
        # --- keypoint encoder
        keyp_conc = torch.cat((keyp_sh.reshape((-1, keyp_sh.shape[1]*keyp_sh.shape[2])), keyp_pred.reshape((-1, keyp_sh.shape[1]*keyp_sh.shape[2]))), axis=1) 
        network_output_keyp = self.linear_keyp(keyp_conc)
        # --- decoder
        x = torch.cat((network_output_seg, network_output_keyp), axis=1)
        y_comb = self.linear_combined(x)
        in_pose_6d = rotmat_to_rot6d(in_pose_3x3.reshape((-1, 3, 3))).reshape((in_pose_3x3.shape[0], -1, 6))
        in_dict = {'pose': in_pose_6d,
                    'trans_notnorm': in_trans_notnorm,
                    'flength_notnorm': in_cam_notnorm, 
                    'betas': in_betas, 
                    'betas_limbs': in_betas_limbs}
        results = {}
        for element in self.output_info:
            # import pdb; pdb.set_trace()

            linear_model = self.output_info_linear_models[element['linear_model_index']]
            y = torch.cat((y_comb, in_dict[element['name']].reshape((-1, element['n']))), axis=1)
            if 'out_shape' in element.keys():
                if element['name'] == 'pose':
                    if self.ref_net_type in ['multrot', 'multrot01', 'multrot01all', 'multrotxx', 'multrot_res34', 'multrot01all_res34']:      # if self.ref_net_type == 'multrot' or self.ref_net_type == 'multrot_res34':
                        #   multiply the rotations with each other -> just predict a correction
                        #   the correction should be initialized as identity
                        # res_pose_out = (linear_model(y)).reshape((-1, element['out_shape'][0], element['out_shape'][1])) + in_dict[element['name']]
                        identity_rot6d = torch.tensor(([1., 0., 0., 1., 0., 0.])).repeat((in_pose_3x3.shape[0]*in_pose_3x3.shape[1], 1)).to(device=device, dtype=dtype)
                        if self.ref_net_type in ['multrot01', 'multrot01all', 'multrot01all_res34']:
                            res_pose_out = identity_rot6d + 0.1*(linear_model(y)).reshape((-1, element['out_shape'][1]))  
                        elif self.ref_net_type == 'multrotxx':
                            res_pose_out = identity_rot6d + 0.0*(linear_model(y)).reshape((-1, element['out_shape'][1]))    
                        else:
                            res_pose_out = identity_rot6d + (linear_model(y)).reshape((-1, element['out_shape'][1]))    
                        res_pose_rotmat = rot6d_to_rotmat(res_pose_out.reshape((-1, 6)))    # (bs*35, 3, 3)     .reshape((batch_size, -1, 3, 3))
                        res_tot_rotmat = torch.bmm(res_pose_rotmat.reshape((-1, 3, 3)), in_pose_3x3.reshape((-1, 3, 3))).reshape((batch_size, -1, 3, 3))   # (bs, 5, 3, 3)
                        results['pose_rotmat'] = res_tot_rotmat
                    elif self.ref_net_type == 'add':
                        res_6d = (linear_model(y)).reshape((-1, element['out_shape'][0], element['out_shape'][1])) + in_dict['pose']
                        results['pose_rotmat'] = rot6d_to_rotmat(res_6d.reshape((-1, 6))).reshape((batch_size, -1, 3, 3)) 
                    else:
                        raise ValueError
                else:
                    if self.ref_net_type in ['multrot01all', 'multrot01all_res34']:
                        results[element['name']] = (0.1*linear_model(y)).reshape((-1, element['out_shape'][0], element['out_shape'][1])) + in_dict[element['name']]
                    else:
                        results[element['name']] = (linear_model(y)).reshape((-1, element['out_shape'][0], element['out_shape'][1])) + in_dict[element['name']]
            else:
                if self.ref_net_type in ['multrot01all', 'multrot01all_res34']:
                    results[element['name']] = 0.1*linear_model(y) + in_dict[element['name']]     
                else:
                    results[element['name']] = linear_model(y) + in_dict[element['name']]  

        # add prediction if ground is flat
        if not self.isflat_type=='inexistent':
            isflat = self.linear_isflat(y_comb)
            results['isflat'] = isflat

        # add graph cnn
        if not self.graphcnn_type=='inexistent':
            ground_contact_downsampled, ground_cntact_all_stages_output = self.graph_cnn(y_comb)
            ground_contact = self.my_custom_smal_dog_mesh.upsample(ground_contact_downsampled.transpose(1,2))
            results['vertexwise_ground_contact'] = ground_contact

        return results




class ModelImageToBreed(nn.Module):
    def __init__(self, smal_model_type, arch='hg8', n_joints=35, n_classes=20, n_partseg=15, n_keyp=20, n_bones=24, n_betas=10, n_betas_limbs=7, n_breeds=121, image_size=256, n_z=512, thr_keyp_sc=None, add_partseg=True):
        super(ModelImageToBreed, self).__init__()
        self.n_classes = n_classes
        self.n_partseg = n_partseg
        self.n_betas = n_betas
        self.n_betas_limbs = n_betas_limbs
        self.n_keyp = n_keyp
        self.n_bones = n_bones
        self.n_breeds = n_breeds
        self.image_size = image_size
        self.upsample_seg = True
        self.threshold_scores = thr_keyp_sc 
        self.n_z = n_z
        self.add_partseg = add_partseg
        self.smal_model_type = smal_model_type
        # ------------------------------ STACKED HOUR GLASS ------------------------------        
        if arch == 'hg8':
            self.stacked_hourglass = hg8(pretrained=False, num_classes=self.n_classes, num_partseg=self.n_partseg, upsample_seg=self.upsample_seg, add_partseg=self.add_partseg)
        else:
            raise Exception('unrecognised model architecture: ' + arch)
        # ------------------------------ SHAPE AND BREED MODEL ------------------------------
        self.breed_model = ModelShapeAndBreed(smal_model_type=self.smal_model_type, n_betas=self.n_betas, n_betas_limbs=self.n_betas_limbs, n_breeds=self.n_breeds, n_z=self.n_z)
    def forward(self, input_img, norm_dict=None, bone_lengths_prepared=None, betas=None):
        batch_size = input_img.shape[0]
        device = input_img.device
        # ------------------------------ STACKED HOUR GLASS ------------------------------
        hourglass_out_dict = self.stacked_hourglass(input_img)
        last_seg = hourglass_out_dict['seg_final']
        last_heatmap = hourglass_out_dict['out_list_kp'][-1] 
        # - prepare keypoints (from heatmap)
        # normalize predictions -> from logits to probability distribution
        # last_heatmap_norm = dsnt.spatial_softmax2d(last_heatmap, temperature=torch.tensor(1))
        # keypoints = dsnt.spatial_expectation2d(last_heatmap_norm, normalized_coordinates=False) + 1   # (bs, 20, 2)
        # keypoints_norm = dsnt.spatial_expectation2d(last_heatmap_norm, normalized_coordinates=True)    # (bs, 20, 2)
        keypoints_norm, scores = get_preds_soft(last_heatmap, return_maxval=True, norm_coords=True)
        if self.threshold_scores is not None:
            scores[scores>self.threshold_scores] = 1.0
            scores[scores<=self.threshold_scores] = 0.0
        # ------------------------------ SHAPE AND BREED MODEL ------------------------------
        # breed_model takes as input the image as well as the predicted segmentation map 
        #     -> we need to split up ModelImageTo3d, such that we can use the silhouette
        resnet_output = self.breed_model(img=input_img, seg_raw=last_seg)
        pred_breed = resnet_output['breeds']       # (bs, n_breeds)
        pred_betas = resnet_output['betas']
        pred_betas_limbs = resnet_output['betas_limbs']
        small_output = {'keypoints_norm': keypoints_norm,
                        'keypoints_scores': scores}
        small_output_reproj = {'betas': pred_betas,
                                'betas_limbs': pred_betas_limbs,
                                'dog_breed': pred_breed}
        return small_output, None, small_output_reproj

class ModelImageTo3d_withshape_withproj(nn.Module):
    def __init__(self, smal_model_type, smal_keyp_conf=None, arch='hg8', num_stage_comb=2, num_stage_heads=1, num_stage_heads_pose=1, trans_sep=False, n_joints=35, n_classes=20, n_partseg=15, n_keyp=20, n_bones=24, n_betas=10, n_betas_limbs=6, n_breeds=121, image_size=256, n_z=512, n_segbps=64*2, thr_keyp_sc=None, add_z_to_3d_input=True, add_segbps_to_3d_input=False, add_partseg=True, silh_no_tail=True, fix_flength=False, render_partseg=False, structure_z_to_betas='default', structure_pose_net='default', nf_version=None, ref_net_type='add', ref_detach_shape=True, graphcnn_type='inexistent', isflat_type='inexistent', shaperef_type='inexistent'):
        super(ModelImageTo3d_withshape_withproj, self).__init__()
        self.n_classes = n_classes
        self.n_partseg = n_partseg
        self.n_betas = n_betas
        self.n_betas_limbs = n_betas_limbs
        self.n_keyp = n_keyp
        self.n_joints = n_joints
        self.n_bones = n_bones
        self.n_breeds = n_breeds
        self.image_size = image_size
        self.threshold_scores = thr_keyp_sc 
        self.upsample_seg = True
        self.silh_no_tail = silh_no_tail
        self.add_z_to_3d_input = add_z_to_3d_input       
        self.add_segbps_to_3d_input = add_segbps_to_3d_input
        self.add_partseg = add_partseg
        self.ref_net_type = ref_net_type
        self.ref_detach_shape = ref_detach_shape
        self.graphcnn_type = graphcnn_type
        self.isflat_type = isflat_type
        self.shaperef_type = shaperef_type
        assert (not self.add_segbps_to_3d_input) or (not self.add_z_to_3d_input)
        self.n_z = n_z   
        if add_segbps_to_3d_input:
            self.n_segbps = n_segbps    # 64
            self.segbps_model = SegBPS()
        else:
            self.n_segbps = 0
        self.fix_flength = fix_flength
        self.render_partseg = render_partseg
        self.structure_z_to_betas = structure_z_to_betas
        self.structure_pose_net = structure_pose_net
        assert self.structure_pose_net in ['default', 'vae', 'normflow']
        self.nf_version = nf_version
        self.smal_model_type = smal_model_type
        assert (smal_keyp_conf is not None)
        self.smal_keyp_conf = smal_keyp_conf
        self.register_buffer('betas_zeros', torch.zeros((1, self.n_betas)))
        self.register_buffer('mean_dog_bone_lengths', torch.tensor(MEAN_DOG_BONE_LENGTHS_NO_RED, dtype=torch.float32))
        p_dropout = 0.2      # 0.5     
        # ------------------------------ SMAL MODEL ------------------------------
        self.smal = SMAL(smal_model_type=self.smal_model_type, template_name='neutral')     
        print('SMAL model type: ' + self.smal.smal_model_type)      
        # New for rendering without tail
        f_np = self.smal.faces.detach().cpu().numpy()
        self.f_no_tail_np = f_np[np.isin(f_np[:,:], VERTEX_IDS_TAIL).sum(axis=1)==0, :]
        # in theory we could optimize for improved shapedirs, but we do not do that
        #   -> would need to implement regularizations 
        #   -> there are better ways than changing the shapedirs
        self.model_learnable_shapedirs = LearnableShapedirs(self.smal.sym_ids_dict, self.smal.shapedirs.detach(), self.n_betas, 10)
        # ------------------------------ STACKED HOUR GLASS ------------------------------        
        if arch == 'hg8':
            self.stacked_hourglass = hg8(pretrained=False, num_classes=self.n_classes, num_partseg=self.n_partseg, upsample_seg=self.upsample_seg, add_partseg=self.add_partseg)
        else:
            raise Exception('unrecognised model architecture: ' + arch)
        # ------------------------------ SHAPE AND BREED MODEL ------------------------------
        self.breed_model = ModelShapeAndBreed(self.smal_model_type, n_betas=self.n_betas, n_betas_limbs=self.n_betas_limbs, n_breeds=self.n_breeds, n_z=self.n_z, structure_z_to_betas=self.structure_z_to_betas)
        # ------------------------------ LINEAR 3D MODEL ------------------------------
        # 3d model -> from image to 3d parameters {2d keypoints from heatmap, pose, trans, flength}
        self.soft_max = torch.nn.Softmax(dim=1)
        input_size = self.n_keyp*3 + self.n_bones
        self.model_3d = LinearModelComplete(linear_size=1024,
                    num_stage_comb=num_stage_comb,
                    num_stage_heads=num_stage_heads,
                    num_stage_heads_pose=num_stage_heads_pose,
                    trans_sep=trans_sep, 
                    p_dropout=p_dropout,        # 0.5, 
                    input_size=input_size,
                    intermediate_size=1024,
                    output_info=None,
                    n_joints=self.n_joints,
                    n_z=self.n_z,
                    add_z_to_3d_input=self.add_z_to_3d_input,
                    n_segbps=self.n_segbps,
                    add_segbps_to_3d_input=self.add_segbps_to_3d_input, 
                    structure_pose_net=self.structure_pose_net,
                    nf_version = self.nf_version)
        # ------------------------------ RENDERING ------------------------------
        self.silh_renderer = SilhRenderer(image_size) 
        # ------------------------------ REFINEMENT -----------------------------
        self.refinement_model = ModelRefinement(n_betas=self.n_betas, n_betas_limbs=self.n_betas_limbs, n_breeds=self.n_breeds, n_keyp=self.n_keyp, n_joints=self.n_joints, ref_net_type=self.ref_net_type, graphcnn_type=self.graphcnn_type, isflat_type=self.isflat_type, shaperef_type=self.shaperef_type)


    def forward(self, input_img, norm_dict=None, bone_lengths_prepared=None, betas=None):
        batch_size = input_img.shape[0]
        device = input_img.device
        # ------------------------------ STACKED HOUR GLASS ------------------------------
        hourglass_out_dict = self.stacked_hourglass(input_img)
        last_seg = hourglass_out_dict['seg_final']
        last_heatmap = hourglass_out_dict['out_list_kp'][-1] 
        # - prepare keypoints (from heatmap)
        # normalize predictions -> from logits to probability distribution
        # last_heatmap_norm = dsnt.spatial_softmax2d(last_heatmap, temperature=torch.tensor(1))
        # keypoints = dsnt.spatial_expectation2d(last_heatmap_norm, normalized_coordinates=False) + 1   # (bs, 20, 2)
        # keypoints_norm = dsnt.spatial_expectation2d(last_heatmap_norm, normalized_coordinates=True)    # (bs, 20, 2)
        keypoints_norm, scores = get_preds_soft(last_heatmap, return_maxval=True, norm_coords=True)
        if self.threshold_scores is not None:
            scores[scores>self.threshold_scores] = 1.0
            scores[scores<=self.threshold_scores] = 0.0
        # ------------------------------ LEARNABLE SHAPE MODEL ------------------------------
        # in our cvpr 2022 paper we do not change the shapedirs
        # learnable_sd_complete has shape (3889, 3, n_sd)
        # learnable_sd_complete_prepared has shape (n_sd, 11667)
        learnable_sd_complete, learnable_sd_complete_prepared = self.model_learnable_shapedirs()
        shapedirs_sel = learnable_sd_complete_prepared        # None
        # ------------------------------ SHAPE AND BREED MODEL ------------------------------
        # breed_model takes as input the image as well as the predicted segmentation map 
        #     -> we need to split up ModelImageTo3d, such that we can use the silhouette
        resnet_output = self.breed_model(img=input_img, seg_raw=last_seg)
        pred_breed = resnet_output['breeds']       # (bs, n_breeds)
        pred_z = resnet_output['z']
        # - prepare shape
        pred_betas = resnet_output['betas']     
        pred_betas_limbs = resnet_output['betas_limbs'] 
        # - calculate bone lengths
        with torch.no_grad():
            use_mean_bone_lengths = False
            if use_mean_bone_lengths:
                bone_lengths_prepared = torch.cat(batch_size*[self.mean_dog_bone_lengths.reshape((1, -1))])
            else:
                assert (bone_lengths_prepared is None)
                bone_lengths_prepared = self.smal.caclulate_bone_lengths(pred_betas, pred_betas_limbs, shapedirs_sel=shapedirs_sel, short=True)
        # ------------------------------ LINEAR 3D MODEL ------------------------------
        # 3d model -> from image to 3d parameters {2d keypoints from heatmap, pose, trans, flength}
        # prepare input for 2d-to-3d network
        keypoints_prepared = torch.cat((keypoints_norm, scores), axis=2)
        if bone_lengths_prepared is None:
            bone_lengths_prepared = torch.cat(batch_size*[self.mean_dog_bone_lengths.reshape((1, -1))])
        # should we add silhouette to 3d input? should we add z?
        if self.add_segbps_to_3d_input:
            seg_raw = last_seg
            seg_prep_bps = self.soft_max(seg_raw)[:, 1, :, :] # class 1 is the dog
            with torch.no_grad():
                seg_prep_np = seg_prep_bps.detach().cpu().numpy()
                bps_output_np = self.segbps_model.calculate_bps_points_batch(seg_prep_np)  # (bs, 64, 2)
                bps_output = torch.tensor(bps_output_np, dtype=torch.float32).to(device).reshape((batch_size, -1))
                bps_output_prep = bps_output * 2. - 1
            input_vec_keyp_bones = torch.cat((keypoints_prepared.reshape((batch_size, -1)), bone_lengths_prepared), axis=1)  
            input_vec = torch.cat((input_vec_keyp_bones, bps_output_prep), dim=1)
        elif self.add_z_to_3d_input:
            # we do not use this in our cvpr 2022 version
            input_vec_keyp_bones = torch.cat((keypoints_prepared.reshape((batch_size, -1)), bone_lengths_prepared), axis=1)  
            input_vec_additional = pred_z       
            input_vec = torch.cat((input_vec_keyp_bones, input_vec_additional), dim=1)
        else:
            input_vec = torch.cat((keypoints_prepared.reshape((batch_size, -1)), bone_lengths_prepared), axis=1)  
        # predict 3d parameters (those are normalized, we need to correct mean and std in a next step)
        output = self.model_3d(input_vec)      
        # add predicted keypoints to the output dict
        output['keypoints_norm'] = keypoints_norm
        output['keypoints_scores'] = scores
        # add predicted segmentation to output dictc
        output['seg_hg'] = hourglass_out_dict['seg_final']
        # - denormalize 3d parameters -> so far predictions were normalized, now we denormalize them again
        pred_trans = output['trans'] * norm_dict['trans_std'][None, :] + norm_dict['trans_mean'][None, :]    # (bs, 3)
        if  self.structure_pose_net == 'default':
            pred_pose_rot6d = output['pose'] + norm_dict['pose_rot6d_mean'][None, :]
        elif self.structure_pose_net == 'normflow':
            pose_rot6d_mean_zeros = torch.zeros_like(norm_dict['pose_rot6d_mean'][None, :])
            pose_rot6d_mean_zeros[:, 0, :] = norm_dict['pose_rot6d_mean'][None, 0, :]
            pred_pose_rot6d = output['pose'] + pose_rot6d_mean_zeros
        else:
            pose_rot6d_mean_zeros = torch.zeros_like(norm_dict['pose_rot6d_mean'][None, :])
            pose_rot6d_mean_zeros[:, 0, :] = norm_dict['pose_rot6d_mean'][None, 0, :]
            pred_pose_rot6d = output['pose'] + pose_rot6d_mean_zeros
        pred_pose_reshx33 = rot6d_to_rotmat(pred_pose_rot6d.reshape((-1, 6)))
        pred_pose = pred_pose_reshx33.reshape((batch_size, -1, 3, 3))
        pred_pose_rot6d = rotmat_to_rot6d(pred_pose_reshx33).reshape((batch_size, -1, 6))

        if self.fix_flength:
            output['flength'] = torch.zeros_like(output['flength'])
            pred_flength = torch.ones_like(output['flength'])*2100  # norm_dict['flength_mean'][None, :]
        else:
            pred_flength_orig = output['flength'] * norm_dict['flength_std'][None, :] + norm_dict['flength_mean'][None, :]   # (bs, 1)
            pred_flength = pred_flength_orig.clone()  # torch.abs(pred_flength_orig)
            pred_flength[pred_flength_orig<=0] = norm_dict['flength_mean'][None, :]

        # ------------------------------ RENDERING ------------------------------
        # get 3d model (SMAL)
        V, keyp_green_3d, _ = self.smal(beta=pred_betas, betas_limbs=pred_betas_limbs, pose=pred_pose, trans=pred_trans, get_skin=True, keyp_conf=self.smal_keyp_conf, shapedirs_sel=shapedirs_sel)
        keyp_3d = keyp_green_3d[:, :self.n_keyp, :]     # (bs, 20, 3)
        # render silhouette
        faces_prep = self.smal.faces.unsqueeze(0).expand((batch_size, -1, -1))
        if not self.silh_no_tail:
            pred_silh_images, pred_keyp = self.silh_renderer(vertices=V, 
                points=keyp_3d, faces=faces_prep, focal_lengths=pred_flength)
        else:
            faces_no_tail_prep = torch.tensor(self.f_no_tail_np).to(device).expand((batch_size, -1, -1))
            pred_silh_images, pred_keyp = self.silh_renderer(vertices=V, 
                points=keyp_3d, faces=faces_no_tail_prep, focal_lengths=pred_flength)
        # get torch 'Meshes'
        torch_meshes = self.silh_renderer.get_torch_meshes(vertices=V, faces=faces_prep) 

        #  render body parts (not part of cvpr 2022 version)
        if self.render_partseg:
            raise NotImplementedError
        else:
            partseg_images = None
            partseg_images_hg = None


        # ------------------------------ REFINEMENT MODEL ------------------------------

        # refinement model
        pred_keyp_norm = (pred_keyp.detach() / (self.image_size - 1) - 0.5)*2
        '''output_ref = self.refinement_model(keypoints_norm.detach(), pred_keyp_norm, \
                            seg_sh_raw=last_seg[:, :, :, :].detach(), seg_pred_prep=pred_silh_images[:, :, :, :].detach()-0.5, \
                            in_pose=output['pose'].detach(), in_trans=output['trans'].detach(), in_cam=output['flength'].detach(), in_betas=pred_betas.detach())'''
        output_ref = self.refinement_model(keypoints_norm.detach(), pred_keyp_norm, \
                            seg_sh_raw=last_seg[:, :, :, :].detach(), seg_pred_prep=pred_silh_images[:, :, :, :].detach()-0.5, \
                            in_pose_3x3=pred_pose.detach(), in_trans_notnorm=output['trans'].detach(), in_cam_notnorm=output['flength'].detach(), in_betas=pred_betas.detach(), in_betas_limbs=pred_betas_limbs.detach())
        # a better alternative would be to submit pred_pose_reshx33



        # nothing changes for betas or shapedirs or z       ##################### should probably not be detached in the end
        if self.shaperef_type == 'inexistent':
            if self.ref_detach_shape:
                output_ref['betas'] = pred_betas.detach()
                output_ref['betas_limbs'] = pred_betas_limbs.detach()
                output_ref['z'] = pred_z.detach()
                output_ref['shapedirs'] = shapedirs_sel.detach()
            else:
                output_ref['betas'] = pred_betas
                output_ref['betas_limbs'] = pred_betas_limbs
                output_ref['z'] = pred_z
                output_ref['shapedirs'] = shapedirs_sel
        else:
            assert ('betas' in output_ref.keys())
            assert ('betas_limbs' in output_ref.keys())
            output_ref['shapedirs'] = shapedirs_sel     


        # we denormalize flength and trans, but pose is handled differently
        if self.fix_flength:
            output_ref['flength_notnorm'] = torch.zeros_like(output['flength'])
            ref_pred_flength = torch.ones_like(output['flength_notnorm'])*2100  # norm_dict['flength_mean'][None, :]
            raise ValueError    # not sure if we want to have a fixed flength in refinement
        else:
            ref_pred_flength_orig = output_ref['flength_notnorm'] * norm_dict['flength_std'][None, :] + norm_dict['flength_mean'][None, :]   # (bs, 1)
            ref_pred_flength = ref_pred_flength_orig.clone()  # torch.abs(pred_flength_orig)
            ref_pred_flength[ref_pred_flength_orig<=0] = norm_dict['flength_mean'][None, :]
        ref_pred_trans = output_ref['trans_notnorm'] * norm_dict['trans_std'][None, :] + norm_dict['trans_mean'][None, :]    # (bs, 3)




        # ref_pred_pose_rot6d = output_ref['pose']
        # ref_pred_pose_reshx33 = rot6d_to_rotmat(output_ref['pose'].reshape((-1, 6))).reshape((batch_size, -1, 3, 3))
        ref_pred_pose_reshx33 = output_ref['pose_rotmat'].reshape((batch_size, -1, 3, 3))
        ref_pred_pose_rot6d = rotmat_to_rot6d(ref_pred_pose_reshx33.reshape((-1, 3, 3))).reshape((batch_size, -1, 6))

        ref_V, ref_keyp_green_3d, _ = self.smal(beta=output_ref['betas'], betas_limbs=output_ref['betas_limbs'], 
                                        pose=ref_pred_pose_reshx33, trans=ref_pred_trans, get_skin=True, keyp_conf=self.smal_keyp_conf, 
                                        shapedirs_sel=output_ref['shapedirs'])
        ref_keyp_3d = ref_keyp_green_3d[:, :self.n_keyp, :]     # (bs, 20, 3)

        if not self.silh_no_tail:
            faces_prep = self.smal.faces.unsqueeze(0).expand((batch_size, -1, -1))
            ref_pred_silh_images, ref_pred_keyp = self.silh_renderer(vertices=ref_V, 
                points=ref_keyp_3d, faces=faces_prep, focal_lengths=ref_pred_flength)
        else:
            faces_no_tail_prep = torch.tensor(self.f_no_tail_np).to(device).expand((batch_size, -1, -1))
            ref_pred_silh_images, ref_pred_keyp = self.silh_renderer(vertices=ref_V, 
                points=ref_keyp_3d, faces=faces_no_tail_prep, focal_lengths=ref_pred_flength)

        output_ref_unnorm = {'vertices_smal': ref_V,
                            'keyp_3d': ref_keyp_3d,
                            'keyp_2d': ref_pred_keyp,
                            'silh': ref_pred_silh_images,
                            'trans': ref_pred_trans,
                            'flength': ref_pred_flength,
                            'betas': output_ref['betas'],
                            'betas_limbs': output_ref['betas_limbs'],
                            # 'z': output_ref['z'],
                            'pose_rot6d': ref_pred_pose_rot6d,   
                            'pose_rotmat':  ref_pred_pose_reshx33} 
                            # 'shapedirs': shapedirs_sel}

        if not self.graphcnn_type == 'inexistent':
            output_ref_unnorm['vertexwise_ground_contact'] = output_ref['vertexwise_ground_contact']
        if not self.isflat_type=='inexistent':
            output_ref_unnorm['isflat'] = output_ref['isflat']
        if self.shaperef_type == 'inexistent':
            output_ref_unnorm['z'] = output_ref['z']

        # REMARK: we will want to have the predicted differences, for pose this would 
        #   be a rotation matrix, ...
        #       -> TODO: adjust output_orig_ref_comparison
        output_orig_ref_comparison = {#'pose': output['pose'].detach(),
                                    #'trans': output['trans'].detach(),
                                    #'flength': output['flength'].detach(),
                                    # 'pose': output['pose'],
                                    'old_pose_rotmat': pred_pose_reshx33,
                                    'old_trans_notnorm': output['trans'],
                                    'old_flength_notnorm': output['flength'],
                                    # 'ref_pose': output_ref['pose'],
                                    'ref_pose_rotmat': ref_pred_pose_reshx33,
                                    'ref_trans_notnorm': output_ref['trans_notnorm'],
                                    'ref_flength_notnorm': output_ref['flength_notnorm']}



        # ------------------------------ PREPARE OUTPUT ------------------------------
        # create output dictionarys
        # output: contains all output from model_image_to_3d
        # output_unnorm: same as output, but normalizations are undone
        # output_reproj: smal output and reprojected keypoints as well as silhouette 
        keypoints_heatmap_256 = (output['keypoints_norm'] / 2. + 0.5) * (self.image_size - 1)
        output_unnorm = {'pose_rotmat': pred_pose,
                        'flength': pred_flength,
                        'trans': pred_trans,
                        'keypoints':keypoints_heatmap_256}
        output_reproj = {'vertices_smal': V,
                        'torch_meshes': torch_meshes,
                        'keyp_3d': keyp_3d,
                        'keyp_2d': pred_keyp,
                        'silh': pred_silh_images,
                        'betas': pred_betas,
                        'betas_limbs': pred_betas_limbs,
                        'pose_rot6d': pred_pose_rot6d,       # used for pose prior...
                        'dog_breed': pred_breed,
                        'shapedirs': shapedirs_sel,
                        'z': pred_z,
                        'flength_unnorm': pred_flength,
                        'flength': output['flength'],
                        'partseg_images_rend': partseg_images,
                        'partseg_images_hg_nograd': partseg_images_hg,
                        'normflow_z': output['normflow_z']}

        return output, output_unnorm, output_reproj, output_ref_unnorm, output_orig_ref_comparison


    def forward_with_multiple_refinements(self, input_img, norm_dict=None, bone_lengths_prepared=None, betas=None):
        
        # import pdb; pdb.set_trace()

        # run normal network part
        output, output_unnorm, output_reproj, output_ref_unnorm, output_orig_ref_comparison = self.forward(input_img, norm_dict=norm_dict, bone_lengths_prepared=bone_lengths_prepared, betas=betas)

        # prepare input for second refinement stage
        batch_size = output['keypoints_norm'].shape[0]
        keypoints_norm = output['keypoints_norm']
        pred_keyp_norm = (output_ref_unnorm['keyp_2d'].detach() / (self.image_size - 1) - 0.5)*2

        last_seg = output['seg_hg']
        pred_silh_images = output_ref_unnorm['silh'].detach() 

        trans_notnorm = output_orig_ref_comparison['ref_trans_notnorm']
        flength_notnorm = output_orig_ref_comparison['ref_flength_notnorm']
        # trans_notnorm = output_orig_ref_comparison['ref_pose_rotmat']
        pred_pose = output_ref_unnorm['pose_rotmat'].reshape((batch_size, -1, 3, 3))

        # run second refinement step
        output_ref_new = self.refinement_model(keypoints_norm.detach(), pred_keyp_norm, \
                            seg_sh_raw=last_seg[:, :, :, :].detach(), seg_pred_prep=pred_silh_images[:, :, :, :].detach()-0.5, \
                            in_pose_3x3=pred_pose.detach(), in_trans_notnorm=trans_notnorm.detach(), in_cam_notnorm=flength_notnorm.detach(), \
                            in_betas=output_ref_unnorm['betas'].detach(), in_betas_limbs=output_ref_unnorm['betas_limbs'].detach()) 
        # output_ref_new = self.refinement_model(keypoints_norm.detach(), pred_keyp_norm, seg_sh_raw=last_seg[:, :, :, :].detach(), seg_pred_prep=pred_silh_images[:, :, :, :].detach()-0.5, in_pose_3x3=pred_pose.detach(), in_trans_notnorm=trans_notnorm.detach(), in_cam_notnorm=flength_notnorm.detach(), in_betas=output_ref_unnorm['betas'].detach(), in_betas_limbs=output_ref_unnorm['betas_limbs'].detach()) 


        # new shape
        if self.shaperef_type == 'inexistent':
            if self.ref_detach_shape:
                output_ref_new['betas'] = output_ref_unnorm['betas'].detach()
                output_ref_new['betas_limbs'] = output_ref_unnorm['betas_limbs'].detach()
                output_ref_new['z'] = output_ref_unnorm['z'].detach()
                output_ref_new['shapedirs'] = output_reproj['shapedirs'].detach()
            else:
                output_ref_new['betas'] = output_ref_unnorm['betas']
                output_ref_new['betas_limbs'] = output_ref_unnorm['betas_limbs']
                output_ref_new['z'] = output_ref_unnorm['z']
                output_ref_new['shapedirs'] = output_reproj['shapedirs']
        else:
            assert ('betas' in output_ref_new.keys())
            assert ('betas_limbs' in output_ref_new.keys())
            output_ref_new['shapedirs'] = output_reproj['shapedirs']    

        # we denormalize flength and trans, but pose is handled differently
        if self.fix_flength:
            raise ValueError    # not sure if we want to have a fixed flength in refinement
        else:
            ref_pred_flength_orig = output_ref_new['flength_notnorm'] * norm_dict['flength_std'][None, :] + norm_dict['flength_mean'][None, :]   # (bs, 1)
            ref_pred_flength = ref_pred_flength_orig.clone()  # torch.abs(pred_flength_orig)
            ref_pred_flength[ref_pred_flength_orig<=0] = norm_dict['flength_mean'][None, :]
        ref_pred_trans = output_ref_new['trans_notnorm'] * norm_dict['trans_std'][None, :] + norm_dict['trans_mean'][None, :]    # (bs, 3)


        ref_pred_pose_reshx33 = output_ref_new['pose_rotmat'].reshape((batch_size, -1, 3, 3))
        ref_pred_pose_rot6d = rotmat_to_rot6d(ref_pred_pose_reshx33.reshape((-1, 3, 3))).reshape((batch_size, -1, 6))

        ref_V, ref_keyp_green_3d, _ = self.smal(beta=output_ref_new['betas'], betas_limbs=output_ref_new['betas_limbs'], 
                                        pose=ref_pred_pose_reshx33, trans=ref_pred_trans, get_skin=True, keyp_conf=self.smal_keyp_conf, 
                                        shapedirs_sel=output_ref_new['shapedirs'])

        # ref_V, ref_keyp_green_3d, _ = self.smal(beta=output_ref_new['betas'], betas_limbs=output_ref_new['betas_limbs'], pose=ref_pred_pose_reshx33, trans=ref_pred_trans, get_skin=True, keyp_conf=self.smal_keyp_conf, shapedirs_sel=output_ref_new['shapedirs'])
        ref_keyp_3d = ref_keyp_green_3d[:, :self.n_keyp, :]     # (bs, 20, 3)

        if not self.silh_no_tail:
            faces_prep = self.smal.faces.unsqueeze(0).expand((batch_size, -1, -1))
            ref_pred_silh_images, ref_pred_keyp = self.silh_renderer(vertices=ref_V, 
                points=ref_keyp_3d, faces=faces_prep, focal_lengths=ref_pred_flength)
        else:
            faces_no_tail_prep = torch.tensor(self.f_no_tail_np).to(device).expand((batch_size, -1, -1))
            ref_pred_silh_images, ref_pred_keyp = self.silh_renderer(vertices=ref_V, 
                points=ref_keyp_3d, faces=faces_no_tail_prep, focal_lengths=ref_pred_flength)

        output_ref_unnorm_new = {'vertices_smal': ref_V,
                            'keyp_3d': ref_keyp_3d,
                            'keyp_2d': ref_pred_keyp,
                            'silh': ref_pred_silh_images,
                            'trans': ref_pred_trans,
                            'flength': ref_pred_flength,
                            'betas': output_ref_new['betas'],
                            'betas_limbs': output_ref_new['betas_limbs'],
                            'pose_rot6d': ref_pred_pose_rot6d,   
                            'pose_rotmat':  ref_pred_pose_reshx33} 

        if not self.graphcnn_type == 'inexistent':
            output_ref_unnorm_new['vertexwise_ground_contact'] = output_ref_new['vertexwise_ground_contact']
        if not self.isflat_type=='inexistent':
            output_ref_unnorm_new['isflat'] = output_ref_new['isflat']
        if self.shaperef_type == 'inexistent':
            output_ref_unnorm_new['z'] = output_ref_new['z']

        output_orig_ref_comparison_new = {'ref_pose_rotmat': ref_pred_pose_reshx33,
                                    'ref_trans_notnorm': output_ref_new['trans_notnorm'],
                                    'ref_flength_notnorm': output_ref_new['flength_notnorm']}

        results = {
            'output': output, 
            'output_unnorm': output_unnorm, 
            'output_reproj':output_reproj, 
            'output_ref_unnorm': output_ref_unnorm, 
            'output_orig_ref_comparison':output_orig_ref_comparison,
            'output_ref_unnorm_new': output_ref_unnorm_new, 
            'output_orig_ref_comparison_new': output_orig_ref_comparison_new}
        return results





















    def render_vis_nograd(self, vertices, focal_lengths, color=0):
        # this function is for visualization only
        # vertices: (bs, n_verts, 3)
        # focal_lengths: (bs, 1)
        # color: integer, either 0 or 1
        # returns a torch tensor of shape (bs, image_size, image_size, 3)
        with torch.no_grad():
            batch_size = vertices.shape[0]
            faces_prep = self.smal.faces.unsqueeze(0).expand((batch_size, -1, -1))
            visualizations = self.silh_renderer.get_visualization_nograd(vertices, 
                faces_prep, focal_lengths, color=color)
        return visualizations