src/combined_model/model_shape_v7.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 UNITY_SMAL_SHAPE_PRIOR_DOGS as SHAPE_PRIOR
from configs.SMAL_configs import MEAN_DOG_BONE_LENGTHS_NO_RED, VERTEX_IDS_TAIL



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, 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
        # 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(SHAPE_PRIOR, '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
        self.betas_mean_np = res['dog_cluster_mean']  
                                        
    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 ModelImageToBreed(nn.Module):
    def __init__(self, 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
        # ------------------------------ 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(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, 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):
        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_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
        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.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(template_name='neutral')       
        # 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(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=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) 

    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
        # - 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='green', 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

        # ------------------------------ 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

    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