src/test_time_optimization/evaluate_ttopt.py

# evaluate test time optimization from refinement
# python src/test_time_optimization/evaluate_ttopt.py --workers 12 --save-images True --config refinement_cfg_test_withvertexwisegc_csaddnonflat.yaml --model-file-complete=cvpr23_dm39dnnv3barcv2b_refwithgcpervertisflat0morestanding0/checkpoint.pth.tar --ttopt-result-name ttoptv6_stanext_v16b 

# python src/test_time_optimization/evaluate_ttopt.py --workers 12 --save-images True --config refinement_cfg_test_withvertexwisegc_csaddnonflat.yaml --model-file-complete=cvpr23_dm39dnnv3barcv2b_refwithgcpervertisflat0morestanding0/checkpoint.pth.tar --ttopt-result-name ttoptv6_stanext_v16 



import argparse
import os.path
import json
import numpy as np
import pickle as pkl
from distutils.util import strtobool
import torch
from torch import nn
import torch.backends.cudnn
from torch.nn import DataParallel
from torch.utils.data import DataLoader
import pytorch3d as p3d
from collections import OrderedDict
import glob
from tqdm import tqdm
from dominate import document
from dominate.tags import *
from PIL import Image
from matplotlib import pyplot as plt
import trimesh
import cv2
import shutil

from pytorch3d.structures import Meshes
from pytorch3d.loss import mesh_edge_loss, mesh_laplacian_smoothing, mesh_normal_consistency

import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))

from combined_model.train_main_image_to_3d_wbr_withref import do_validation_epoch
# from combined_model.model_shape_v7 import ModelImageTo3d_withshape_withproj  
# from combined_model.model_shape_v7_withref import ModelImageTo3d_withshape_withproj 
from combined_model.model_shape_v7_withref_withgraphcnn import ModelImageTo3d_withshape_withproj 

from combined_model.loss_image_to_3d_withbreedrel import Loss
from combined_model.loss_image_to_3d_refinement import LossRef
from configs.barc_cfg_defaults import get_cfg_defaults, update_cfg_global_with_yaml, get_cfg_global_updated

from lifting_to_3d.utils.geometry_utils import rot6d_to_rotmat, rotmat_to_rot6d  # , batch_rot2aa, geodesic_loss_R


# from test_time_optimization.utils_ttopt import get_evaluation_dataset, get_norm_dict
from stacked_hourglass.datasets.utils_dataset_selection import get_evaluation_dataset, get_norm_dict

from test_time_optimization.bite_inference_model_for_ttopt import BITEInferenceModel
from smal_pytorch.smal_model.smal_torch_new import SMAL
from configs.SMAL_configs import SMAL_MODEL_CONFIG
from smal_pytorch.renderer.differentiable_renderer import SilhRenderer
from test_time_optimization.utils.utils_ttopt import reset_loss_values, get_optimed_pose_with_glob

from combined_model.loss_utils.loss_utils import leg_sideway_error, leg_torsion_error, tail_sideway_error, tail_torsion_error, spine_torsion_error, spine_sideway_error
from combined_model.loss_utils.loss_utils_gc import LossGConMesh, calculate_plane_errors_batch
from combined_model.loss_utils.loss_arap import Arap_Loss
from combined_model.loss_utils.loss_laplacian_mesh_comparison import LaplacianCTF     # (coarse to fine animal)
from graph_networks import graphcmr     # .utils_mesh import Mesh
from stacked_hourglass.utils.visualization import save_input_image_with_keypoints, save_input_image

from metrics.metrics import Metrics
from configs.SMAL_configs import EVAL_KEYPOINTS, KEYPOINT_GROUPS


ROOT_LOSS_WEIGH_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/configs/ttopt_loss_weights/'



def main(args):

    # load configs
    #   step 1: load default configs
    #   step 2: load updates from .yaml file
    path_config = os.path.join(get_cfg_defaults().barc_dir, 'src', 'configs', args.config)
    update_cfg_global_with_yaml(path_config)
    cfg = get_cfg_global_updated()

    pck_thresh = 0.15
    print('pck_thresh: ' + str(pck_thresh))




    ROOT_IN_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/results/results_ttopt/' + args.ttopt_result_name + '/'      # ttoptv6_debug_x8/'    
    ROOT_IN_PATH_DETAIL = ROOT_IN_PATH + 'details/'

    ROOT_OUT_PATH = ROOT_IN_PATH + 'evaluation/'
    if not os.path.exists(ROOT_OUT_PATH): os.makedirs(ROOT_OUT_PATH)









    
    # NEW!!!
    logscale_part_list = ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'] 
    # logscale_part_list = ['front_legs_l', 'front_legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l', 'back_legs_l', 'back_legs_f'] 


    # Select the hardware device to use for training.
    if torch.cuda.is_available() and cfg.device=='cuda':
        device = torch.device('cuda', torch.cuda.current_device())
        torch.backends.cudnn.benchmark = False      # True
    else:
        device = torch.device('cpu')

    print('structure_pose_net: ' + cfg.params.STRUCTURE_POSE_NET)
    print('refinement network type: ' + cfg.params.REF_NET_TYPE)
    print('smal_model_type: ' + cfg.smal.SMAL_MODEL_TYPE)

    path_model_file_complete = os.path.join(cfg.paths.ROOT_CHECKPOINT_PATH, args.model_file_complete) 

    # Disable gradient calculations.
    # torch.set_grad_enabled(False)


    # prepare dataset and dataset loadr
    val_dataset, val_loader, len_val_dataset, test_name_list, stanext_data_info, stanext_acc_joints = get_evaluation_dataset(cfg.data.DATASET, cfg.data.VAL_OPT, cfg.data.V12, cfg.optim.BATCH_SIZE, args.workers)
    len_data = len_val_dataset
    # summarize information for normalization 
    norm_dict = get_norm_dict(stanext_data_info, device)

    # prepare complete model
    bite_model = BITEInferenceModel(cfg, path_model_file_complete, norm_dict)
    # smal_model_type = bite_model.complete_model.smal.smal_model_type
    smal_model_type = bite_model.smal_model_type
    smal = SMAL(smal_model_type=smal_model_type, template_name='neutral', logscale_part_list=logscale_part_list).to(device)    
    silh_renderer = SilhRenderer(image_size=256).to(device)    



    # ----------------------------------------------------------------------------------

    summary = {}
    summary['pck'] = np.zeros((len_data))
    summary['pck_by_part'] = {group:np.zeros((len_data)) for group in KEYPOINT_GROUPS}
    summary['acc_sil_2d'] = np.zeros(len_data)











    # Put the model in training mode.
    # model.train()
    # prepare progress bar
    iterable = enumerate(val_loader)
    progress = None
    if True:        # not quiet:
        progress = tqdm(iterable, desc='Train', total=len(val_loader), ascii=True, leave=False)
        iterable = progress
    ind_img_tot = 0
    # prepare variables, put them on the right device

    my_step = 0
    batch_size = cfg.optim.BATCH_SIZE

    for index, (input, target_dict) in iterable:
        for key in target_dict.keys(): 
            if key == 'breed_index':
                target_dict[key] = target_dict[key].long().to(device)
            elif key in ['index', 'pts', 'tpts', 'target_weight', 'silh', 'silh_distmat_tofg', 'silh_distmat_tobg', 'sim_breed_index', 'img_border_mask']:
                target_dict[key] = target_dict[key].float().to(device)
            elif key == 'has_seg':
                target_dict[key] = target_dict[key].to(device)
            else:
                pass
        input = input.float().to(device)



        # get starting values for the optimization 
        #   -> here from barc, but could also be saved and loaded
        preds_dict = bite_model.get_all_results(input)
        res_normal_and_ref = bite_model.get_selected_results(preds_dict=preds_dict, result_networks=['normal', 'ref'])
        res = bite_model.get_selected_results(preds_dict=preds_dict, result_networks=['ref'])['ref']

        # --------------------------------------------------------------------

        # ind_img = 0

        batch_verts_smal = []
        batch_faces_prep = []
        batch_optimed_camera_flength = []



        for ind_img in range(input.shape[0]):
            name = (test_name_list[target_dict['index'][ind_img].long()]).replace('/', '__').split('.')[0]

            print('ind_img_tot: ' + str(ind_img_tot) + '   -> ' + name)
            ind_img_tot += 1

            e_name = 'e000'     # 'e300'

            npy_file = ROOT_IN_PATH_DETAIL + name + '_flength_' + e_name +'.npy'
            flength = np.load(npy_file)
            optimed_camera_flength = torch.tensor(flength, device=device)

            obj_file = ROOT_IN_PATH + name + '_res_' + e_name +'.obj'

            verts, faces, aux = p3d.io.load_obj(obj_file)
            verts_smal = verts[None, ...].to(device)
            faces_prep = faces.verts_idx[None, ...].to(device)
            batch_verts_smal.append(verts_smal)
            batch_faces_prep.append(faces_prep)
            batch_optimed_camera_flength.append(optimed_camera_flength)


        # import pdb; pdb.set_trace()

        verts_smal = torch.cat(batch_verts_smal, dim=0)
        faces_prep = torch.cat(batch_faces_prep, dim=0)
        optimed_camera_flength = torch.cat(batch_optimed_camera_flength, dim=0)

        # get keypoint locations from mesh vertices
        keyp_3d = smal.get_joints_from_verts(verts_smal, keyp_conf='olive')
        

        # render silhouette and keypoints
        pred_silh_images, pred_keyp_raw = silh_renderer(vertices=verts_smal, points=keyp_3d, faces=faces_prep, focal_lengths=optimed_camera_flength)
        pred_keyp = pred_keyp_raw[:, :24, :]



        # --------------- calculate iou and pck values --------------------

        gt_keypoints_256 = target_dict['tpts'][:, :, :2] / 64. * (256. - 1)
        gt_keypoints = torch.cat((gt_keypoints_256, target_dict['tpts'][:, :, 2:3]), dim=2)    
        # prepare silhouette for IoU calculation - predicted as well as ground truth
        has_seg = target_dict['has_seg']
        img_border_mask = target_dict['img_border_mask'][:, 0, :, :]
        gtseg = target_dict['silh']
        synth_silhouettes = pred_silh_images[:, 0, :, :]   # pred_silh[:, 0, :, :]       # output_reproj['silh']
        synth_silhouettes[synth_silhouettes>0.5] = 1
        synth_silhouettes[synth_silhouettes<0.5] = 0
        # calculate PCK as well as IoU (similar to WLDO)
        preds = {}
        preds['acc_PCK'] = Metrics.PCK(
            pred_keyp, gt_keypoints, 
            gtseg, has_seg, idxs=EVAL_KEYPOINTS,
            thresh_range=[pck_thresh],       # [0.15],
        )
        preds['acc_IOU'] = Metrics.IOU(
            synth_silhouettes, gtseg, 
            img_border_mask, mask=has_seg
        )
        for group, group_kps in KEYPOINT_GROUPS.items():
            preds[f'{group}_PCK'] = Metrics.PCK(
                pred_keyp, gt_keypoints, gtseg, has_seg, 
                thresh_range=[pck_thresh],       # [0.15],
                idxs=group_kps
            )

        curr_batch_size = pred_keyp.shape[0]
        if not (preds['acc_PCK'].data.cpu().numpy().shape == (summary['pck'][my_step * batch_size:my_step * batch_size + curr_batch_size]).shape):
            import pdb; pdb.set_trace()
        summary['pck'][my_step * batch_size:my_step * batch_size + curr_batch_size] = preds['acc_PCK'].data.cpu().numpy()
        summary['acc_sil_2d'][my_step * batch_size:my_step * batch_size + curr_batch_size] = preds['acc_IOU'].data.cpu().numpy()
        for part in summary['pck_by_part']:
            summary['pck_by_part'][part][my_step * batch_size:my_step * batch_size + curr_batch_size] = preds[f'{part}_PCK'].data.cpu().numpy()




        my_step += 1





    # import pdb; pdb.set_trace()






    iou = np.nanmean(summary['acc_sil_2d'])
    pck = np.nanmean(summary['pck'])
    pck_legs = np.nanmean(summary['pck_by_part']['legs'])
    pck_tail = np.nanmean(summary['pck_by_part']['tail'])
    pck_ears = np.nanmean(summary['pck_by_part']['ears'])
    pck_face = np.nanmean(summary['pck_by_part']['face'])
    print('------------------------------------------------')
    print("iou:         {:.2f}".format(iou*100))
    print('                                                ')
    print("pck:         {:.2f}".format(pck*100))
    print('                                                ')
    print("pck_legs:    {:.2f}".format(pck_legs*100))
    print("pck_tail:    {:.2f}".format(pck_tail*100))
    print("pck_ears:    {:.2f}".format(pck_ears*100))
    print("pck_face:    {:.2f}".format(pck_face*100))
    print('------------------------------------------------')
    # save results in a .txt file
    with open(ROOT_OUT_PATH + "a_evaluation_" + e_name + ".txt", "w") as text_file:
        print("iou:         {:.2f}".format(iou*100), file=text_file)
        print("pck:         {:.2f}".format(pck*100), file=text_file)
        print("pck_legs:    {:.2f}".format(pck_legs*100), file=text_file)
        print("pck_tail:    {:.2f}".format(pck_tail*100), file=text_file)
        print("pck_ears:    {:.2f}".format(pck_ears*100), file=text_file)
        print("pck_face:    {:.2f}".format(pck_face*100), file=text_file)




































if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Evaluate a stacked hourglass model.')
    parser.add_argument('--model-file-complete', default='', type=str, metavar='PATH',
                        help='path to saved model weights')
    parser.add_argument('--ttopt-result-name', default='', type=str, metavar='PATH',
                        help='path to saved ttopt results')
    parser.add_argument('-cg', '--config', default='barc_cfg_test.yaml', type=str, metavar='PATH',
                        help='name of config file (default: barc_cfg_test.yaml within src/configs folder)')
    parser.add_argument('--save-images', default='True', type=lambda x: bool(strtobool(x)),
                        help='bool indicating if images should be saved')
    parser.add_argument('--workers', default=4, type=int, metavar='N',
                        help='number of data loading workers')
    parser.add_argument('--metrics', '-m', metavar='METRICS', default='all',
                        choices=['all', None],
                        help='model architecture')
    main(parser.parse_args())