initial commit with code and data

LICENSE

@@ -0,0 +1,60 @@
+License
+Software Copyright License for non-commercial scientific research purposes
+Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use BITE data, model and software, (the "Data & Software"), including 3D meshes, images, videos, textures, software, scripts, and animations. By downloading and/or using the Data & Software (including downloading, cloning, installing, and any other use of the corresponding github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Data & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License
+
+Ownership / Licensees
+The Software and the associated materials has been developed at the
+
+Max Planck Institute for Intelligent Systems 
+and 
+ETH Zurich 
+
+Any copyright or patent right is owned by and proprietary material of the
+
+Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”)
+
+hereinafter the “Licensor”.
+
+License Grant
+Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right:
+
+To install the Data & Software on computers owned, leased or otherwise controlled by you and/or your organization;
+To use the Data & Software for the sole purpose of performing non-commercial scientific research, non-commercial education, or non-commercial artistic projects;
+Any other use, in particular any use for commercial, pornographic, military, or surveillance, purposes is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artifacts for commercial purposes. The Data & Software may not be used to create fake, libelous, misleading, or defamatory content of any kind excluding analyses in peer-reviewed scientific research. The Data & Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission.
+
+The Data & Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Software to train methods/algorithms/neural networks/etc. for commercial, pornographic, military, surveillance, or defamatory use of any kind. By downloading the Data & Software, you agree not to reverse engineer it.
+
+No Distribution
+The Data & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only.
+
+Disclaimer of Representations and Warranties
+You expressly acknowledge and agree that the Data & Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Data & Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE DATA & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Data & Software, (ii) that the use of the Data & Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Data & Software will not cause any damage of any kind to you or a third party.
+
+Limitation of Liability
+Because this Data & Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage.
+Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded.
+Patent claims generated through the usage of the Data & Software cannot be directed towards the copyright holders.
+The Data & Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Data & Software and is not responsible for any problems such modifications cause.
+
+No Maintenance Services
+You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Data & Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Data & Software at any time.
+
+Defects of the Data & Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication.
+
+Publications using the Data & Software
+You acknowledge that the Data & Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Data & Software.
+
+Citation:
+
+
+@inproceedings{BITE:2023,
+    title = {BITE: Beyond priors for Improved Three-D dog pose Estimation},
+    author = {Rueegg, Nadine and Tripathi, Shashank and Schindler, Konrad and Black, Michael J. and Zuffi, Silvia},
+    booktitle = {under review},
+    year = {2023}
+  url = {https://bite.is.tue.mpg.de}
+}
+Commercial licensing opportunities
+For commercial uses of the Data & Software, please send email to ps-license@tue.mpg.de
+
+This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention.

README.md

@@ -1,12 +1,9 @@
----
-title: Bite Gradio
-emoji: 👀
-colorFrom: blue
-colorTo: pink
+title: BITE
+emoji: 🐩  🐶  🐕
+colorFrom: pink
+colorTo: green
 sdk: gradio
-sdk_version: 3.35.2
-app_file: app.py
+sdk_version: 3.0.2
+app_file: ./scripts/gradio_demo.py
 pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+python_version: 3.7.6

packages.txt

@@ -0,0 +1,8 @@
+libgl1
+unzip
+ffmpeg 
+libsm6 
+libxext6
+libgl1-mesa-dri
+libegl1-mesa
+libgbm1 

requirements.txt

@@ -0,0 +1,15 @@
+torch==1.6.0
+torchvision==0.7.0
+pytorch3d==0.2.5
+kornia==0.4.0 
+matplotlib
+opencv-python
+trimesh
+scipy
+chumpy
+pymp
+importlib-resources
+pycocotools
+openpyxl
+dominate
+git+https://github.com/runa91/FrEIA.git

scripts/gradio_demo.py

@@ -0,0 +1,672 @@
+
+# aenv_new_icon_2
+
+# was used for ttoptv6_sketchfab_v16: python src/test_time_optimization/ttopt_fromref_v6_sketchfab.py --workers 12 --save-images True --config refinement_cfg_visualization_withgc_withvertexwisegc_isflat.yaml --model-file-complete=cvpr23_dm39dnnv3barcv2b_refwithgcpervertisflat0morestanding0/checkpoint.pth.tar --sketchfab 1
+
+# for stanext images:
+#   python scripts/gradio.py --workers 12 --config refinement_cfg_test_withvertexwisegc_csaddnonflat.yaml --model-file-complete=cvpr23_dm39dnnv3barcv2b_refwithgcpervertisflat0morestanding0/checkpoint.pth.tar -s ttopt_vtest1
+# for all images from the folder datasets/test_image_crops:
+#   python scripts/gradio.py --workers 12 --config refinement_cfg_test_withvertexwisegc_csaddnonflat_crops.yaml --model-file-complete=cvpr23_dm39dnnv3barcv2b_refwithgcpervertisflat0morestanding0/checkpoint.pth.tar -s ttopt_vtest2
+
+'''import os
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  
+os.environ["CUDA_VISIBLE_DEVICES"]="0"
+try:
+    # os.system("pip install --upgrade  torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html")
+    os.system("pip install --upgrade  torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/cu101/torch_stable.html")
+except Exception as e:
+    print(e)'''
+
+import argparse
+import os.path
+import json
+import numpy as np
+import pickle as pkl
+import csv
+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
+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
+import random
+import gradio as gr
+
+import torchvision
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+import torchvision.transforms as T
+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_withref_withgraphcnn import ModelImageTo3d_withshape_withproj 
+
+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  
+from stacked_hourglass.datasets.utils_dataset_selection import get_evaluation_dataset, get_sketchfab_evaluation_dataset, get_crop_evaluation_dataset, get_norm_dict, get_single_crop_dataset_from_image
+
+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
+
+random.seed(0)
+
+print(
+    "torch: ", torch.__version__,
+    "\ntorchvision: ", torchvision.__version__,
+)
+
+
+def get_prediction(model, img_path_or_img, confidence=0.5):
+    """
+    see https://haochen23.github.io/2020/04/object-detection-faster-rcnn.html#.YsMCm4TP3-g
+    get_prediction
+        parameters:
+        - img_path - path of the input image
+        - confidence - threshold value for prediction score
+        method:
+        - Image is obtained from the image path
+        - the image is converted to image tensor using PyTorch's Transforms
+        - image is passed through the model to get the predictions
+        - class, box coordinates are obtained, but only prediction score > threshold
+            are chosen.
+    """
+    if isinstance(img_path_or_img, str):
+        img = Image.open(img_path_or_img).convert('RGB')
+    else:
+        img = img_path_or_img
+    transform = T.Compose([T.ToTensor()])
+    img = transform(img)
+    pred = model([img])
+    # pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]
+    pred_class = list(pred[0]['labels'].numpy())
+    pred_boxes = [[(int(i[0]), int(i[1])), (int(i[2]), int(i[3]))] for i in list(pred[0]['boxes'].detach().numpy())]
+    pred_score = list(pred[0]['scores'].detach().numpy())
+    try:
+        pred_t = [pred_score.index(x) for x in pred_score if x>confidence][-1]
+        pred_boxes = pred_boxes[:pred_t+1]
+        pred_class = pred_class[:pred_t+1]
+        return pred_boxes, pred_class, pred_score
+    except:
+        print('no bounding box with a score that is high enough found! -> work on full image')
+        return None, None, None
+
+
+def detect_object(model, img_path_or_img, confidence=0.5, rect_th=2, text_size=0.5, text_th=1):
+    """
+    see https://haochen23.github.io/2020/04/object-detection-faster-rcnn.html#.YsMCm4TP3-g
+    object_detection_api
+        parameters:
+        - img_path_or_img - path of the input image
+        - confidence - threshold value for prediction score
+        - rect_th - thickness of bounding box
+        - text_size - size of the class label text
+        - text_th - thichness of the text
+        method:
+        - prediction is obtained from get_prediction method
+        - for each prediction, bounding box is drawn and text is written 
+            with opencv
+        - the final image is displayed
+    """
+    boxes, pred_cls, pred_scores = get_prediction(model, img_path_or_img, confidence)
+    if isinstance(img_path_or_img, str):
+        img = cv2.imread(img_path_or_img)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    else:
+        img = img_path_or_img
+    is_first = True
+    bbox = None
+    if boxes is not None:
+        for i in range(len(boxes)):
+            cls = pred_cls[i]
+            if cls == 18 and bbox is None:
+                cv2.rectangle(img, boxes[i][0], boxes[i][1],color=(0, 255, 0), thickness=rect_th)
+                # cv2.putText(img, pred_cls[i], boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)
+                # cv2.putText(img, str(pred_scores[i]), boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)
+                bbox = boxes[i]
+    return img, bbox
+
+
+# -------------------------------------------------------------------------------------------------------------------- #
+model_bbox = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+model_bbox.eval()
+
+def run_bbox_inference(input_image):
+    # load configs
+    cfg = get_cfg_global_updated()
+    out_path = os.path.join(cfg.paths.ROOT_OUT_PATH, 'gradio_examples', 'test2.png')
+    img, bbox = detect_object(model=model_bbox, img_path_or_img=input_image, confidence=0.5)
+    fig = plt.figure()   #  plt.figure(figsize=(20,30))
+    plt.imsave(out_path, img)
+    return img, bbox
+
+
+
+# -------------------------------------------------------------------------------------------------------------------- #
+# python scripts/gradio.py --workers 12 --config refinement_cfg_test_withvertexwisegc_csaddnonflat.yaml --model-file-complete=cvpr23_dm39dnnv3barcv2b_refwithgcpervertisflat0morestanding0/checkpoint.pth.tar
+args_config = "refinement_cfg_test_withvertexwisegc_csaddnonflat.yaml"
+args_model_file_complete = "cvpr23_dm39dnnv3barcv2b_refwithgcpervertisflat0morestanding0/checkpoint.pth.tar"
+args_suffix = "ttopt_v0"
+args_loss_weight_ttopt_path = "bite_loss_weights_ttopt.json"
+args_workers = 12
+# -------------------------------------------------------------------------------------------------------------------- #
+
+
+
+# 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()
+
+# define path to load the trained model
+path_model_file_complete = os.path.join(cfg.paths.ROOT_CHECKPOINT_PATH, args_model_file_complete) 
+
+# define and create paths to save results
+out_sub_name = cfg.data.VAL_OPT + '_' + cfg.data.DATASET + '_' + args_suffix + '/'
+root_out_path = os.path.join(os.path.dirname(path_model_file_complete).replace(cfg.paths.ROOT_CHECKPOINT_PATH, cfg.paths.ROOT_OUT_PATH + 'results_gradio/'), out_sub_name)
+root_out_path_details = root_out_path + 'details/'
+if not os.path.exists(root_out_path): os.makedirs(root_out_path)
+if not os.path.exists(root_out_path_details): os.makedirs(root_out_path_details)
+print('root_out_path: ' + root_out_path)
+
+# other paths
+root_data_path = os.path.join(os.path.dirname(__file__), '../', 'data')
+# downsampling as used in graph neural network
+root_smal_downsampling = os.path.join(root_data_path, 'graphcmr_data')   
+# remeshing as used for ground contact
+remeshing_path = os.path.join(root_data_path, 'smal_data_remeshed', 'uniform_surface_sampling', 'my_smpl_39dogsnorm_Jr_4_dog_remesh4000_info.pkl')
+
+loss_weight_path = os.path.join(os.path.dirname(__file__), '../', 'src', 'configs', 'ttopt_loss_weights', args_loss_weight_ttopt_path)  
+print(loss_weight_path)
+
+
+# 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)
+
+# prepare complete model
+norm_dict = get_norm_dict(data_info=None, device=device)
+bite_model = BITEInferenceModel(cfg, path_model_file_complete, norm_dict)
+smal_model_type = bite_model.smal_model_type
+logscale_part_list = SMAL_MODEL_CONFIG[smal_model_type]['logscale_part_list']       # ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'] 
+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)    
+
+# load loss modules -> not necessary!
+# loss_module = Loss(smal_model_type=cfg.smal.SMAL_MODEL_TYPE, data_info=StanExt.DATA_INFO, nf_version=cfg.params.NF_VERSION).to(device)    
+# loss_module_ref = LossRef(smal_model_type=cfg.smal.SMAL_MODEL_TYPE, data_info=StanExt.DATA_INFO, nf_version=cfg.params.NF_VERSION).to(device)    
+
+# remeshing utils
+with open(remeshing_path, 'rb') as fp: 
+    remeshing_dict = pkl.load(fp)
+remeshing_relevant_faces = torch.tensor(remeshing_dict['smal_faces'][remeshing_dict['faceid_closest']], dtype=torch.long, device=device)
+remeshing_relevant_barys = torch.tensor(remeshing_dict['barys_closest'], dtype=torch.float32, device=device)
+
+
+
+
+# create path for output files
+save_imgs_path = os.path.join(cfg.paths.ROOT_OUT_PATH, 'gradio_examples')
+if not os.path.exists(save_imgs_path):
+    os.makedirs(save_imgs_path)
+
+
+
+
+
+def run_bite_inference(input_image, bbox=None):
+
+    with open(loss_weight_path, 'r') as j: 
+        losses = json.loads(j.read())
+    shutil.copyfile(loss_weight_path, root_out_path_details + os.path.basename(loss_weight_path))
+    print(losses)
+
+    # prepare dataset and dataset loader
+    val_dataset, val_loader, len_val_dataset, test_name_list, stanext_data_info, stanext_acc_joints = get_single_crop_dataset_from_image(input_image, bbox=bbox)
+
+    # summarize information for normalization 
+    norm_dict = get_norm_dict(stanext_data_info, device)
+    # get keypoint weights
+    keypoint_weights = torch.tensor(stanext_data_info.keypoint_weights, dtype=torch.float)[None, :].to(device) 
+
+
+    # prepare progress bar
+    iterable = enumerate(val_loader) # the length of this iterator should be 1
+    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
+
+    for i, (input, target_dict) in iterable:
+        batch_size = input.shape[0]
+        # prepare variables, put them on the right device
+        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 
+        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']
+        bs = res['pose_rotmat'].shape[0]
+        all_pose_6d = rotmat_to_rot6d(res['pose_rotmat'][:, None, 1:, :, :].clone().reshape((-1, 3, 3))).reshape((bs, -1, 6))       # [bs, 34, 6]
+        all_orient_6d = rotmat_to_rot6d(res['pose_rotmat'][:, None, :1, :, :].clone().reshape((-1, 3, 3))).reshape((bs, -1, 6))     # [bs, 1, 6]
+        
+
+    ind_img = 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
+    batch_size = 1
+
+    # save initial visualizations
+    # save the image with keypoints as predicted by the stacked hourglass
+    pred_unp_prep = torch.cat((res['hg_keyp_256'][ind_img, :, :].detach(), res['hg_keyp_scores'][ind_img, :, :]), 1)
+    inp_img = input[ind_img, :, :, :].detach().clone()
+    out_path = root_out_path +  name + '_hg_key.png'
+    save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_path, threshold=0.01, print_scores=True, ratio_in_out=1.0)    # threshold=0.3
+    # save the input image
+    img_inp = input[ind_img, :, :, :].clone()
+    for t, m, s in zip(img_inp, stanext_data_info.rgb_mean, stanext_data_info.rgb_stddev): t.add_(m)       # inverse to transforms.color_normalize()
+    img_inp = img_inp.detach().cpu().numpy().transpose(1, 2, 0)  
+    img_init = Image.fromarray(np.uint8(255*img_inp)).convert('RGB') 
+    img_init.save(root_out_path_details + name + '_img_ainit.png')
+    # save ground truth silhouette (for visualization only, it is not used during the optimization)
+    target_img_silh = Image.fromarray(np.uint8(255*target_dict['silh'][ind_img, :, :].detach().cpu().numpy())).convert('RGB')
+    target_img_silh.save(root_out_path_details +  name + '_target_silh.png')
+    # save the silhouette as predicted by the stacked hourglass
+    hg_img_silh = Image.fromarray(np.uint8(255*res['hg_silh_prep'][ind_img, :, :].detach().cpu().numpy())).convert('RGB')
+    hg_img_silh.save(root_out_path +  name + '_hg_silh.png')
+
+    # initialize the variables over which we want to optimize
+    optimed_pose_6d = all_pose_6d[ind_img, None, :, :].to(device).clone().detach().requires_grad_(True)
+    optimed_orient_6d = all_orient_6d[ind_img, None, :, :].to(device).clone().detach().requires_grad_(True)  # [1, 1, 6]
+    optimed_betas = res['betas'][ind_img, None, :].to(device).clone().detach().requires_grad_(True)   # [1,30]
+    optimed_trans_xy = res['trans'][ind_img, None, :2].to(device).clone().detach().requires_grad_(True)
+    optimed_trans_z =res['trans'][ind_img, None, 2:3].to(device).clone().detach().requires_grad_(True)
+    optimed_camera_flength = res['flength'][ind_img, None, :].to(device).clone().detach().requires_grad_(True)  # [1,1]
+    n_vert_comp = 2*smal.n_center + 3*smal.n_left
+    optimed_vert_off_compact = torch.tensor(np.zeros((batch_size, n_vert_comp)), dtype=torch.float,
+                                device=device,
+                                requires_grad=True) 
+    assert len(logscale_part_list) == 7
+    new_betas_limb_lengths = res['betas_limbs'][ind_img, None, :]
+    optimed_betas_limbs = new_betas_limb_lengths.to(device).clone().detach().requires_grad_(True)  # [1,7]
+
+    # define the optimizers 
+    optimizer = torch.optim.SGD(
+        # [optimed_pose, optimed_trans_xy, optimed_betas, optimed_betas_limbs, optimed_orient, optimed_vert_off_compact],
+        [optimed_camera_flength, optimed_trans_z, optimed_trans_xy, optimed_pose_6d, optimed_orient_6d, optimed_betas, optimed_betas_limbs],
+        lr=5*1e-4,        # 1e-3,
+        momentum=0.9)
+    optimizer_vshift = torch.optim.SGD(
+        [optimed_camera_flength, optimed_trans_z, optimed_trans_xy, optimed_pose_6d, optimed_orient_6d, optimed_betas, optimed_betas_limbs, optimed_vert_off_compact],
+        lr=1e-4,  # 1e-4,
+        momentum=0.9)
+    nopose_optimizer = torch.optim.SGD(
+        # [optimed_pose, optimed_trans_xy, optimed_betas, optimed_betas_limbs, optimed_orient, optimed_vert_off_compact],
+        [optimed_camera_flength, optimed_trans_z, optimed_trans_xy, optimed_orient_6d, optimed_betas, optimed_betas_limbs],
+        lr=5*1e-4,        # 1e-3,
+        momentum=0.9)
+    nopose_optimizer_vshift = torch.optim.SGD(
+        [optimed_camera_flength, optimed_trans_z, optimed_trans_xy, optimed_orient_6d, optimed_betas, optimed_betas_limbs, optimed_vert_off_compact],
+        lr=1e-4,  # 1e-4,
+        momentum=0.9)
+    # define schedulers
+    patience = 5
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer,
+        mode='min',
+        factor=0.5,
+        verbose=0,
+        min_lr=1e-5,
+        patience=patience)
+    scheduler_vshift = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer_vshift,
+        mode='min',
+        factor=0.5,
+        verbose=0,
+        min_lr=1e-5,
+        patience=patience)
+
+    # set all loss values to 0
+    losses = reset_loss_values(losses)
+
+    # prepare all the target labels: keypoints, silhouette, ground contact, ...
+    with torch.no_grad():
+        thr_kp = 0.2
+        kp_weights = res['hg_keyp_scores']
+        kp_weights[res['hg_keyp_scores']<thr_kp] = 0
+        weights_resh = kp_weights[ind_img, None, :, :].reshape((-1))  # target_dict['tpts'][:, :, 2].reshape((-1)) 
+        keyp_w_resh = keypoint_weights.repeat((batch_size, 1)).reshape((-1))
+        # prepare predicted ground contact labels
+        sm = nn.Softmax(dim=1)
+        target_gc_class = sm(res['vertexwise_ground_contact'][ind_img, :, :])[None, :, 1]       # values between 0 and 1
+        target_gc_class_remeshed = torch.einsum('ij,aij->ai', remeshing_relevant_barys, target_gc_class[:, remeshing_relevant_faces].to(device=device, dtype=torch.float32))
+        target_gc_class_remeshed_prep = torch.round(target_gc_class_remeshed).to(torch.long)
+        vert_colors = np.repeat(255*target_gc_class.detach().cpu().numpy()[0, :, None], 3, 1)
+        vert_colors[:, 2] = 255
+        faces_prep = smal.faces.unsqueeze(0).expand((batch_size, -1, -1))
+        # prepare target silhouette and keypoints, from stacked hourglass predictions
+        target_hg_silh = res['hg_silh_prep'][ind_img, :, :].detach()
+        target_kp_resh = res['hg_keyp_256'][ind_img, None, :, :].reshape((-1, 2)).detach()
+        # find out if ground contact constraints should be used for the image at hand
+        # print('is flat: ' + str(res['isflat_prep'][ind_img]))
+        if res['isflat_prep'][ind_img] >= 0.5: # threshold should probably be set higher
+            isflat = [True]
+        else:
+            isflat = [False] 
+        if target_gc_class_remeshed_prep.sum() > 3:
+            istouching = [True]
+        else:
+            istouching = [False]
+        ignore_pose_optimization = False
+
+
+    ##########################################################################################################
+    # start optimizing for this image
+    n_iter = 301    # how many iterations are desired? (+1)
+    loop = tqdm(range(n_iter))
+    per_loop_lst = []
+    list_error_procrustes = []
+    for i in loop:
+        # for the first 150 iterations steps we don't allow vertex shifts
+        if i == 0:          
+            current_i = 0
+            if ignore_pose_optimization:
+                current_optimizer = nopose_optimizer
+            else:
+                current_optimizer = optimizer
+            current_scheduler = scheduler
+            current_weight_name = 'weight' 
+        # after 150 iteration steps we start with vertex shifts    
+        elif i == 150:      
+            current_i = 0
+            if ignore_pose_optimization:
+                current_optimizer = nopose_optimizer_vshift
+            else:
+                current_optimizer = optimizer_vshift
+            current_scheduler = scheduler_vshift
+            current_weight_name = 'weight_vshift'    
+            # set up arap loss
+            if losses["arap"]['weight_vshift'] > 0.0:
+                with torch.no_grad():
+                    torch_mesh_comparison = Meshes(smal_verts.detach(), faces_prep.detach())
+                arap_loss = Arap_Loss(meshes=torch_mesh_comparison, device=device)  
+            #  is there a laplacian loss similar as in coarse-to-fine?
+            if losses["lapctf"]['weight_vshift'] > 0.0:
+                torch_verts_comparison = smal_verts.detach().clone()
+                smal_model_type_downsampling = '39dogs_norm'
+                smal_downsampling_npz_name = 'mesh_downsampling_' + os.path.basename(SMAL_MODEL_CONFIG[smal_model_type_downsampling]['smal_model_path']).replace('.pkl', '_template.npz')
+                smal_downsampling_npz_path = os.path.join(root_smal_downsampling, smal_downsampling_npz_name)  
+                data = np.load(smal_downsampling_npz_path, encoding='latin1', allow_pickle=True)
+                adjmat = data['A'][0]
+                laplacian_ctf = LaplacianCTF(adjmat, device=device)
+        else:
+            pass
+
+
+        current_optimizer.zero_grad()
+
+        # get 3d smal model
+        optimed_pose_with_glob = get_optimed_pose_with_glob(optimed_orient_6d, optimed_pose_6d)
+        optimed_trans = torch.cat((optimed_trans_xy, optimed_trans_z), dim=1)
+        smal_verts, keyp_3d, _ = smal(beta=optimed_betas, betas_limbs=optimed_betas_limbs, pose=optimed_pose_with_glob, vert_off_compact=optimed_vert_off_compact, trans=optimed_trans, keyp_conf='olive', get_skin=True)
+
+        # render silhouette and keypoints
+        pred_silh_images, pred_keyp_raw = silh_renderer(vertices=smal_verts, points=keyp_3d, faces=faces_prep, focal_lengths=optimed_camera_flength)
+        pred_keyp = pred_keyp_raw[:, :24, :]
+
+        # save silhouette reprojection visualization
+        if i==0:
+            img_silh = Image.fromarray(np.uint8(255*pred_silh_images[0, 0, :, :].detach().cpu().numpy())).convert('RGB')
+            img_silh.save(root_out_path_details +  name + '_silh_ainit.png')
+            my_mesh_tri = trimesh.Trimesh(vertices=smal_verts[0, ...].detach().cpu().numpy(), faces=faces_prep[0, ...].detach().cpu().numpy(), process=False,  maintain_order=True)
+            my_mesh_tri.export(root_out_path_details +  name + '_res_ainit.obj')
+
+        # silhouette loss
+        diff_silh = torch.abs(pred_silh_images[0, 0, :, :] - target_hg_silh)
+        losses['silhouette']['value'] = diff_silh.mean()
+
+        # keypoint_loss
+        output_kp_resh = (pred_keyp[0, :, :]).reshape((-1, 2))    
+        losses['keyp']['value'] = ((((output_kp_resh - target_kp_resh)[weights_resh>0]**2).sum(axis=1).sqrt() * \
+            weights_resh[weights_resh>0])*keyp_w_resh[weights_resh>0]).sum() / \
+            max((weights_resh[weights_resh>0]*keyp_w_resh[weights_resh>0]).sum(), 1e-5)
+        # losses['keyp']['value'] = ((((output_kp_resh - target_kp_resh)[weights_resh>0]**2).sum(axis=1).sqrt()*weights_resh[weights_resh>0])*keyp_w_resh[weights_resh>0]).sum() / max((weights_resh[weights_resh>0]*keyp_w_resh[weights_resh>0]).sum(), 1e-5)
+
+        # pose priors on refined pose
+        losses['pose_legs_side']['value'] = leg_sideway_error(optimed_pose_with_glob)
+        losses['pose_legs_tors']['value'] = leg_torsion_error(optimed_pose_with_glob)
+        losses['pose_tail_side']['value'] = tail_sideway_error(optimed_pose_with_glob)
+        losses['pose_tail_tors']['value'] = tail_torsion_error(optimed_pose_with_glob)
+        losses['pose_spine_side']['value'] = spine_sideway_error(optimed_pose_with_glob)
+        losses['pose_spine_tors']['value'] = spine_torsion_error(optimed_pose_with_glob)
+
+        # ground contact loss
+        sel_verts = torch.index_select(smal_verts, dim=1, index=remeshing_relevant_faces.reshape((-1))).reshape((batch_size, remeshing_relevant_faces.shape[0], 3, 3))
+        verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, sel_verts)
+
+        # gc_errors_plane, gc_errors_under_plane = calculate_plane_errors_batch(verts_remeshed, target_gc_class_remeshed_prep, target_dict['has_gc'], target_dict['has_gc_is_touching'])
+        gc_errors_plane, gc_errors_under_plane = calculate_plane_errors_batch(verts_remeshed, target_gc_class_remeshed_prep, isflat, istouching)
+
+        losses['gc_plane']['value'] = torch.mean(gc_errors_plane) 
+        losses['gc_belowplane']['value'] = torch.mean(gc_errors_under_plane)
+
+        # edge length of the predicted mesh
+        if (losses["edge"][current_weight_name] + losses["normal"][ current_weight_name] + losses["laplacian"][ current_weight_name]) > 0:
+            torch_mesh = Meshes(smal_verts, faces_prep.detach())
+            losses["edge"]['value'] = mesh_edge_loss(torch_mesh)
+            # mesh normal consistency
+            losses["normal"]['value'] = mesh_normal_consistency(torch_mesh)
+            # mesh laplacian smoothing
+            losses["laplacian"]['value'] = mesh_laplacian_smoothing(torch_mesh, method="uniform")
+
+        # arap loss
+        if losses["arap"][current_weight_name] > 0.0:
+            torch_mesh = Meshes(smal_verts, faces_prep.detach())
+            losses["arap"]['value'] =  arap_loss(torch_mesh)
+
+        # laplacian loss for comparison (from coarse-to-fine paper)
+        if losses["lapctf"][current_weight_name] > 0.0:
+            verts_refine = smal_verts
+            loss_almost_arap, loss_smooth = laplacian_ctf(verts_refine, torch_verts_comparison)
+            losses["lapctf"]['value'] =  loss_almost_arap
+
+        # Weighted sum of the losses
+        total_loss = 0.0 
+        for k in ['keyp', 'silhouette', 'pose_legs_side', 'pose_legs_tors', 'pose_tail_side', 'pose_tail_tors', 'pose_spine_tors', 'pose_spine_side', 'gc_plane', 'gc_belowplane', 'edge', 'normal', 'laplacian', 'arap', 'lapctf']:
+            if losses[k][current_weight_name] > 0.0:
+                total_loss += losses[k]['value'] * losses[k][current_weight_name]
+
+        # calculate gradient and make optimization step
+        total_loss.backward(retain_graph=True)  #  
+        current_optimizer.step()
+        current_scheduler.step(total_loss)
+        loop.set_description(f"Body Fitting = {total_loss.item():.3f}")
+
+        # save the result three times (0, 150, 300)
+        if i % 150 == 0:    
+            # save silhouette image
+            img_silh = Image.fromarray(np.uint8(255*pred_silh_images[0, 0, :, :].detach().cpu().numpy())).convert('RGB')
+            img_silh.save(root_out_path_details +  name + '_silh_e' + format(i, '03d') + '.png')
+            # save image overlay
+            visualizations = silh_renderer.get_visualization_nograd(smal_verts, faces_prep, optimed_camera_flength, color=0)
+            pred_tex = visualizations[0, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+            # out_path = root_out_path_details +  name + '_tex_pred_e' + format(i, '03d') + '.png'
+            # plt.imsave(out_path, pred_tex)
+            input_image_np = img_inp.copy()
+            im_masked = cv2.addWeighted(input_image_np,0.2,pred_tex,0.8,0)
+            pred_tex_max = np.max(pred_tex, axis=2)
+            im_masked[pred_tex_max<0.01, :] = input_image_np[pred_tex_max<0.01, :]
+            out_path = root_out_path +  name + '_comp_pred_e' + format(i, '03d') + '.png'
+            plt.imsave(out_path, im_masked)
+            # save mesh
+            my_mesh_tri = trimesh.Trimesh(vertices=smal_verts[0, ...].detach().cpu().numpy(), faces=faces_prep[0, ...].detach().cpu().numpy(), process=False,  maintain_order=True)
+            my_mesh_tri.visual.vertex_colors = vert_colors
+            my_mesh_tri.export(root_out_path +  name + '_res_e' + format(i, '03d') + '.obj')
+            # save focal length (together with the mesh this is enough to create an overlay in blender)
+            out_file_flength = root_out_path_details +  name + '_flength_e' + format(i, '03d') # + '.npz'
+            np.save(out_file_flength, optimed_camera_flength.detach().cpu().numpy())
+        current_i += 1
+
+    # prepare output mesh
+    mesh = my_mesh_tri  # all_results[0]['mesh_posed']
+    mesh.apply_transform([[-1, 0, 0, 0],
+                            [0, -1, 0, 0],
+                            [0, 0, 1, 1],
+                            [0, 0, 0, 1]])
+    result_path = os.path.join(save_imgs_path, test_name_list[0] + '_z')
+    mesh.export(file_obj=result_path + '.glb')
+    result_gltf = result_path + '.glb'
+    return result_gltf
+
+
+
+
+
+# -------------------------------------------------------------------------------------------------------------------- #
+
+
+def run_complete_inference(img_path_or_img, crop_choice):
+    # depending on crop_choice: run faster r-cnn or take the input image directly
+    if crop_choice == "input image is cropped":
+        if isinstance(img_path_or_img, str):
+            img = cv2.imread(img_path_or_img)
+            output_interm_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        else:
+            output_interm_image = img_path_or_img
+        output_interm_bbox = None
+    else:
+        output_interm_image, output_interm_bbox = run_bbox_inference(img_path_or_img.copy())
+    # run barc inference
+    result_gltf = run_bite_inference(img_path_or_img, output_interm_bbox)
+    # add white border to image for nicer alignment
+    output_interm_image_vis = np.concatenate((255*np.ones_like(output_interm_image), output_interm_image, 255*np.ones_like(output_interm_image)), axis=1)
+    return [result_gltf, result_gltf, output_interm_image_vis]
+
+
+
+
+########################################################################################################################
+
+# see: https://huggingface.co/spaces/radames/PIFu-Clothed-Human-Digitization/blob/main/PIFu/spaces.py
+
+description = '''
+# BITE
+
+#### Project Page
+* https://bite.is.tue.mpg.de/
+
+#### Description
+This is a demo for BITE (*B*eyond Priors for *I*mproved *T*hree-{D} Dog Pose *E*stimation). 
+You can either submit a cropped image or choose the option to run a pretrained Faster R-CNN in order to obtain a bounding box.  
+Please have a look at the examples below.
+<details>
+
+<summary>More</summary>
+
+#### Citation
+
+```
+@inproceedings{bite2023rueegg,
+        title = {{BITE}: Beyond Priors for Improved Three-{D} Dog Pose Estimation},
+        author = {R\"uegg, Nadine and Tripathi, Shashank and Schindler, Konrad and Black, Michael J. and Zuffi, Silvia},
+        booktitle = {IEEE/CVF Conf.~on Computer Vision and Pattern Recognition (CVPR)},
+        pages = {8867-8876},
+        year = {2023},
+}
+```
+
+#### Image Sources
+* Stanford extra image dataset
+* Images from google search engine
+    * https://www.dogtrainingnation.com/wp-content/uploads/2015/02/keep-dog-training-sessions-short.jpg
+    * https://thumbs.dreamstime.com/b/hund-und-seine-neue-hundeh%C3%BCtte-36757551.jpg
+    * https://www.mydearwhippet.com/wp-content/uploads/2021/04/whippet-temperament-2.jpg
+    * https://media.istockphoto.com/photos/ibizan-hound-at-the-shore-in-winter-picture-id1092705644?k=20&m=1092705644&s=612x612&w=0&h=ppwg92s9jI8GWnk22SOR_DWWNP8b2IUmLXSQmVey5Ss=
+
+
+</details>
+'''
+
+
+
+
+
+
+example_images = sorted(glob.glob(os.path.join(os.path.dirname(__file__), '../', 'datasets', 'test_image_crops', '*.jpg')) + glob.glob(os.path.join(os.path.dirname(__file__), '../', 'datasets', 'test_image_crops', '*.png')))  
+random.shuffle(example_images)
+# example_images.reverse()
+# examples = [[img, "input image is cropped"] for img in example_images]
+examples = []
+for img in example_images:
+    if os.path.basename(img)[:2] == 'z_':
+        examples.append([img, "use Faster R-CNN to get a bounding box"])
+    else:
+        examples.append([img, "input image is cropped"])
+
+demo = gr.Interface(
+    fn=run_complete_inference,
+    description=description,
+    # inputs=gr.Image(type="filepath", label="Input Image"),
+    inputs=[gr.Image(label="Input Image"),
+        gr.Radio(["input image is cropped", "use Faster R-CNN to get a bounding box"], value="use Faster R-CNN to get a bounding box", label="Crop Choice"),
+    ],
+    outputs=[
+        gr.Model3D(
+            clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model"),
+        gr.File(label="Download 3D Model"),
+        gr.Image(label="Bounding Box (Faster R-CNN prediction)"),
+
+    ],
+    examples=examples,
+    thumbnail="bite_thumbnail.png",
+    allow_flagging="never",
+    cache_examples=True,
+    examples_per_page=14,
+)
+
+demo.launch(share=True)

src/bps_2d/bps_for_segmentation.py

@@ -0,0 +1,114 @@
+
+# code idea from https://github.com/sergeyprokudin/bps
+
+import os
+import numpy as np
+from PIL import Image
+import time
+import scipy
+import scipy.spatial
+import pymp
+
+
+#####################
+QUERY_POINTS = np.asarray([30, 34, 31, 55, 29, 84, 35, 108, 34, 145, 29, 171, 27, 
+    196, 29, 228, 58, 35, 61, 55, 57, 83, 56, 109, 63, 148, 58, 164, 57, 197, 60, 
+    227, 81, 26, 87, 58, 85, 87, 89, 117, 86, 142, 89, 172, 84, 197, 88, 227, 113, 
+    32, 116, 58, 112, 88, 118, 113, 109, 147, 114, 173, 119, 201, 113, 229, 139, 
+    29, 141, 59, 142, 93, 139, 117, 146, 147, 141, 173, 142, 201, 143, 227, 170, 
+    26, 173, 59, 166, 90, 174, 117, 176, 141, 169, 175, 167, 198, 172, 227, 198, 
+    30, 195, 59, 204, 85, 198, 116, 195, 140, 198, 175, 194, 193, 199, 227, 221, 
+    26, 223, 57, 227, 83, 227, 113, 227, 140, 226, 173, 230, 196, 228, 229]).reshape((64, 2))
+#####################
+
+class SegBPS():
+
+    def __init__(self, query_points=QUERY_POINTS, size=256):
+        self.size = size
+        self.query_points = query_points
+        row, col = np.indices((self.size, self.size))
+        self.indices_rc = np.stack((row, col), axis=2)   # (256, 256, 2)
+        self.pts_aranged = np.arange(64)
+        return
+
+    def _do_kdtree(self, combined_x_y_arrays, points):
+        # see https://stackoverflow.com/questions/10818546/finding-index-of-nearest-
+        #   point-in-numpy-arrays-of-x-and-y-coordinates
+        mytree = scipy.spatial.cKDTree(combined_x_y_arrays)
+        dist, indexes = mytree.query(points)
+        return indexes
+
+    def calculate_bps_points(self, seg, thr=0.5, vis=False, out_path=None):
+        # seg: input segmentation image of shape (256, 256) with values between 0 and 1
+        query_val = seg[self.query_points[:, 0], self.query_points[:, 1]]
+        pts_fg = self.pts_aranged[query_val>=thr]
+        pts_bg = self.pts_aranged[query_val<thr]
+        candidate_inds_bg = self.indices_rc[seg<thr]
+        candidate_inds_fg = self.indices_rc[seg>=thr]
+        if candidate_inds_bg.shape[0] == 0:
+            candidate_inds_bg = np.ones((1, 2)) * 128        #  np.zeros((1, 2))
+        if candidate_inds_fg.shape[0] == 0:
+            candidate_inds_fg = np.ones((1, 2)) * 128        #  np.zeros((1, 2))
+        # calculate nearest points
+        all_nearest_points = np.zeros((64, 2))
+        all_nearest_points[pts_fg, :] = candidate_inds_bg[self._do_kdtree(candidate_inds_bg, self.query_points[pts_fg, :]), :]
+        all_nearest_points[pts_bg, :] = candidate_inds_fg[self._do_kdtree(candidate_inds_fg, self.query_points[pts_bg, :]), :]
+        all_nearest_points_01 = all_nearest_points / 255.
+        if vis:
+            self.visualize_result(seg, all_nearest_points, out_path=out_path)
+        return all_nearest_points_01
+
+    def calculate_bps_points_batch(self, seg_batch, thr=0.5, vis=False, out_path=None):
+        # seg_batch: input segmentation image of shape (bs, 256, 256) with values between 0 and 1
+        bs = seg_batch.shape[0]
+        all_nearest_points_01_batch = np.zeros((bs, self.query_points.shape[0], 2))
+        for ind in range(0, bs):         # 0.25
+            seg = seg_batch[ind, :, :]
+            all_nearest_points_01 = self.calculate_bps_points(seg, thr=thr, vis=vis, out_path=out_path)
+            all_nearest_points_01_batch[ind, :, :] = all_nearest_points_01
+        return all_nearest_points_01_batch
+
+    def visualize_result(self, seg, all_nearest_points, out_path=None):
+        import matplotlib as mpl
+        mpl.use('Agg')
+        import matplotlib.pyplot as plt
+        # img: (256, 256, 3) 
+        img = (np.stack((seg, seg, seg), axis=2) * 155).astype(np.int)
+        if out_path is None:
+            ind_img = 0
+            out_path = '../test_img' + str(ind_img) + '.png'
+        fig, ax = plt.subplots()
+        plt.imshow(img)     
+        plt.gca().set_axis_off()
+        plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
+        plt.margins(0,0)
+        ratio_in_out = 1    # 255
+        for idx, (y, x) in enumerate(self.query_points):
+            x = int(x*ratio_in_out)     
+            y = int(y*ratio_in_out)     
+            plt.scatter([x], [y], marker="x", s=50)
+            x2 = int(all_nearest_points[idx, 1])
+            y2 = int(all_nearest_points[idx, 0])
+            plt.scatter([x2], [y2], marker="o", s=50)
+            plt.plot([x, x2], [y, y2])
+        plt.savefig(out_path, bbox_inches='tight', pad_inches=0)
+        plt.close()
+        return
+
+
+
+
+
+if __name__ == "__main__":
+    ind_img = 2 # 4
+    path_seg_top = '...../pytorch-stacked-hourglass/results/dogs_hg8_ks_24_v1/test/'
+    path_seg = os.path.join(path_seg_top, 'seg_big_' + str(ind_img) + '.png')
+    img = np.asarray(Image.open(path_seg))
+    # min is 0.004, max is 0.9
+    # low values are background, high values are foreground
+    seg = img[:, :, 1] / 255.
+    # calculate points
+    bps = SegBPS()
+    bps.calculate_bps_points(seg, thr=0.5, vis=False, out_path=None)
+
+

src/combined_model/helper.py

@@ -0,0 +1,207 @@
+
+import torch
+import torch.nn as nn
+import torch.backends.cudnn
+import torch.nn.parallel
+from tqdm import tqdm
+import os
+import pathlib
+from matplotlib import pyplot as plt
+import cv2
+import numpy as np
+import torch
+import trimesh
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
+from stacked_hourglass.utils.evaluation import accuracy, AverageMeter, final_preds, get_preds, get_preds_soft
+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
+
+
+# GOAL: have all the functions from the validation and visual epoch together
+
+
+'''
+save_imgs_path = ...
+prefix = ''
+input   # this is the image
+data_info
+target_dict
+render_all
+model
+
+
+vertices_smal = output_reproj['vertices_smal']
+flength = output_unnorm['flength']
+hg_keyp_norm = output['keypoints_norm']
+hg_keyp_scores = output['keypoints_scores']
+betas = output_reproj['betas']
+betas_limbs = output_reproj['betas_limbs']
+zz = output_reproj['z']
+pose_rotmat = output_unnorm['pose_rotmat']
+trans = output_unnorm['trans']
+pred_keyp = output_reproj['keyp_2d']
+pred_silh = output_reproj['silh']
+'''
+
+#################################################
+
+def eval_save_visualizations_and_meshes(model, input, data_info, target_dict, test_name_list, vertices_smal, hg_keyp_norm, hg_keyp_scores, zz, betas, betas_limbs, pose_rotmat, trans, flength, pred_keyp, pred_silh, save_imgs_path, prefix, index, render_all=False):
+    device = input.device
+    curr_batch_size = input.shape[0]
+    # render predicted 3d models
+    visualizations = model.render_vis_nograd(vertices=vertices_smal,
+                                            focal_lengths=flength,
+                                            color=0)        # color=2)
+    for ind_img in range(len(target_dict['index'])):    
+        try: 
+            # import pdb; pdb.set_trace()
+            if test_name_list is not None:
+                img_name = test_name_list[int(target_dict['index'][ind_img].cpu().detach().numpy())].replace('/', '_')
+                img_name = img_name.split('.')[0]
+            else:
+                img_name = str(index) + '_' + str(ind_img)
+            # save image with predicted keypoints
+            out_path = save_imgs_path + '/keypoints_pred_' + img_name + '.png'
+            pred_unp = (hg_keyp_norm[ind_img, :, :] + 1.) / 2 * (data_info.image_size - 1)
+            pred_unp_maxval = hg_keyp_scores[ind_img, :, :]
+            pred_unp_prep = torch.cat((pred_unp, pred_unp_maxval), 1)
+            inp_img = input[ind_img, :, :, :].detach().clone()
+            save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_path, threshold=0.1, print_scores=True, ratio_in_out=1.0)    # threshold=0.3
+            # save predicted 3d model (front view)
+            pred_tex = visualizations[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+            pred_tex_max = np.max(pred_tex, axis=2)
+            out_path = save_imgs_path + '/' + prefix + 'tex_pred_' + img_name + '.png'
+            plt.imsave(out_path, pred_tex)
+            input_image = input[ind_img, :, :, :].detach().clone()
+            for t, m, s in zip(input_image, data_info.rgb_mean, data_info.rgb_stddev): t.add_(m)
+            input_image_np = input_image.detach().cpu().numpy().transpose(1, 2, 0) 
+            im_masked = cv2.addWeighted(input_image_np,0.2,pred_tex,0.8,0)
+            im_masked[pred_tex_max<0.01, :] = input_image_np[pred_tex_max<0.01, :]
+            out_path = save_imgs_path + '/' + prefix + 'comp_pred_' + img_name + '.png'
+            plt.imsave(out_path, im_masked)
+            # save predicted 3d model (side view)
+            vertices_cent = vertices_smal - vertices_smal.mean(dim=1)[:, None, :]
+            roll = np.pi / 2 * torch.ones(1).float().to(device)
+            pitch = np.pi / 2 * torch.ones(1).float().to(device)
+            tensor_0 = torch.zeros(1).float().to(device)
+            tensor_1 = torch.ones(1).float().to(device)
+            RX = torch.stack([torch.stack([tensor_1, tensor_0, tensor_0]), torch.stack([tensor_0, torch.cos(roll), -torch.sin(roll)]),torch.stack([tensor_0, torch.sin(roll), torch.cos(roll)])]).reshape(3,3)
+            RY = torch.stack([
+                torch.stack([torch.cos(pitch), tensor_0, torch.sin(pitch)]),
+                torch.stack([tensor_0, tensor_1, tensor_0]),
+                torch.stack([-torch.sin(pitch), tensor_0, torch.cos(pitch)])]).reshape(3,3)
+            vertices_rot = (torch.matmul(RY, vertices_cent.reshape((-1, 3))[:, :, None])).reshape((curr_batch_size, -1, 3))
+            vertices_rot[:, :, 2] = vertices_rot[:, :, 2] + torch.ones_like(vertices_rot[:, :, 2]) * 20     # 18     # *16
+
+            visualizations_rot = model.render_vis_nograd(vertices=vertices_rot,
+                                                    focal_lengths=flength,
+                                                    color=0)        # 2)
+            pred_tex = visualizations_rot[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+            pred_tex_max = np.max(pred_tex, axis=2)
+            out_path = save_imgs_path + '/' + prefix + 'rot_tex_pred_' + img_name + '.png'
+            plt.imsave(out_path, pred_tex)
+            if render_all:
+                # save input image
+                inp_img = input[ind_img, :, :, :].detach().clone()
+                out_path = save_imgs_path + '/image_' + img_name + '.png'
+                save_input_image(inp_img, out_path)
+                # save mesh
+                V_posed = vertices_smal[ind_img, :, :].detach().cpu().numpy()
+                Faces = model.smal.f
+                mesh_posed = trimesh.Trimesh(vertices=V_posed, faces=Faces, process=False,  maintain_order=True)
+                mesh_posed.export(save_imgs_path + '/' + prefix + 'mesh_posed_' + img_name + '.obj')
+        except: 
+            print('dont save an image')
+
+############
+
+def eval_prepare_pck_and_iou(model, input, data_info, target_dict, test_name_list, vertices_smal, hg_keyp_norm, hg_keyp_scores, zz, betas, betas_limbs, pose_rotmat, trans, flength, pred_keyp, pred_silh, save_imgs_path, prefix, index, pck_thresh, progress=None, skip_pck_and_iou=False):
+    preds = {}
+    preds['betas'] = betas.cpu().detach().numpy()
+    preds['betas_limbs'] = betas_limbs.cpu().detach().numpy()
+    preds['z'] = zz.cpu().detach().numpy()
+    preds['pose_rotmat'] = pose_rotmat.cpu().detach().numpy()
+    preds['flength'] = flength.cpu().detach().numpy()
+    preds['trans'] = trans.cpu().detach().numpy()
+    preds['breed_index'] = target_dict['breed_index'].cpu().detach().numpy().reshape((-1))
+    img_names = []
+    for ind_img2 in range(0, betas.shape[0]):
+        if test_name_list is not None:
+            img_name2 = test_name_list[int(target_dict['index'][ind_img2].cpu().detach().numpy())].replace('/', '_')
+            img_name2 = img_name2.split('.')[0]
+        else:
+            img_name2 = str(index) + '_' + str(ind_img2)
+        img_names.append(img_name2)
+    preds['image_names'] = img_names
+    if not skip_pck_and_iou:
+        # prepare keypoints for PCK calculation - predicted as well as ground truth
+        # pred_keyp = output_reproj['keyp_2d']   # 256
+        gt_keypoints_256 = target_dict['tpts'][:, :, :2] / 64. * (256. - 1)
+        # gt_keypoints_norm = gt_keypoints_256 / 256 / 0.5 - 1
+        gt_keypoints = torch.cat((gt_keypoints_256, target_dict['tpts'][:, :, 2:3]), dim=2)     # gt_keypoints_norm
+        # 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[:, 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['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
+            )
+    return preds
+
+
+# preds['acc_PCK'] = Metrics.PCK(pred_keyp, gt_keypoints, gtseg, has_seg, idxs=EVAL_KEYPOINTS, thresh_range=[pck_thresh])
+# preds['acc_IOU'] = Metrics.IOU(synth_silhouettes, gtseg, img_border_mask, mask=has_seg)
+#############################
+
+def eval_add_preds_to_summary(summary, preds, my_step, batch_size, curr_batch_size, skip_pck_and_iou=False):
+    if not skip_pck_and_iou:
+        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()
+    summary['betas'][my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['betas'] 
+    summary['betas_limbs'][my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['betas_limbs'] 
+    summary['z'][my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['z'] 
+    summary['pose_rotmat'][my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['pose_rotmat'] 
+    summary['flength'][my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['flength'] 
+    summary['trans'][my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['trans'] 
+    summary['breed_indices'][my_step * batch_size:my_step * batch_size + curr_batch_size] = preds['breed_index']
+    summary['image_names'].extend(preds['image_names'])
+    return
+
+
+def get_triangle_faces_from_pyvista_poly(poly):
+    """Fetch all triangle faces."""
+    stream = poly.faces
+    tris = []
+    i = 0
+    while i < len(stream):
+        n = stream[i]
+        if n != 3:
+            i += n + 1
+            continue
+        stop = i + n + 1
+        tris.append(stream[i+1:stop])
+        i = stop
+    return np.array(tris)

src/combined_model/helper3.py

@@ -0,0 +1,17 @@
+
+import numpy as np 
+
+def get_triangle_faces_from_pyvista_poly(poly):
+    """Fetch all triangle faces."""
+    stream = poly.faces
+    tris = []
+    i = 0
+    while i < len(stream):
+        n = stream[i]
+        if n != 3:
+            i += n + 1
+            continue
+        stop = i + n + 1
+        tris.append(stream[i+1:stop])
+        i = stop
+    return np.array(tris)

src/combined_model/loss_image_to_3d_refinement.py

@@ -0,0 +1,216 @@
+
+
+import torch
+import numpy as np
+import pickle as pkl
+
+import os
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
+# from priors.pose_prior_35 import Prior
+# from priors.tiger_pose_prior.tiger_pose_prior import GaussianMixturePrior
+from priors.normalizing_flow_prior.normalizing_flow_prior import NormalizingFlowPrior
+from priors.shape_prior import ShapePrior
+from lifting_to_3d.utils.geometry_utils import rot6d_to_rotmat, batch_rot2aa, geodesic_loss_R
+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 priors.shape_prior import ShapePrior
+from configs.SMAL_configs import SMAL_MODEL_CONFIG
+
+from priors.helper_3dcgmodel_loss import load_dog_betas_for_3dcgmodel_loss
+
+
+class LossRef(torch.nn.Module):
+    def __init__(self, smal_model_type, data_info, nf_version=None):
+        super(LossRef, self).__init__()
+        self.criterion_regr = torch.nn.MSELoss()        # takes the mean   
+        self.criterion_class = torch.nn.CrossEntropyLoss()
+
+        class_weights_isflat = torch.tensor([12, 2])
+        self.criterion_class_isflat = torch.nn.CrossEntropyLoss(weight=class_weights_isflat)
+        self.criterion_l1 = torch.nn.L1Loss()
+        self.geodesic_loss = geodesic_loss_R(reduction='mean')
+        self.gc_loss_on_mesh = LossGConMesh()
+        self.data_info = data_info   
+        self.smal_model_type = smal_model_type
+        self.register_buffer('keypoint_weights', torch.tensor(data_info.keypoint_weights)[None, :])
+        # if nf_version is not None:
+        #     self.normalizing_flow_pose_prior = NormalizingFlowPrior(nf_version=nf_version)
+
+        self.smal_model_data_path = SMAL_MODEL_CONFIG[self.smal_model_type]['smal_model_data_path']
+        self.shape_prior = ShapePrior(self.smal_model_data_path) # here we just need mean and cov        
+
+        remeshing_path = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/smal_data_remeshed/uniform_surface_sampling/my_smpl_39dogsnorm_Jr_4_dog_remesh4000_info.pkl'
+        with open(remeshing_path, 'rb') as fp: 
+            self.remeshing_dict = pkl.load(fp)
+        self.remeshing_relevant_faces = torch.tensor(self.remeshing_dict['smal_faces'][self.remeshing_dict['faceid_closest']], dtype=torch.long)
+        self.remeshing_relevant_barys = torch.tensor(self.remeshing_dict['barys_closest'], dtype=torch.float32)
+
+
+
+        # load 3d data for the unity dogs (an optional shape prior for 11 breeds)
+        self.unity_smal_shape_prior_dogs = SMAL_MODEL_CONFIG[self.smal_model_type]['unity_smal_shape_prior_dogs']
+        if self.unity_smal_shape_prior_dogs is not None:
+            self.dog_betas_unity = load_dog_betas_for_3dcgmodel_loss(self.unity_smal_shape_prior_dogs, self.smal_model_type)
+        else:
+            self.dog_betas_unity = None
+
+
+
+
+
+
+
+    def forward(self, output_ref, output_ref_comp, target_dict, weight_dict_ref):
+        # output_reproj: ['vertices_smal', 'keyp_3d', 'keyp_2d', 'silh_image']
+        # target_dict: ['index', 'center', 'scale', 'pts', 'tpts', 'target_weight']
+        batch_size = output_ref['keyp_2d'].shape[0]
+        loss_dict_temp = {}
+
+        # loss on reprojected keypoints 
+        output_kp_resh = (output_ref['keyp_2d']).reshape((-1, 2))    
+        target_kp_resh = (target_dict['tpts'][:, :, :2] / 64. * (256. - 1)).reshape((-1, 2))
+        weights_resh = target_dict['tpts'][:, :, 2].reshape((-1)) 
+        keyp_w_resh = self.keypoint_weights.repeat((batch_size, 1)).reshape((-1))
+        loss_dict_temp['keyp_ref'] = ((((output_kp_resh - target_kp_resh)[weights_resh>0]**2).sum(axis=1).sqrt()*weights_resh[weights_resh>0])*keyp_w_resh[weights_resh>0]).sum() / \
+            max((weights_resh[weights_resh>0]*keyp_w_resh[weights_resh>0]).sum(), 1e-5)
+
+        # loss on reprojected silhouette
+        assert output_ref['silh'].shape == (target_dict['silh'][:, None, :, :]).shape
+        silh_loss_type = 'default'
+        if silh_loss_type == 'default':
+            with torch.no_grad():
+                thr_silh = 20
+                diff = torch.norm(output_kp_resh - target_kp_resh, dim=1)
+                diff_x = diff.reshape((batch_size, -1))
+                weights_resh_x = weights_resh.reshape((batch_size, -1))
+                unweighted_kp_mean_dist = (diff_x * weights_resh_x).sum(dim=1) / ((weights_resh_x).sum(dim=1)+1e-6)
+            loss_silh_bs = ((output_ref['silh'] - target_dict['silh'][:, None, :, :]) ** 2).sum(axis=3).sum(axis=2).sum(axis=1) / (output_ref['silh'].shape[2]*output_ref['silh'].shape[3])
+            loss_dict_temp['silh_ref'] = loss_silh_bs[unweighted_kp_mean_dist<thr_silh].sum() / batch_size
+        else:
+            print('silh_loss_type: ' + silh_loss_type)
+            raise ValueError
+
+        # regularization: losses on difference between previous prediction and refinement
+        loss_dict_temp['reg_trans'] = self.criterion_l1(output_ref_comp['ref_trans_notnorm'], output_ref_comp['old_trans_notnorm'].detach()) * 3
+        loss_dict_temp['reg_flength'] = self.criterion_l1(output_ref_comp['ref_flength_notnorm'], output_ref_comp['old_flength_notnorm'].detach()) * 1
+        loss_dict_temp['reg_pose'] = self.geodesic_loss(output_ref_comp['ref_pose_rotmat'], output_ref_comp['old_pose_rotmat'].detach()) * 35 * 6
+
+        # pose priors on refined pose
+        loss_dict_temp['pose_legs_side'] = leg_sideway_error(output_ref['pose_rotmat'])
+        loss_dict_temp['pose_legs_tors'] = leg_torsion_error(output_ref['pose_rotmat'])
+        loss_dict_temp['pose_tail_side'] = tail_sideway_error(output_ref['pose_rotmat'])
+        loss_dict_temp['pose_tail_tors'] = tail_torsion_error(output_ref['pose_rotmat'])
+        loss_dict_temp['pose_spine_side'] = spine_sideway_error(output_ref['pose_rotmat'])
+        loss_dict_temp['pose_spine_tors'] = spine_torsion_error(output_ref['pose_rotmat'])
+
+        # loss to predict ground contact per vertex
+        # import pdb; pdb.set_trace()
+        if 'gc_vertexwise' in weight_dict_ref.keys():
+            # import pdb; pdb.set_trace()
+            device = output_ref['vertexwise_ground_contact'].device
+            pred_gc = output_ref['vertexwise_ground_contact']
+            loss_dict_temp['gc_vertexwise'] = self.gc_loss_on_mesh(pred_gc, target_dict['gc'].to(device=device, dtype=torch.long), target_dict['has_gc'], loss_type_gcmesh='ce')
+
+        keep_smal_mesh = False 
+        if 'gc_plane' in weight_dict_ref.keys():
+            if weight_dict_ref['gc_plane'] > 0:
+                if keep_smal_mesh:
+                    target_gc_class = target_dict['gc'][:, :, 0]
+                    gc_errors_plane = calculate_plane_errors_batch(output_ref['vertices_smal'], target_gc_class, target_dict['has_gc'], target_dict['has_gc_is_touching'])
+                    loss_dict_temp['gc_plane'] = torch.mean(gc_errors_plane)
+                else:   # use a uniformly sampled mesh
+                    target_gc_class = target_dict['gc'][:, :, 0]
+                    device = output_ref['vertices_smal'].device
+                    remeshing_relevant_faces = self.remeshing_relevant_faces.to(device)
+                    remeshing_relevant_barys = self.remeshing_relevant_barys.to(device)
+
+                    bs = output_ref['vertices_smal'].shape[0]
+                    # verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, output_ref['vertices_smal'][:, self.remeshing_relevant_faces])
+                    # sel_verts_comparison = output_ref['vertices_smal'][:, self.remeshing_relevant_faces]
+                    # verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, sel_verts_comparison)
+                    sel_verts = torch.index_select(output_ref['vertices_smal'], dim=1, index=remeshing_relevant_faces.reshape((-1))).reshape((bs, remeshing_relevant_faces.shape[0], 3, 3))
+                    verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, sel_verts)
+                    target_gc_class_remeshed = torch.einsum('ij,aij->ai', remeshing_relevant_barys, target_gc_class[:, self.remeshing_relevant_faces].to(device=device, dtype=torch.float32))
+                    target_gc_class_remeshed_prep = torch.round(target_gc_class_remeshed).to(torch.long)
+                    gc_errors_plane, gc_errors_under_plane = calculate_plane_errors_batch(verts_remeshed, target_gc_class_remeshed_prep, target_dict['has_gc'], target_dict['has_gc_is_touching'])
+                    loss_dict_temp['gc_plane'] = torch.mean(gc_errors_plane) 
+                    loss_dict_temp['gc_blowplane'] = torch.mean(gc_errors_under_plane)
+
+        # error on classification if the ground plane is flat
+        if 'gc_isflat' in weight_dict_ref.keys():
+            # import pdb; pdb.set_trace()
+            self.criterion_class_isflat.to(device)
+            loss_dict_temp['gc_isflat'] = self.criterion_class(output_ref['isflat'], target_dict['isflat'].to(device))
+
+        # if we refine the shape WITHIN the refinement newtork (shaperef_type is not inexistent)
+        # shape regularization
+        #   'smal': loss on betas (pca coefficients), betas should be close to 0
+        #   'limbs...' loss on selected betas_limbs
+        device = output_ref_comp['ref_trans_notnorm'].device
+        loss_shape_weighted_list = [torch.zeros((1), device=device).mean()]  
+        if 'shape_options' in weight_dict_ref.keys():
+            for ind_sp, sp in enumerate(weight_dict_ref['shape_options']):
+                weight_sp = weight_dict_ref['shape'][ind_sp]
+                # self.logscale_part_list = ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'] 
+                if sp == 'smal':
+                    loss_shape_tmp = self.shape_prior(output_ref['betas'])
+                elif sp == 'limbs':
+                    loss_shape_tmp = torch.mean((output_ref['betas_limbs'])**2)  
+                elif sp == 'limbs7':
+                    limb_coeffs_list = [0.01, 1, 0.1, 1, 1, 0.1, 2]
+                    limb_coeffs = torch.tensor(limb_coeffs_list).to(torch.float32).to(target_dict['tpts'].device)   
+                    loss_shape_tmp = torch.mean((output_ref['betas_limbs'] * limb_coeffs[None, :])**2)            
+                else:
+                    raise NotImplementedError
+                loss_shape_weighted_list.append(weight_sp * loss_shape_tmp)
+        loss_shape_weighted = torch.stack((loss_shape_weighted_list)).sum()
+
+
+
+
+
+        # 3D loss for dogs for which we have a unity model or toy figure
+        loss_dict_temp['models3d'] = torch.zeros((1), device=device).mean().to(output_ref['betas'].device)
+        if 'models3d' in weight_dict_ref.keys():
+            if weight_dict_ref['models3d'] > 0:
+                assert (self.dog_betas_unity is not None)
+                if weight_dict_ref['models3d'] > 0:
+                    for ind_dog in range(target_dict['breed_index'].shape[0]):
+                        breed_index = np.asscalar(target_dict['breed_index'][ind_dog].detach().cpu().numpy())
+                        if breed_index in self.dog_betas_unity.keys():
+                            betas_target = self.dog_betas_unity[breed_index][:output_ref['betas'].shape[1]].to(output_ref['betas'].device)
+                            betas_output = output_ref['betas'][ind_dog, :]
+                            betas_limbs_output = output_ref['betas_limbs'][ind_dog, :]
+                            loss_dict_temp['models3d'] += ((betas_limbs_output**2).sum() + ((betas_output-betas_target)**2).sum()) / (output_ref['betas'].shape[1] + output_ref['betas_limbs'].shape[1])
+            else:
+                weight_dict_ref['models3d'] = 0.0
+        else:
+            weight_dict_ref['models3d'] = 0.0
+
+
+
+
+
+
+
+
+
+
+
+        # weight the losses
+        loss = torch.zeros((1)).mean().to(device=output_ref['keyp_2d'].device, dtype=output_ref['keyp_2d'].dtype)
+        loss_dict = {}
+        for loss_name in weight_dict_ref.keys():
+            if not loss_name in ['shape', 'shape_options']:
+                if weight_dict_ref[loss_name] > 0:
+                    loss_weighted = loss_dict_temp[loss_name] * weight_dict_ref[loss_name]
+                    loss_dict[loss_name] = loss_weighted.item()
+                    loss += loss_weighted
+        loss += loss_shape_weighted
+        loss_dict['loss'] = loss.item()
+
+        return loss, loss_dict
+
+

src/combined_model/loss_image_to_3d_withbreedrel.py

@@ -0,0 +1,342 @@
+
+
+import torch
+import numpy as np
+import pickle as pkl
+
+import os
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
+# from priors.pose_prior_35 import Prior
+# from priors.tiger_pose_prior.tiger_pose_prior import GaussianMixturePrior
+from priors.normalizing_flow_prior.normalizing_flow_prior import NormalizingFlowPrior
+from priors.shape_prior import ShapePrior
+from lifting_to_3d.utils.geometry_utils import rot6d_to_rotmat, batch_rot2aa
+# from configs.SMAL_configs import SMAL_MODEL_DATA_PATH, UNITY_SMAL_SHAPE_PRIOR_DOGS, SMAL_MODEL_TYPE  
+from configs.SMAL_configs import SMAL_MODEL_CONFIG
+
+from priors.helper_3dcgmodel_loss import load_dog_betas_for_3dcgmodel_loss
+from combined_model.loss_utils.loss_utils_gc import calculate_plane_errors_batch
+
+
+
+class Loss(torch.nn.Module):
+    def __init__(self, smal_model_type, data_info, nf_version=None):
+        super(Loss, self).__init__()
+        self.criterion_regr = torch.nn.MSELoss()        # takes the mean   
+        self.criterion_class = torch.nn.CrossEntropyLoss()
+        self.data_info = data_info   
+        self.register_buffer('keypoint_weights', torch.tensor(data_info.keypoint_weights)[None, :])
+        self.l_anchor = None
+        self.l_pos = None
+        self.l_neg = None
+        self.smal_model_type = smal_model_type
+        self.smal_model_data_path = SMAL_MODEL_CONFIG[self.smal_model_type]['smal_model_data_path']
+        self.unity_smal_shape_prior_dogs = SMAL_MODEL_CONFIG[self.smal_model_type]['unity_smal_shape_prior_dogs']
+
+        if nf_version is not None:
+            self.normalizing_flow_pose_prior = NormalizingFlowPrior(nf_version=nf_version)
+        self.shape_prior = ShapePrior(self.smal_model_data_path) # here we just need mean and cov
+        self.criterion_triplet = torch.nn.TripletMarginLoss(margin=1)
+
+        # load 3d data for the unity dogs (an optional shape prior for 11 breeds)
+        if self.unity_smal_shape_prior_dogs is not None:
+            self.dog_betas_unity = load_dog_betas_for_3dcgmodel_loss(self.unity_smal_shape_prior_dogs, self.smal_model_type)
+        else:
+            self.dog_betas_unity = None
+
+        remeshing_path = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/smal_data_remeshed/uniform_surface_sampling/my_smpl_39dogsnorm_Jr_4_dog_remesh4000_info.pkl'
+        with open(remeshing_path, 'rb') as fp: 
+            self.remeshing_dict = pkl.load(fp)
+        self.remeshing_relevant_faces = torch.tensor(self.remeshing_dict['smal_faces'][self.remeshing_dict['faceid_closest']], dtype=torch.long)
+        self.remeshing_relevant_barys = torch.tensor(self.remeshing_dict['barys_closest'], dtype=torch.float32)
+
+
+    def prepare_anchor_pos_neg(self, batch_size, device):
+        l0 = np.arange(0, batch_size, 2)
+        l_anchor = []
+        l_pos = []
+        l_neg = []
+        for ind in l0:
+            xx = set(np.arange(0, batch_size))
+            xx.discard(ind)
+            xx.discard(ind+1)
+            for ind2 in xx:
+                if ind2 % 2 == 0:
+                    l_anchor.append(ind)
+                    l_pos.append(ind + 1)
+                else:
+                    l_anchor.append(ind + 1)
+                    l_pos.append(ind)
+                l_neg.append(ind2)
+        self.l_anchor = torch.Tensor(l_anchor).to(torch.int64).to(device)
+        self.l_pos = torch.Tensor(l_pos).to(torch.int64).to(device)
+        self.l_neg = torch.Tensor(l_neg).to(torch.int64).to(device)
+        return
+
+
+    def forward(self, output_reproj, target_dict, weight_dict=None):
+
+        # output_reproj: ['vertices_smal', 'keyp_3d', 'keyp_2d', 'silh_image']
+        # target_dict: ['index', 'center', 'scale', 'pts', 'tpts', 'target_weight']
+        batch_size = output_reproj['keyp_2d'].shape[0]
+        device = output_reproj['keyp_2d'].device
+
+        # loss on reprojected keypoints 
+        output_kp_resh = (output_reproj['keyp_2d']).reshape((-1, 2))    
+        target_kp_resh = (target_dict['tpts'][:, :, :2] / 64. * (256. - 1)).reshape((-1, 2))
+        weights_resh = target_dict['tpts'][:, :, 2].reshape((-1)) 
+        keyp_w_resh = self.keypoint_weights.repeat((batch_size, 1)).reshape((-1))
+        loss_keyp = ((((output_kp_resh - target_kp_resh)[weights_resh>0]**2).sum(axis=1).sqrt()*weights_resh[weights_resh>0])*keyp_w_resh[weights_resh>0]).sum() / \
+            max((weights_resh[weights_resh>0]*keyp_w_resh[weights_resh>0]).sum(), 1e-5)
+
+        # loss on reprojected silhouette
+        assert output_reproj['silh'].shape == (target_dict['silh'][:, None, :, :]).shape
+        silh_loss_type = 'default'
+        if silh_loss_type == 'default':
+            with torch.no_grad():
+                thr_silh = 20
+                diff = torch.norm(output_kp_resh - target_kp_resh, dim=1)
+                diff_x = diff.reshape((batch_size, -1))
+                weights_resh_x = weights_resh.reshape((batch_size, -1))
+                unweighted_kp_mean_dist = (diff_x * weights_resh_x).sum(dim=1) / ((weights_resh_x).sum(dim=1)+1e-6)
+            loss_silh_bs = ((output_reproj['silh'] - target_dict['silh'][:, None, :, :]) ** 2).sum(axis=3).sum(axis=2).sum(axis=1) / (output_reproj['silh'].shape[2]*output_reproj['silh'].shape[3])
+            loss_silh = loss_silh_bs[unweighted_kp_mean_dist<thr_silh].sum() / batch_size
+        else:
+            print('silh_loss_type: ' + silh_loss_type)
+            raise ValueError
+
+        # shape regularization
+        #   'smal': loss on betas (pca coefficients), betas should be close to 0
+        #   'limbs...' loss on selected betas_limbs
+        loss_shape_weighted_list = [torch.zeros((1), device=device).mean().to(output_reproj['keyp_2d'].device)]  
+        for ind_sp, sp in enumerate(weight_dict['shape_options']):
+            weight_sp = weight_dict['shape'][ind_sp]
+            # self.logscale_part_list = ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'] 
+            if sp == 'smal':
+                loss_shape_tmp = self.shape_prior(output_reproj['betas'])
+            elif sp == 'limbs':
+                loss_shape_tmp = torch.mean((output_reproj['betas_limbs'])**2)  
+            elif sp == 'limbs7':
+                limb_coeffs_list = [0.01, 1, 0.1, 1, 1, 0.1, 2]
+                limb_coeffs = torch.tensor(limb_coeffs_list).to(torch.float32).to(target_dict['tpts'].device)   
+                loss_shape_tmp = torch.mean((output_reproj['betas_limbs'] * limb_coeffs[None, :])**2)            
+            else:
+                raise NotImplementedError
+            loss_shape_weighted_list.append(weight_sp * loss_shape_tmp)
+        loss_shape_weighted = torch.stack((loss_shape_weighted_list)).sum()
+
+        # 3D loss for dogs for which we have a unity model or toy figure
+        loss_models3d = torch.zeros((1), device=device).mean().to(output_reproj['betas'].device)
+        if 'models3d' in weight_dict.keys():
+            if weight_dict['models3d'] > 0:
+                assert (self.dog_betas_unity is not None)
+                if weight_dict['models3d'] > 0:
+                    for ind_dog in range(target_dict['breed_index'].shape[0]):
+                        breed_index = np.asscalar(target_dict['breed_index'][ind_dog].detach().cpu().numpy())
+                        if breed_index in self.dog_betas_unity.keys():
+                            betas_target = self.dog_betas_unity[breed_index][:output_reproj['betas'].shape[1]].to(output_reproj['betas'].device)
+                            betas_output = output_reproj['betas'][ind_dog, :]
+                            betas_limbs_output = output_reproj['betas_limbs'][ind_dog, :]
+                            loss_models3d += ((betas_limbs_output**2).sum() + ((betas_output-betas_target)**2).sum()) / (output_reproj['betas'].shape[1] + output_reproj['betas_limbs'].shape[1])
+            else:
+                weight_dict['models3d'] = 0.0
+        else:
+            weight_dict['models3d'] = 0.0
+
+        # shape resularization loss on shapedirs
+        #   -> in the current version shapedirs are kept fixed, so we don't need those losses
+        if weight_dict['shapedirs'] > 0:
+            raise NotImplementedError  
+        else:
+            loss_shapedirs = torch.zeros((1), device=device).mean().to(output_reproj['betas'].device)
+
+        # prior on back joints (not used in cvpr 2022 paper)
+        #   -> elementwise MSE loss on all 6 coefficients of 6d rotation representation
+        if 'pose_0' in weight_dict.keys(): 
+            if weight_dict['pose_0'] > 0:
+                pred_pose_rot6d = output_reproj['pose_rot6d']
+                w_rj_np = np.zeros((pred_pose_rot6d.shape[1]))
+                w_rj_np[[2, 3, 4, 5]] = 1.0         # back
+                w_rj = torch.tensor(w_rj_np).to(torch.float32).to(pred_pose_rot6d.device)     
+                zero_rot = torch.tensor([1, 0, 0, 1, 0, 0]).to(pred_pose_rot6d.device).to(torch.float32)[None, None, :].repeat((batch_size, pred_pose_rot6d.shape[1], 1))
+                loss_pose = self.criterion_regr(pred_pose_rot6d*w_rj[None, :, None], zero_rot*w_rj[None, :, None])
+            else:
+                loss_pose = torch.zeros((1), device=device).mean()
+
+        # pose prior 
+        #   -> we did experiment with different pose priors, for example:
+        #       * similart to SMALify (https://github.com/benjiebob/SMALify/blob/master/smal_fitter/smal_fitter.py, 
+        #         https://github.com/benjiebob/SMALify/blob/master/smal_fitter/priors/pose_prior_35.py)
+        #       * vae 
+        #       * normalizing flow pose prior
+        #   -> our cvpr 2022 paper uses the normalizing flow pose prior as implemented below
+        if 'poseprior' in weight_dict.keys():
+            if weight_dict['poseprior'] > 0:
+                pred_pose_rot6d = output_reproj['pose_rot6d']
+                pred_pose = rot6d_to_rotmat(pred_pose_rot6d.reshape((-1, 6))).reshape((batch_size, -1, 3, 3))
+                if 'normalizing_flow_tiger' in weight_dict['poseprior_options']:
+                    if output_reproj['normflow_z'] is not None:
+                        loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss_from_z(output_reproj['normflow_z'], type='square')
+                    else:
+                        loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss(pred_pose_rot6d, type='square')
+                elif 'normalizing_flow_tiger_logprob' in weight_dict['poseprior_options']:
+                    if output_reproj['normflow_z'] is not None:
+                        loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss_from_z(output_reproj['normflow_z'], type='neg_log_prob')
+                    else:
+                        loss_poseprior = self.normalizing_flow_pose_prior.calculate_loss(pred_pose_rot6d, type='neg_log_prob')
+                else:
+                    raise NotImplementedError
+            else:
+                loss_poseprior = torch.zeros((1), device=device).mean()
+        else:
+            weight_dict['poseprior'] = 0
+            loss_poseprior = torch.zeros((1), device=device).mean()
+
+        # add a prior which penalizes side-movement angles for legs 
+        if 'poselegssidemovement' in weight_dict.keys():
+            if weight_dict['poselegssidemovement'] > 0:
+                use_pose_legs_side_loss = True
+            else:
+                use_pose_legs_side_loss = False
+        else:
+            use_pose_legs_side_loss = False
+        if use_pose_legs_side_loss:
+            leg_indices_right = np.asarray([7, 8, 9, 10, 17, 18, 19, 20])      # front, back
+            leg_indices_left = np.asarray([11, 12, 13, 14, 21, 22, 23, 24])     # front, back
+            vec = torch.zeros((3, 1)).to(device=pred_pose.device, dtype=pred_pose.dtype)
+            vec[2] = -1
+            x0_rotmat = pred_pose   
+            x0_rotmat_legs_left = x0_rotmat[:, leg_indices_left, :, :]
+            x0_rotmat_legs_right = x0_rotmat[:, leg_indices_right, :, :]
+            x0_legs_left = x0_rotmat_legs_left.reshape((-1, 3, 3))@vec
+            x0_legs_right = x0_rotmat_legs_right.reshape((-1, 3, 3))@vec
+            eps=0       # 1e-7
+            # use the component of the vector which points to the side
+            loss_poselegssidemovement = (x0_legs_left[:, 1]**2).mean() + (x0_legs_right[:, 1]**2).mean()
+        else:
+            loss_poselegssidemovement = torch.zeros((1), device=device).mean()
+            weight_dict['poselegssidemovement'] = 0
+
+        # dog breed classification loss
+        dog_breed_gt = target_dict['breed_index']
+        dog_breed_pred = output_reproj['dog_breed']
+        loss_class = self.criterion_class(dog_breed_pred, dog_breed_gt)
+
+        # dog breed relationship loss
+        #   -> we did experiment with many other options, but none was significantly better 
+        if '4' in weight_dict['breed_options']:      # we have pairs of dogs of the same breed 
+            if weight_dict['breed'] > 0:
+                assert output_reproj['dog_breed'].shape[0] == 12
+                # assert weight_dict['breed'] > 0
+                z = output_reproj['z']   
+                # go through all pairs and compare them to each other sample
+                if self.l_anchor is None:
+                    self.prepare_anchor_pos_neg(batch_size, z.device)
+                anchor = torch.index_select(z, 0, self.l_anchor)
+                positive = torch.index_select(z, 0, self.l_pos)
+                negative = torch.index_select(z, 0, self.l_neg)
+                loss_breed = self.criterion_triplet(anchor, positive, negative)
+            else:
+                loss_breed = torch.zeros((1), device=device).mean()
+        else:
+            loss_breed = torch.zeros((1), device=device).mean()
+
+        # regularizarion for focal length
+        loss_flength_near_mean = torch.mean(output_reproj['flength']**2)
+        loss_flength = loss_flength_near_mean
+
+        # bodypart segmentation loss
+        if 'partseg' in weight_dict.keys():
+            if weight_dict['partseg'] > 0:
+                raise NotImplementedError
+            else:
+                loss_partseg = torch.zeros((1), device=device).mean()
+        else:
+            weight_dict['partseg'] = 0
+            loss_partseg = torch.zeros((1), device=device).mean()
+
+
+        # NEW: ground contact loss for main network
+        keep_smal_mesh = False 
+        if 'gc_plane' in weight_dict.keys():
+            if weight_dict['gc_plane'] > 0:
+                if keep_smal_mesh:
+                    target_gc_class = target_dict['gc'][:, :, 0]
+                    gc_errors_plane = calculate_plane_errors_batch(output_reproj['vertices_smal'], target_gc_class, target_dict['has_gc'], target_dict['has_gc_is_touching'])
+                    loss_gc_plane = torch.mean(gc_errors_plane)
+                else:   # use a uniformly sampled mesh
+                    target_gc_class = target_dict['gc'][:, :, 0]
+                    device = output_reproj['vertices_smal'].device
+                    remeshing_relevant_faces = self.remeshing_relevant_faces.to(device)
+                    remeshing_relevant_barys = self.remeshing_relevant_barys.to(device)
+
+                    bs = output_reproj['vertices_smal'].shape[0]
+                    # verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, output_reproj['vertices_smal'][:, self.remeshing_relevant_faces])
+                    # sel_verts_comparison = output_reproj['vertices_smal'][:, self.remeshing_relevant_faces]
+                    # verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, sel_verts_comparison)
+                    sel_verts = torch.index_select(output_reproj['vertices_smal'], dim=1, index=remeshing_relevant_faces.reshape((-1))).reshape((bs, remeshing_relevant_faces.shape[0], 3, 3))
+                    verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, sel_verts)
+                    target_gc_class_remeshed = torch.einsum('ij,aij->ai', remeshing_relevant_barys, target_gc_class[:, self.remeshing_relevant_faces].to(device=device, dtype=torch.float32))
+                    target_gc_class_remeshed_prep = torch.round(target_gc_class_remeshed).to(torch.long)
+                    gc_errors_plane, gc_errors_under_plane = calculate_plane_errors_batch(verts_remeshed, target_gc_class_remeshed_prep, target_dict['has_gc'], target_dict['has_gc_is_touching'])
+                    loss_gc_plane = torch.mean(gc_errors_plane) 
+                    loss_gc_belowplane = torch.mean(gc_errors_under_plane)
+                    # loss_dict_temp['gc_plane'] = torch.mean(gc_errors_plane)
+            else:
+                loss_gc_plane = torch.zeros((1), device=device).mean()
+                loss_gc_belowplane = torch.zeros((1), device=device).mean()
+        else:
+            loss_gc_plane = torch.zeros((1), device=device).mean()
+            loss_gc_belowplane = torch.zeros((1), device=device).mean()
+            weight_dict['gc_plane'] = 0
+            weight_dict['gc_belowplane'] = 0
+
+
+
+        # weight and combine losses
+        loss_keyp_weighted = loss_keyp * weight_dict['keyp']
+        loss_silh_weighted = loss_silh * weight_dict['silh']
+        loss_shapedirs_weighted = loss_shapedirs * weight_dict['shapedirs']
+        loss_pose_weighted = loss_pose * weight_dict['pose_0']
+        loss_class_weighted = loss_class * weight_dict['class']
+        loss_breed_weighted = loss_breed * weight_dict['breed']
+        loss_flength_weighted = loss_flength * weight_dict['flength']
+        loss_poseprior_weighted = loss_poseprior * weight_dict['poseprior']
+        loss_partseg_weighted = loss_partseg * weight_dict['partseg']
+        loss_models3d_weighted = loss_models3d * weight_dict['models3d']
+        loss_poselegssidemovement_weighted = loss_poselegssidemovement * weight_dict['poselegssidemovement']
+
+        loss_gc_plane_weighted = loss_gc_plane * weight_dict['gc_plane']
+        loss_gc_belowplane_weighted = loss_gc_belowplane * weight_dict['gc_belowplane']
+
+
+        ####################################################################################################
+        loss = loss_keyp_weighted + loss_silh_weighted + loss_shape_weighted + loss_pose_weighted + loss_class_weighted + \
+                loss_shapedirs_weighted + loss_breed_weighted + loss_flength_weighted + loss_poseprior_weighted + \
+                loss_partseg_weighted + loss_models3d_weighted + loss_poselegssidemovement_weighted + \
+                loss_gc_plane_weighted + loss_gc_belowplane_weighted
+        ####################################################################################################
+        
+        loss_dict = {'loss': loss.item(), 
+                    'loss_keyp_weighted': loss_keyp_weighted.item(), \
+                    'loss_silh_weighted': loss_silh_weighted.item(), \
+                    'loss_shape_weighted': loss_shape_weighted.item(), \
+                    'loss_shapedirs_weighted': loss_shapedirs_weighted.item(), \
+                    'loss_pose0_weighted': loss_pose_weighted.item(), \
+                    'loss_class_weighted': loss_class_weighted.item(), \
+                    'loss_breed_weighted': loss_breed_weighted.item(), \
+                    'loss_flength_weighted': loss_flength_weighted.item(), \
+                    'loss_poseprior_weighted': loss_poseprior_weighted.item(), \
+                    'loss_partseg_weighted': loss_partseg_weighted.item(), \
+                    'loss_models3d_weighted': loss_models3d_weighted.item(), \
+                    'loss_poselegssidemovement_weighted': loss_poselegssidemovement_weighted.item(), \
+                    'loss_gc_plane_weighted': loss_gc_plane_weighted.item(), \
+                    'loss_gc_belowplane_weighted': loss_gc_belowplane_weighted.item()
+                    }
+                    
+        return loss, loss_dict
+
+
+
+

src/combined_model/loss_utils/loss_arap.py

@@ -0,0 +1,153 @@
+import torch
+
+# code from https://raw.githubusercontent.com/yufu-wang/aves/main/optimization/loss_arap.py
+
+
+class Arap_Loss():
+    '''
+    Pytorch implementaion: As-rigid-as-possible loss class
+    
+    '''
+
+    def __init__(self, meshes, device='cpu', vertex_w=None):
+
+        with torch.no_grad():       # new nadine 
+
+            self.device = device
+            self.bn = len(meshes)
+
+            # get lapacian cotangent matrix
+            L = self.get_laplacian_cot(meshes)
+            self.wij = L.values().clone()
+            self.wij[self.wij<0] = 0.
+
+            # get ajacency matrix
+            V = meshes.num_verts_per_mesh().sum()
+            edges_packed = meshes.edges_packed()
+            e0, e1 = edges_packed.unbind(1)
+            idx01 = torch.stack([e0, e1], dim=1)
+            idx10 = torch.stack([e1, e0], dim=1)
+            idx = torch.cat([idx01, idx10], dim=0).t()
+
+            ones = torch.ones(idx.shape[1], dtype=torch.float32).to(device)
+            A = torch.sparse.FloatTensor(idx, ones, (V, V))
+            self.deg = torch.sparse.sum(A, dim=1).to_dense().long()
+            self.idx = self.sort_idx(idx)
+
+            # get edges of default mesh
+            self.eij = self.get_edges(meshes)
+
+            # get per vertex regularization strength
+            self.vertex_w = vertex_w
+
+
+    def __call__(self, new_meshes):
+        new_meshes._compute_packed()
+        
+        optimal_R = self.step_1(new_meshes)
+        arap_loss = self.step_2(optimal_R, new_meshes)
+        return arap_loss
+
+
+    def step_1(self, new_meshes):
+        bn = self.bn
+        eij = self.eij.view(bn, -1, 3).cpu()
+
+        with torch.no_grad():
+            eij_ = self.get_edges(new_meshes)
+
+            eij_ = eij_.view(bn, -1, 3).cpu()
+            wij = self.wij.view(bn, -1).cpu()
+
+            deg_1 = self.deg.view(bn, -1)[0].cpu()  # assuming same topology
+            S = torch.zeros([bn, len(deg_1), 3, 3])
+            for i in range(len(deg_1)):
+                start, end = deg_1[:i].sum(), deg_1[:i+1].sum()
+                P  = eij[:, start : end]
+                P_ = eij_[:, start : end]
+                D = wij[:, start : end]
+                D = torch.diag_embed(D)
+                S[:, i] = P.transpose(-2,-1) @ D @ P_
+                
+            S = S.view(-1, 3, 3)
+
+            u, _, v = torch.svd(S)
+            R = v @ u.transpose(-2, -1)
+            det = torch.det(R)
+
+            u[det<0, :, -1] *= -1
+            R = v @ u.transpose(-2, -1)
+            R = R.to(self.device)
+
+        return R
+
+
+    def step_2(self, R, new_meshes):
+        R = torch.repeat_interleave(R, self.deg, dim=0)
+        Reij = R @ self.eij.unsqueeze(2)
+        Reij = Reij.squeeze()
+        
+        eij_ = self.get_edges(new_meshes)
+        arap_loss = self.wij * (eij_ - Reij).norm(dim=1)
+
+        if self.vertex_w is not None:
+            vertex_w = torch.repeat_interleave(self.vertex_w, self.deg, dim=0)
+            arap_loss = arap_loss * vertex_w
+
+        arap_loss = arap_loss.sum() / self.bn
+
+        return arap_loss
+
+
+    def get_edges(self, meshes):
+        verts_packed = meshes.verts_packed()
+        vi = torch.repeat_interleave(verts_packed, self.deg, dim=0)
+        vj = verts_packed[self.idx[1]]
+        eij = vi - vj
+        return eij
+
+
+    def sort_idx(self, idx):
+        _, order = (idx[0] + idx[1]*1e-6).sort()
+
+        return idx[:, order]
+
+
+    def get_laplacian_cot(self, meshes):
+        '''
+        Routine modified from :
+        pytorch3d/loss/mesh_laplacian_smoothing.py
+        '''
+        verts_packed = meshes.verts_packed()
+        faces_packed = meshes.faces_packed()
+        V, F = verts_packed.shape[0], faces_packed.shape[0]
+
+        face_verts = verts_packed[faces_packed]
+        v0, v1, v2 = face_verts[:,0], face_verts[:,1], face_verts[:,2]
+
+        A = (v1-v2).norm(dim=1)
+        B = (v0-v2).norm(dim=1)
+        C = (v0-v1).norm(dim=1)
+
+        s = 0.5 * (A+B+C)
+        area = (s * (s - A) * (s - B) * (s - C)).clamp_(min=1e-12).sqrt()
+
+        A2, B2, C2 = A * A, B * B, C * C
+        cota = (B2 + C2 - A2) / area
+        cotb = (A2 + C2 - B2) / area
+        cotc = (A2 + B2 - C2) / area
+        cot = torch.stack([cota, cotb, cotc], dim=1)
+        cot /= 4.0
+
+        ii = faces_packed[:, [1,2,0]]  
+        jj = faces_packed[:, [2,0,1]]
+        idx = torch.stack([ii, jj], dim=0).view(2, F*3)
+        L = torch.sparse.FloatTensor(idx, cot.view(-1), (V, V))
+        L += L.t()
+        L = L.coalesce()
+        L /= 2.0  # normalized according to arap paper
+
+        return L
+
+
+

src/combined_model/loss_utils/loss_laplacian_mesh_comparison.py

@@ -0,0 +1,45 @@
+# code from: https://github.com/chaneyddtt/Coarse-to-fine-3D-Animal/blob/main/util/loss_utils.py
+
+import numpy as np
+import torch 
+
+
+# Laplacian loss, calculate the Laplacian coordiante of both coarse and refined vertices and then compare the difference
+class LaplacianCTF(torch.nn.Module):
+    def __init__(self, adjmat, device):
+        '''
+        Args:
+            adjmat: adjacency matrix of the input graph data
+            device: specify device for training
+        '''
+        super(LaplacianCTF, self).__init__()
+        adjmat.data = np.ones_like(adjmat.data)
+        adjmat = torch.from_numpy(adjmat.todense()).float()
+        dg = torch.sum(adjmat, dim=-1)
+        dg_m = torch.diag(dg)
+        ls = dg_m - adjmat
+        self.ls = ls.unsqueeze(0).to(device)  # Should be normalized by the diagonal elements according to
+                                              # the origial definition, this one also works fine.
+
+    def forward(self, verts_pred, verts_gt, smooth=False):
+        verts_pred = torch.matmul(self.ls, verts_pred)
+        verts_gt = torch.matmul(self.ls, verts_gt)
+        loss = torch.norm(verts_pred - verts_gt, dim=-1).mean()
+        if smooth:
+            loss_smooth = torch.norm(torch.matmul(self.ls, verts_pred), dim=-1).mean()
+            return loss, loss_smooth
+        return loss, None
+
+
+
+
+# 
+# read the adjacency matrix, which will used in the Laplacian regularizer
+# data = np.load('./data/mesh_down_sampling_4.npz', encoding='latin1', allow_pickle=True)
+# adjmat = data['A'][0]
+# laplacianloss = Laplacian(adjmat, device)
+# 
+# verts_clone = verts.detach().clone()
+# loss_arap, loss_smooth = laplacianloss(verts_refine, verts_clone)
+# loss_arap = args.w_arap * loss_arap
+#

src/combined_model/loss_utils/loss_sdf.py

@@ -0,0 +1,122 @@
+
+# code from: https://github.com/chaneyddtt/Coarse-to-fine-3D-Animal/blob/main/util/loss_sdf.py
+
+import torch
+import numpy as np
+from scipy.ndimage import distance_transform_edt as distance
+from skimage import segmentation as skimage_seg
+import matplotlib.pyplot as plt
+
+
+def dice_loss(score, target):
+    # implemented from paper https://arxiv.org/pdf/1606.04797.pdf
+    target = target.float()
+    smooth = 1e-5
+    intersect = torch.sum(score * target)
+    y_sum = torch.sum(target * target)
+    z_sum = torch.sum(score * score)
+    loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
+    loss = 1 - loss
+    return loss
+
+
+class tversky_loss(torch.nn.Module):
+    # implemented from https://arxiv.org/pdf/1706.05721.pdf
+    def __init__(self, alpha, beta):
+        '''
+        Args:
+            alpha: coefficient for false positive prediction
+            beta: coefficient for false negtive prediction
+        '''
+        super(tversky_loss, self).__init__()
+        self.alpha = alpha
+        self.beta = beta
+
+    def __call__(self, score, target):
+        target = target.float()
+        smooth = 1e-5
+        tp = torch.sum(score * target)
+        fn = torch.sum(target * (1 - score))
+        fp = torch.sum((1-target) * score)
+        loss = (tp + smooth) / (tp + self.alpha * fp + self.beta * fn + smooth)
+        loss = 1 - loss
+        return loss
+
+
+def compute_sdf1_1(img_gt, out_shape):
+    """
+    compute the normalized signed distance map of binary mask
+    input: segmentation, shape = (batch_size, x, y, z)
+    output: the Signed Distance Map (SDM)
+    sdf(x) = 0; x in segmentation boundary
+             -inf|x-y|; x in segmentation
+             +inf|x-y|; x out of segmentation
+    normalize sdf to [-1, 1]
+    """
+
+    img_gt = img_gt.astype(np.uint8)
+
+    normalized_sdf = np.zeros(out_shape)
+
+    for b in range(out_shape[0]): # batch size
+            # ignore background
+        for c in range(1, out_shape[1]):
+            posmask = img_gt[b]
+            negmask = 1-posmask
+            posdis = distance(posmask)
+            negdis = distance(negmask)
+            boundary = skimage_seg.find_boundaries(posmask, mode='inner').astype(np.uint8)
+            sdf = (negdis-np.min(negdis))/(np.max(negdis)-np.min(negdis)) - (posdis-np.min(posdis))/(np.max(posdis)-np.min(posdis))
+            sdf[boundary==1] = 0
+            normalized_sdf[b][c] = sdf
+            assert np.min(sdf) == -1.0, print(np.min(posdis), np.min(negdis), np.max(posdis), np.max(negdis))
+            assert np.max(sdf) ==  1.0, print(np.min(posdis), np.min(negdis), np.max(posdis), np.max(negdis))
+
+    return normalized_sdf
+
+
+def compute_sdf(img_gt, out_shape):
+    """
+    compute the signed distance map of binary mask
+    input: segmentation, shape = (batch_size, x, y, z)
+    output: the Signed Distance Map (SDM)
+    sdf(x) = 0; x in segmentation boundary
+             -inf|x-y|; x in segmentation
+             +inf|x-y|; x out of segmentation
+    """
+
+    img_gt = img_gt.astype(np.uint8)
+
+    gt_sdf = np.zeros(out_shape)
+    debug = False
+    for b in range(out_shape[0]): # batch size
+        for c in range(0, out_shape[1]):
+            posmask = img_gt[b]
+            negmask = 1-posmask
+            posdis = distance(posmask)
+            negdis = distance(negmask)
+            boundary = skimage_seg.find_boundaries(posmask, mode='inner').astype(np.uint8)
+            sdf = negdis - posdis
+            sdf[boundary==1] = 0
+            gt_sdf[b][c] = sdf
+            if debug:
+                plt.figure()
+                plt.subplot(1, 2, 1), plt.imshow(img_gt[b, 0, :, :]), plt.colorbar()
+                plt.subplot(1, 2, 2), plt.imshow(gt_sdf[b, 0, :, :]), plt.colorbar()
+                plt.show()
+
+    return gt_sdf
+
+
+def boundary_loss(output, gt):
+    """
+    compute boundary loss for binary segmentation
+    input: outputs_soft: softmax results,  shape=(b,2,x,y,z)
+           gt_sdf: sdf of ground truth (can be original or normalized sdf); shape=(b,2,x,y,z)
+    output: boundary_loss; sclar
+    adopted from http://proceedings.mlr.press/v102/kervadec19a/kervadec19a.pdf
+    """
+    multipled = torch.einsum('bcxy, bcxy->bcxy', output, gt)
+    bd_loss = multipled.mean()
+
+    return bd_loss

src/combined_model/loss_utils/loss_utils.py

@@ -0,0 +1,191 @@
+
+import torch
+import numpy as np
+
+
+'''
+def keyp_rep_error_l1(smpl_keyp_2d, keyp_hourglass, keyp_hourglass_scores, thr_kp=0.3):
+    # step 1: make sure that the hg prediction and barc are close
+    with torch.no_grad():
+        kp_weights = keyp_hourglass_scores
+        kp_weights[keyp_hourglass_scores<thr_kp] = 0
+    loss_keyp_rep = torch.mean((torch.abs((smpl_keyp_2d - keyp_hourglass)/512)).sum(dim=2)*kp_weights[:, :, 0])
+    return loss_keyp_rep
+
+def keyp_rep_error(smpl_keyp_2d, keyp_hourglass, keyp_hourglass_scores, thr_kp=0.3):
+    # step 1: make sure that the hg prediction and barc are close
+    with torch.no_grad():
+        kp_weights = keyp_hourglass_scores
+        kp_weights[keyp_hourglass_scores<thr_kp] = 0
+    # losses['kp_reproj']['value'] = torch.mean((((smpl_keyp_2d - keyp_reproj_init)/512)**2).sum(dim=2)*kp_weights[:, :, 0])
+    loss_keyp_rep = torch.mean((((smpl_keyp_2d - keyp_hourglass)/512)**2).sum(dim=2)*kp_weights[:, :, 0])
+    return loss_keyp_rep
+'''
+
+def leg_sideway_error(optimed_pose_with_glob):
+    assert optimed_pose_with_glob.shape[1] == 35
+    leg_indices_right = np.asarray([7, 8, 9, 10, 17, 18, 19, 20])      # front, back
+    leg_indices_left = np.asarray([11, 12, 13, 14, 21, 22, 23, 24])     # front, back
+    # leg_indices_right = np.asarray([8, 9, 10, 18, 19, 20])      # front, back
+    # leg_indices_left = np.asarray([12, 13, 14, 22, 23, 24])     # front, back
+    x0_rotmat = optimed_pose_with_glob   # (1, 35, 3, 3)
+    x0_rotmat_legs_left = x0_rotmat[:, leg_indices_left, :, :]
+    x0_rotmat_legs_right = x0_rotmat[:, leg_indices_right, :, :]
+    vec = torch.zeros((3, 1)).to(device=optimed_pose_with_glob.device, dtype=optimed_pose_with_glob.dtype)
+    vec[2] = -1
+    x0_legs_left = x0_rotmat_legs_left.reshape((-1, 3, 3))@vec
+    x0_legs_right = x0_rotmat_legs_right.reshape((-1, 3, 3))@vec
+    loss_pose_legs_side = (x0_legs_left[:, 1]**2).mean() + (x0_legs_right[:, 1]**2).mean()
+    return loss_pose_legs_side
+
+
+def leg_torsion_error(optimed_pose_with_glob):
+    leg_indices_right = np.asarray([7, 8, 9, 10, 17, 18, 19, 20])      # front, back
+    leg_indices_left = np.asarray([11, 12, 13, 14, 21, 22, 23, 24])     # front, back
+    x0_rotmat = optimed_pose_with_glob   # (1, 35, 3, 3)
+    x0_rotmat_legs_left = x0_rotmat[:, leg_indices_left, :, :]
+    x0_rotmat_legs_right = x0_rotmat[:, leg_indices_right, :, :]
+    vec_x = torch.zeros((3, 1)).to(device=optimed_pose_with_glob.device, dtype=optimed_pose_with_glob.dtype)
+    vec_x[0] = 1      # in x direction
+    x_x_legs_left = x0_rotmat_legs_left.reshape((-1, 3, 3))@vec_x
+    x_x_legs_right = x0_rotmat_legs_right.reshape((-1, 3, 3))@vec_x
+    loss_pose_legs_torsion = (x_x_legs_left[:, 1]**2).mean() + (x_x_legs_right[:, 1]**2).mean()
+    return loss_pose_legs_torsion
+
+
+def frontleg_walkingdir_error(optimed_pose_with_glob):
+    # this prior should only be used for standing poses!
+    leg_indices_right = np.asarray([7, 8, 9, 10])      # front, back
+    leg_indices_left = np.asarray([11, 12, 13, 14])     # front, back
+    relevant_back_indices = np.asarray([1, 2, 3, 4, 5, 6])      # np.asarray([6])             # back joint in the front
+    x0_rotmat = optimed_pose_with_glob   # (1, 35, 3, 3)
+    x0_rotmat_legs_left = x0_rotmat[:, leg_indices_left, :, :]
+    x0_rotmat_legs_right = x0_rotmat[:, leg_indices_right, :, :]
+    x0_rotmat_back = x0_rotmat[:, relevant_back_indices, :, :]
+    vec = torch.zeros((3, 1)).to(device=optimed_pose_with_glob.device, dtype=optimed_pose_with_glob.dtype)
+    vec[2] = -1     # vector down
+    x0_legs_left = x0_rotmat_legs_left.reshape((-1, 3, 3))@vec
+    x0_legs_right = x0_rotmat_legs_right.reshape((-1, 3, 3))@vec
+    x0_back = x0_rotmat_back.reshape((-1, 3, 3))@vec
+    loss_pose_legs_side = (x0_legs_left[:, 0]**2).mean() + (x0_legs_right[:, 0]**2).mean() + (x0_back[:, 0]**2).mean()  # penalize movement to front
+    return loss_pose_legs_side
+
+
+def tail_sideway_error(optimed_pose_with_glob):
+    tail_indices = np.asarray([25, 26, 27, 28, 29, 30, 31])      
+    x0_rotmat = optimed_pose_with_glob   # (1, 35, 3, 3)
+    x0_rotmat_tail = x0_rotmat[:, tail_indices, :, :]
+    vec = torch.zeros((3, 1)).to(device=optimed_pose_with_glob.device, dtype=optimed_pose_with_glob.dtype)
+    '''vec[2] = -1    
+    x0_tail = x0_rotmat_tail.reshape((-1, 3, 3))@vec
+    loss_pose_tail_side = (x0_tail[:, 1]**2).mean()'''
+    vec[0] = -1    
+    x0_tail = x0_rotmat_tail.reshape((-1, 3, 3))@vec
+    loss_pose_tail_side = (x0_tail[:, 1]**2).mean()
+    return loss_pose_tail_side
+
+
+def tail_torsion_error(optimed_pose_with_glob):
+    tail_indices = np.asarray([25, 26, 27, 28, 29, 30, 31])      
+    x0_rotmat = optimed_pose_with_glob   # (1, 35, 3, 3)
+    x0_rotmat_tail = x0_rotmat[:, tail_indices, :, :]
+    vec_x = torch.zeros((3, 1)).to(device=optimed_pose_with_glob.device, dtype=optimed_pose_with_glob.dtype)
+    '''vec_x[0] = 1      # in x direction
+    x_x_tail = x0_rotmat_tail.reshape((-1, 3, 3))@vec_x
+    loss_pose_tail_torsion = (x_x_tail[:, 1]**2).mean()'''
+    vec_x[2] = 1      # in y direction
+    x_x_tail = x0_rotmat_tail.reshape((-1, 3, 3))@vec_x
+    loss_pose_tail_torsion = (x_x_tail[:, 1]**2).mean()
+    return loss_pose_tail_torsion
+
+
+def spine_sideway_error(optimed_pose_with_glob):
+    tail_indices = np.asarray([1, 2, 3, 4, 5, 6])   # was wrong      
+    x0_rotmat = optimed_pose_with_glob   # (1, 35, 3, 3)
+    x0_rotmat_tail = x0_rotmat[:, tail_indices, :, :]
+    vec = torch.zeros((3, 1)).to(device=optimed_pose_with_glob.device, dtype=optimed_pose_with_glob.dtype)
+    vec[0] = -1    
+    x0_tail = x0_rotmat_tail.reshape((-1, 3, 3))@vec
+    loss_pose_tail_side = (x0_tail[:, 1]**2).mean()
+    return loss_pose_tail_side
+
+
+def spine_torsion_error(optimed_pose_with_glob):
+    tail_indices = np.asarray([1, 2, 3, 4, 5, 6])      
+    x0_rotmat = optimed_pose_with_glob   # (1, 35, 3, 3)
+    x0_rotmat_tail = x0_rotmat[:, tail_indices, :, :]
+    vec_x = torch.zeros((3, 1)).to(device=optimed_pose_with_glob.device, dtype=optimed_pose_with_glob.dtype)
+    vec_x[2] = 1    # vec_x[0] = 1      # in z direction
+    x_x_tail = x0_rotmat_tail.reshape((-1, 3, 3))@vec_x
+    loss_pose_tail_torsion = (x_x_tail[:, 1]**2).mean()     # (x_x_tail[:, 1]**2).mean()
+    return loss_pose_tail_torsion
+
+
+def fit_plane(points_npx3):
+    # remarks:
+    #   visualization of the plane: debug_code/curve_fitting_v2.py
+    #   theory: https://www.ltu.se/cms_fs/1.51590!/svd-fitting.pdf
+    #   remark: torch.svd is depreciated
+    # new plane equation:
+    #   a(x−x0)+b(y−y0)+c(z−z0)=0
+    #   ax+by+cz=d with  d=ax0+by0+cz0
+    #   z = (d-ax-by)/c
+    #   here:
+    #   a, b, c describe the plane normal 
+    #   d can be calculated (from a, b, c, x0, y0, z0)
+    #   (x0, y0, z0) are the coordinates of a point on the 
+    #     plane, for example points_centroid
+    #   (x, y, z) are the coordinates of a query point on the plane 
+    # 
+    # points_npx3: (n_points, 3) 
+    # REMARK: this loss is not yet for batches!
+    # import pdb; pdb.set_trace()
+    # print('this loss is not yet for batches!')
+    assert (points_npx3.ndim == 2)
+    assert (points_npx3.shape[1] == 3)
+    points = torch.transpose(points_npx3, 0, 1)       # (3, n_points)
+    points_centroid = torch.mean(points, dim=1)
+    input_svd = points - points_centroid[:, None] 
+    U_svd, sigma_svd, V_svd = torch.svd(input_svd, compute_uv=True)
+    plane_normal = U_svd[:, 2]
+    plane_squaredsumofdists = sigma_svd[2]
+    error = plane_squaredsumofdists
+    return points_centroid, plane_normal, error
+
+
+def paws_to_groundplane_error(vertices, return_details=False):
+    # list of feet vertices (some of them)
+    #   remark: we did annotate left indices and find the right insices using sym_ids_dict
+    # REMARK: this loss is not yet for batches!
+    # import pdb; pdb.set_trace()
+    # print('this loss is not yet for batches!')
+    list_back_left = [1524, 1517, 1512, 1671, 1678, 1664, 1956, 1680, 1685, 1602, 1953, 1569]
+    list_front_left = [1331, 1327, 1332, 1764, 1767, 1747, 1779, 1789, 1944, 1339, 1323, 1420]
+    list_back_right = [3476, 3469, 3464, 3623, 3630, 3616, 3838, 3632, 3637, 3554, 3835, 3521]
+    list_front_right = [3283, 3279, 3284, 3715, 3718, 3698, 3730, 3740, 3826, 3291, 3275, 3372]
+    assert vertices.shape[0] == 3889
+    assert vertices.shape[1] == 3
+    all_paw_vert_idxs = list_back_left + list_front_left + list_back_right + list_front_right
+    verts_paws = vertices[all_paw_vert_idxs, :]
+    plane_centroid, plane_normal, error = fit_plane(verts_paws)
+    if return_details:
+        return plane_centroid, plane_normal, error
+    else:
+        return error
+
+def groundcontact_error(vertices, gclabels, return_details=False):
+    # import pdb; pdb.set_trace()
+    # REMARK: this loss is not yet for batches!
+    import pdb; pdb.set_trace()
+    print('this loss is not yet for batches!')
+    assert vertices.shape[0] == 3889
+    assert vertices.shape[1] == 3
+    verts_gc = vertices[gclabels, :]
+    plane_centroid, plane_normal, error = fit_plane(verts_gc)
+    if return_details:
+        return plane_centroid, plane_normal, error
+    else:
+        return error
+
+
+

src/combined_model/loss_utils/loss_utils_gc.py

@@ -0,0 +1,179 @@
+
+import torch
+
+
+
+
+
+class LossGConMesh(torch.nn.Module):
+    def __init__(self , n_verts=3889):
+        super(LossGConMesh, self).__init__()
+        self.n_verts = n_verts
+        self.criterion_class = torch.nn.CrossEntropyLoss(reduction='mean')
+
+    def forward(self, pred_gc, target_gc, has_gc, loss_type_gcmesh='ce'):
+        # pred_gc has shape (bs, n_verts, 2)
+        # target_gc has shape (bs, n_verts, 3)
+        #   with [first: no-contact=0 contact=1     
+        #        second: index of closest vertex with opposite label     
+        #        third: dist to that closest vertex]
+        target_gc_class = target_gc[:, :, 0]
+        target_gc_nearoppvert_ind = target_gc[:, :, 1]
+        target_gc_nearoppvert_dist = target_gc[:, :, 2]
+        # bs = pred_gc.shape[0]
+        bs = has_gc.sum()
+
+        if loss_type_gcmesh == 'ce':        # cross entropy
+            # import pdb; pdb.set_trace()
+
+            # classification_loss = self.criterion_class(pred_gc.reshape((bs*self.n_verts, 2)), target_gc_class.reshape((bs*self.n_verts)))
+            classification_loss = self.criterion_class(pred_gc[has_gc==True, ...].reshape((bs*self.n_verts, 2)), target_gc_class[has_gc==True, ...].reshape((bs*self.n_verts)))
+            loss = classification_loss
+        else:
+            raise ValueError
+
+        return loss
+
+
+
+
+
+
+def calculate_plane_errors_batch(vertices, target_gc_class, target_has_gc, has_gc_is_touching, return_error_under_plane=True):            
+    # remarks:
+    #   visualization of the plane: debug_code/curve_fitting_v2.py
+    #   theory: https://www.ltu.se/cms_fs/1.51590!/svd-fitting.pdf
+    #   remark: torch.svd is depreciated
+    # new plane equation:
+    #   a(x−x0)+b(y−y0)+c(z−z0)=0
+    #   ax+by+cz=d with  d=ax0+by0+cz0
+    #   z = (d-ax-by)/c
+    #   here:
+    #   a, b, c describe the plane normal 
+    #   d can be calculated (from a, b, c, x0, y0, z0)
+    #   (x0, y0, z0) are the coordinates of a point on the 
+    #     plane, for example points_centroid
+    #   (x, y, z) are the coordinates of a query point on the plane 
+    # 
+    # input:
+    #   vertices: (bs, 3889, 3)
+    #   target_gc_class: (bs, 3889)
+    # 
+    bs = vertices.shape[0]
+    error_list = []
+    error_under_plane_list = []
+
+    for ind_b in range(bs):
+        if target_has_gc[ind_b] == 1 and has_gc_is_touching[ind_b] == 1:
+            try:
+                points_npx3 = vertices[ind_b, target_gc_class[ind_b, :]==1, :]
+                points = torch.transpose(points_npx3, 0, 1)       # (3, n_points)
+                points_centroid = torch.mean(points, dim=1)
+                input_svd = points - points_centroid[:, None] 
+                # U_svd, sigma_svd, V_svd = torch.svd(input_svd, compute_uv=True)
+                # plane_normal = U_svd[:, 2]
+                # _, sigma_svd, _ = torch.svd(input_svd, compute_uv=False)
+                # _, sigma_svd, _ = torch.svd(input_svd, compute_uv=True)
+                U_svd, sigma_svd, V_svd = torch.svd(input_svd, compute_uv=True)
+                plane_squaredsumofdists = sigma_svd[2]
+                error_list.append(plane_squaredsumofdists)
+
+                if return_error_under_plane:
+                    # plane information
+                    # plane_centroid = points_centroid
+                    plane_normal = U_svd[:, 2]
+
+                    # non-plane points
+                    nonplane_points_npx3 = vertices[ind_b, target_gc_class[ind_b, :]==0, :]     # (n_points_3)
+                    nonplane_points = torch.transpose(nonplane_points_npx3, 0, 1)       # (3, n_points)
+                    nonplane_points_centered = nonplane_points - points_centroid[:, None] 
+
+                    nonplane_points_projected = torch.matmul(plane_normal[None, :], nonplane_points_centered)   # plane normal already has length 1
+                    
+                    if nonplane_points_projected.sum() > 0:
+                        # bug corrected 07.11.22
+                        # error_under_plane = nonplane_points_projected[nonplane_points_projected<0].sum() / 100
+                        error_under_plane = - nonplane_points_projected[nonplane_points_projected<0].sum() / 100
+                    else:
+                        error_under_plane = nonplane_points_projected[nonplane_points_projected>0].sum() / 100
+                    error_under_plane_list.append(error_under_plane)
+            except:
+                print('was not able to calculate plane error for this image')
+                error_list.append(torch.zeros((1), dtype=vertices.dtype, device=vertices.device)[0])
+                error_under_plane_list.append(torch.zeros((1), dtype=vertices.dtype, device=vertices.device)[0])           
+        else:
+            error_list.append(torch.zeros((1), dtype=vertices.dtype, device=vertices.device)[0])
+            error_under_plane_list.append(torch.zeros((1), dtype=vertices.dtype, device=vertices.device)[0])
+    errors = torch.stack(error_list, dim=0)
+    errors_under_plane = torch.stack(error_under_plane_list, dim=0)
+
+    if return_error_under_plane:
+        return errors, errors_under_plane
+    else:
+        return errors
+
+
+
+# def calculate_vertex_wise_labeling_error():
+    # vertexwise_ground_contact
+
+
+
+
+
+
+
+
+
+
+
+
+
+'''
+
+def paws_to_groundplane_error_batch(vertices, return_details=False):
+    # list of feet vertices (some of them)
+    #   remark: we did annotate left indices and find the right insices using sym_ids_dict
+    # REMARK: this loss is not yet for batches!
+    import pdb; pdb.set_trace()
+    print('this loss is not yet for batches!')
+    list_back_left = [1524, 1517, 1512, 1671, 1678, 1664, 1956, 1680, 1685, 1602, 1953, 1569]
+    list_front_left = [1331, 1327, 1332, 1764, 1767, 1747, 1779, 1789, 1944, 1339, 1323, 1420]
+    list_back_right = [3476, 3469, 3464, 3623, 3630, 3616, 3838, 3632, 3637, 3554, 3835, 3521]
+    list_front_right = [3283, 3279, 3284, 3715, 3718, 3698, 3730, 3740, 3826, 3291, 3275, 3372]
+    assert vertices.shape[0] == 3889
+    assert vertices.shape[1] == 3
+    all_paw_vert_idxs = list_back_left + list_front_left + list_back_right + list_front_right
+    verts_paws = vertices[all_paw_vert_idxs, :]
+    plane_centroid, plane_normal, error = fit_plane_batch(verts_paws)
+    if return_details:
+        return plane_centroid, plane_normal, error
+    else:
+        return error
+
+def paws_to_groundplane_error_batch_new(vertices, return_details=False):
+    # list of feet vertices (some of them)
+    #   remark: we did annotate left indices and find the right insices using sym_ids_dict
+    # REMARK: this loss is not yet for batches!
+    import pdb; pdb.set_trace()
+    print('this loss is not yet for batches!')
+    list_back_left = [1524, 1517, 1512, 1671, 1678, 1664, 1956, 1680, 1685, 1602, 1953, 1569]
+    list_front_left = [1331, 1327, 1332, 1764, 1767, 1747, 1779, 1789, 1944, 1339, 1323, 1420]
+    list_back_right = [3476, 3469, 3464, 3623, 3630, 3616, 3838, 3632, 3637, 3554, 3835, 3521]
+    list_front_right = [3283, 3279, 3284, 3715, 3718, 3698, 3730, 3740, 3826, 3291, 3275, 3372]
+    assert vertices.shape[0] == 3889
+    assert vertices.shape[1] == 3
+    all_paw_vert_idxs = list_back_left + list_front_left + list_back_right + list_front_right
+    verts_paws = vertices[all_paw_vert_idxs, :]
+    plane_centroid, plane_normal, error = fit_plane_batch(verts_paws)
+    print('this loss is not yet for batches!')
+    points = torch.transpose(points_npx3, 0, 1)       # (3, n_points)
+    points_centroid = torch.mean(points, dim=1)
+    input_svd = points - points_centroid[:, None] 
+    U_svd, sigma_svd, V_svd = torch.svd(input_svd, compute_uv=True)
+    plane_normal = U_svd[:, 2]
+    plane_squaredsumofdists = sigma_svd[2]
+    error = plane_squaredsumofdists
+    print('error: ' + str(error.item()))
+    return error
+'''

src/combined_model/model_shape_v7_withref_withgraphcnn.py

@@ -0,0 +1,927 @@
+
+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 = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/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
+

src/combined_model/train_main_image_to_3d_wbr_withref.py

@@ -0,0 +1,955 @@
+
+import torch
+import torch.nn as nn
+import torch.backends.cudnn
+import torch.nn.parallel
+from tqdm import tqdm
+import os
+import pathlib
+from matplotlib import pyplot as plt
+import cv2
+import numpy as np
+import torch
+import trimesh
+import pickle as pkl
+import csv
+from scipy.spatial.transform import Rotation as R_sc
+
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
+from stacked_hourglass.utils.evaluation import accuracy, AverageMeter, final_preds, get_preds, get_preds_soft
+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, SMAL_KEYPOINT_NAMES_FOR_3D_EVAL, SMAL_KEYPOINT_INDICES_FOR_3D_EVAL, SMAL_KEYPOINT_WHICHTOUSE_FOR_3D_EVAL
+from combined_model.helper import eval_save_visualizations_and_meshes, eval_prepare_pck_and_iou, eval_add_preds_to_summary
+
+from smal_pytorch.smal_model.smal_torch_new import SMAL     # for gc visualization
+from src.combined_model.loss_utils.loss_utils import fit_plane
+# from src.evaluation.sketchfab_evaluation.alignment_utils.calculate_v2v_error_release import compute_similarity_transform
+# from src.evaluation.sketchfab_evaluation.alignment_utils.calculate_alignment_error import calculate_alignemnt_errors
+
+# ---------------------------------------------------------------------------------------------------------------------------
+def do_training_epoch(train_loader, model, loss_module, loss_module_ref, device, data_info, optimiser, quiet=False, acc_joints=None, weight_dict=None, weight_dict_ref=None):
+    losses = AverageMeter()
+    losses_keyp = AverageMeter()
+    losses_silh = AverageMeter()
+    losses_shape = AverageMeter()
+    losses_pose = AverageMeter()
+    losses_class = AverageMeter()
+    losses_breed = AverageMeter()
+    losses_partseg = AverageMeter()
+    losses_ref_keyp = AverageMeter()
+    losses_ref_silh = AverageMeter()
+    losses_ref_pose = AverageMeter()
+    losses_ref_reg = AverageMeter()
+    accuracies = AverageMeter()
+    # Put the model in training mode.
+    model.train()
+    # prepare progress bar
+    iterable = enumerate(train_loader)
+    progress = None
+    if not quiet:
+        progress = tqdm(iterable, desc='Train', total=len(train_loader), ascii=True, leave=False)
+        iterable = progress
+    # information for normalization 
+    norm_dict = {
+        'pose_rot6d_mean': torch.from_numpy(data_info.pose_rot6d_mean).float().to(device),
+        'trans_mean': torch.from_numpy(data_info.trans_mean).float().to(device),
+        'trans_std': torch.from_numpy(data_info.trans_std).float().to(device),
+        'flength_mean': torch.from_numpy(data_info.flength_mean).float().to(device),
+        'flength_std': torch.from_numpy(data_info.flength_std).float().to(device)}
+    # prepare variables, put them on the right device
+    for i, (input, target_dict) in iterable:
+        batch_size = input.shape[0]
+        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 in ['has_seg', 'gc']:
+                target_dict[key] = target_dict[key].to(device)
+            else:
+                pass
+        input = input.float().to(device)
+
+        # ----------------------- do training step -----------------------
+        assert model.training, 'model must be in training mode.'
+        with torch.enable_grad():
+            # ----- forward pass -----  
+            output, output_unnorm, output_reproj, output_ref, output_ref_comp = model(input, norm_dict=norm_dict)        
+            # ----- loss -----
+            # --- from main network
+            loss, loss_dict = loss_module(output_reproj=output_reproj, 
+                target_dict=target_dict, 
+                weight_dict=weight_dict)
+            # ---from refinement network
+            loss_ref, loss_dict_ref = loss_module_ref(output_ref=output_ref, 
+                output_ref_comp=output_ref_comp,
+                target_dict=target_dict, 
+                weight_dict_ref=weight_dict_ref)
+            loss_total = loss + loss_ref
+            # ----- backward pass and parameter update -----
+            optimiser.zero_grad()
+            loss_total.backward()
+            optimiser.step()
+        # ----------------------------------------------------------------
+
+        # prepare losses for progress bar
+        bs_fake = 1     # batch_size
+        losses.update(loss_dict['loss'] + loss_dict_ref['loss'], bs_fake)
+        losses_keyp.update(loss_dict['loss_keyp_weighted'], bs_fake)
+        losses_silh.update(loss_dict['loss_silh_weighted'], bs_fake)
+        losses_shape.update(loss_dict['loss_shape_weighted'], bs_fake)
+        losses_pose.update(loss_dict['loss_poseprior_weighted'], bs_fake)   
+        losses_class.update(loss_dict['loss_class_weighted'], bs_fake)
+        losses_breed.update(loss_dict['loss_breed_weighted'], bs_fake)
+        losses_partseg.update(loss_dict['loss_partseg_weighted'], bs_fake)
+        losses_ref_keyp.update(loss_dict_ref['keyp_ref'], bs_fake)
+        losses_ref_silh.update(loss_dict_ref['silh_ref'], bs_fake)
+        loss_ref_pose = 0
+        for l_name in ['pose_legs_side', 'pose_legs_tors', 'pose_tail_side', 'pose_tail_tors', 'pose_spine_side', 'pose_spine_tors']:
+            if l_name in loss_dict_ref.keys():
+                loss_ref_pose += loss_dict_ref[l_name]
+        losses_ref_pose.update(loss_ref_pose, bs_fake)
+        loss_ref_reg = 0
+        for l_name in ['reg_trans', 'reg_flength', 'reg_pose']:
+            if l_name in loss_dict_ref.keys():
+                loss_ref_reg += loss_dict_ref[l_name]
+        losses_ref_reg.update(loss_ref_reg, bs_fake)
+        acc = - loss_dict['loss_keyp_weighted']     # this will be used to keep track of the 'best model'
+        accuracies.update(acc, bs_fake)
+        # Show losses as part of the progress bar.
+        if progress is not None:
+            my_string = 'Loss: {loss:0.4f}, loss_keyp: {loss_keyp:0.4f}, loss_silh: {loss_silh:0.4f}, loss_partseg: {loss_partseg:0.4f}, loss_shape: {loss_shape:0.4f}, loss_pose: {loss_pose:0.4f}, loss_class: {loss_class:0.4f}, loss_breed: {loss_breed:0.4f}, loss_ref_keyp: {loss_ref_keyp:0.4f}, loss_ref_silh: {loss_ref_silh:0.4f}, loss_ref_pose: {loss_ref_pose:0.4f}, loss_ref_reg: {loss_ref_reg:0.4f}'.format(
+                loss=losses.avg,
+                loss_keyp=losses_keyp.avg,
+                loss_silh=losses_silh.avg,
+                loss_shape=losses_shape.avg,
+                loss_pose=losses_pose.avg,
+                loss_class=losses_class.avg,
+                loss_breed=losses_breed.avg,
+                loss_partseg=losses_partseg.avg,
+                loss_ref_keyp=losses_ref_keyp.avg,
+                loss_ref_silh=losses_ref_silh.avg,
+                loss_ref_pose=losses_ref_pose.avg,
+                loss_ref_reg=losses_ref_reg.avg)
+            my_string_short = 'Loss: {loss:0.4f}, loss_keyp: {loss_keyp:0.4f}, loss_silh: {loss_silh:0.4f}, loss_ref_keyp: {loss_ref_keyp:0.4f}, loss_ref_silh: {loss_ref_silh:0.4f}, loss_ref_pose: {loss_ref_pose:0.4f}, loss_ref_reg: {loss_ref_reg:0.4f}'.format(
+                loss=losses.avg,
+                loss_keyp=losses_keyp.avg,
+                loss_silh=losses_silh.avg,
+                loss_ref_keyp=losses_ref_keyp.avg,
+                loss_ref_silh=losses_ref_silh.avg,
+                loss_ref_pose=losses_ref_pose.avg,
+                loss_ref_reg=losses_ref_reg.avg)
+            progress.set_postfix_str(my_string_short)
+
+    return my_string, accuracies.avg       
+
+
+# ---------------------------------------------------------------------------------------------------------------------------
+def do_validation_epoch(val_loader, model, loss_module, loss_module_ref, device, data_info, flip=False, quiet=False, acc_joints=None, save_imgs_path=None, weight_dict=None, weight_dict_ref=None, metrics=None, val_opt='default', test_name_list=None, render_all=False, pck_thresh=0.15, len_dataset=None):
+    losses = AverageMeter()
+    losses_keyp = AverageMeter()
+    losses_silh = AverageMeter()
+    losses_shape = AverageMeter()
+    losses_pose = AverageMeter()
+    losses_class = AverageMeter()
+    losses_breed = AverageMeter()
+    losses_partseg = AverageMeter()
+    losses_ref_keyp = AverageMeter()
+    losses_ref_silh = AverageMeter()
+    losses_ref_pose = AverageMeter()
+    losses_ref_reg = AverageMeter()
+    accuracies = AverageMeter()
+    if save_imgs_path is not None:
+        pathlib.Path(save_imgs_path).mkdir(parents=True, exist_ok=True) 
+    # Put the model in evaluation mode.
+    model.eval()
+    # prepare progress bar
+    iterable = enumerate(val_loader)
+    progress = None
+    if not quiet:
+        progress = tqdm(iterable, desc='Valid', total=len(val_loader), ascii=True, leave=False)
+        iterable = progress
+    # summarize information for normalization 
+    norm_dict = {
+        'pose_rot6d_mean': torch.from_numpy(data_info.pose_rot6d_mean).float().to(device),
+        'trans_mean': torch.from_numpy(data_info.trans_mean).float().to(device),
+        'trans_std': torch.from_numpy(data_info.trans_std).float().to(device),
+        'flength_mean': torch.from_numpy(data_info.flength_mean).float().to(device),
+        'flength_std': torch.from_numpy(data_info.flength_std).float().to(device)}
+    batch_size = val_loader.batch_size
+
+    return_mesh_with_gt_groundplane = True
+    if return_mesh_with_gt_groundplane:
+        remeshing_path = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/smal_data_remeshed/uniform_surface_sampling/my_smpl_39dogsnorm_Jr_4_dog_remesh4000_info.pkl'
+        with open(remeshing_path, 'rb') as fp: 
+            remeshing_dict = pkl.load(fp)
+        remeshing_relevant_faces = torch.tensor(remeshing_dict['smal_faces'][remeshing_dict['faceid_closest']], dtype=torch.long, device=device)
+        remeshing_relevant_barys = torch.tensor(remeshing_dict['barys_closest'], dtype=torch.float32, device=device)
+
+
+    # from smal_pytorch.smal_model.smal_torch_new import SMAL
+    print('start: load smal default model (barc), but only for vertices')
+    smal = SMAL()
+    print('end: load smal default model (barc), but only for vertices')
+    smal_template_verts = smal.v_template.detach().cpu().numpy()
+    smal_faces = smal.faces.detach().cpu().numpy()
+
+    
+    my_step = 0
+    for index, (input, target_dict) in iterable:
+
+        # prepare variables, put them on the right device
+        curr_batch_size = input.shape[0]
+        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 in ['has_seg', 'gc']:
+                target_dict[key] = target_dict[key].to(device)
+            else:
+                pass
+        input = input.float().to(device)
+
+        # ----------------------- do validation step -----------------------
+        with torch.no_grad():
+            # ----- forward pass -----  
+            # output: (['pose', 'flength', 'trans', 'keypoints_norm', 'keypoints_scores'])
+            # output_unnorm: (['pose_rotmat', 'flength', 'trans', 'keypoints'])
+            # output_reproj: (['vertices_smal', 'torch_meshes', 'keyp_3d', 'keyp_2d', 'silh', 'betas', 'pose_rot6d', 'dog_breed', 'shapedirs', 'z', 'flength_unnorm', 'flength'])
+            # target_dict: (['index', 'center', 'scale', 'pts', 'tpts', 'target_weight', 'breed_index', 'sim_breed_index', 'ind_dataset', 'silh'])
+            output, output_unnorm, output_reproj, output_ref, output_ref_comp = model(input, norm_dict=norm_dict)        
+            # ----- loss -----
+            if metrics == 'no_loss':
+                # --- from main network
+                loss, loss_dict = loss_module(output_reproj=output_reproj, 
+                    target_dict=target_dict, 
+                    weight_dict=weight_dict)
+                # ---from refinement network
+                loss_ref, loss_dict_ref = loss_module_ref(output_ref=output_ref, 
+                    output_ref_comp=output_ref_comp,
+                    target_dict=target_dict, 
+                    weight_dict_ref=weight_dict_ref)
+                loss_total = loss + loss_ref
+
+        # ----------------------------------------------------------------
+
+
+        for result_network in ['normal', 'ref']:
+            # variabled that are not refined
+            hg_keyp_norm = output['keypoints_norm']
+            hg_keyp_scores = output['keypoints_scores']
+            betas = output_reproj['betas']
+            betas_limbs = output_reproj['betas_limbs']
+            zz = output_reproj['z']
+            if result_network == 'normal':
+                # STEP 1: normal network
+                vertices_smal = output_reproj['vertices_smal']
+                flength = output_unnorm['flength']
+                pose_rotmat = output_unnorm['pose_rotmat']
+                trans = output_unnorm['trans']
+                pred_keyp = output_reproj['keyp_2d']
+                pred_silh = output_reproj['silh']
+                prefix = 'normal_'
+            else:
+                # STEP 1: refinement network
+                vertices_smal = output_ref['vertices_smal']
+                flength = output_ref['flength']
+                pose_rotmat = output_ref['pose_rotmat']
+                trans = output_ref['trans']
+                pred_keyp = output_ref['keyp_2d']
+                pred_silh = output_ref['silh']
+                prefix = 'ref_'
+                if return_mesh_with_gt_groundplane and 'gc' in target_dict.keys():
+                    bs = vertices_smal.shape[0]
+                    target_gc_class = target_dict['gc'][:, :, 0]
+                    sel_verts = torch.index_select(output_ref['vertices_smal'], dim=1, index=remeshing_relevant_faces.reshape((-1))).reshape((bs, remeshing_relevant_faces.shape[0], 3, 3))
+                    verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, sel_verts)
+                    target_gc_class_remeshed = torch.einsum('ij,aij->ai', remeshing_relevant_barys, target_gc_class[:, remeshing_relevant_faces].to(device=device, dtype=torch.float32))
+                    target_gc_class_remeshed_prep = torch.round(target_gc_class_remeshed).to(torch.long)
+
+
+
+
+
+                # import pdb; pdb.set_trace()
+
+                # new for vertex wise ground contact
+                if (not model.graphcnn_type == 'inexistent') and (save_imgs_path is not None):
+                    # import pdb; pdb.set_trace()
+
+                    sm = torch.nn.Softmax(dim=2)
+                    ground_contact_probs = sm(output_ref['vertexwise_ground_contact'])
+
+                    for ind_img in range(ground_contact_probs.shape[0]):
+                        # ind_img = 0
+                        if test_name_list is not None:
+                            img_name = test_name_list[int(target_dict['index'][ind_img].cpu().detach().numpy())].replace('/', '_')
+                            img_name = img_name.split('.')[0]
+                        else:
+                            img_name = str(index) + '_' + str(ind_img)
+                        out_path_gcmesh = save_imgs_path + '/' + prefix + 'gcmesh_' + img_name + '.obj'
+
+                        gc_prob = ground_contact_probs[ind_img, :, 1]   # contact probability
+                        vert_colors = np.repeat(255*gc_prob.detach().cpu().numpy()[:, None], 3, 1)
+                        my_mesh = trimesh.Trimesh(vertices=smal_template_verts, faces=smal_faces, process=False,  maintain_order=True)
+                        my_mesh.visual.vertex_colors = vert_colors
+                        save_gc_mesh = True # False
+                        if save_gc_mesh:
+                            my_mesh.export(out_path_gcmesh)
+            
+                        '''
+                        input_image = input[ind_img, :, :, :].detach().clone()
+                        for t, m, s in zip(input_image, data_info.rgb_mean,data_info.rgb_stddev): t.add_(m)
+                        input_image_np = input_image.detach().cpu().numpy().transpose(1, 2, 0) 
+                        out_path = save_debug_path + 'b' + str(ind_img) +'_input.png'
+                        plt.imsave(out_path, input_image_np)
+                        '''
+
+                        # -------------------------------------
+
+                        # import pdb; pdb.set_trace()
+
+
+                        '''
+                        target_gc_class = target_dict['gc'][ind_img, :, 0]
+
+                        current_vertices_smal = vertices_smal[ind_img, :, :]
+
+                        points_centroid, plane_normal, error = fit_plane(current_vertices_smal[target_gc_class==1, :])
+                        '''
+
+                        # calculate ground plane
+                        #   (see /is/cluster/work/nrueegg/icon_pifu_related/ICON/debug_code/curve_fitting_v2.py)
+                        if return_mesh_with_gt_groundplane and 'gc' in target_dict.keys():
+
+                            current_verts_remeshed = verts_remeshed[ind_img, :, :]
+                            current_target_gc_class_remeshed_prep = target_gc_class_remeshed_prep[ind_img, ...]
+
+                            if current_target_gc_class_remeshed_prep.sum() > 3:
+                                points_on_plane = current_verts_remeshed[current_target_gc_class_remeshed_prep==1, :]
+                                data_centroid, plane_normal, error = fit_plane(points_on_plane)
+                                nonplane_points_centered = current_verts_remeshed[current_target_gc_class_remeshed_prep==0, :] - data_centroid[None, :]
+                                nonplane_points_projected = torch.matmul(plane_normal[None, :], nonplane_points_centered.transpose(0,1))
+
+                                if nonplane_points_projected.sum() > 0: # plane normal points towards the animal
+                                    plane_normal = plane_normal.detach().cpu().numpy()
+                                else:
+                                    plane_normal = - plane_normal.detach().cpu().numpy()
+                                data_centroid = data_centroid.detach().cpu().numpy()
+
+
+
+                                # import pdb; pdb.set_trace()
+
+
+                                desired_plane_normal_vector = np.asarray([[0, -1, 0]])
+                                # new approach: use cross product
+                                rotation_axis = np.cross(plane_normal, desired_plane_normal_vector)     #  np.cross(plane_normal, desired_plane_normal_vector)
+                                lengt_rotation_axis = np.linalg.norm(rotation_axis)     # = sin(alpha)      (because vectors have unit length)
+                                angle = np.sin(lengt_rotation_axis)
+                                rot = R_sc.from_rotvec(angle * rotation_axis * 1/lengt_rotation_axis)
+                                rot_mat = rot[0].as_matrix()
+                                rot_upsidedown = R_sc.from_rotvec(np.pi * np.asarray([[1, 0, 0]]))
+                                # rot_upsidedown[0].apply(rot[0].apply(plane_normal))
+                                current_vertices_smal = vertices_smal[ind_img, :, :].detach().cpu().numpy()
+                                new_smal_vertices = rot_upsidedown[0].apply(rot[0].apply(current_vertices_smal - data_centroid[None, :]))
+                                my_mesh = trimesh.Trimesh(vertices=new_smal_vertices, faces=smal_faces, process=False,  maintain_order=True)
+                                vert_colors[:, 2] = 255
+                                my_mesh.visual.vertex_colors = vert_colors
+                                out_path_gc_rotated = save_imgs_path + '/' + prefix + 'gc_rotated_' + img_name + '_new.obj'
+                                my_mesh.export(out_path_gc_rotated)
+
+
+
+
+
+
+                                '''# rot = R_sc.align_vectors(plane_normal.reshape((1, -1)), desired_plane_normal_vector)
+                                desired_plane_normal_vector = np.asarray([[0, 1, 0]])
+
+                                rot = R_sc.align_vectors(desired_plane_normal_vector, plane_normal.reshape((1, -1)))        # inv
+                                rot_mat = rot[0].as_matrix()
+
+
+                                current_vertices_smal = vertices_smal[ind_img, :, :].detach().cpu().numpy()
+                                new_smal_vertices = rot[0].apply((current_vertices_smal - data_centroid[None, :]))
+
+                                my_mesh = trimesh.Trimesh(vertices=new_smal_vertices, faces=smal_faces, process=False,  maintain_order=True)
+                                my_mesh.visual.vertex_colors = vert_colors
+                                out_path_gc_rotated = save_imgs_path + '/' + prefix + 'gc_rotated_' + img_name + '_y.obj'
+                                my_mesh.export(out_path_gc_rotated)
+                                '''
+
+
+
+
+
+
+
+
+
+                        # ----
+
+
+                        # -------------------------------------
+
+
+
+
+            if index == 0:
+                if len_dataset is None:
+                    len_data = val_loader.batch_size * len(val_loader)  # 1703
+                else:
+                    len_data = len_dataset
+                if metrics == 'all' or metrics == 'no_loss':
+                    if result_network == 'normal':
+                        summaries = {'normal': dict(), 'ref': dict()}
+                        summary = summaries['normal']
+                    else:
+                        summary = summaries['ref']
+                    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)
+                    summary['betas'] = np.zeros((len_data,betas.shape[1]))
+                    summary['betas_limbs'] = np.zeros((len_data, betas_limbs.shape[1]))
+                    summary['z'] = np.zeros((len_data, zz.shape[1]))
+                    summary['pose_rotmat'] = np.zeros((len_data, pose_rotmat.shape[1], 3, 3))
+                    summary['flength'] = np.zeros((len_data, flength.shape[1]))
+                    summary['trans'] = np.zeros((len_data, trans.shape[1]))
+                    summary['breed_indices'] = np.zeros((len_data))
+                    summary['image_names'] = []        # len_data * [None]
+            else:
+                if result_network == 'normal':
+                    summary = summaries['normal']
+                else:
+                    summary = summaries['ref']
+
+            if save_imgs_path is not None:
+                eval_save_visualizations_and_meshes(model, input, data_info, target_dict, test_name_list, vertices_smal, hg_keyp_norm, hg_keyp_scores, zz, betas, betas_limbs, pose_rotmat, trans, flength, pred_keyp, pred_silh, save_imgs_path, prefix, index, render_all=render_all)
+
+            if metrics == 'all' or metrics == 'no_loss':
+                preds = eval_prepare_pck_and_iou(model, input, data_info, target_dict, test_name_list, vertices_smal, hg_keyp_norm, hg_keyp_scores, zz, betas, betas_limbs, pose_rotmat, trans, flength, pred_keyp, pred_silh, save_imgs_path, prefix, index, pck_thresh, progress=progress)
+                # add results for all images in this batch to lists
+                curr_batch_size = pred_keyp.shape[0]
+                eval_add_preds_to_summary(summary, preds, my_step, batch_size, curr_batch_size)
+            else:
+                # measure accuracy and record loss
+                bs_fake = 1     # batch_size
+        # import pdb; pdb.set_trace()
+
+
+        # save_imgs_path + '/' + prefix + 'rot_tex_pred_' + img_name + '.png'
+        # import pdb; pdb.set_trace()
+        '''
+        for ind_img in range(len(target_dict['index'])): 
+            try: 
+                if test_name_list is not None:
+                    img_name = test_name_list[int(target_dict['index'][ind_img].cpu().detach().numpy())].replace('/', '_')
+                    img_name = img_name.split('.')[0]
+                else:
+                    img_name = str(index) + '_' + str(ind_img)
+                all_image_names = ['keypoints_pred_' + img_name + '.png',  'normal_comp_pred_' + img_name + '.png', 'normal_rot_tex_pred_' + img_name + '.png',  'ref_comp_pred_' + img_name + '.png', 'ref_rot_tex_pred_' + img_name + '.png']
+                all_saved_images = []
+                for sub_img_name in all_image_names:
+                    saved_img = cv2.imread(save_imgs_path + '/' + sub_img_name)
+                    if not (saved_img.shape[0] == 256 and saved_img.shape[1] == 256):
+                        saved_img = cv2.resize(saved_img, (256, 256)) 
+                    all_saved_images.append(saved_img)
+                final_image = np.concatenate(all_saved_images, axis=1)
+                save_imgs_path_sum = save_imgs_path.replace('test_', 'summary_test_')
+                if not os.path.exists(save_imgs_path_sum): os.makedirs(save_imgs_path_sum)
+                final_image_path = save_imgs_path_sum +  '/summary_' + img_name + '.png'
+                cv2.imwrite(final_image_path, final_image)
+            except: 
+                print('dont save a summary image')
+        '''
+                
+        
+        bs_fake = 1
+        if metrics == 'all' or metrics == 'no_loss':
+            # update progress bar
+            if progress is not None:
+                '''my_string = "PCK: {0:.2f}, IOU: {1:.2f}".format(
+                    pck[:(my_step * batch_size + curr_batch_size)].mean(),
+                    acc_sil_2d[:(my_step * batch_size + curr_batch_size)].mean())'''
+                my_string = "normal_PCK: {0:.2f}, normal_IOU: {1:.2f}, ref_PCK: {2:.2f}, ref_IOU: {3:.2f}".format(
+                    summaries['normal']['pck'][:(my_step * batch_size + curr_batch_size)].mean(),
+                    summaries['normal']['acc_sil_2d'][:(my_step * batch_size + curr_batch_size)].mean(),
+                    summaries['ref']['pck'][:(my_step * batch_size + curr_batch_size)].mean(),
+                    summaries['ref']['acc_sil_2d'][:(my_step * batch_size + curr_batch_size)].mean())
+                progress.set_postfix_str(my_string)
+        else:
+            losses.update(loss_dict['loss'] + loss_dict_ref['loss'], bs_fake)
+            losses_keyp.update(loss_dict['loss_keyp_weighted'], bs_fake)
+            losses_silh.update(loss_dict['loss_silh_weighted'], bs_fake)
+            losses_shape.update(loss_dict['loss_shape_weighted'], bs_fake)
+            losses_pose.update(loss_dict['loss_poseprior_weighted'], bs_fake)   
+            losses_class.update(loss_dict['loss_class_weighted'], bs_fake)
+            losses_breed.update(loss_dict['loss_breed_weighted'], bs_fake)
+            losses_partseg.update(loss_dict['loss_partseg_weighted'], bs_fake)
+            losses_ref_keyp.update(loss_dict_ref['keyp_ref'], bs_fake)
+            losses_ref_silh.update(loss_dict_ref['silh_ref'], bs_fake)
+            loss_ref_pose = 0
+            for l_name in ['pose_legs_side', 'pose_legs_tors', 'pose_tail_side', 'pose_tail_tors', 'pose_spine_side', 'pose_spine_tors']: 
+                loss_ref_pose += loss_dict_ref[l_name]
+            losses_ref_pose.update(loss_ref_pose, bs_fake)
+            loss_ref_reg = 0
+            for l_name in ['reg_trans', 'reg_flength', 'reg_pose']:
+                loss_ref_reg += loss_dict_ref[l_name]
+            losses_ref_reg.update(loss_ref_reg, bs_fake)
+            acc = - loss_dict['loss_keyp_weighted']     # this will be used to keep track of the 'best model'
+            accuracies.update(acc, bs_fake)
+            # Show losses as part of the progress bar.
+            if progress is not None:
+                my_string = 'Loss: {loss:0.4f}, loss_keyp: {loss_keyp:0.4f}, loss_silh: {loss_silh:0.4f}, loss_partseg: {loss_partseg:0.4f}, loss_shape: {loss_shape:0.4f}, loss_pose: {loss_pose:0.4f}, loss_class: {loss_class:0.4f}, loss_breed: {loss_breed:0.4f}, loss_ref_keyp: {loss_ref_keyp:0.4f}, loss_ref_silh: {loss_ref_silh:0.4f}, loss_ref_pose: {loss_ref_pose:0.4f}, loss_ref_reg: {loss_ref_reg:0.4f}'.format(
+                    loss=losses.avg,
+                    loss_keyp=losses_keyp.avg,
+                    loss_silh=losses_silh.avg,
+                    loss_shape=losses_shape.avg,
+                    loss_pose=losses_pose.avg,
+                    loss_class=losses_class.avg,
+                    loss_breed=losses_breed.avg,
+                    loss_partseg=losses_partseg.avg,
+                    loss_ref_keyp=losses_ref_keyp.avg,
+                    loss_ref_silh=losses_ref_silh.avg,
+                    loss_ref_pose=losses_ref_pose.avg,
+                    loss_ref_reg=losses_ref_reg.avg)
+                my_string_short = 'Loss: {loss:0.4f}, loss_keyp: {loss_keyp:0.4f}, loss_silh: {loss_silh:0.4f}, loss_ref_keyp: {loss_ref_keyp:0.4f}, loss_ref_silh: {loss_ref_silh:0.4f}, loss_ref_pose: {loss_ref_pose:0.4f}, loss_ref_reg: {loss_ref_reg:0.4f}'.format(
+                    loss=losses.avg,
+                    loss_keyp=losses_keyp.avg,
+                    loss_silh=losses_silh.avg,
+                    loss_ref_keyp=losses_ref_keyp.avg,
+                    loss_ref_silh=losses_ref_silh.avg,
+                    loss_ref_pose=losses_ref_pose.avg,
+                    loss_ref_reg=losses_ref_reg.avg)
+                progress.set_postfix_str(my_string_short)
+        my_step += 1
+    if metrics == 'all':
+        return my_string, summaries     # summary    
+    elif metrics == 'no_loss':
+        return my_string, np.average(np.asarray(summaries['ref']['acc_sil_2d']))     # np.average(np.asarray(summary['acc_sil_2d']))
+    else:
+        return my_string, accuracies.avg       
+
+
+# ---------------------------------------------------------------------------------------------------------------------------
+def do_visual_epoch(val_loader, model, device, data_info, flip=False, quiet=False, acc_joints=None, save_imgs_path=None, weight_dict=None, weight_dict_ref=None, metrics=None, val_opt='default', test_name_list=None, render_all=False, pck_thresh=0.15, return_results=False, len_dataset=None):
+    if save_imgs_path is not None:
+        pathlib.Path(save_imgs_path).mkdir(parents=True, exist_ok=True) 
+    all_results = []
+
+    # Put the model in evaluation mode.
+    model.eval()
+
+    iterable = enumerate(val_loader)
+
+    # information for normalization 
+    norm_dict = {
+        'pose_rot6d_mean': torch.from_numpy(data_info.pose_rot6d_mean).float().to(device),
+        'trans_mean': torch.from_numpy(data_info.trans_mean).float().to(device),
+        'trans_std': torch.from_numpy(data_info.trans_std).float().to(device),
+        'flength_mean': torch.from_numpy(data_info.flength_mean).float().to(device),
+        'flength_std': torch.from_numpy(data_info.flength_std).float().to(device)}
+
+
+    return_mesh_with_gt_groundplane = True
+    if return_mesh_with_gt_groundplane:
+        remeshing_path = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/smal_data_remeshed/uniform_surface_sampling/my_smpl_39dogsnorm_Jr_4_dog_remesh4000_info.pkl'
+        with open(remeshing_path, 'rb') as fp: 
+            remeshing_dict = pkl.load(fp)
+        remeshing_relevant_faces = torch.tensor(remeshing_dict['smal_faces'][remeshing_dict['faceid_closest']], dtype=torch.long, device=device)
+        remeshing_relevant_barys = torch.tensor(remeshing_dict['barys_closest'], dtype=torch.float32, device=device)
+
+    # from smal_pytorch.smal_model.smal_torch_new import SMAL
+    print('start: load smal default model (barc), but only for vertices')
+    smal = SMAL()
+    print('end: load smal default model (barc), but only for vertices')
+    smal_template_verts = smal.v_template.detach().cpu().numpy()
+    smal_faces = smal.faces.detach().cpu().numpy()
+
+    file_alignment_errors = open(save_imgs_path + '/a_ref_procrustes_alignmnet_errors.txt', 'a') # append mode
+    file_alignment_errors.write(" -----------  start evaluation  ------------- \n ")
+
+    csv_file_alignment_errors = open(save_imgs_path + '/a_ref_procrustes_alignmnet_errors.csv', 'w') # write mode
+    fieldnames = ['name', 'error']
+    writer = csv.DictWriter(csv_file_alignment_errors, fieldnames=fieldnames)
+    writer.writeheader()
+
+    my_step = 0
+    for index, (input, target_dict) in iterable:
+        batch_size = input.shape[0]
+        input = input.float().to(device)
+        partial_results = {}
+
+        # ----------------------- do visualization step -----------------------
+        with torch.no_grad():
+            output, output_unnorm, output_reproj, output_ref, output_ref_comp = model(input, norm_dict=norm_dict)        
+
+
+        # import pdb; pdb.set_trace()
+
+
+        sm = torch.nn.Softmax(dim=2)
+        ground_contact_probs = sm(output_ref['vertexwise_ground_contact'])
+
+        for result_network in ['normal', 'ref']:
+            # variabled that are not refined
+            hg_keyp_norm = output['keypoints_norm']
+            hg_keyp_scores = output['keypoints_scores']
+            betas = output_reproj['betas']
+            betas_limbs = output_reproj['betas_limbs']
+            zz = output_reproj['z']
+            if result_network == 'normal':
+                # STEP 1: normal network
+                vertices_smal = output_reproj['vertices_smal']
+                flength = output_unnorm['flength']
+                pose_rotmat = output_unnorm['pose_rotmat']
+                trans = output_unnorm['trans']
+                pred_keyp = output_reproj['keyp_2d']
+                pred_silh = output_reproj['silh']
+                prefix = 'normal_'
+            else:
+                # STEP 1: refinement network
+                vertices_smal = output_ref['vertices_smal']
+                flength = output_ref['flength']
+                pose_rotmat = output_ref['pose_rotmat']
+                trans = output_ref['trans']
+                pred_keyp = output_ref['keyp_2d']
+                pred_silh = output_ref['silh']
+                prefix = 'ref_'
+                
+                bs = vertices_smal.shape[0]
+                # target_gc_class = target_dict['gc'][:, :, 0]
+                target_gc_class = torch.round(ground_contact_probs).long()[:, :, 1]
+                sel_verts = torch.index_select(output_ref['vertices_smal'], dim=1, index=remeshing_relevant_faces.reshape((-1))).reshape((bs, remeshing_relevant_faces.shape[0], 3, 3))
+                verts_remeshed = torch.einsum('ij,aijk->aik', remeshing_relevant_barys, sel_verts)
+                target_gc_class_remeshed = torch.einsum('ij,aij->ai', remeshing_relevant_barys, target_gc_class[:, remeshing_relevant_faces].to(device=device, dtype=torch.float32))
+                target_gc_class_remeshed_prep = torch.round(target_gc_class_remeshed).to(torch.long)
+
+
+
+
+                # index = i  
+                # ind_img = 0
+                for ind_img in range(batch_size): #  range(min(12, batch_size)):     # range(12):    # [0]:  #range(0, batch_size):
+
+                    # ind_img = 0
+                    if test_name_list is not None:
+                        img_name = test_name_list[int(target_dict['index'][ind_img].cpu().detach().numpy())].replace('/', '_')
+                        img_name = img_name.split('.')[0]
+                    else:
+                        img_name = str(index) + '_' + str(ind_img)
+                    out_path_gcmesh = save_imgs_path + '/' + prefix + 'gcmesh_' + img_name + '.obj'
+
+                    gc_prob = ground_contact_probs[ind_img, :, 1]   # contact probability
+                    vert_colors = np.repeat(255*gc_prob.detach().cpu().numpy()[:, None], 3, 1)
+                    my_mesh = trimesh.Trimesh(vertices=smal_template_verts, faces=smal_faces, process=False,  maintain_order=True)
+                    my_mesh.visual.vertex_colors = vert_colors
+                    save_gc_mesh = False
+                    if save_gc_mesh:
+                        my_mesh.export(out_path_gcmesh)            
+        
+                    current_verts_remeshed = verts_remeshed[ind_img, :, :]
+                    current_target_gc_class_remeshed_prep = target_gc_class_remeshed_prep[ind_img, ...]
+
+                    if current_target_gc_class_remeshed_prep.sum() > 3:
+                        points_on_plane = current_verts_remeshed[current_target_gc_class_remeshed_prep==1, :]
+                        data_centroid, plane_normal, error = fit_plane(points_on_plane)
+                        nonplane_points_centered = current_verts_remeshed[current_target_gc_class_remeshed_prep==0, :] - data_centroid[None, :]
+                        nonplane_points_projected = torch.matmul(plane_normal[None, :], nonplane_points_centered.transpose(0,1))
+
+                        if nonplane_points_projected.sum() > 0: # plane normal points towards the animal
+                            plane_normal = plane_normal.detach().cpu().numpy()
+                        else:
+                            plane_normal = - plane_normal.detach().cpu().numpy()
+                        data_centroid = data_centroid.detach().cpu().numpy()
+
+
+
+                        # import pdb; pdb.set_trace()
+
+
+                        desired_plane_normal_vector = np.asarray([[0, -1, 0]])
+                        # new approach: use cross product
+                        rotation_axis = np.cross(plane_normal, desired_plane_normal_vector)     #  np.cross(plane_normal, desired_plane_normal_vector)
+                        lengt_rotation_axis = np.linalg.norm(rotation_axis)     # = sin(alpha)      (because vectors have unit length)
+                        angle = np.sin(lengt_rotation_axis)
+                        rot = R_sc.from_rotvec(angle * rotation_axis * 1/lengt_rotation_axis)
+                        rot_mat = rot[0].as_matrix()
+                        rot_upsidedown = R_sc.from_rotvec(np.pi * np.asarray([[1, 0, 0]]))
+                        # rot_upsidedown[0].apply(rot[0].apply(plane_normal))
+                        current_vertices_smal = vertices_smal[ind_img, :, :].detach().cpu().numpy()
+                        new_smal_vertices = rot_upsidedown[0].apply(rot[0].apply(current_vertices_smal - data_centroid[None, :]))
+                        my_mesh = trimesh.Trimesh(vertices=new_smal_vertices, faces=smal_faces, process=False,  maintain_order=True)
+                        vert_colors[:, 2] = 255
+                        my_mesh.visual.vertex_colors = vert_colors
+                        out_path_gc_rotated = save_imgs_path + '/' + prefix + 'gc_rotated_' + img_name + '_new.obj'
+                        my_mesh.export(out_path_gc_rotated)
+
+
+
+                        '''
+                        import pdb; pdb.set_trace()
+
+                        from src.evaluation.registration import preprocess_point_cloud, o3d_ransac, draw_registration_result
+                        import open3d as o3d
+                        import copy
+
+
+                        mesh_gt_path = target_dict['mesh_path'][ind_img]
+                        mesh_gt = o3d.io.read_triangle_mesh(mesh_gt_path)
+
+                        mesh_gt_verts = np.asarray(mesh_gt.vertices)
+                        mesh_gt_faces = np.asarray(mesh_gt.triangles)
+                        diag_gt = np.sqrt(sum((mesh_gt_verts.max(axis=0) - mesh_gt_verts.min(axis=0))**2))
+
+                        mesh_pred_verts = np.asarray(new_smal_vertices)
+                        mesh_pred_faces = np.asarray(smal_faces)
+                        diag_pred = np.sqrt(sum((mesh_pred_verts.max(axis=0) - mesh_pred_verts.min(axis=0))**2))
+                        mesh_pred = o3d.geometry.TriangleMesh()
+                        mesh_pred.vertices = o3d.utility.Vector3dVector(mesh_pred_verts)
+                        mesh_pred.triangles = o3d.utility.Vector3iVector(mesh_pred_faces)
+
+                        # center the predicted mesh around 0
+                        trans = - mesh_pred_verts.mean(axis=0)
+                        mesh_pred_verts_new = mesh_pred_verts + trans
+                        # change the size of the predicted mesh
+                        mesh_pred_verts_new = mesh_pred_verts_new * diag_gt / diag_pred
+
+                        # transform the predicted mesh (rough alignment)
+                        mesh_pred_new = copy.deepcopy(mesh_pred)
+                        mesh_pred_new.vertices = o3d.utility.Vector3dVector(np.asarray(mesh_pred_verts_new))    # normals should not have changed
+                        voxel_size = 0.01       # 0.5
+                        distance_threshold = 0.015  # 0.005 #  0.02   # 1.0
+                        result, src_down, src_fpfh, dst_down, dst_fpfh = o3d_ransac(mesh_pred_new, mesh_gt, voxel_size=voxel_size, distance_threshold=distance_threshold, return_all=True)
+                        transform = result.transformation
+                        mesh_pred_transf = copy.deepcopy(mesh_pred_new).transform(transform)
+
+                        out_path_pred_transf = save_imgs_path + '/' + prefix + 'alignment_initial_' + img_name + '.obj'
+                        o3d.io.write_triangle_mesh(out_path_pred_transf, mesh_pred_transf)
+
+                        # img_name_part = img_name.split(img_name.split('_')[-1] + '_')[0]
+                        # out_path_gt = save_imgs_path + '/' + prefix + 'ground_truth_' + img_name_part + '.obj'
+                        # o3d.io.write_triangle_mesh(out_path_gt, mesh_gt)
+
+                        
+                        trans_init = transform
+                        threshold = 0.02        #  0.1  # 0.02
+
+                        n_points = 10000
+                        src = mesh_pred_new.sample_points_uniformly(number_of_points=n_points)
+                        dst = mesh_gt.sample_points_uniformly(number_of_points=n_points)
+
+                        # reg_p2p = o3d.pipelines.registration.registration_icp(src_down, dst_down, threshold, trans_init, o3d.pipelines.registration.TransformationEstimationPointToPoint(), o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=2000))
+                        reg_p2p = o3d.pipelines.registration.registration_icp(src, dst, threshold, trans_init, o3d.pipelines.registration.TransformationEstimationPointToPoint(), o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=2000))
+
+                        # mesh_pred_transf_refined = copy.deepcopy(mesh_pred_new).transform(reg_p2p.transformation)
+                        # out_path_pred_transf_refined =  save_imgs_path + '/' + prefix + 'alignment_final_' + img_name + '.obj'
+                        # o3d.io.write_triangle_mesh(out_path_pred_transf_refined, mesh_pred_transf_refined)
+
+
+                        aligned_mesh_final = trimesh.Trimesh(mesh_pred_new.vertices, mesh_pred_new.triangles, vertex_colors=[0, 255, 0])
+                        gt_mesh = trimesh.Trimesh(mesh_gt.vertices, mesh_gt.triangles, vertex_colors=[255, 0, 0])
+                        scene = trimesh.Scene([aligned_mesh_final, gt_mesh])
+                        out_path_alignment_with_gt =  save_imgs_path + '/' + prefix + 'alignment_with_gt_' + img_name + '.obj'
+
+                        scene.export(out_path_alignment_with_gt)
+                        '''
+
+                        # import pdb; pdb.set_trace()
+
+
+                        # SMAL_KEYPOINT_NAMES_FOR_3D_EVAL     # 17 keypoints
+                        # prepare target
+                        target_keyp_isvalid = target_dict['keypoints_3d'][ind_img, :, 3].detach().cpu().numpy()
+                        keyp_to_use = (np.asarray(SMAL_KEYPOINT_WHICHTOUSE_FOR_3D_EVAL)==1)*(target_keyp_isvalid==1)
+                        target_keyp_raw = target_dict['keypoints_3d'][ind_img, :, :3].detach().cpu().numpy()
+                        target_keypoints = target_keyp_raw[keyp_to_use, :]
+                        target_pointcloud = target_dict['pointcloud_points'][ind_img, :, :].detach().cpu().numpy()
+                        # prepare prediction
+                        pred_keypoints_raw = output_ref['vertices_smal'][ind_img, SMAL_KEYPOINT_INDICES_FOR_3D_EVAL, :].detach().cpu().numpy()
+                        pred_keypoints = pred_keypoints_raw[keyp_to_use, :]
+                        pred_pointcloud =  verts_remeshed[ind_img, :, :].detach().cpu().numpy()
+
+
+
+
+                        '''
+                        pred_keypoints_transf, pred_pointcloud_transf, procrustes_params = compute_similarity_transform(pred_keypoints, target_keypoints, num_joints=None, verts=pred_pointcloud)
+                        pa_error = np.sqrt(np.sum((target_keypoints - pred_keypoints_transf) ** 2, axis=1))
+                        error_procrustes = np.mean(pa_error)
+
+
+                        col_target = np.zeros((target_pointcloud.shape[0], 3), dtype=np.uint8)
+                        col_target[:, 0] = 255
+                        col_pred = np.zeros((pred_pointcloud_transf.shape[0], 3), dtype=np.uint8)
+                        col_pred[:, 1] = 255
+                        pc = trimesh.points.PointCloud(np.concatenate((target_pointcloud, pred_pointcloud_transf)), colors=np.concatenate((col_target, col_pred)))
+                        out_path_pc = save_imgs_path + '/' + prefix + 'pointclouds_aligned_' + img_name + '.obj'
+                        pc.export(out_path_pc)
+
+                        print(target_dict['mesh_path'][ind_img])
+                        print(error_procrustes)
+                        file_alignment_errors.write(target_dict['mesh_path'][ind_img] + '\n')
+                        file_alignment_errors.write('error: ' + str(error_procrustes) + ' \n')
+
+                        writer.writerow({'name': (target_dict['mesh_path'][ind_img]).split('/')[-1], 'error': str(error_procrustes)})
+
+                        # import pdb; pdb.set_trace()
+                        # alignment_dict = calculate_alignemnt_errors(output_ref['vertices_smal'][ind_img, :, :], target_dict['keypoints_3d'][ind_img, :, :], target_dict['pointcloud_points'][ind_img, :, :])
+                        # file_alignment_errors.write('error: ' + str(alignment_dict['error_procrustes']) + ' \n')
+                        '''
+
+
+
+
+
+
+            if index == 0:
+                if len_dataset is None:
+                    len_data = val_loader.batch_size * len(val_loader)  # 1703
+                else:
+                    len_data = len_dataset
+                if result_network == 'normal':
+                    summaries = {'normal': dict(), 'ref': dict()}
+                    summary = summaries['normal']
+                else:
+                    summary = summaries['ref']
+                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)
+                summary['betas'] = np.zeros((len_data,betas.shape[1]))
+                summary['betas_limbs'] = np.zeros((len_data, betas_limbs.shape[1]))
+                summary['z'] = np.zeros((len_data, zz.shape[1]))
+                summary['pose_rotmat'] = np.zeros((len_data, pose_rotmat.shape[1], 3, 3))
+                summary['flength'] = np.zeros((len_data, flength.shape[1]))
+                summary['trans'] = np.zeros((len_data, trans.shape[1]))
+                summary['breed_indices'] = np.zeros((len_data))
+                summary['image_names'] = []        # len_data * [None]
+                # ['vertices_smal'] = np.zeros((len_data, vertices_smal.shape[1], 3))
+            else:
+                if result_network == 'normal':
+                    summary = summaries['normal']
+                else:
+                    summary = summaries['ref']
+            
+            
+            # import pdb; pdb.set_trace()
+
+
+            eval_save_visualizations_and_meshes(model, input, data_info, target_dict, test_name_list, vertices_smal, hg_keyp_norm, hg_keyp_scores, zz, betas, betas_limbs, pose_rotmat, trans, flength, pred_keyp, pred_silh, save_imgs_path, prefix, index, render_all=render_all)
+
+
+            preds = eval_prepare_pck_and_iou(model, input, data_info, target_dict, test_name_list, vertices_smal, hg_keyp_norm, hg_keyp_scores, zz, betas, betas_limbs, pose_rotmat, trans, flength, pred_keyp, pred_silh, save_imgs_path, prefix, index, pck_thresh=None, skip_pck_and_iou=True)
+            # add results for all images in this batch to lists
+            curr_batch_size = pred_keyp.shape[0]
+            eval_add_preds_to_summary(summary, preds, my_step, batch_size, curr_batch_size, skip_pck_and_iou=True)
+
+            # summary['vertices_smal'][my_step * batch_size:my_step * batch_size + curr_batch_size] = vertices_smal.detach().cpu().numpy()
+
+
+
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            '''
+            try: 
+                if test_name_list is not None:
+                    img_name = test_name_list[int(target_dict['index'][ind_img].cpu().detach().numpy())].replace('/', '_')
+                    img_name = img_name.split('.')[0]
+                else:
+                    img_name = str(index) + '_' + str(ind_img)
+                partial_results['img_name'] = img_name
+                visualizations = model.render_vis_nograd(vertices=output_reproj['vertices_smal'],
+                                                        focal_lengths=output_unnorm['flength'],
+                                                        color=0)    # 2)
+                # save image with predicted keypoints
+                pred_unp = (output['keypoints_norm'][ind_img, :, :] + 1.) / 2 * (data_info.image_size - 1)
+                pred_unp_maxval = output['keypoints_scores'][ind_img, :, :]
+                pred_unp_prep = torch.cat((pred_unp, pred_unp_maxval), 1)
+                inp_img = input[ind_img, :, :, :].detach().clone()
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/keypoints_pred_' + img_name + '.png'
+                    save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_path, threshold=0.1, print_scores=True, ratio_in_out=1.0)    # threshold=0.3
+                # save predicted 3d model
+                #   (1) front view
+                pred_tex = visualizations[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+                pred_tex_max = np.max(pred_tex, axis=2)
+                partial_results['tex_pred'] = pred_tex
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/tex_pred_' + img_name + '.png'
+                    plt.imsave(out_path, pred_tex)
+                input_image = input[ind_img, :, :, :].detach().clone()
+                for t, m, s in zip(input_image, data_info.rgb_mean, data_info.rgb_stddev): t.add_(m)
+                input_image_np = input_image.detach().cpu().numpy().transpose(1, 2, 0) 
+                im_masked = cv2.addWeighted(input_image_np,0.2,pred_tex,0.8,0)
+                im_masked[pred_tex_max<0.01, :] = input_image_np[pred_tex_max<0.01, :]
+                partial_results['comp_pred'] = im_masked
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/comp_pred_' + img_name + '.png'
+                    plt.imsave(out_path, im_masked)
+                #   (2) side view
+                vertices_cent = output_reproj['vertices_smal'] - output_reproj['vertices_smal'].mean(dim=1)[:, None, :]
+                roll = np.pi / 2 * torch.ones(1).float().to(device)
+                pitch = np.pi / 2 * torch.ones(1).float().to(device)
+                tensor_0 = torch.zeros(1).float().to(device)
+                tensor_1 = torch.ones(1).float().to(device)
+                RX = torch.stack([torch.stack([tensor_1, tensor_0, tensor_0]), torch.stack([tensor_0, torch.cos(roll), -torch.sin(roll)]),torch.stack([tensor_0, torch.sin(roll), torch.cos(roll)])]).reshape(3,3)
+                RY = torch.stack([
+                    torch.stack([torch.cos(pitch), tensor_0, torch.sin(pitch)]),
+                    torch.stack([tensor_0, tensor_1, tensor_0]),
+                    torch.stack([-torch.sin(pitch), tensor_0, torch.cos(pitch)])]).reshape(3,3)
+                vertices_rot = (torch.matmul(RY, vertices_cent.reshape((-1, 3))[:, :, None])).reshape((batch_size, -1, 3))
+                vertices_rot[:, :, 2] = vertices_rot[:, :, 2] + torch.ones_like(vertices_rot[:, :, 2]) * 20     # 18     # *16
+                visualizations_rot = model.render_vis_nograd(vertices=vertices_rot,
+                                                        focal_lengths=output_unnorm['flength'],
+                                                        color=0)    # 2)
+                pred_tex = visualizations_rot[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+                pred_tex_max = np.max(pred_tex, axis=2)
+                partial_results['rot_tex_pred'] = pred_tex
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/rot_tex_pred_' + img_name + '.png'
+                    plt.imsave(out_path, pred_tex)
+                render_all = True
+                if render_all:
+                    # save input image 
+                    inp_img = input[ind_img, :, :, :].detach().clone()
+                    if save_imgs_path is not None:
+                        out_path = save_imgs_path + '/image_' + img_name + '.png'
+                        save_input_image(inp_img, out_path)
+                    # save posed mesh
+                    V_posed = output_reproj['vertices_smal'][ind_img, :, :].detach().cpu().numpy()
+                    Faces = model.smal.f
+                    mesh_posed = trimesh.Trimesh(vertices=V_posed, faces=Faces, process=False,  maintain_order=True)
+                    partial_results['mesh_posed'] = mesh_posed
+                    if save_imgs_path is not None:
+                        mesh_posed.export(save_imgs_path + '/mesh_posed_' + img_name + '.obj')
+            except:
+                print('pass...')
+            all_results.append(partial_results)
+            '''
+
+        my_step += 1
+
+
+    file_alignment_errors.close()
+    csv_file_alignment_errors.close()
+
+
+    if return_results:
+        return all_results
+    else:
+        return summaries

src/combined_model/train_main_image_to_3d_withbreedrel.py

@@ -0,0 +1,496 @@
+
+import torch
+import torch.nn as nn
+import torch.backends.cudnn
+import torch.nn.parallel
+from tqdm import tqdm
+import os
+import pathlib
+from matplotlib import pyplot as plt
+import cv2
+import numpy as np
+import torch
+import trimesh
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
+from stacked_hourglass.utils.evaluation import accuracy, AverageMeter, final_preds, get_preds, get_preds_soft
+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
+
+
+# ---------------------------------------------------------------------------------------------------------------------------
+def do_training_epoch(train_loader, model, loss_module, device, data_info, optimiser, quiet=False, acc_joints=None, weight_dict=None):
+    losses = AverageMeter()
+    losses_keyp = AverageMeter()
+    losses_silh = AverageMeter()
+    losses_shape = AverageMeter()
+    losses_pose = AverageMeter()
+    losses_class = AverageMeter()
+    losses_breed = AverageMeter()
+    losses_partseg = AverageMeter()
+    accuracies = AverageMeter()
+    # Put the model in training mode.
+    model.train()
+    # prepare progress bar
+    iterable = enumerate(train_loader)
+    progress = None
+    if not quiet:
+        progress = tqdm(iterable, desc='Train', total=len(train_loader), ascii=True, leave=False)
+        iterable = progress
+    # information for normalization 
+    norm_dict = {
+        'pose_rot6d_mean': torch.from_numpy(data_info.pose_rot6d_mean).float().to(device),
+        'trans_mean': torch.from_numpy(data_info.trans_mean).float().to(device),
+        'trans_std': torch.from_numpy(data_info.trans_std).float().to(device),
+        'flength_mean': torch.from_numpy(data_info.flength_mean).float().to(device),
+        'flength_std': torch.from_numpy(data_info.flength_std).float().to(device)}
+    # prepare variables, put them on the right device
+    for i, (input, target_dict) in iterable:
+        batch_size = input.shape[0]
+        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 in ['has_seg', 'gc']:
+                target_dict[key] = target_dict[key].to(device)
+            else:
+                pass
+        input = input.float().to(device)
+
+        # ----------------------- do training step -----------------------
+        assert model.training, 'model must be in training mode.'
+        with torch.enable_grad():
+            # ----- forward pass -----  
+            output, output_unnorm, output_reproj = model(input, norm_dict=norm_dict)        
+            # ----- loss -----
+            loss, loss_dict = loss_module(output_reproj=output_reproj, 
+                target_dict=target_dict, 
+                weight_dict=weight_dict)
+            # ----- backward pass and parameter update -----
+            optimiser.zero_grad()
+            loss.backward()
+            optimiser.step()
+        # ----------------------------------------------------------------
+
+        # prepare losses for progress bar
+        bs_fake = 1     # batch_size
+        losses.update(loss_dict['loss'], bs_fake)
+        losses_keyp.update(loss_dict['loss_keyp_weighted'], bs_fake)
+        losses_silh.update(loss_dict['loss_silh_weighted'], bs_fake)
+        losses_shape.update(loss_dict['loss_shape_weighted'], bs_fake)
+        losses_pose.update(loss_dict['loss_poseprior_weighted'], bs_fake)   
+        losses_class.update(loss_dict['loss_class_weighted'], bs_fake)
+        losses_breed.update(loss_dict['loss_breed_weighted'], bs_fake)
+        losses_partseg.update(loss_dict['loss_partseg_weighted'], bs_fake)
+        acc = - loss_dict['loss_keyp_weighted']     # this will be used to keep track of the 'best model'
+        accuracies.update(acc, bs_fake)
+        # Show losses as part of the progress bar.
+        if progress is not None:
+            my_string = 'Loss: {loss:0.4f}, loss_keyp: {loss_keyp:0.4f}, loss_silh: {loss_silh:0.4f}, loss_partseg: {loss_partseg:0.4f}, loss_shape: {loss_shape:0.4f}, loss_pose: {loss_pose:0.4f}, loss_class: {loss_class:0.4f}, loss_breed: {loss_breed:0.4f}'.format(
+                loss=losses.avg,
+                loss_keyp=losses_keyp.avg,
+                loss_silh=losses_silh.avg,
+                loss_shape=losses_shape.avg,
+                loss_pose=losses_pose.avg,
+                loss_class=losses_class.avg,
+                loss_breed=losses_breed.avg,
+                loss_partseg=losses_partseg.avg
+            )
+            progress.set_postfix_str(my_string)
+
+    return my_string, accuracies.avg       
+
+
+# ---------------------------------------------------------------------------------------------------------------------------
+def do_validation_epoch(val_loader, model, loss_module, device, data_info, flip=False, quiet=False, acc_joints=None, save_imgs_path=None, weight_dict=None, metrics=None, val_opt='default', test_name_list=None, render_all=False, pck_thresh=0.15, len_dataset=None):
+    losses = AverageMeter()
+    losses_keyp = AverageMeter()
+    losses_silh = AverageMeter()
+    losses_shape = AverageMeter()
+    losses_pose = AverageMeter()
+    losses_class = AverageMeter()
+    losses_breed = AverageMeter()
+    losses_partseg = AverageMeter()
+    accuracies = AverageMeter()
+    if save_imgs_path is not None:
+        pathlib.Path(save_imgs_path).mkdir(parents=True, exist_ok=True) 
+    # Put the model in evaluation mode.
+    model.eval()
+    # prepare progress bar
+    iterable = enumerate(val_loader)
+    progress = None
+    if not quiet:
+        progress = tqdm(iterable, desc='Valid', total=len(val_loader), ascii=True, leave=False)
+        iterable = progress
+    # summarize information for normalization 
+    norm_dict = {
+        'pose_rot6d_mean': torch.from_numpy(data_info.pose_rot6d_mean).float().to(device),
+        'trans_mean': torch.from_numpy(data_info.trans_mean).float().to(device),
+        'trans_std': torch.from_numpy(data_info.trans_std).float().to(device),
+        'flength_mean': torch.from_numpy(data_info.flength_mean).float().to(device),
+        'flength_std': torch.from_numpy(data_info.flength_std).float().to(device)}
+    batch_size = val_loader.batch_size
+    # prepare variables, put them on the right device
+    my_step = 0
+    for i, (input, target_dict) in iterable:
+        curr_batch_size = input.shape[0]
+        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 in ['has_seg', 'gc']:
+                target_dict[key] = target_dict[key].to(device)
+            else:
+                pass
+        input = input.float().to(device)
+
+        # ----------------------- do validation step -----------------------
+        with torch.no_grad():
+            # ----- forward pass -----  
+            # output: (['pose', 'flength', 'trans', 'keypoints_norm', 'keypoints_scores'])
+            # output_unnorm: (['pose_rotmat', 'flength', 'trans', 'keypoints'])
+            # output_reproj: (['vertices_smal', 'torch_meshes', 'keyp_3d', 'keyp_2d', 'silh', 'betas', 'pose_rot6d', 'dog_breed', 'shapedirs', 'z', 'flength_unnorm', 'flength'])
+            # target_dict: (['index', 'center', 'scale', 'pts', 'tpts', 'target_weight', 'breed_index', 'sim_breed_index', 'ind_dataset', 'silh'])
+            output, output_unnorm, output_reproj = model(input, norm_dict=norm_dict)        
+            # ----- loss -----
+            if metrics == 'no_loss':
+                loss, loss_dict = loss_module(output_reproj=output_reproj, 
+                    target_dict=target_dict, 
+                    weight_dict=weight_dict)
+        # ----------------------------------------------------------------
+
+        if i == 0:
+            if len_dataset is None:
+                len_data = val_loader.batch_size * len(val_loader)  # 1703
+            else:
+                len_data = len_dataset
+            if metrics == 'all' or metrics == 'no_loss':
+                pck = np.zeros((len_data))
+                pck_by_part = {group:np.zeros((len_data)) for group in KEYPOINT_GROUPS}
+                acc_sil_2d = np.zeros(len_data)
+
+                all_betas = np.zeros((len_data, output_reproj['betas'].shape[1]))
+                all_betas_limbs = np.zeros((len_data, output_reproj['betas_limbs'].shape[1]))
+                all_z = np.zeros((len_data, output_reproj['z'].shape[1]))
+                all_pose_rotmat = np.zeros((len_data, output_unnorm['pose_rotmat'].shape[1], 3, 3))
+                all_flength = np.zeros((len_data, output_unnorm['flength'].shape[1]))
+                all_trans = np.zeros((len_data, output_unnorm['trans'].shape[1]))
+                all_breed_indices = np.zeros((len_data))
+                all_image_names = []        # len_data * [None]
+
+        index = i  
+        ind_img = 0
+        if save_imgs_path is not None:
+            # render predicted 3d models
+            visualizations = model.render_vis_nograd(vertices=output_reproj['vertices_smal'],
+                                                    focal_lengths=output_unnorm['flength'],
+                                                    color=0)        # color=2)
+            for ind_img in range(len(target_dict['index'])):    
+                try: 
+                    if test_name_list is not None:
+                        img_name = test_name_list[int(target_dict['index'][ind_img].cpu().detach().numpy())].replace('/', '_')
+                        img_name = img_name.split('.')[0]
+                    else:
+                        img_name = str(index) + '_' + str(ind_img)
+                    # save image with predicted keypoints
+                    out_path = save_imgs_path + '/keypoints_pred_' + img_name + '.png'
+                    pred_unp = (output['keypoints_norm'][ind_img, :, :] + 1.) / 2 * (data_info.image_size - 1)
+                    pred_unp_maxval = output['keypoints_scores'][ind_img, :, :]
+                    pred_unp_prep = torch.cat((pred_unp, pred_unp_maxval), 1)
+                    inp_img = input[ind_img, :, :, :].detach().clone()
+                    save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_path, threshold=0.1, print_scores=True, ratio_in_out=1.0)    # threshold=0.3
+                    # save predicted 3d model (front view)
+                    pred_tex = visualizations[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+                    pred_tex_max = np.max(pred_tex, axis=2)
+                    out_path = save_imgs_path + '/tex_pred_' + img_name + '.png'
+                    plt.imsave(out_path, pred_tex)
+                    input_image = input[ind_img, :, :, :].detach().clone()
+                    for t, m, s in zip(input_image, data_info.rgb_mean, data_info.rgb_stddev): t.add_(m)
+                    input_image_np = input_image.detach().cpu().numpy().transpose(1, 2, 0) 
+                    im_masked = cv2.addWeighted(input_image_np,0.2,pred_tex,0.8,0)
+                    im_masked[pred_tex_max<0.01, :] = input_image_np[pred_tex_max<0.01, :]
+                    out_path = save_imgs_path + '/comp_pred_' + img_name + '.png'
+                    plt.imsave(out_path, im_masked)
+                    # save predicted 3d model (side view)
+                    vertices_cent = output_reproj['vertices_smal'] - output_reproj['vertices_smal'].mean(dim=1)[:, None, :]
+                    roll = np.pi / 2 * torch.ones(1).float().to(device)
+                    pitch = np.pi / 2 * torch.ones(1).float().to(device)
+                    tensor_0 = torch.zeros(1).float().to(device)
+                    tensor_1 = torch.ones(1).float().to(device)
+                    RX = torch.stack([torch.stack([tensor_1, tensor_0, tensor_0]), torch.stack([tensor_0, torch.cos(roll), -torch.sin(roll)]),torch.stack([tensor_0, torch.sin(roll), torch.cos(roll)])]).reshape(3,3)
+                    RY = torch.stack([
+                        torch.stack([torch.cos(pitch), tensor_0, torch.sin(pitch)]),
+                        torch.stack([tensor_0, tensor_1, tensor_0]),
+                        torch.stack([-torch.sin(pitch), tensor_0, torch.cos(pitch)])]).reshape(3,3)
+                    vertices_rot = (torch.matmul(RY, vertices_cent.reshape((-1, 3))[:, :, None])).reshape((curr_batch_size, -1, 3))
+                    vertices_rot[:, :, 2] = vertices_rot[:, :, 2] + torch.ones_like(vertices_rot[:, :, 2]) * 20     # 18     # *16
+
+                    visualizations_rot = model.render_vis_nograd(vertices=vertices_rot,
+                                                            focal_lengths=output_unnorm['flength'],
+                                                            color=0)        # 2)
+                    pred_tex = visualizations_rot[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+                    pred_tex_max = np.max(pred_tex, axis=2)
+                    out_path = save_imgs_path + '/rot_tex_pred_' + img_name + '.png'
+                    plt.imsave(out_path, pred_tex)
+                    if render_all:
+                        # save input image
+                        inp_img = input[ind_img, :, :, :].detach().clone()
+                        out_path = save_imgs_path + '/image_' + img_name + '.png'
+                        save_input_image(inp_img, out_path)
+                        # save mesh
+                        V_posed = output_reproj['vertices_smal'][ind_img, :, :].detach().cpu().numpy()
+                        Faces = model.smal.f
+                        mesh_posed = trimesh.Trimesh(vertices=V_posed, faces=Faces, process=False,  maintain_order=True)
+                        mesh_posed.export(save_imgs_path + '/mesh_posed_' + img_name + '.obj')
+                except: 
+                    print('dont save an image')
+
+        if metrics == 'all' or metrics == 'no_loss':
+            # prepare a dictionary with all the predicted results
+            preds = {}
+            preds['betas'] = output_reproj['betas'].cpu().detach().numpy()
+            preds['betas_limbs'] = output_reproj['betas_limbs'].cpu().detach().numpy()
+            preds['z'] = output_reproj['z'].cpu().detach().numpy()
+            preds['pose_rotmat'] = output_unnorm['pose_rotmat'].cpu().detach().numpy()
+            preds['flength'] = output_unnorm['flength'].cpu().detach().numpy()
+            preds['trans'] = output_unnorm['trans'].cpu().detach().numpy()
+            preds['breed_index'] = target_dict['breed_index'].cpu().detach().numpy().reshape((-1))
+            img_names = []
+            for ind_img2 in range(0, output_reproj['betas'].shape[0]):
+                if test_name_list is not None:
+                    img_name2 = test_name_list[int(target_dict['index'][ind_img2].cpu().detach().numpy())].replace('/', '_')
+                    img_name2 = img_name2.split('.')[0]
+                else:
+                    img_name2 = str(index) + '_' + str(ind_img2)
+                img_names.append(img_name2)
+            preds['image_names'] = img_names
+            # prepare keypoints for PCK calculation - predicted as well as ground truth
+            pred_keypoints_norm = output['keypoints_norm']   # -1 to 1
+            pred_keypoints_256 = output_reproj['keyp_2d']
+            pred_keypoints = pred_keypoints_256
+            gt_keypoints_256 = target_dict['tpts'][:, :, :2] / 64. * (256. - 1)
+            gt_keypoints_norm = gt_keypoints_256 / 256 / 0.5 - 1
+            gt_keypoints = torch.cat((gt_keypoints_256, target_dict['tpts'][:, :, 2:3]), dim=2)     # gt_keypoints_norm
+            # 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 = output_reproj['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['acc_PCK'] = Metrics.PCK(
+                pred_keypoints, 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_keypoints, gt_keypoints, gtseg, has_seg, 
+                    thresh_range=[pck_thresh],       # [0.15],
+                    idxs=group_kps
+                )
+            # add results for all images in this batch to lists
+            curr_batch_size = pred_keypoints_256.shape[0]
+            if not (preds['acc_PCK'].data.cpu().numpy().shape == (pck[my_step * batch_size:my_step * batch_size + curr_batch_size]).shape):
+                import pdb; pdb.set_trace()
+            pck[my_step * batch_size:my_step * batch_size + curr_batch_size] = preds['acc_PCK'].data.cpu().numpy()
+            acc_sil_2d[my_step * batch_size:my_step * batch_size + curr_batch_size] = preds['acc_IOU'].data.cpu().numpy()
+            for part in pck_by_part:
+                pck_by_part[part][my_step * batch_size:my_step * batch_size + curr_batch_size] = preds[f'{part}_PCK'].data.cpu().numpy()
+            all_betas[my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['betas'] 
+            all_betas_limbs[my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['betas_limbs'] 
+            all_z[my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['z'] 
+            all_pose_rotmat[my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['pose_rotmat'] 
+            all_flength[my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['flength'] 
+            all_trans[my_step * batch_size:my_step * batch_size + curr_batch_size, ...] = preds['trans'] 
+            all_breed_indices[my_step * batch_size:my_step * batch_size + curr_batch_size] = preds['breed_index']
+            all_image_names.extend(preds['image_names'])
+            # update progress bar
+            if progress is not None:
+                my_string = "PCK: {0:.2f}, IOU: {1:.2f}".format(
+                    pck[:(my_step * batch_size + curr_batch_size)].mean(),
+                    acc_sil_2d[:(my_step * batch_size + curr_batch_size)].mean())
+                progress.set_postfix_str(my_string)
+        else:
+            # measure accuracy and record loss
+            bs_fake = 1     # batch_size
+            losses.update(loss_dict['loss'], bs_fake)
+            losses_keyp.update(loss_dict['loss_keyp_weighted'], bs_fake)
+            losses_silh.update(loss_dict['loss_silh_weighted'], bs_fake)
+            losses_shape.update(loss_dict['loss_shape_weighted'], bs_fake)
+            losses_pose.update(loss_dict['loss_poseprior_weighted'], bs_fake)  
+            losses_class.update(loss_dict['loss_class_weighted'], bs_fake)
+            losses_breed.update(loss_dict['loss_breed_weighted'], bs_fake)
+            losses_partseg.update(loss_dict['loss_partseg_weighted'], bs_fake)
+            acc = - loss_dict['loss_keyp_weighted']     # this will be used to keep track of the 'best model'
+            accuracies.update(acc, bs_fake)
+            # Show losses as part of the progress bar.
+            if progress is not None:
+                my_string = 'Loss: {loss:0.4f}, loss_keyp: {loss_keyp:0.4f}, loss_silh: {loss_silh:0.4f}, loss_partseg: {loss_partseg:0.4f}, loss_shape: {loss_shape:0.4f}, loss_pose: {loss_pose:0.4f}, loss_class: {loss_class:0.4f}, loss_breed: {loss_breed:0.4f}'.format(
+                    loss=losses.avg,
+                    loss_keyp=losses_keyp.avg,
+                    loss_silh=losses_silh.avg,
+                    loss_shape=losses_shape.avg,
+                    loss_pose=losses_pose.avg,
+                    loss_class=losses_class.avg,
+                    loss_breed=losses_breed.avg,
+                    loss_partseg=losses_partseg.avg
+                )
+                progress.set_postfix_str(my_string)
+        my_step += 1
+    if metrics == 'all':
+        summary = {'pck': pck, 'acc_sil_2d': acc_sil_2d, 'pck_by_part':pck_by_part,
+                    'betas': all_betas, 'betas_limbs': all_betas_limbs, 'z': all_z, 'pose_rotmat': all_pose_rotmat,
+                    'flenght': all_flength, 'trans': all_trans, 'image_names': all_image_names, 'breed_indices': all_breed_indices}
+        return my_string, summary    
+    elif metrics == 'no_loss':
+        return my_string, np.average(np.asarray(acc_sil_2d))
+    else:
+        return my_string, accuracies.avg       
+
+
+# ---------------------------------------------------------------------------------------------------------------------------
+def do_visual_epoch(val_loader, model, device, data_info, flip=False, quiet=False, acc_joints=None, save_imgs_path=None, weight_dict=None, metrics=None, val_opt='default', test_name_list=None, render_all=False, pck_thresh=0.15, return_results=False):
+    if save_imgs_path is not None:
+        pathlib.Path(save_imgs_path).mkdir(parents=True, exist_ok=True) 
+    all_results = []
+
+    # Put the model in evaluation mode.
+    model.eval()
+
+    iterable = enumerate(val_loader)
+
+    # information for normalization 
+    norm_dict = {
+        'pose_rot6d_mean': torch.from_numpy(data_info.pose_rot6d_mean).float().to(device),
+        'trans_mean': torch.from_numpy(data_info.trans_mean).float().to(device),
+        'trans_std': torch.from_numpy(data_info.trans_std).float().to(device),
+        'flength_mean': torch.from_numpy(data_info.flength_mean).float().to(device),
+        'flength_std': torch.from_numpy(data_info.flength_std).float().to(device)}
+
+    '''
+    return_mesh_with_gt_groundplane = True
+    if return_mesh_with_gt_groundplane:
+        remeshing_path = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/smal_data_remeshed/uniform_surface_sampling/my_smpl_39dogsnorm_Jr_4_dog_remesh4000_info.pkl'
+        with open(remeshing_path, 'rb') as fp: 
+            remeshing_dict = pkl.load(fp)
+        remeshing_relevant_faces = torch.tensor(remeshing_dict['smal_faces'][remeshing_dict['faceid_closest']], dtype=torch.long, device=device)
+        remeshing_relevant_barys = torch.tensor(remeshing_dict['barys_closest'], dtype=torch.float32, device=device)
+
+    # from smal_pytorch.smal_model.smal_torch_new import SMAL
+    print('start: load smal default model (barc), but only for vertices')
+    smal = SMAL()
+    print('end: load smal default model (barc), but only for vertices')
+    smal_template_verts = smal.v_template.detach().cpu().numpy()
+    smal_faces = smal.faces.detach().cpu().numpy()
+    
+    file_alignment_errors = open(save_imgs_path + '/a_ref_procrustes_alignmnet_errors.txt', 'a') # append mode
+    file_alignment_errors.write(" -----------  start evaluation  ------------- \n ")
+
+    csv_file_alignment_errors = open(save_imgs_path + '/a_ref_procrustes_alignmnet_errors.csv', 'w') # write mode
+    fieldnames = ['name', 'error']
+    writer = csv.DictWriter(csv_file_alignment_errors, fieldnames=fieldnames)
+    writer.writeheader()
+    '''
+
+    my_step = 0
+    for i, (input, target_dict) in iterable:
+        batch_size = input.shape[0]
+        input = input.float().to(device)
+        partial_results = {}
+
+        # ----------------------- do visualization step -----------------------
+        with torch.no_grad():
+            output, output_unnorm, output_reproj = model(input, norm_dict=norm_dict)        
+
+        index = i  
+        ind_img = 0
+        for ind_img in range(batch_size): #  range(min(12, batch_size)):     # range(12):    # [0]:  #range(0, batch_size):
+
+            try: 
+                if test_name_list is not None:
+                    img_name = test_name_list[int(target_dict['index'][ind_img].cpu().detach().numpy())].replace('/', '_')
+                    img_name = img_name.split('.')[0]
+                else:
+                    img_name = str(index) + '_' + str(ind_img)
+                partial_results['img_name'] = img_name
+                visualizations = model.render_vis_nograd(vertices=output_reproj['vertices_smal'],
+                                                        focal_lengths=output_unnorm['flength'],
+                                                        color=0)    # 2)
+                # save image with predicted keypoints
+                pred_unp = (output['keypoints_norm'][ind_img, :, :] + 1.) / 2 * (data_info.image_size - 1)
+                pred_unp_maxval = output['keypoints_scores'][ind_img, :, :]
+                pred_unp_prep = torch.cat((pred_unp, pred_unp_maxval), 1)
+                inp_img = input[ind_img, :, :, :].detach().clone()
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/keypoints_pred_' + img_name + '.png'
+                    save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_path, threshold=0.1, print_scores=True, ratio_in_out=1.0)    # threshold=0.3
+                # save predicted 3d model
+                #   (1) front view
+                pred_tex = visualizations[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+                pred_tex_max = np.max(pred_tex, axis=2)
+                partial_results['tex_pred'] = pred_tex
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/tex_pred_' + img_name + '.png'
+                    plt.imsave(out_path, pred_tex)
+                input_image = input[ind_img, :, :, :].detach().clone()
+                for t, m, s in zip(input_image, data_info.rgb_mean, data_info.rgb_stddev): t.add_(m)
+                input_image_np = input_image.detach().cpu().numpy().transpose(1, 2, 0) 
+                im_masked = cv2.addWeighted(input_image_np,0.2,pred_tex,0.8,0)
+                im_masked[pred_tex_max<0.01, :] = input_image_np[pred_tex_max<0.01, :]
+                partial_results['comp_pred'] = im_masked
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/comp_pred_' + img_name + '.png'
+                    plt.imsave(out_path, im_masked)
+                #   (2) side view
+                vertices_cent = output_reproj['vertices_smal'] - output_reproj['vertices_smal'].mean(dim=1)[:, None, :]
+                roll = np.pi / 2 * torch.ones(1).float().to(device)
+                pitch = np.pi / 2 * torch.ones(1).float().to(device)
+                tensor_0 = torch.zeros(1).float().to(device)
+                tensor_1 = torch.ones(1).float().to(device)
+                RX = torch.stack([torch.stack([tensor_1, tensor_0, tensor_0]), torch.stack([tensor_0, torch.cos(roll), -torch.sin(roll)]),torch.stack([tensor_0, torch.sin(roll), torch.cos(roll)])]).reshape(3,3)
+                RY = torch.stack([
+                    torch.stack([torch.cos(pitch), tensor_0, torch.sin(pitch)]),
+                    torch.stack([tensor_0, tensor_1, tensor_0]),
+                    torch.stack([-torch.sin(pitch), tensor_0, torch.cos(pitch)])]).reshape(3,3)
+                vertices_rot = (torch.matmul(RY, vertices_cent.reshape((-1, 3))[:, :, None])).reshape((batch_size, -1, 3))
+                vertices_rot[:, :, 2] = vertices_rot[:, :, 2] + torch.ones_like(vertices_rot[:, :, 2]) * 20     # 18     # *16
+                visualizations_rot = model.render_vis_nograd(vertices=vertices_rot,
+                                                        focal_lengths=output_unnorm['flength'],
+                                                        color=0)    # 2)
+                pred_tex = visualizations_rot[ind_img, :, :, :].permute((1, 2, 0)).cpu().detach().numpy() / 256
+                pred_tex_max = np.max(pred_tex, axis=2)
+                partial_results['rot_tex_pred'] = pred_tex
+                if save_imgs_path is not None:
+                    out_path = save_imgs_path + '/rot_tex_pred_' + img_name + '.png'
+                    plt.imsave(out_path, pred_tex)
+                render_all = True
+                if render_all:
+                    # save input image 
+                    inp_img = input[ind_img, :, :, :].detach().clone()
+                    if save_imgs_path is not None:
+                        out_path = save_imgs_path + '/image_' + img_name + '.png'
+                        save_input_image(inp_img, out_path)
+                    # save posed mesh
+                    V_posed = output_reproj['vertices_smal'][ind_img, :, :].detach().cpu().numpy()
+                    Faces = model.smal.f
+                    mesh_posed = trimesh.Trimesh(vertices=V_posed, faces=Faces, process=False,  maintain_order=True)
+                    partial_results['mesh_posed'] = mesh_posed
+                    if save_imgs_path is not None:
+                        mesh_posed.export(save_imgs_path + '/mesh_posed_' + img_name + '.obj')
+            except:
+                print('pass...')
+            all_results.append(partial_results)
+    if return_results:
+        return all_results
+    else:
+        return

src/configs/SMAL_configs.py

@@ -0,0 +1,230 @@
+
+
+import numpy as np
+import os
+import sys
+
+
+# SMAL_DATA_DIR = '/is/cluster/work/nrueegg/dog_project/pytorch-dogs-inference/src/smal_pytorch/smpl_models/'
+# SMAL_DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'smal_pytorch', 'smal_data')
+SMAL_DATA_DIR = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'smal_data')
+
+# we replace the old SMAL model by a more dog specific model (see BARC cvpr 2022 paper)
+# our model has several differences compared to the original SMAL model, some of them are:
+#   - the PCA shape space is recalculated (from partially new data and weighted)  
+#   - coefficients for limb length changes are allowed (similar to WLDO, we did borrow some of their code)
+#   - all dogs have a core of approximately the same length
+#   - dogs are centered in their root joint (which is close to the tail base)
+#       -> like this the root rotations is always around this joint AND (0, 0, 0)
+#       -> before this it would happen that the animal 'slips' from the image middle to the side when rotating it. Now
+#          'trans' also defines the center of the rotation
+#   - we correct the back joint locations such that all those joints are more aligned
+
+# 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'] 
+
+SMAL_MODEL_CONFIG = {
+  'barc': {
+    'smal_model_type': 'barc',
+    'smal_model_path': os.path.join(SMAL_DATA_DIR, 'my_smpl_SMBLD_nbj_v3.pkl'),
+    'smal_model_data_path': os.path.join(SMAL_DATA_DIR, 'my_smpl_data_SMBLD_v3.pkl'),        
+    'unity_smal_shape_prior_dogs': os.path.join(SMAL_DATA_DIR, 'my_smpl_data_SMBLD_v3.pkl'),
+    'logscale_part_list': ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'],
+  },
+  '39dogs_diffsize': {
+    'smal_model_type': '39dogs_diffsize',
+    'smal_model_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_00791_nadine_Jr_4_dog.pkl'),
+    'smal_model_data_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_00791_nadine_Jr_4_dog.pkl'),       
+    'unity_smal_shape_prior_dogs': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_00791_nadine_Jr_4_dog.pkl'),
+    'logscale_part_list': ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'],
+  },
+  '39dogs_norm': {
+    'smal_model_type': '39dogs_norm',
+    'smal_model_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_39dogsnorm_Jr_4_dog.pkl'),
+    'smal_model_data_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_Jr_4_dog.pkl'),       
+    'unity_smal_shape_prior_dogs': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_Jr_4_dog.pkl'), 
+    'logscale_part_list': ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'],
+  },
+  '39dogs_norm_9ll': {    # 9 limb length parameters
+    'smal_model_type': '39dogs_norm_9ll',
+    'smal_model_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_39dogsnorm_Jr_4_dog.pkl'),
+    'smal_model_data_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_Jr_4_dog.pkl'),       
+    'unity_smal_shape_prior_dogs': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_Jr_4_dog.pkl'), 
+    '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'],
+  },
+  '39dogs_norm_newv2': {  # front and back legs of equal lengths
+    'smal_model_type': '39dogs_norm_newv2',
+    'smal_model_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_39dogsnorm_newv2_dog.pkl'),
+    'smal_model_data_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_newv2_dog.pkl'),       
+    'unity_smal_shape_prior_dogs': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_newv2_dog.pkl'), 
+    'logscale_part_list': ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'],
+  },
+  '39dogs_norm_newv3': {  # pca on dame AND different front and back legs lengths
+    'smal_model_type': '39dogs_norm_newv3',
+    'smal_model_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_39dogsnorm_newv3_dog.pkl'),
+    'smal_model_data_path': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_newv3_dog.pkl'),       
+    'unity_smal_shape_prior_dogs': os.path.join(SMAL_DATA_DIR, 'new_dog_models', 'my_smpl_data_39dogsnorm_newv3_dog.pkl'), 
+    'logscale_part_list': ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'],
+  },
+}
+
+
+SYMMETRY_INDS_FILE = os.path.join(SMAL_DATA_DIR, 'symmetry_inds.json')
+
+mean_dog_bone_lengths_txt = os.path.join(SMAL_DATA_DIR, 'mean_dog_bone_lengths.txt')
+
+# some vertex indices, (from silvia zuffi´s code, create_projected_images_cats.py)
+KEY_VIDS = np.array(([1068, 1080, 1029, 1226],  # left eye
+       [2660, 3030, 2675, 3038],                # right eye
+       [910],                                   # mouth low
+       [360, 1203, 1235, 1230],                 # front left leg, low
+       [3188, 3156, 2327, 3183],                # front right leg, low
+       [1976, 1974, 1980,  856],                # back left leg, low
+       [3854, 2820, 3852, 3858],                # back right leg, low
+       [452, 1811],                             # tail start
+       [416, 235, 182],                         # front left leg, top
+       [2156, 2382, 2203],                      # front right leg, top
+       [829],                                   # back left leg, top
+       [2793],                                  # back right leg, top
+       [60, 114, 186,  59],                     # throat, close to base of neck
+       [2091, 2037, 2036, 2160],                # withers (a bit lower than in reality)
+       [384,  799, 1169,  431],                 # front left leg, middle
+       [2351, 2763, 2397, 3127],                # front right leg, middle
+       [221, 104],                              # back left leg, middle
+       [2754, 2192],                            # back right leg, middle
+       [191, 1158, 3116, 2165],                 # neck
+       [28],                                    # Tail tip
+       [542],                                   # Left Ear
+       [2507],                                  # Right Ear
+       [1039, 1845, 1846, 1870, 1879, 1919, 2997, 3761, 3762], # nose tip
+       [0, 464, 465, 726, 1824, 2429, 2430, 2690]), dtype=object) # half tail
+
+# the following vertices are used for visibility only: if one of the vertices is visible, 
+# then we assume that the joint is visible!  There is some noise, but we don't care, as this is 
+# for generation of the synthetic dataset only    
+KEY_VIDS_VISIBILITY_ONLY = np.array(([1068, 1080, 1029, 1226, 645], # left eye
+       [2660, 3030, 2675, 3038, 2567],                              # right eye
+       [910, 11, 5],                                                # mouth low
+       [360, 1203, 1235, 1230, 298, 408, 303, 293, 384],            # front left leg, low
+       [3188, 3156, 2327, 3183, 2261, 2271, 2573, 2265],            # front right leg, low
+       [1976, 1974, 1980,  856, 559, 851, 556],                     # back left leg, low
+       [3854, 2820, 3852, 3858, 2524, 2522, 2815, 2072],            # back right leg, low
+       [452, 1811, 63, 194, 52, 370, 64],                           # tail start
+       [416, 235, 182, 440, 8, 80, 73, 112],                        # front left leg, top
+       [2156, 2382, 2203, 2050, 2052, 2406, 3],                     # front right leg, top
+       [829, 219, 218, 173, 17, 7, 279],                            # back left leg, top
+       [2793, 582, 140, 87, 2188, 2147, 2063],                      # back right leg, top
+       [60, 114, 186,  59, 878, 130, 189, 45],                      # throat, close to base of neck
+       [2091, 2037, 2036, 2160, 190, 2164],                         # withers (a bit lower than in reality)
+       [384,  799, 1169,  431, 321, 314, 437, 310, 323],            # front left leg, middle
+       [2351, 2763, 2397, 3127, 2278, 2285, 2282, 2275, 2359],      # front right leg, middle
+       [221, 104, 105, 97, 103],                                    # back left leg, middle
+       [2754, 2192, 2080, 2251, 2075, 2074],                        # back right leg, middle
+       [191, 1158, 3116, 2165, 154, 653, 133, 339],                 # neck
+       [28, 474, 475, 731, 24],                                     # Tail tip
+       [542, 147, 509, 200, 522],                                   # Left Ear
+       [2507,2174, 2122, 2126, 2474],                               # Right Ear
+       [1039, 1845, 1846, 1870, 1879, 1919, 2997, 3761, 3762],      # nose tip
+       [0, 464, 465, 726, 1824, 2429, 2430, 2690]), dtype=object)   # half tail
+
+# Keypoint indices for 3d sketchfab evaluation
+SMAL_KEYPOINT_NAMES_FOR_3D_EVAL = ['right_front_paw','right_front_elbow','right_back_paw','right_back_hock','right_ear_top','right_ear_bottom','right_eye', \
+                                  'left_front_paw','left_front_elbow','left_back_paw','left_back_hock','left_ear_top','left_ear_bottom','left_eye', \
+                                  'nose','tail_start','tail_end']
+SMAL_KEYPOINT_INDICES_FOR_3D_EVAL = [2577,	2361,	2820,	2085,	2125,	2453,	2668,	613,	394,	855,	786,	149,	486,	1079,	1845,	1820,	28]
+SMAL_KEYPOINT_WHICHTOUSE_FOR_3D_EVAL = [1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0]    # [1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0]
+
+
+
+
+# see: https://github.com/benjiebob/SMALify/blob/master/config.py
+# JOINT DEFINITIONS - based on SMAL joints and additional {eyes, ear tips, chin and nose}
+TORSO_JOINTS = [2, 5, 8, 11, 12, 23]
+CANONICAL_MODEL_JOINTS = [
+  10, 9, 8, # upper_left [paw, middle, top]
+  20, 19, 18, # lower_left [paw, middle, top]
+  14, 13, 12, # upper_right [paw, middle, top]
+  24, 23, 22, # lower_right [paw, middle, top]
+  25, 31, # tail [start, end]
+  33, 34, # ear base [left, right]
+  35, 36, # nose, chin
+  38, 37, # ear tip [left, right]
+  39, 40, # eyes [left, right]
+  6, 11, # withers, throat (throat is inaccurate and withers also)
+  28] # tail middle
+  # old:   15, 15, # withers, throat (TODO: Labelled same as throat for now), throat 
+
+CANONICAL_MODEL_JOINTS_REFINED = [
+  41, 9, 8, # upper_left [paw, middle, top]
+  43, 19, 18, # lower_left [paw, middle, top]
+  42, 13, 12, # upper_right [paw, middle, top]
+  44, 23, 22, # lower_right [paw, middle, top]
+  25, 31, # tail [start, end]
+  33, 34, # ear base [left, right]
+  35, 36, # nose, chin
+  38, 37, # ear tip [left, right]
+  39, 40, # eyes [left, right]
+  46, 45, # withers, throat
+  28] # tail middle
+
+# the following list gives the indices of the KEY_VIDS_JOINTS that must be taken in order 
+# to judge if the CANONICAL_MODEL_JOINTS are visible - those are all approximations!
+CMJ_VISIBILITY_IN_KEY_VIDS = [
+    3, 14, 8,       # left front leg
+    5, 16, 10,      # left rear leg
+    4, 15, 9,       # right front leg
+    6, 17, 11,      # right rear leg
+    7, 19,          # tail front, tail back
+    20, 21,         # ear base (but can not be found in blue, se we take the tip)
+    2, 2,           # mouth  (was: 22, 2)
+    20, 21,         # ear tips
+    1, 0,           # eyes
+    18,             # withers, not sure where this point is
+    12,             # throat
+    23,             # mid tail
+    ]
+
+# define which bone lengths are used as input to the 2d-to-3d network
+IDXS_BONES_NO_REDUNDANCY = [6,7,8,9,16,17,18,19,32,1,2,3,4,5,14,15,24,25,26,27,28,29,30,31]
+# load bone lengths of the mean dog (already filtered)
+mean_dog_bone_lengths = []
+with open(mean_dog_bone_lengths_txt, 'r') as f:
+    for line in f:
+       mean_dog_bone_lengths.append(float(line.split('\n')[0]))
+MEAN_DOG_BONE_LENGTHS_NO_RED = np.asarray(mean_dog_bone_lengths)[IDXS_BONES_NO_REDUNDANCY]        # (24, )
+
+# Body part segmentation:
+#   the body can be segmented based on the bones and for the new dog model also based on the new shapedirs
+#   axis_horizontal = self.shapedirs[2, :].reshape((-1, 3))[:, 0]
+#   all_indices =  np.arange(3889)
+#   tail_indices = all_indices[axis_horizontal.detach().cpu().numpy() < 0.0]
+VERTEX_IDS_TAIL = [   0,    4,    9,   10,   24,   25,   28,  453,  454,  456,  457,
+        458,  459,  460,  461,  462,  463,  464,  465,  466,  467,  468,
+        469,  470,  471,  472,  473,  474,  475,  724,  725,  726,  727,
+        728,  729,  730,  731,  813,  975,  976,  977, 1109, 1110, 1111,
+       1811, 1813, 1819, 1820, 1821, 1822, 1823, 1824, 1825, 1826, 1827,
+       1828, 1835, 1836, 1960, 1961, 1962, 1963, 1964, 1965, 1966, 1967,
+       1968, 1969, 2418, 2419, 2421, 2422, 2423, 2424, 2425, 2426, 2427,
+       2428, 2429, 2430, 2431, 2432, 2433, 2434, 2435, 2436, 2437, 2438,
+       2439, 2440, 2688, 2689, 2690, 2691, 2692, 2693, 2694, 2695, 2777,
+       3067, 3068, 3069, 3842, 3843, 3844, 3845, 3846, 3847]
+
+# same as in https://github.com/benjiebob/WLDO/blob/master/global_utils/config.py
+EVAL_KEYPOINTS = [
+  0, 1, 2, # left front
+  3, 4, 5, # left rear
+  6, 7, 8, # right front
+  9, 10, 11, # right rear
+  12, 13, # tail start -> end
+  14, 15, # left ear, right ear
+  16, 17, # nose, chin
+  18, 19] # left tip, right tip
+
+KEYPOINT_GROUPS = {
+  'legs': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], # legs
+  'tail': [12, 13], # tail
+  'ears': [14, 15, 18, 19], # ears
+  'face': [16, 17] # face
+}
+
+

src/configs/anipose_data_info.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from typing import List
+import json
+import numpy as np
+import os
+
+STATISTICS_DATA_DIR = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'statistics')
+STATISTICS_PATH = os.path.join(STATISTICS_DATA_DIR, 'statistics_modified_v1.json')
+
+@dataclass
+class DataInfo:
+    rgb_mean: List[float]
+    rgb_stddev: List[float]
+    joint_names: List[str]
+    hflip_indices: List[int]
+    n_joints: int
+    n_keyp: int
+    n_bones: int
+    n_betas: int
+    image_size: int
+    trans_mean: np.ndarray
+    trans_std: np.ndarray
+    flength_mean: np.ndarray
+    flength_std: np.ndarray
+    pose_rot6d_mean: np.ndarray
+    keypoint_weights: List[float]
+
+# SMAL samples 3d statistics
+#   statistics like mean values were calculated once when the project was started and they were not changed afterwards anymore
+def load_statistics(statistics_path):
+    with open(statistics_path) as f:
+        statistics = json.load(f)
+    '''new_pose_mean = [[[np.round(val, 2) for val in sublst] for sublst in sublst_big] for sublst_big in statistics['pose_mean']]
+    statistics['pose_mean'] = new_pose_mean
+    j_out = json.dumps(statistics, indent=4)    #, sort_keys=True)
+    with open(self.statistics_path, 'w') as file: file.write(j_out)'''
+    new_statistics = {'trans_mean': np.asarray(statistics['trans_mean']),
+                    'trans_std': np.asarray(statistics['trans_std']),       
+                    'flength_mean': np.asarray(statistics['flength_mean']),
+                    'flength_std': np.asarray(statistics['flength_std']),  
+                    'pose_mean': np.asarray(statistics['pose_mean']),
+                    }
+    new_statistics['pose_rot6d_mean'] = new_statistics['pose_mean'][:, :, :2].reshape((-1, 6))
+    return new_statistics
+STATISTICS = load_statistics(STATISTICS_PATH)
+
+AniPose_JOINT_NAMES_swapped = [
+    'L_F_Paw', 'L_F_Knee', 'L_F_Elbow', 
+    'L_B_Paw', 'L_B_Knee', 'L_B_Elbow', 
+    'R_F_Paw', 'R_F_Knee', 'R_F_Elbow',
+    'R_B_Paw', 'R_B_Knee', 'R_B_Elbow', 
+    'TailBase', '_Tail_end_', 'L_EarBase', 'R_EarBase',
+    'Nose', '_Chin_', '_Left_ear_tip_', '_Right_ear_tip_',
+    'L_Eye', 'R_Eye', 'Withers', 'Throat']
+
+KEYPOINT_WEIGHTS = [3, 2, 2, 3, 2, 2, 3, 2, 2, 3, 2, 2, 3, 3, 2, 2, 3, 1, 2, 2]      
+
+COMPLETE_DATA_INFO = DataInfo(
+    rgb_mean=[0.4404, 0.4440, 0.4327],      # not sure
+    rgb_stddev=[0.2458, 0.2410, 0.2468],    # not sure
+    joint_names=AniPose_JOINT_NAMES_swapped,        # AniPose_JOINT_NAMES,
+    hflip_indices=[6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5, 12, 13, 15, 14, 16, 17, 19, 18, 21, 20, 22, 23],
+    n_joints = 35,
+    n_keyp = 24,    # 20,    # 25,
+    n_bones = 24,
+    n_betas = 30,       # 10,
+    image_size = 256,
+    trans_mean = STATISTICS['trans_mean'],
+    trans_std = STATISTICS['trans_std'],
+    flength_mean = STATISTICS['flength_mean'],
+    flength_std = STATISTICS['flength_std'],
+    pose_rot6d_mean = STATISTICS['pose_rot6d_mean'],
+    keypoint_weights = KEYPOINT_WEIGHTS
+    )

src/configs/barc_cfg_defaults.py

@@ -0,0 +1,121 @@
+
+from yacs.config import CfgNode as CN
+import argparse
+import yaml
+import os
+
+abs_barc_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..',))
+
+_C = CN()
+_C.barc_dir = abs_barc_dir
+_C.device = 'cuda'
+
+## path settings
+_C.paths = CN()
+_C.paths.ROOT_OUT_PATH = abs_barc_dir + '/results/'
+_C.paths.ROOT_CHECKPOINT_PATH = abs_barc_dir + '/checkpoint/'
+_C.paths.MODELPATH_NORMFLOW = abs_barc_dir + '/checkpoint/barc_normflow_pret/rgbddog_v3_model.pt'
+
+## parameter settings
+_C.params = CN()
+_C.params.ARCH = 'hg8'    
+_C.params.STRUCTURE_POSE_NET = 'normflow'     # 'default'   # 'vae' 
+_C.params.NF_VERSION = 3
+_C.params.N_JOINTS = 35   
+_C.params.N_KEYP = 24      #20    
+_C.params.N_SEG = 2
+_C.params.N_PARTSEG = 15
+_C.params.UPSAMPLE_SEG = True
+_C.params.ADD_PARTSEG = True   # partseg: for the CVPR paper this part of the network exists, but is not trained (no part labels in StanExt)
+_C.params.N_BETAS = 30    # 10
+_C.params.N_BETAS_LIMBS = 7
+_C.params.N_BONES = 24
+_C.params.N_BREEDS = 121      # 120 breeds plus background
+_C.params.IMG_SIZE = 256
+_C.params.SILH_NO_TAIL = False
+_C.params.KP_THRESHOLD = None    
+_C.params.ADD_Z_TO_3D_INPUT = False   
+_C.params.N_SEGBPS = 64*2
+_C.params.ADD_SEGBPS_TO_3D_INPUT = True
+_C.params.FIX_FLENGTH = False   
+_C.params.RENDER_ALL = True
+_C.params.VLIN = 2    
+_C.params.STRUCTURE_Z_TO_B = 'lin'
+_C.params.N_Z_FREE = 64
+_C.params.PCK_THRESH = 0.15   
+_C.params.REF_NET_TYPE = 'add' # refinement network type
+_C.params.REF_DETACH_SHAPE = True
+_C.params.GRAPHCNN_TYPE = 'inexistent'
+_C.params.ISFLAT_TYPE = 'inexistent'
+_C.params.SHAPEREF_TYPE = 'inexistent'
+
+## SMAL settings
+_C.smal = CN()
+_C.smal.SMAL_MODEL_TYPE = 'barc'    
+_C.smal.SMAL_KEYP_CONF = 'green'    
+
+## optimization settings
+_C.optim = CN()
+_C.optim.LR = 5e-4
+_C.optim.SCHEDULE = [150, 175, 200]
+_C.optim.GAMMA = 0.1
+_C.optim.MOMENTUM = 0
+_C.optim.WEIGHT_DECAY = 0
+_C.optim.EPOCHS = 220
+_C.optim.BATCH_SIZE = 12       # keep 12 (needs to be an even number, as we have a custom data sampler)
+_C.optim.TRAIN_PARTS = 'all_without_shapedirs'
+
+## dataset settings
+_C.data = CN()
+_C.data.DATASET = 'stanext24'
+_C.data.V12 = True     
+_C.data.SHORTEN_VAL_DATASET_TO = None        
+_C.data.VAL_OPT = 'val'
+_C.data.VAL_METRICS = 'no_loss'
+
+# ---------------------------------------
+def update_dependent_vars(cfg):    
+    cfg.params.N_CLASSES = cfg.params.N_KEYP + cfg.params.N_SEG
+    if cfg.params.VLIN == 0: 
+        cfg.params.NUM_STAGE_COMB = 2
+        cfg.params.NUM_STAGE_HEADS = 1  
+        cfg.params.NUM_STAGE_HEADS_POSE = 1
+        cfg.params.TRANS_SEP = False
+    elif cfg.params.VLIN == 1:
+        cfg.params.NUM_STAGE_COMB = 3              
+        cfg.params.NUM_STAGE_HEADS = 1             
+        cfg.params.NUM_STAGE_HEADS_POSE = 2        
+        cfg.params.TRANS_SEP = False
+    elif cfg.params.VLIN == 2:
+        cfg.params.NUM_STAGE_COMB = 3              
+        cfg.params.NUM_STAGE_HEADS = 1             
+        cfg.params.NUM_STAGE_HEADS_POSE = 2        
+        cfg.params.TRANS_SEP = True
+    else:
+        raise NotImplementedError
+    if cfg.params.STRUCTURE_Z_TO_B == '1dconv':
+        cfg.params.N_Z = cfg.params.N_BETAS + cfg.params.N_BETAS_LIMBS
+    else:
+        cfg.params.N_Z = cfg.params.N_Z_FREE
+    return
+
+
+update_dependent_vars(_C)
+global _cfg_global 
+_cfg_global = _C.clone()
+
+
+def get_cfg_defaults():
+    # Get a yacs CfgNode object with default values as defined within this file.
+    # Return a clone so that the defaults will not be altered.
+    return _C.clone()
+
+def update_cfg_global_with_yaml(cfg_yaml_file):    
+    _cfg_global.merge_from_file(cfg_yaml_file)
+    update_dependent_vars(_cfg_global)
+    return 
+
+def get_cfg_global_updated():
+    # return _cfg_global.clone()
+    return _cfg_global
+

src/configs/barc_cfg_train.yaml

@@ -0,0 +1,24 @@
+
+paths:
+  ROOT_OUT_PATH: './results/'
+  ROOT_CHECKPOINT_PATH: './checkpoint/'
+  MODELPATH_NORMFLOW: './checkpoint/barc_normflow_pret/rgbddog_v3_model.pt'
+
+smal:
+  SMAL_MODEL_TYPE: '39dogs_norm_newv3'  # '39dogs_diffsize'   # 'barc'
+  SMAL_KEYP_CONF: 'olive'   # 'green'
+
+optim:
+  LR: 5e-4
+  SCHEDULE: [150, 175, 200]
+  GAMMA: 0.1
+  MOMENTUM: 0
+  WEIGHT_DECAY: 0
+  EPOCHS: 220
+  BATCH_SIZE: 12       # keep 12 (needs to be an even number, as we have a custom data sampler)
+  TRAIN_PARTS: 'all_without_shapedirs'
+
+data:
+  DATASET: 'stanext24'
+  SHORTEN_VAL_DATASET_TO: 600  # this is faster as we do not evaluate on the whole validation set  
+  VAL_OPT: 'val'

src/configs/barc_loss_weights_allzeros.json

@@ -0,0 +1,30 @@
+
+
+
+{
+    "breed_options": [
+        "4"             
+    ],
+    "breed": 0.0,
+    "class": 0.0,
+    "models3d": 0.0,
+    "keyp": 0.0,
+    "silh": 0.0,
+    "shape_options": [
+        "smal",
+        "limbs7"
+    ],
+    "shape": [
+        0,
+        0
+    ],
+    "poseprior_options": [
+        "normalizing_flow_tiger_logprob"
+    ],
+    "poseprior": 0.0,
+    "poselegssidemovement": 0.0,
+    "flength": 0.0,
+    "partseg": 0,
+    "shapedirs": 0,
+    "pose_0": 0.0
+}

src/configs/barc_loss_weights_with3dcgloss_higherbetaloss_v2_dm39dnnv3v2.json

@@ -0,0 +1,30 @@
+
+
+
+{
+    "breed_options": [
+        "4"             
+    ],
+    "breed": 5.0,
+    "class": 5.0,
+    "models3d": 0.1,
+    "keyp": 0.2,
+    "silh": 50.0,
+    "shape_options": [
+        "smal",
+        "limbs7"
+    ],
+    "shape": [
+        0.1,
+        1.0
+    ],
+    "poseprior_options": [
+        "normalizing_flow_tiger_logprob"
+    ],
+    "poseprior": 0.1,
+    "poselegssidemovement": 10.0,
+    "flength": 1.0,
+    "partseg": 0,
+    "shapedirs": 0,
+    "pose_0": 0.0
+}

src/configs/data_info.py

@@ -0,0 +1,115 @@
+from dataclasses import dataclass
+from typing import List
+import json
+import numpy as np
+import os
+import sys
+
+STATISTICS_DATA_DIR = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'statistics')
+STATISTICS_PATH = os.path.join(STATISTICS_DATA_DIR, 'statistics_modified_v1.json')
+
+@dataclass
+class DataInfo:
+    rgb_mean: List[float]
+    rgb_stddev: List[float]
+    joint_names: List[str]
+    hflip_indices: List[int]
+    n_joints: int
+    n_keyp: int
+    n_bones: int
+    n_betas: int
+    image_size: int
+    trans_mean: np.ndarray
+    trans_std: np.ndarray
+    flength_mean: np.ndarray
+    flength_std: np.ndarray
+    pose_rot6d_mean: np.ndarray
+    keypoint_weights: List[float]
+
+# SMAL samples 3d statistics
+#   statistics like mean values were calculated once when the project was started and they were not changed afterwards anymore
+def load_statistics(statistics_path):
+    with open(statistics_path) as f:
+        statistics = json.load(f)
+    '''new_pose_mean = [[[np.round(val, 2) for val in sublst] for sublst in sublst_big] for sublst_big in statistics['pose_mean']]
+    statistics['pose_mean'] = new_pose_mean
+    j_out = json.dumps(statistics, indent=4)    #, sort_keys=True)
+    with open(self.statistics_path, 'w') as file: file.write(j_out)'''
+    new_statistics = {'trans_mean': np.asarray(statistics['trans_mean']),
+                    'trans_std': np.asarray(statistics['trans_std']),       
+                    'flength_mean': np.asarray(statistics['flength_mean']),
+                    'flength_std': np.asarray(statistics['flength_std']),  
+                    'pose_mean': np.asarray(statistics['pose_mean']),
+                    }
+    new_statistics['pose_rot6d_mean'] = new_statistics['pose_mean'][:, :, :2].reshape((-1, 6))
+    return new_statistics
+STATISTICS = load_statistics(STATISTICS_PATH)
+
+
+############################################################################
+# for StanExt (original number of keypoints, 20 not 24)
+
+# for keypoint names see: https://github.com/benjiebob/StanfordExtra/blob/master/keypoint_definitions.csv
+StanExt_JOINT_NAMES = [
+    'Left_front_leg_paw', 'Left_front_leg_middle_joint', 'Left_front_leg_top', 
+    'Left_rear_leg_paw', 'Left_rear_leg_middle_joint', 'Left_rear_leg_top',
+    'Right_front_leg_paw', 'Right_front_leg_middle_joint', 'Right_front_leg_top',
+    'Right_rear_leg_paw', 'Right_rear_leg_middle_joint', 'Right_rear_leg_top',
+    'Tail_start', 'Tail_end', 'Base_of_left_ear', 'Base_of_right_ear', 
+    'Nose', 'Chin', 'Left_ear_tip', 'Right_ear_tip']
+
+KEYPOINT_WEIGHTS = [3, 2, 2, 3, 2, 2, 3, 2, 2, 3, 2, 2, 3, 3, 2, 2, 3, 1, 2, 2]      
+
+COMPLETE_DATA_INFO = DataInfo(
+    rgb_mean=[0.4404, 0.4440, 0.4327],      # not sure
+    rgb_stddev=[0.2458, 0.2410, 0.2468],    # not sure
+    joint_names=StanExt_JOINT_NAMES,
+    hflip_indices=[6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5, 12, 13, 15, 14, 16, 17, 19, 18],
+    n_joints = 35,
+    n_keyp = 20,    # 25,
+    n_bones = 24,
+    n_betas = 30,       # 10,
+    image_size = 256,
+    trans_mean = STATISTICS['trans_mean'],
+    trans_std = STATISTICS['trans_std'],
+    flength_mean = STATISTICS['flength_mean'],
+    flength_std = STATISTICS['flength_std'],
+    pose_rot6d_mean = STATISTICS['pose_rot6d_mean'],
+    keypoint_weights = KEYPOINT_WEIGHTS
+    )
+
+
+############################################################################
+# new for StanExt24
+
+# ..., 'Left_eye', 'Right_eye', 'Withers', 'Throat']       # the last 4 keypoints are in the animal_pose dataset, but not StanfordExtra
+StanExt_JOINT_NAMES_24 = [
+    'Left_front_leg_paw', 'Left_front_leg_middle_joint', 'Left_front_leg_top', 
+    'Left_rear_leg_paw', 'Left_rear_leg_middle_joint', 'Left_rear_leg_top',
+    'Right_front_leg_paw', 'Right_front_leg_middle_joint', 'Right_front_leg_top',
+    'Right_rear_leg_paw', 'Right_rear_leg_middle_joint', 'Right_rear_leg_top',
+    'Tail_start', 'Tail_end', 'Base_of_left_ear', 'Base_of_right_ear', 
+    'Nose', 'Chin', 'Left_ear_tip', 'Right_ear_tip',
+    'Left_eye', 'Right_eye', 'Withers', 'Throat']
+
+KEYPOINT_WEIGHTS_24 = [3, 2, 2, 3, 2, 2, 3, 2, 2, 3, 2, 2, 3, 3, 2, 2, 3, 1, 2, 2, 1, 1, 0, 0]     
+
+COMPLETE_DATA_INFO_24 = DataInfo(
+    rgb_mean=[0.4404, 0.4440, 0.4327],      # not sure
+    rgb_stddev=[0.2458, 0.2410, 0.2468],    # not sure
+    joint_names=StanExt_JOINT_NAMES_24,
+    hflip_indices=[6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5, 12, 13, 15, 14, 16, 17, 19, 18, 21, 20, 22, 23],
+    n_joints = 35,
+    n_keyp = 24,    # 20,    # 25,
+    n_bones = 24,
+    n_betas = 30,       # 10,
+    image_size = 256,
+    trans_mean = STATISTICS['trans_mean'],
+    trans_std = STATISTICS['trans_std'],
+    flength_mean = STATISTICS['flength_mean'],
+    flength_std = STATISTICS['flength_std'],
+    pose_rot6d_mean = STATISTICS['pose_rot6d_mean'],
+    keypoint_weights = KEYPOINT_WEIGHTS_24
+    )
+
+

src/configs/dataset_path_configs.py

@@ -0,0 +1,21 @@
+
+
+import numpy as np
+import os
+import sys
+
+abs_barc_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..',))
+
+# stanext dataset
+# (1) path to stanext dataset
+STAN_V12_ROOT_DIR = '/ps/scratch/nrueegg/new_projects/Animals/data/dog_datasets/Stanford_Dogs_Dataset' + '/StanfordExtra_V12/'
+IMG_V12_DIR = os.path.join(STAN_V12_ROOT_DIR, 'StanExtV12_Images')		      
+JSON_V12_DIR = os.path.join(STAN_V12_ROOT_DIR, 'labels', "StanfordExtra_v12.json")    
+STAN_V12_TRAIN_LIST_DIR = os.path.join(STAN_V12_ROOT_DIR, 'labels', 'train_stanford_StanfordExtra_v12.npy')  
+STAN_V12_VAL_LIST_DIR = os.path.join(STAN_V12_ROOT_DIR, 'labels', 'val_stanford_StanfordExtra_v12.npy')  
+STAN_V12_TEST_LIST_DIR = os.path.join(STAN_V12_ROOT_DIR, 'labels', 'test_stanford_StanfordExtra_v12.npy')  
+# (2) path to related data such as breed indices and prepared predictions for withers, throat and eye keypoints 
+STANEXT_RELATED_DATA_ROOT_DIR = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'stanext_related_data')
+
+# test image crop dataset
+TEST_IMAGE_CROP_ROOT_DIR = os.path.join(os.path.dirname(__file__), '..', '..', 'datasets', 'test_image_crops') 

src/configs/dog_breeds/dog_breed_class.py

@@ -0,0 +1,170 @@
+
+import os
+import warnings
+warnings.filterwarnings("ignore", category=DeprecationWarning) 
+import pandas as pd
+import difflib
+import json
+import pickle as pkl
+import csv
+import numpy as np
+
+
+# ----------------------------------------------------------------------------------------------------------------- #
+class DogBreed(object):
+    def __init__(self, abbrev, name_akc=None, name_stanext=None, name_xlsx=None, path_akc=None, path_stanext=None, ind_in_xlsx=None, ind_in_xlsx_matrix=None, ind_in_stanext=None, clade=None):
+        self._abbrev = abbrev 
+        self._name_xlsx = name_xlsx 
+        self._name_akc = name_akc
+        self._name_stanext = name_stanext
+        self._path_stanext = path_stanext
+        self._additional_names = set()
+        if self._name_akc is not None:
+            self.add_akc_info(name_akc, path_akc)
+        if self._name_stanext is not None:
+            self.add_stanext_info(name_stanext, path_stanext, ind_in_stanext)
+        if self._name_xlsx is not None:
+            self.add_xlsx_info(name_xlsx, ind_in_xlsx, ind_in_xlsx_matrix, clade)
+    def add_xlsx_info(self, name_xlsx, ind_in_xlsx, ind_in_xlsx_matrix, clade):
+        assert (name_xlsx is not None) and (ind_in_xlsx is not None) and (ind_in_xlsx_matrix is not None) and (clade is not None)
+        self._name_xlsx = name_xlsx
+        self._ind_in_xlsx = ind_in_xlsx
+        self._ind_in_xlsx_matrix = ind_in_xlsx_matrix
+        self._clade = clade
+    def add_stanext_info(self, name_stanext, path_stanext, ind_in_stanext):
+        assert (name_stanext is not None) and (path_stanext is not None) and (ind_in_stanext is not None)
+        self._name_stanext = name_stanext
+        self._path_stanext = path_stanext
+        self._ind_in_stanext = ind_in_stanext
+    def add_akc_info(self, name_akc, path_akc):
+        assert (name_akc is not None) and (path_akc is not None)
+        self._name_akc = name_akc
+        self._path_akc = path_akc
+    def add_additional_names(self, name_list):
+        self._additional_names = self._additional_names.union(set(name_list)) 
+    def add_text_info(self, text_height, text_weight, text_life_exp):
+        self._text_height = text_height
+        self._text_weight = text_weight
+        self._text_life_exp = text_life_exp
+    def get_datasets(self):
+        # all datasets in which this breed is found
+        datasets = set()
+        if self._name_akc is not None:
+            datasets.add('akc')
+        if self._name_stanext is not None:
+            datasets.add('stanext')
+        if self._name_xlsx is not None:
+            datasets.add('xlsx')
+        return datasets
+    def get_names(self):
+        # set of names for this breed
+        names = {self._abbrev, self._name_akc, self._name_stanext, self._name_xlsx, self._path_stanext}.union(self._additional_names)
+        names.discard(None)
+        return names
+    def get_names_as_pointing_dict(self):
+        # each name points to the abbreviation
+        names = self.get_names()
+        my_dict = {}
+        for name in names:
+            my_dict[name] = self._abbrev
+        return my_dict
+    def print_overview(self):
+        # print important information to get an overview of the class instance
+        if self._name_akc is not None:
+            name = self._name_akc
+        elif self._name_xlsx is not None:
+            name = self._name_xlsx
+        else:
+            name = self._name_stanext
+        print('----------------------------------------------------')
+        print('----- dog breed: ' + name )
+        print('----------------------------------------------------')
+        print('[names]')
+        print(self.get_names())
+        print('[datasets]')
+        print(self.get_datasets())
+        # see https://stackoverflow.com/questions/9058305/getting-attributes-of-a-class
+        print('[instance attributes]')
+        for attribute, value in self.__dict__.items():
+            print(attribute, '=', value)
+    def use_dict_to_save_class_instance(self):
+        my_dict = {}
+        for attribute, value in self.__dict__.items():
+            my_dict[attribute] = value
+        return my_dict
+    def use_dict_to_load_class_instance(self, my_dict):
+        for attribute, value in my_dict.items():
+            setattr(self, attribute, value)
+        return 
+
+# ----------------------------------------------------------------------------------------------------------------- #
+def get_name_list_from_summary(summary):
+    name_from_abbrev_dict = {}
+    for breed in summary.values():
+        abbrev = breed._abbrev
+        all_names = breed.get_names()
+        name_from_abbrev_dict[abbrev] = list(all_names)
+    return name_from_abbrev_dict
+def get_partial_summary(summary, part):
+    assert part in ['xlsx', 'akc', 'stanext']
+    partial_summary = {}
+    for key, value in summary.items():
+        if (part == 'xlsx' and value._name_xlsx is not None) \
+            or (part == 'akc' and value._name_akc is not None) \
+            or (part == 'stanext' and value._name_stanext is not None):
+            partial_summary[key] = value
+    return partial_summary
+def get_akc_but_not_stanext_partial_summary(summary):
+    partial_summary = {}
+    for key, value in summary.items():
+        if value._name_akc is not None:
+            if value._name_stanext is None:
+                partial_summary[key] = value
+    return partial_summary    
+
+# ----------------------------------------------------------------------------------------------------------------- #
+def main_load_dog_breed_classes(path_complete_abbrev_dict_v1, path_complete_summary_breeds_v1):
+    with open(path_complete_abbrev_dict_v1, 'rb') as file:
+        complete_abbrev_dict = pkl.load(file)
+    with open(path_complete_summary_breeds_v1, 'rb') as file: 
+        complete_summary_breeds_attributes_only = pkl.load(file)
+    
+    complete_summary_breeds = {}
+    for key, value in complete_summary_breeds_attributes_only.items():
+        attributes_only = complete_summary_breeds_attributes_only[key]
+        complete_summary_breeds[key] = DogBreed(abbrev=attributes_only['_abbrev'])
+        complete_summary_breeds[key].use_dict_to_load_class_instance(attributes_only)
+    return complete_abbrev_dict, complete_summary_breeds
+
+
+# ----------------------------------------------------------------------------------------------------------------- #
+def load_similarity_matrix_raw(xlsx_path):
+    # --- LOAD EXCEL FILE FROM DOG BREED PAPER
+    xlsx = pd.read_excel(xlsx_path)
+    # create an array
+    abbrev_indices = {}
+    matrix_raw = np.zeros((168, 168))
+    for ind in range(1, 169):
+        abbrev = xlsx[xlsx.columns[2]][ind]
+        abbrev_indices[abbrev] = ind-1
+    for ind_col in range(0, 168):
+        for ind_row in range(0, 168):
+            matrix_raw[ind_col, ind_row] = float(xlsx[xlsx.columns[3+ind_col]][1+ind_row])
+    return matrix_raw, abbrev_indices
+
+
+
+# ----------------------------------------------------------------------------------------------------------------- #
+# ----------------------------------------------------------------------------------------------------------------- #
+# load the (in advance created) final dict of dog breed classes
+ROOT_PATH_BREED_DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', '..', 'data', 'breed_data')
+path_complete_abbrev_dict_v1 = os.path.join(ROOT_PATH_BREED_DATA, 'complete_abbrev_dict_v2.pkl')
+path_complete_summary_breeds_v1 = os.path.join(ROOT_PATH_BREED_DATA, 'complete_summary_breeds_v2.pkl')
+COMPLETE_ABBREV_DICT, COMPLETE_SUMMARY_BREEDS = main_load_dog_breed_classes(path_complete_abbrev_dict_v1, path_complete_summary_breeds_v1)
+# load similarity matrix, data from: 
+#   Parker H. G., Dreger D. L., Rimbault M., Davis B. W., Mullen A. B., Carpintero-Ramirez G., and Ostrander E. A.
+#   Genomic analyses reveal the influence of geographic origin, migration, and hybridization on modern dog breed 
+#   development. Cell Reports, 4(19):697–708, 2017.
+xlsx_path = os.path.join(ROOT_PATH_BREED_DATA, 'NIHMS866262-supplement-2.xlsx')
+SIM_MATRIX_RAW, SIM_ABBREV_INDICES = load_similarity_matrix_raw(xlsx_path)
+

src/configs/refinement_cfg_test_withvertexwisegc_csaddnonflat.yaml

@@ -0,0 +1,23 @@
+
+paths:
+  ROOT_OUT_PATH: './results/'
+  ROOT_CHECKPOINT_PATH: './checkpoint/'
+  MODELPATH_NORMFLOW: './checkpoint/barc_normflow_pret/rgbddog_v3_model.pt'
+
+smal:
+  SMAL_MODEL_TYPE: '39dogs_norm_newv3'  # '39dogs_norm'  # '39dogs_diffsize'   # 'barc'
+  SMAL_KEYP_CONF: 'olive'       # 'green'
+
+optim:
+  BATCH_SIZE: 12       
+
+params:
+  REF_NET_TYPE: 'multrot01all_res34' # 'multrot01all_res34'    # 'multrot01all'    # 'multrot01'  # 'multrot' # 'add'
+  REF_DETACH_SHAPE: True
+  GRAPHCNN_TYPE: 'multistage_simple'  # 'inexistent'
+  SHAPEREF_TYPE: 'inexistent'   # 'linear'  # 'inexistent'
+  ISFLAT_TYPE: 'linear'     # 'inexistent'  # 'inexistent'
+
+data:
+  DATASET: 'stanext24'
+  VAL_OPT: 'test'  # 'val'

src/configs/refinement_cfg_test_withvertexwisegc_csaddnonflat_crops.yaml

@@ -0,0 +1,23 @@
+
+paths:
+  ROOT_OUT_PATH: './results/'
+  ROOT_CHECKPOINT_PATH: './checkpoint/'
+  MODELPATH_NORMFLOW: './checkpoint/barc_normflow_pret/rgbddog_v3_model.pt'
+
+smal:
+  SMAL_MODEL_TYPE: '39dogs_norm_newv3'  # '39dogs_norm'  # '39dogs_diffsize'   # 'barc'
+  SMAL_KEYP_CONF: 'olive'       # 'green'
+
+optim:
+  BATCH_SIZE: 12       
+
+params:
+  REF_NET_TYPE: 'multrot01all_res34' # 'multrot01all_res34'    # 'multrot01all'    # 'multrot01'  # 'multrot' # 'add'
+  REF_DETACH_SHAPE: True
+  GRAPHCNN_TYPE: 'multistage_simple'  # 'inexistent'
+  SHAPEREF_TYPE: 'inexistent'   # 'linear'  # 'inexistent'
+  ISFLAT_TYPE: 'linear'     # 'inexistent'  # 'inexistent'
+
+data:
+  DATASET: 'ImgCropList'
+  VAL_OPT: 'test'  # 'val'

src/configs/refinement_cfg_train_withvertexwisegc_isflat_csmorestanding.yaml

@@ -0,0 +1,31 @@
+
+paths:
+  ROOT_OUT_PATH: './results/'
+  ROOT_CHECKPOINT_PATH: './checkpoint/'
+  MODELPATH_NORMFLOW: './checkpoint/barc_normflow_pret/rgbddog_v3_model.pt'
+
+smal:
+  SMAL_MODEL_TYPE: '39dogs_norm_newv3'    # '39dogs_norm'  # '39dogs_diffsize'  # 'barc'
+  SMAL_KEYP_CONF: 'olive'   # 'green'
+
+optim:
+  LR: 5e-5    # 5e-7  # (new) 5e-6   # 5e-5    #  5e-5  # 5e-4
+  SCHEDULE: [150, 175, 200]   # [220, 270]   # [150, 175, 200]
+  GAMMA: 0.1
+  MOMENTUM: 0
+  WEIGHT_DECAY: 0
+  EPOCHS: 220   # 300
+  BATCH_SIZE: 14  # 12       # keep 12 (needs to be an even number, as we have a custom data sampler)
+  TRAIN_PARTS: 'refinement_model'   # 'refinement_model_and_shape' # 'refinement_model'
+
+params:
+  REF_NET_TYPE: 'multrot01all_res34' # 'multrot01all_res34'    # 'multrot01all'    # 'multrot01'   # 'multrot01'   # 'multrot01'  # 'multrot'   # 'multrot_res34' # 'multrot' # 'add'
+  REF_DETACH_SHAPE: True
+  GRAPHCNN_TYPE: 'multistage_simple'  # 'inexistent'
+  SHAPEREF_TYPE: 'inexistent'   # 'linear'  # 'inexistent'
+  ISFLAT_TYPE: 'linear'     # 'inexistent'  # 'inexistent'
+
+data:
+  DATASET: 'stanext24_withgc_csaddnonflatmorestanding'  #  'stanext24_withgc_csaddnonflat'  # 'stanext24_withgc_cs0'
+  SHORTEN_VAL_DATASET_TO: 600  # this is faster as we do not evaluate on the whole validation set  
+  VAL_OPT: 'val'

src/configs/refinement_loss_weights_withgc_withvertexwise_addnonflat.json

@@ -0,0 +1,20 @@
+
+
+
+{
+    "keyp_ref": 0.2,
+    "silh_ref": 50.0,
+    "pose_legs_side": 1.0,
+    "pose_legs_tors": 1.0,
+    "pose_tail_side": 0.0,
+    "pose_tail_tors": 0.0,
+    "pose_spine_side": 0.0,
+    "pose_spine_tors": 0.0,
+    "reg_trans": 0.0,
+    "reg_flength": 0.0,
+    "reg_pose": 0.0,
+    "gc_plane": 5.0,     
+    "gc_blowplane": 5.0,        
+    "gc_vertexwise": 10.0, 
+    "gc_isflat": 0.5
+}

src/configs/ttopt_loss_weights/bite_loss_weights_ttopt.json

@@ -0,0 +1,77 @@
+{
+    "silhouette": {     
+        "weight": 40.0,   
+        "weight_vshift": 20.0,
+        "value": 0.0
+    },
+    "keyp":{                   
+        "weight": 0.2,     
+        "weight_vshift": 0.01,
+        "value": 0.0  
+    },
+    "pose_legs_side":{                  
+        "weight": 1.0,      
+        "weight_vshift": 1.0,
+        "value": 0.0  
+    },
+    "pose_legs_tors":{                  
+        "weight": 10.0,     
+        "weight_vshift": 10.0,
+        "value": 0.0  
+    },
+    "pose_tail_side":{                  
+        "weight": 1,      
+        "weight_vshift": 1,     
+        "value": 0.0  
+    },
+    "pose_tail_tors":{                   
+        "weight": 10.0,     
+        "weight_vshift": 10.0,
+        "value": 0.0  
+    },
+    "pose_spine_side":{                  
+        "weight": 0.0,      
+        "weight_vshift": 0.0,
+        "value": 0.0  
+    },
+    "pose_spine_tors":{                  
+        "weight": 0.0,      
+        "weight_vshift": 0.0,
+        "value": 0.0  
+    },
+    "gc_plane":{                   
+        "weight": 10.0,      
+        "weight_vshift": 20.0,
+        "value": 0.0  
+    },
+    "gc_belowplane":{                   
+        "weight": 10.0,      
+        "weight_vshift": 20.0,
+        "value": 0.0  
+    },
+    "lapctf": {
+        "weight": 0.0,
+        "weight_vshift": 10.0,    
+        "value": 0.0
+    },
+    "arap": {
+        "weight": 0.0,
+        "weight_vshift": 0.0,    
+        "value": 0.0
+    },
+    "edge": {
+        "weight": 0.0,
+        "weight_vshift": 10.0,  
+        "value": 0.0
+    },
+    "normal": {
+        "weight": 0.0,      
+        "weight_vshift": 1.0,   
+        "value": 0.0
+    },
+    "laplacian": {
+        "weight": 0.0,
+        "weight_vshift": 0.0,   
+        "value": 0.0
+    }
+}

src/configs/ttopt_loss_weights/ttopt_loss_weights_v2c_withlapcft_v2.json

@@ -0,0 +1,77 @@
+{
+    "silhouette": {     
+        "weight": 40.0,   
+        "weight_vshift": 20.0,
+        "value": 0.0
+    },
+    "keyp":{                   
+        "weight": 0.2,     
+        "weight_vshift": 0.01,
+        "value": 0.0  
+    },
+    "pose_legs_side":{                  
+        "weight": 1.0,      
+        "weight_vshift": 1.0,
+        "value": 0.0  
+    },
+    "pose_legs_tors":{                  
+        "weight": 10.0,     
+        "weight_vshift": 10.0,
+        "value": 0.0  
+    },
+    "pose_tail_side":{                  
+        "weight": 1,      
+        "weight_vshift": 1,     
+        "value": 0.0  
+    },
+    "pose_tail_tors":{                   
+        "weight": 10.0,     
+        "weight_vshift": 10.0,
+        "value": 0.0  
+    },
+    "pose_spine_side":{                  
+        "weight": 0.0,      
+        "weight_vshift": 0.0,
+        "value": 0.0  
+    },
+    "pose_spine_tors":{                  
+        "weight": 0.0,      
+        "weight_vshift": 0.0,
+        "value": 0.0  
+    },
+    "gc_plane":{                   
+        "weight": 10.0,      
+        "weight_vshift": 20.0,
+        "value": 0.0  
+    },
+    "gc_belowplane":{                   
+        "weight": 10.0,      
+        "weight_vshift": 20.0,
+        "value": 0.0  
+    },
+    "lapctf": {
+        "weight": 0.0,
+        "weight_vshift": 10.0,    
+        "value": 0.0
+    },
+    "arap": {
+        "weight": 0.0,
+        "weight_vshift": 0.0,    
+        "value": 0.0
+    },
+    "edge": {
+        "weight": 0.0,
+        "weight_vshift": 10.0,  
+        "value": 0.0
+    },
+    "normal": {
+        "weight": 0.0,      
+        "weight_vshift": 1.0,   
+        "value": 0.0
+    },
+    "laplacian": {
+        "weight": 0.0,
+        "weight_vshift": 0.0,   
+        "value": 0.0
+    }
+}

src/graph_networks/graphcmr/get_downsampled_mesh_npz.py

@@ -0,0 +1,84 @@
+
+# try to use aenv_conda3  (maybe also export PYOPENGL_PLATFORM=osmesa)
+# python src/graph_networks/graphcmr/get_downsampled_mesh_npz.py
+
+# see https://github.com/nkolot/GraphCMR/issues/35
+
+
+from __future__ import print_function
+# import mesh_sampling
+from psbody.mesh import Mesh, MeshViewer, MeshViewers
+import numpy as np
+import json
+import os
+import copy
+import argparse
+import pickle
+import time
+import sys
+import trimesh
+
+
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "../../../../"))
+from barc_for_bite.src.graph_networks.graphcmr.pytorch_coma_mesh_operations import generate_transform_matrices
+from barc_for_bite.src.configs.SMAL_configs import SMAL_MODEL_CONFIG
+from barc_for_bite.src.smal_pytorch.smal_model.smal_torch_new import SMAL
+# smal_model_path = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/smal_data/new_dog_models/my_smpl_00791_nadine_Jr_4_dog.pkl'
+
+
+SMAL_MODEL_TYPE = '39dogs_diffsize'        # '39dogs_diffsize'     # '39dogs_norm'  # 'barc'
+smal_model_path = SMAL_MODEL_CONFIG[SMAL_MODEL_TYPE]['smal_model_path']
+
+# data_path_root = "/is/cluster/work/nrueegg/icon_pifu_related/ICON/lib/graph_networks/graphcmr/data/"
+data_path_root = "/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/graphcmr/data/"
+
+smal_dog_model_name = os.path.basename(smal_model_path).split('.pkl')[0]    # 'my_smpl_SMBLD_nbj_v3'
+suffix = "_template"
+template_obj_path = data_path_root + smal_dog_model_name + suffix + ".obj"
+
+print("Loading smal .. ")
+print(SMAL_MODEL_TYPE)
+print(smal_model_path)
+
+smal = SMAL(smal_model_type=SMAL_MODEL_TYPE, template_name='neutral')
+smal_verts = smal.v_template.detach().cpu().numpy()     # (3889, 3)
+smal_faces = smal.f                                     # (7774, 3)
+smal_trimesh = trimesh.base.Trimesh(vertices=smal_verts, faces=smal_faces, process=False,  maintain_order=True)
+smal_trimesh.export(file_obj=template_obj_path)  # file_type='obj')
+
+
+print("Loading data .. ")
+reference_mesh_file = template_obj_path # 'data/barc_neutral_vertices.obj'      # 'data/smpl_neutral_vertices.obj'
+reference_mesh = Mesh(filename=reference_mesh_file)
+
+# ds_factors = [4, 4]     # ds_factors = [4,1]	# Sampling factor of the mesh at each stage of sampling
+ds_factors = [4, 4, 4, 4]
+print("Generating Transform Matrices ..")
+
+
+# Generates adjecency matrices A, downsampling matrices D, and upsamling matrices U by sampling
+# the mesh 4 times. Each time the mesh is sampled by a factor of 4
+
+# M,A,D,U = mesh_sampling.generate_transform_matrices(reference_mesh, ds_factors)
+M,A,D,U = generate_transform_matrices(reference_mesh, ds_factors)
+
+# REMARK: there is a warning:
+#   lib/graph_networks/graphcmr/../../../lib/graph_networks/graphcmr/pytorch_coma_mesh_operations.py:237: FutureWarning: `rcond` parameter will 
+#   change to the default of machine precision times ``max(M, N)`` where M and N are the input matrix dimensions.
+#   To use the future default and silence this warning we advise to pass `rcond=None`, to keep using the old, explicitly pass `rcond=-1`.
+
+
+print(type(A))
+np.savez(data_path_root + 'mesh_downsampling_' + smal_dog_model_name + suffix + '.npz', A = A, D = D, U = U)
+np.savez(data_path_root + 'meshes/' + 'mesh_downsampling_meshes' + smal_dog_model_name + suffix + '.npz', M = M)
+
+for ind_m, my_mesh in enumerate(M):
+    new_suffix = '_template_downsampled' + str(ind_m)
+    my_mesh_tri = trimesh.Trimesh(vertices=my_mesh.v, faces=my_mesh.f, process=False,  maintain_order=True)
+    my_mesh_tri.export(data_path_root + 'meshes/' + 'mesh_downsampling_meshes' + smal_dog_model_name + new_suffix + '.obj')
+
+
+
+
+

src/graph_networks/graphcmr/graph_cnn.py

@@ -0,0 +1,53 @@
+"""
+code from https://raw.githubusercontent.com/nkolot/GraphCMR/master/models/graph_cnn.py
+This file contains the Definition of GraphCNN
+GraphCNN includes ResNet50 as a submodule
+"""
+from __future__ import division
+
+import torch
+import torch.nn as nn
+
+from .graph_layers import GraphResBlock, GraphLinear
+from .resnet import resnet50
+
+class GraphCNN(nn.Module):
+    
+    def __init__(self, A, ref_vertices, num_layers=5, num_channels=512):
+        super(GraphCNN, self).__init__()
+        self.A = A
+        self.ref_vertices = ref_vertices
+        self.resnet = resnet50(pretrained=True)
+        layers = [GraphLinear(3 + 2048, 2 * num_channels)]
+        layers.append(GraphResBlock(2 * num_channels, num_channels, A))
+        for i in range(num_layers):
+            layers.append(GraphResBlock(num_channels, num_channels, A))
+        self.shape = nn.Sequential(GraphResBlock(num_channels, 64, A),
+                                   GraphResBlock(64, 32, A),
+                                   nn.GroupNorm(32 // 8, 32),
+                                   nn.ReLU(inplace=True),
+                                   GraphLinear(32, 3))
+        self.gc = nn.Sequential(*layers)
+        self.camera_fc = nn.Sequential(nn.GroupNorm(num_channels // 8, num_channels),
+                                      nn.ReLU(inplace=True),
+                                      GraphLinear(num_channels, 1),
+                                      nn.ReLU(inplace=True),
+                                      nn.Linear(A.shape[0], 3))
+
+    def forward(self, image):
+        """Forward pass
+        Inputs:
+            image: size = (B, 3, 224, 224)
+        Returns:
+            Regressed (subsampled) non-parametric shape: size = (B, 1723, 3)
+            Weak-perspective camera: size = (B, 3)
+        """
+        batch_size = image.shape[0]
+        ref_vertices = self.ref_vertices[None, :, :].expand(batch_size, -1, -1)
+        image_resnet = self.resnet(image)
+        image_enc = image_resnet.view(batch_size, 2048, 1).expand(-1, -1, ref_vertices.shape[-1])
+        x = torch.cat([ref_vertices, image_enc], dim=1)
+        x = self.gc(x)
+        shape = self.shape(x)
+        camera = self.camera_fc(x).view(batch_size, 3)
+        return shape, camera

src/graph_networks/graphcmr/graph_cnn_groundcontact.py

@@ -0,0 +1,101 @@
+"""
+code from https://raw.githubusercontent.com/nkolot/GraphCMR/master/models/graph_cnn.py
+This file contains the Definition of GraphCNN
+GraphCNN includes ResNet50 as a submodule
+"""
+from __future__ import division
+
+import torch
+import torch.nn as nn
+
+# from .resnet import resnet50
+import torchvision.models as models
+
+
+import os
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', '..'))
+from src.graph_networks.graphcmr.utils_mesh import Mesh
+from src.graph_networks.graphcmr.graph_layers import GraphResBlock, GraphLinear
+
+
+class GraphCNN(nn.Module):
+    
+    def __init__(self, A, ref_vertices, n_resnet_in, n_resnet_out, num_layers=5, num_channels=512):
+        super(GraphCNN, self).__init__()
+        self.A = A
+        self.ref_vertices = ref_vertices
+        # self.resnet = resnet50(pretrained=True)
+        #   -> within the GraphCMR network they ignore the last fully connected layer
+        # replace the first layer
+        self.resnet = models.resnet34(pretrained=False)  
+        n_in = 3 + 1
+        self.resnet.conv1 = nn.Conv2d(n_resnet_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_resnet_out) 
+
+
+        layers = [GraphLinear(3 + n_resnet_out, 2 * num_channels)]  # [GraphLinear(3 + 2048, 2 * num_channels)]
+        layers.append(GraphResBlock(2 * num_channels, num_channels, A))
+        for i in range(num_layers):
+            layers.append(GraphResBlock(num_channels, num_channels, A))
+        self.n_out_gc = 2       # two labels per vertex  
+        self.gc = nn.Sequential(GraphResBlock(num_channels, 64, A),
+                                   GraphResBlock(64, 32, A),
+                                   nn.GroupNorm(32 // 8, 32),
+                                   nn.ReLU(inplace=True),
+                                   GraphLinear(32, self.n_out_gc))
+        self.gcnn = nn.Sequential(*layers)
+        self.n_out_flatground = 1
+        self.flat_ground = nn.Sequential(nn.GroupNorm(num_channels // 8, num_channels),
+                                      nn.ReLU(inplace=True),
+                                      GraphLinear(num_channels, 1),
+                                      nn.ReLU(inplace=True),
+                                      nn.Linear(A.shape[0], self.n_out_flatground))
+
+    def forward(self, image):
+        """Forward pass
+        Inputs:
+            image: size = (B, 3, 256, 256)
+        Returns:
+            Regressed (subsampled) non-parametric shape: size = (B, 1723, 3)
+            Weak-perspective camera: size = (B, 3)
+        """
+        # import pdb; pdb.set_trace()
+
+        batch_size = image.shape[0]
+        ref_vertices = self.ref_vertices[None, :, :].expand(batch_size, -1, -1)     # (bs, 3, 973)
+        image_resnet = self.resnet(image)       # (bs, 512)
+        image_enc = image_resnet.view(batch_size, -1, 1).expand(-1, -1, ref_vertices.shape[-1]) # (bs, 512, 973)
+        x = torch.cat([ref_vertices, image_enc], dim=1)
+        x = self.gcnn(x)        # (bs, 512, 973)
+        ground_contact = self.gc(x)      # (bs, 2, 973)
+        ground_flatness = self.flat_ground(x).view(batch_size, self.n_out_flatground)    # (bs, 1)
+        return ground_contact, ground_flatness
+
+
+
+
+# how to use it:
+#
+# from src.graph_networks.graphcmr.utils_mesh import Mesh 
+#
+# create Mesh object
+# self.mesh = Mesh()
+# self.faces = self.mesh.faces.to(self.device)
+#
+# create GraphCNN
+# self.graph_cnn = GraphCNN(self.mesh.adjmat,
+#                     self.mesh.ref_vertices.t(),
+#                     num_channels=self.options.num_channels,
+#                     num_layers=self.options.num_layers
+#                     ).to(self.device)
+# ------------
+#
+# Feed image in the GraphCNN
+# Returns subsampled mesh and camera parameters
+# pred_vertices_sub, pred_camera = self.graph_cnn(images)
+# 
+# Upsample mesh in the original size
+# pred_vertices = self.mesh.upsample(pred_vertices_sub.transpose(1,2))
+# 

src/graph_networks/graphcmr/graph_cnn_groundcontact_multistage.py

@@ -0,0 +1,174 @@
+"""
+code from 
+    https://raw.githubusercontent.com/nkolot/GraphCMR/master/models/graph_cnn.py
+     https://github.com/chaneyddtt/Coarse-to-fine-3D-Animal/blob/main/model/graph_hg.py
+This file contains the Definition of GraphCNN
+GraphCNN includes ResNet50 as a submodule
+"""
+from __future__ import division
+
+import torch
+import torch.nn as nn
+
+# from .resnet import resnet50
+import torchvision.models as models
+
+
+import os
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', '..'))
+from src.graph_networks.graphcmr.utils_mesh import Mesh
+from src.graph_networks.graphcmr.graph_layers import GraphResBlock, GraphLinear
+
+
+class GraphCNNMS(nn.Module):
+    
+    def __init__(self, mesh, num_downsample=0, num_layers=5, n_resnet_out=256, num_channels=256):
+        '''
+        Args:
+            mesh: mesh data that store the adjacency matrix
+            num_channels: number of channels of GCN
+            num_downsample: number of downsampling of the input mesh
+        '''
+        
+        super(GraphCNNMS, self).__init__()
+
+        self.A = mesh._A[num_downsample:] # get the correct adjacency matrix because the input might be downsampled
+        # self.num_layers = len(self.A) - 1
+        self.num_layers = num_layers
+        assert self.num_layers <= len(self.A) - 1
+        print("Number of downsampling layer: {}".format(self.num_layers))
+        self.num_downsample = num_downsample
+        self.n_resnet_out = n_resnet_out
+        
+        
+        '''
+        self.use_pret_res = use_pret_res
+        # self.resnet = resnet50(pretrained=True)
+        #   -> within the GraphCMR network they ignore the last fully connected layer
+        # replace the first layer
+        self.resnet = models.resnet34(pretrained=self.use_pret_res)  
+        if (self.use_pret_res) and (n_resnet_in == 3):
+            print('use full pretrained resnet including first layer!')
+        else:
+            self.resnet.conv1 = nn.Conv2d(n_resnet_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_resnet_out) 
+        '''
+
+        self.lin1 = GraphLinear(3 + n_resnet_out, 2 * num_channels)
+        self.res1 = GraphResBlock(2 * num_channels, num_channels, self.A[0])
+        encode_layers = []
+        decode_layers = []
+
+        for i in range(self.num_layers + 1):    # range(len(self.A)):
+            encode_layers.append(GraphResBlock(num_channels, num_channels, self.A[i]))
+
+            decode_layers.append(GraphResBlock((i+1)*num_channels, (i+1)*num_channels,
+                                                   self.A[self.num_layers - i]))
+            current_channels = (i+1)*num_channels
+            # number of channels for the input is different because of the concatenation operation
+        self.n_out_gc = 2       # two labels per vertex  
+        self.gc  = nn.Sequential(GraphResBlock(current_channels, 64, self.A[0]),
+                                   GraphResBlock(64, 32, self.A[0]),
+                                   nn.GroupNorm(32 // 8, 32),
+                                   nn.ReLU(inplace=True),
+                                   GraphLinear(32, self.n_out_gc))
+
+        '''
+        self.n_out_flatground = 2
+        self.flat_ground = nn.Sequential(nn.GroupNorm(current_channels // 8, current_channels),
+                                      nn.ReLU(inplace=True),
+                                      GraphLinear(current_channels, 1),
+                                      nn.ReLU(inplace=True),
+                                      nn.Linear(A.shape[0], self.n_out_flatground))
+        '''
+
+        self.encoder = nn.Sequential(*encode_layers)
+        self.decoder = nn.Sequential(*decode_layers)
+        self.mesh = mesh
+
+
+
+
+    def forward(self, image_enc):
+        """Forward pass
+        Inputs:
+            image_enc: size = (B, self.n_resnet_out) 
+        Returns:
+            Regressed (subsampled) non-parametric shape: size = (B, 1723, 3)
+            Weak-perspective camera: size = (B, 3)
+        """
+        # import pdb; pdb.set_trace()
+
+        batch_size = image_enc.shape[0]
+        # ref_vertices = (self.mesh.get_ref_vertices(n=self.num_downsample).t())[None, :, :].expand(batch_size, -1, -1)  # (bs, 3, 973)
+        ref_vertices = (self.mesh.ref_vertices.t())[None, :, :].expand(batch_size, -1, -1)  # (bs, 3, 973)
+        '''image_resnet = self.resnet(image)       # (bs, 512)'''
+        image_enc_prep = image_enc.view(batch_size, -1, 1).expand(-1, -1, ref_vertices.shape[-1]) # (bs, 512, 973)
+
+        # prepare network input
+        #   -> for each node we feed the location of the vertex in the template mesh and an image encoding
+        x = torch.cat([ref_vertices, image_enc_prep], dim=1)
+        x = self.lin1(x)
+        x = self.res1(x)
+        x_ = [x]
+        output_list = []
+        for i in range(self.num_layers + 1):
+            if i == self.num_layers:
+                x = self.encoder[i](x)
+            else:
+                x = self.encoder[i](x)
+                x = self.mesh.downsample(x.transpose(1, 2), n1=self.num_downsample+i, n2=self.num_downsample+i+1)
+                x = x.transpose(1, 2)
+                if i < self.num_layers-1:
+                    x_.append(x)
+        for i in range(self.num_layers + 1):
+            if i == self.num_layers:
+                x = self.decoder[i](x)
+                output_list.append(x)
+            else:
+                x = self.decoder[i](x)
+                output_list.append(x)
+                x = self.mesh.upsample(x.transpose(1, 2), n1=self.num_layers-i+self.num_downsample,
+                                       n2=self.num_layers-i-1+self.num_downsample)
+                x = x.transpose(1, 2)
+                x = torch.cat([x, x_[self.num_layers-i-1]], dim=1) # skip connection between encoder and decoder
+
+        ground_contact = self.gc(x)
+
+        '''
+        ground_flatness = self.flat_ground(x).view(batch_size, self.n_out_flatground)    # (bs, 1)
+        '''
+
+        return ground_contact, output_list       # , ground_flatness
+
+
+
+
+
+
+
+# how to use it:
+#
+# from src.graph_networks.graphcmr.utils_mesh import Mesh 
+#
+# create Mesh object
+# self.mesh = Mesh()
+# self.faces = self.mesh.faces.to(self.device)
+#
+# create GraphCNN
+# self.graph_cnn = GraphCNN(self.mesh.adjmat,
+#                     self.mesh.ref_vertices.t(),
+#                     num_channels=self.options.num_channels,
+#                     num_layers=self.options.num_layers
+#                     ).to(self.device)
+# ------------
+#
+# Feed image in the GraphCNN
+# Returns subsampled mesh and camera parameters
+# pred_vertices_sub, pred_camera = self.graph_cnn(images)
+# 
+# Upsample mesh in the original size
+# pred_vertices = self.mesh.upsample(pred_vertices_sub.transpose(1,2))
+# 

src/graph_networks/graphcmr/graph_cnn_groundcontact_multistage_includingresnet.py

@@ -0,0 +1,170 @@
+"""
+code from 
+    https://raw.githubusercontent.com/nkolot/GraphCMR/master/models/graph_cnn.py
+     https://github.com/chaneyddtt/Coarse-to-fine-3D-Animal/blob/main/model/graph_hg.py
+This file contains the Definition of GraphCNN
+GraphCNN includes ResNet50 as a submodule
+"""
+from __future__ import division
+
+import torch
+import torch.nn as nn
+
+# from .resnet import resnet50
+import torchvision.models as models
+
+
+import os
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', '..'))
+from src.graph_networks.graphcmr.utils_mesh import Mesh
+from src.graph_networks.graphcmr.graph_layers import GraphResBlock, GraphLinear
+
+
+class GraphCNNMS(nn.Module):
+    
+    def __init__(self, mesh, num_downsample=0, num_layers=5, n_resnet_in=3, n_resnet_out=256, num_channels=256, use_pret_res=False):
+        '''
+        Args:
+            mesh: mesh data that store the adjacency matrix
+            num_channels: number of channels of GCN
+            num_downsample: number of downsampling of the input mesh
+        '''
+        
+        super(GraphCNNMS, self).__init__()
+
+        self.A = mesh._A[num_downsample:] # get the correct adjacency matrix because the input might be downsampled
+        # self.num_layers = len(self.A) - 1
+        self.num_layers = num_layers
+        assert self.num_layers <= len(self.A) - 1
+        print("Number of downsampling layer: {}".format(self.num_layers))
+        self.num_downsample = num_downsample
+        self.use_pret_res = use_pret_res
+
+        # self.resnet = resnet50(pretrained=True)
+        #   -> within the GraphCMR network they ignore the last fully connected layer
+        # replace the first layer
+        self.resnet = models.resnet34(pretrained=self.use_pret_res)  
+        if (self.use_pret_res) and (n_resnet_in == 3):
+            print('use full pretrained resnet including first layer!')
+        else:
+            self.resnet.conv1 = nn.Conv2d(n_resnet_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_resnet_out) 
+
+        self.lin1 = GraphLinear(3 + n_resnet_out, 2 * num_channels)
+        self.res1 = GraphResBlock(2 * num_channels, num_channels, self.A[0])
+        encode_layers = []
+        decode_layers = []
+
+        for i in range(self.num_layers + 1):    # range(len(self.A)):
+            encode_layers.append(GraphResBlock(num_channels, num_channels, self.A[i]))
+
+            decode_layers.append(GraphResBlock((i+1)*num_channels, (i+1)*num_channels,
+                                                   self.A[self.num_layers - i]))
+            current_channels = (i+1)*num_channels
+            # number of channels for the input is different because of the concatenation operation
+        self.n_out_gc = 2       # two labels per vertex  
+        self.gc  = nn.Sequential(GraphResBlock(current_channels, 64, self.A[0]),
+                                   GraphResBlock(64, 32, self.A[0]),
+                                   nn.GroupNorm(32 // 8, 32),
+                                   nn.ReLU(inplace=True),
+                                   GraphLinear(32, self.n_out_gc))
+
+        '''
+        self.n_out_flatground = 2
+        self.flat_ground = nn.Sequential(nn.GroupNorm(current_channels // 8, current_channels),
+                                      nn.ReLU(inplace=True),
+                                      GraphLinear(current_channels, 1),
+                                      nn.ReLU(inplace=True),
+                                      nn.Linear(A.shape[0], self.n_out_flatground))
+        '''
+
+        self.encoder = nn.Sequential(*encode_layers)
+        self.decoder = nn.Sequential(*decode_layers)
+        self.mesh = mesh
+
+
+
+
+    def forward(self, image):
+        """Forward pass
+        Inputs:
+            image: size = (B, 3, 256, 256)
+        Returns:
+            Regressed (subsampled) non-parametric shape: size = (B, 1723, 3)
+            Weak-perspective camera: size = (B, 3)
+        """
+        # import pdb; pdb.set_trace()
+
+        batch_size = image.shape[0]
+        # ref_vertices = (self.mesh.get_ref_vertices(n=self.num_downsample).t())[None, :, :].expand(batch_size, -1, -1)  # (bs, 3, 973)
+        ref_vertices = (self.mesh.ref_vertices.t())[None, :, :].expand(batch_size, -1, -1)  # (bs, 3, 973)
+        image_resnet = self.resnet(image)       # (bs, 512)
+        image_enc = image_resnet.view(batch_size, -1, 1).expand(-1, -1, ref_vertices.shape[-1]) # (bs, 512, 973)
+
+        # prepare network input
+        #   -> for each node we feed the location of the vertex in the template mesh and an image encoding
+        x = torch.cat([ref_vertices, image_enc], dim=1)
+        x = self.lin1(x)
+        x = self.res1(x)
+        x_ = [x]
+        output_list = []
+        for i in range(self.num_layers + 1):
+            if i == self.num_layers:
+                x = self.encoder[i](x)
+            else:
+                x = self.encoder[i](x)
+                x = self.mesh.downsample(x.transpose(1, 2), n1=self.num_downsample+i, n2=self.num_downsample+i+1)
+                x = x.transpose(1, 2)
+                if i < self.num_layers-1:
+                    x_.append(x)
+        for i in range(self.num_layers + 1):
+            if i == self.num_layers:
+                x = self.decoder[i](x)
+                output_list.append(x)
+            else:
+                x = self.decoder[i](x)
+                output_list.append(x)
+                x = self.mesh.upsample(x.transpose(1, 2), n1=self.num_layers-i+self.num_downsample,
+                                       n2=self.num_layers-i-1+self.num_downsample)
+                x = x.transpose(1, 2)
+                x = torch.cat([x, x_[self.num_layers-i-1]], dim=1) # skip connection between encoder and decoder
+
+        ground_contact = self.gc(x)
+
+        '''
+        ground_flatness = self.flat_ground(x).view(batch_size, self.n_out_flatground)    # (bs, 1)
+        '''
+
+        return ground_contact, output_list       # , ground_flatness
+
+
+
+
+
+
+
+# how to use it:
+#
+# from src.graph_networks.graphcmr.utils_mesh import Mesh 
+#
+# create Mesh object
+# self.mesh = Mesh()
+# self.faces = self.mesh.faces.to(self.device)
+#
+# create GraphCNN
+# self.graph_cnn = GraphCNN(self.mesh.adjmat,
+#                     self.mesh.ref_vertices.t(),
+#                     num_channels=self.options.num_channels,
+#                     num_layers=self.options.num_layers
+#                     ).to(self.device)
+# ------------
+#
+# Feed image in the GraphCNN
+# Returns subsampled mesh and camera parameters
+# pred_vertices_sub, pred_camera = self.graph_cnn(images)
+# 
+# Upsample mesh in the original size
+# pred_vertices = self.mesh.upsample(pred_vertices_sub.transpose(1,2))
+# 

src/graph_networks/graphcmr/graph_layers.py

@@ -0,0 +1,125 @@
+"""
+code from https://github.com/nkolot/GraphCMR/blob/master/models/graph_layers.py
+This file contains definitions of layers used to build the GraphCNN
+"""
+from __future__ import division
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+class GraphConvolution(nn.Module):
+    """Simple GCN layer, similar to https://arxiv.org/abs/1609.02907."""
+    def __init__(self, in_features, out_features, adjmat, bias=True):
+        super(GraphConvolution, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.adjmat = adjmat
+        self.weight = nn.Parameter(torch.FloatTensor(in_features, out_features))
+        if bias:
+            self.bias = nn.Parameter(torch.FloatTensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        stdv = 6. / math.sqrt(self.weight.size(0) + self.weight.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, x):
+        if x.ndimension() == 2:
+            support = torch.matmul(x, self.weight)
+            output = torch.matmul(self.adjmat, support)
+            if self.bias is not None:
+                output = output + self.bias
+            return output
+        else:
+            output = []
+            for i in range(x.shape[0]):
+                support = torch.matmul(x[i], self.weight)
+                # output.append(torch.matmul(self.adjmat, support))
+                output.append(spmm(self.adjmat, support))
+            output = torch.stack(output, dim=0)
+            if self.bias is not None:
+                output = output + self.bias
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+class GraphLinear(nn.Module):
+    """
+    Generalization of 1x1 convolutions on Graphs
+    """
+    def __init__(self, in_channels, out_channels):
+        super(GraphLinear, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.W = nn.Parameter(torch.FloatTensor(out_channels, in_channels))
+        self.b = nn.Parameter(torch.FloatTensor(out_channels))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        w_stdv = 1 / (self.in_channels * self.out_channels)
+        self.W.data.uniform_(-w_stdv, w_stdv)
+        self.b.data.uniform_(-w_stdv, w_stdv)
+
+    def forward(self, x):
+        return torch.matmul(self.W[None, :], x) + self.b[None, :, None]
+
+class GraphResBlock(nn.Module):
+    """
+    Graph Residual Block similar to the Bottleneck Residual Block in ResNet
+    """
+
+    def __init__(self, in_channels, out_channels, A):
+        super(GraphResBlock, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.lin1 = GraphLinear(in_channels, out_channels // 2)
+        self.conv = GraphConvolution(out_channels // 2, out_channels // 2, A)
+        self.lin2 = GraphLinear(out_channels // 2, out_channels)
+        self.skip_conv = GraphLinear(in_channels, out_channels)
+        self.pre_norm = nn.GroupNorm(in_channels // 8, in_channels)
+        self.norm1 = nn.GroupNorm((out_channels // 2) // 8, (out_channels // 2))
+        self.norm2 = nn.GroupNorm((out_channels // 2) // 8, (out_channels // 2))
+
+    def forward(self, x):
+        y = F.relu(self.pre_norm(x))
+        y = self.lin1(y)
+
+        y = F.relu(self.norm1(y))
+        y = self.conv(y.transpose(1,2)).transpose(1,2)
+
+        y = F.relu(self.norm2(y))
+        y = self.lin2(y)
+        if self.in_channels != self.out_channels:
+            x = self.skip_conv(x)
+        return x+y
+
+class SparseMM(torch.autograd.Function):
+    """Redefine sparse @ dense matrix multiplication to enable backpropagation.
+    The builtin matrix multiplication operation does not support backpropagation in some cases.
+    """
+    @staticmethod
+    def forward(ctx, sparse, dense):
+        ctx.req_grad = dense.requires_grad
+        ctx.save_for_backward(sparse)
+        return torch.matmul(sparse, dense)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = None
+        sparse, = ctx.saved_tensors
+        if ctx.req_grad:
+            grad_input = torch.matmul(sparse.t(), grad_output)
+        return None, grad_input
+
+def spmm(sparse, dense):
+    return SparseMM.apply(sparse, dense)

src/graph_networks/graphcmr/graphcnn_coarse_to_fine_animal_pose.py

@@ -0,0 +1,97 @@
+
+"""
+code from: https://github.com/chaneyddtt/Coarse-to-fine-3D-Animal/blob/main/model/graph_hg.py
+This file contains the Definition of GraphCNN
+GraphCNN includes ResNet50 as a submodule
+"""
+from __future__ import division
+
+import torch
+import torch.nn as nn
+
+from model.networks.graph_layers import GraphResBlock, GraphLinear
+from smal.mesh import Mesh
+from smal.smal_torch import SMAL
+
+# encoder-decoder structured GCN with skip connections
+class GraphCNN_hg(nn.Module):
+
+    def __init__(self, mesh,  num_channels=256, local_feat=False, num_downsample=0):
+        '''
+        Args:
+            mesh: mesh data that store the adjacency matrix
+            num_channels: number of channels of GCN
+            local_feat: whether use local feature for refinement
+            num_downsample: number of downsampling of the input mesh
+        '''
+        super(GraphCNN_hg, self).__init__()
+        self.A = mesh._A[num_downsample:] # get the correct adjacency matrix because the input might be downsampled
+        self.num_layers = len(self.A) - 1
+        print("Number of downsampling layer: {}".format(self.num_layers))
+        self.num_downsample = num_downsample
+        if local_feat:
+            self.lin1 = GraphLinear(3 + 2048 + 3840, 2 * num_channels)
+        else:
+            self.lin1 = GraphLinear(3 + 2048, 2 * num_channels)
+        self.res1 = GraphResBlock(2 * num_channels, num_channels, self.A[0])
+        encode_layers = []
+        decode_layers = []
+
+        for i in range(len(self.A)):
+            encode_layers.append(GraphResBlock(num_channels, num_channels, self.A[i]))
+
+            decode_layers.append(GraphResBlock((i+1)*num_channels, (i+1)*num_channels,
+                                                   self.A[len(self.A) - i - 1]))
+            current_channels = (i+1)*num_channels
+            # number of channels for the input is different because of the concatenation operation
+        self.shape = nn.Sequential(GraphResBlock(current_channels, 64, self.A[0]),
+                                   GraphResBlock(64, 32, self.A[0]),
+                                   nn.GroupNorm(32 // 8, 32),
+                                   nn.ReLU(inplace=True),
+                                   GraphLinear(32, 3))
+
+        self.encoder = nn.Sequential(*encode_layers)
+        self.decoder = nn.Sequential(*decode_layers)
+        self.mesh = mesh
+
+    def forward(self, verts_c, img_fea_global, img_fea_multiscale=None, points_local=None):
+        '''
+        Args:
+            verts_c: vertices from the coarse estimation
+            img_fea_global: global feature for mesh refinement
+            img_fea_multiscale: multi-scale feature from the encoder, used for local feature extraction
+            points_local: 2D keypoint for local feature extraction
+        Returns: refined mesh
+        '''
+        batch_size = img_fea_global.shape[0]
+        ref_vertices = verts_c.transpose(1, 2)
+        image_enc = img_fea_global.view(batch_size, 2048, 1).expand(-1, -1, ref_vertices.shape[-1])
+        if points_local is not None:
+            feat_local = torch.nn.functional.grid_sample(img_fea_multiscale, points_local)
+            x = torch.cat([ref_vertices, image_enc, feat_local.squeeze(2)], dim=1)
+        else:
+            x = torch.cat([ref_vertices, image_enc], dim=1)
+        x = self.lin1(x)
+        x = self.res1(x)
+        x_ = [x]
+        for i in range(self.num_layers + 1):
+            if i == self.num_layers:
+                x = self.encoder[i](x)
+            else:
+                x = self.encoder[i](x)
+                x = self.mesh.downsample(x.transpose(1, 2), n1=self.num_downsample+i, n2=self.num_downsample+i+1)
+                x = x.transpose(1, 2)
+                if i < self.num_layers-1:
+                    x_.append(x)
+        for i in range(self.num_layers + 1):
+            if i == self.num_layers:
+                x = self.decoder[i](x)
+            else:
+                x = self.decoder[i](x)
+                x = self.mesh.upsample(x.transpose(1, 2), n1=self.num_layers-i+self.num_downsample,
+                                       n2=self.num_layers-i-1+self.num_downsample)
+                x = x.transpose(1, 2)
+                x = torch.cat([x, x_[self.num_layers-i-1]], dim=1) # skip connection between encoder and decoder
+
+        shape = self.shape(x)
+        return shape

src/graph_networks/graphcmr/my_remarks.txt

@@ -0,0 +1,11 @@
+
+this folder contains code from https://github.com/nkolot/GraphCMR/tree/master/models
+
+
+other (newer) networks operating on meshes such as SMAL would be:
+    https://github.com/microsoft/MeshTransformer
+    https://github.com/microsoft/MeshGraphormer
+    
+see also:
+    https://arxiv.org/pdf/2112.01554.pdf, page 13
+        (Neural Head Avatars from Monocular RGB Videos)

src/graph_networks/graphcmr/pytorch_coma_mesh_operations.py

@@ -0,0 +1,282 @@
+# code from https://github.com/pixelite1201/pytorch_coma/blob/master/mesh_operations.py
+
+import math
+import heapq
+import numpy as np
+import scipy.sparse as sp
+from psbody.mesh import Mesh
+
+def row(A):
+    return A.reshape((1, -1))
+
+def col(A):
+    return A.reshape((-1, 1))
+
+def get_vert_connectivity(mesh_v, mesh_f):
+    """Returns a sparse matrix (of size #verts x #verts) where each nonzero
+    element indicates a neighborhood relation. For example, if there is a
+    nonzero element in position (15,12), that means vertex 15 is connected
+    by an edge to vertex 12."""
+
+    vpv = sp.csc_matrix((len(mesh_v),len(mesh_v)))
+
+    # for each column in the faces...
+    for i in range(3):
+        IS = mesh_f[:,i]
+        JS = mesh_f[:,(i+1)%3]
+        data = np.ones(len(IS))
+        ij = np.vstack((row(IS.flatten()), row(JS.flatten())))
+        mtx = sp.csc_matrix((data, ij), shape=vpv.shape)
+        vpv = vpv + mtx + mtx.T
+
+    return vpv
+
+def get_vertices_per_edge(mesh_v, mesh_f):
+    """Returns an Ex2 array of adjacencies between vertices, where
+    each element in the array is a vertex index. Each edge is included
+    only once. If output of get_faces_per_edge is provided, this is used to
+    avoid call to get_vert_connectivity()"""
+
+    vc = sp.coo_matrix(get_vert_connectivity(mesh_v, mesh_f))
+    result = np.hstack((col(vc.row), col(vc.col)))
+    result = result[result[:,0] < result[:,1]] # for uniqueness
+
+    return result
+
+
+def vertex_quadrics(mesh):
+    """Computes a quadric for each vertex in the Mesh.
+    see also:
+        https://www.cs.cmu.edu/~./garland/Papers/quadrics.pdf
+        https://users.csc.calpoly.edu/~zwood/teaching/csc570/final06/jseeba/
+    Returns:
+       v_quadrics: an (N x 4 x 4) array, where N is # vertices.
+    """
+
+    # Allocate quadrics
+    v_quadrics = np.zeros((len(mesh.v), 4, 4,))
+
+    # For each face...
+    for f_idx in range(len(mesh.f)):
+
+        # Compute normalized plane equation for that face
+        vert_idxs = mesh.f[f_idx]
+        verts = np.hstack((mesh.v[vert_idxs], np.array([1, 1, 1]).reshape(-1, 1)))
+        u, s, v = np.linalg.svd(verts)
+        eq = v[-1, :].reshape(-1, 1)
+        eq = eq / (np.linalg.norm(eq[0:3]))
+
+        # Add the outer product of the plane equation to the
+        # quadrics of the vertices for this face
+        for k in range(3):
+            v_quadrics[mesh.f[f_idx, k], :, :] += np.outer(eq, eq)
+
+    return v_quadrics
+
+def _get_sparse_transform(faces, num_original_verts):
+    verts_left = np.unique(faces.flatten())
+    IS = np.arange(len(verts_left))
+    JS = verts_left
+    data = np.ones(len(JS))
+
+    mp = np.arange(0, np.max(faces.flatten()) + 1)
+    mp[JS] = IS
+    new_faces = mp[faces.copy().flatten()].reshape((-1, 3))
+
+    ij = np.vstack((IS.flatten(), JS.flatten()))
+    mtx = sp.csc_matrix((data, ij), shape=(len(verts_left) , num_original_verts ))
+
+    return (new_faces, mtx)
+
+def qslim_decimator_transformer(mesh, factor=None, n_verts_desired=None):
+    """Return a simplified version of this mesh.
+
+    A Qslim-style approach is used here.
+
+    :param factor: fraction of the original vertices to retain
+    :param n_verts_desired: number of the original vertices to retain
+    :returns: new_faces: An Fx3 array of faces, mtx: Transformation matrix
+    """
+
+    if factor is None and n_verts_desired is None:
+        raise Exception('Need either factor or n_verts_desired.')
+
+    if n_verts_desired is None:
+        n_verts_desired = math.ceil(len(mesh.v) * factor)
+
+    Qv = vertex_quadrics(mesh)
+
+    # fill out a sparse matrix indicating vertex-vertex adjacency
+    # from psbody.mesh.topology.connectivity import get_vertices_per_edge
+    vert_adj = get_vertices_per_edge(mesh.v, mesh.f)
+    # vert_adj = sp.lil_matrix((len(mesh.v), len(mesh.v)))
+    # for f_idx in range(len(mesh.f)):
+    #     vert_adj[mesh.f[f_idx], mesh.f[f_idx]] = 1
+
+    vert_adj = sp.csc_matrix((vert_adj[:, 0] * 0 + 1, (vert_adj[:, 0], vert_adj[:, 1])), shape=(len(mesh.v), len(mesh.v)))
+    vert_adj = vert_adj + vert_adj.T
+    vert_adj = vert_adj.tocoo()
+
+    def collapse_cost(Qv, r, c, v):
+        Qsum = Qv[r, :, :] + Qv[c, :, :]
+        p1 = np.vstack((v[r].reshape(-1, 1), np.array([1]).reshape(-1, 1)))
+        p2 = np.vstack((v[c].reshape(-1, 1), np.array([1]).reshape(-1, 1)))
+
+        destroy_c_cost = p1.T.dot(Qsum).dot(p1)
+        destroy_r_cost = p2.T.dot(Qsum).dot(p2)
+        result = {
+            'destroy_c_cost': destroy_c_cost,
+            'destroy_r_cost': destroy_r_cost,
+            'collapse_cost': min([destroy_c_cost, destroy_r_cost]),
+            'Qsum': Qsum}
+        return result
+
+    # construct a queue of edges with costs
+    queue = []
+    for k in range(vert_adj.nnz):
+        r = vert_adj.row[k]
+        c = vert_adj.col[k]
+
+        if r > c:
+            continue
+
+        cost = collapse_cost(Qv, r, c, mesh.v)['collapse_cost']
+        heapq.heappush(queue, (cost, (r, c)))
+
+    # decimate
+    collapse_list = []
+    nverts_total = len(mesh.v)
+    faces = mesh.f.copy()
+    while nverts_total > n_verts_desired:
+        e = heapq.heappop(queue)
+        r = e[1][0]
+        c = e[1][1]
+        if r == c:
+            continue
+
+        cost = collapse_cost(Qv, r, c, mesh.v)
+        if cost['collapse_cost'] > e[0]:
+            heapq.heappush(queue, (cost['collapse_cost'], e[1]))
+            # print 'found outdated cost, %.2f < %.2f' % (e[0], cost['collapse_cost'])
+            continue
+        else:
+
+            # update old vert idxs to new one,
+            # in queue and in face list
+            if cost['destroy_c_cost'] < cost['destroy_r_cost']:
+                to_destroy = c
+                to_keep = r
+            else:
+                to_destroy = r
+                to_keep = c
+
+            collapse_list.append([to_keep, to_destroy])
+
+            # in our face array, replace "to_destroy" vertidx with "to_keep" vertidx
+            np.place(faces, faces == to_destroy, to_keep)
+
+            # same for queue
+            which1 = [idx for idx in range(len(queue)) if queue[idx][1][0] == to_destroy]
+            which2 = [idx for idx in range(len(queue)) if queue[idx][1][1] == to_destroy]
+            for k in which1:
+                queue[k] = (queue[k][0], (to_keep, queue[k][1][1]))
+            for k in which2:
+                queue[k] = (queue[k][0], (queue[k][1][0], to_keep))
+
+            Qv[r, :, :] = cost['Qsum']
+            Qv[c, :, :] = cost['Qsum']
+
+            a = faces[:, 0] == faces[:, 1]
+            b = faces[:, 1] == faces[:, 2]
+            c = faces[:, 2] == faces[:, 0]
+
+            # remove degenerate faces
+            def logical_or3(x, y, z):
+                return np.logical_or(x, np.logical_or(y, z))
+
+            faces_to_keep = np.logical_not(logical_or3(a, b, c))
+            faces = faces[faces_to_keep, :].copy()
+
+        nverts_total = (len(np.unique(faces.flatten())))
+
+    new_faces, mtx = _get_sparse_transform(faces, len(mesh.v))
+    return new_faces, mtx
+
+
+def setup_deformation_transfer(source, target, use_normals=False):
+    rows = np.zeros(3 * target.v.shape[0])
+    cols = np.zeros(3 * target.v.shape[0])
+    coeffs_v = np.zeros(3 * target.v.shape[0])
+    coeffs_n = np.zeros(3 * target.v.shape[0])
+
+    nearest_faces, nearest_parts, nearest_vertices = source.compute_aabb_tree().nearest(target.v, True)
+    nearest_faces = nearest_faces.ravel().astype(np.int64)
+    nearest_parts = nearest_parts.ravel().astype(np.int64)
+    nearest_vertices = nearest_vertices.ravel()
+
+    for i in range(target.v.shape[0]):
+        # Closest triangle index
+        f_id = nearest_faces[i]
+        # Closest triangle vertex ids
+        nearest_f = source.f[f_id]
+
+        # Closest surface point
+        nearest_v = nearest_vertices[3 * i:3 * i + 3]
+        # Distance vector to the closest surface point
+        dist_vec = target.v[i] - nearest_v
+
+        rows[3 * i:3 * i + 3] = i * np.ones(3)
+        cols[3 * i:3 * i + 3] = nearest_f
+
+        n_id = nearest_parts[i]
+        if n_id == 0:
+            # Closest surface point in triangle
+            A = np.vstack((source.v[nearest_f])).T
+            coeffs_v[3 * i:3 * i + 3] = np.linalg.lstsq(A, nearest_v)[0]
+        elif n_id > 0 and n_id <= 3:
+            # Closest surface point on edge
+            A = np.vstack((source.v[nearest_f[n_id - 1]], source.v[nearest_f[n_id % 3]])).T
+            tmp_coeffs = np.linalg.lstsq(A, target.v[i])[0]
+            coeffs_v[3 * i + n_id - 1] = tmp_coeffs[0]
+            coeffs_v[3 * i + n_id % 3] = tmp_coeffs[1]
+        else:
+            # Closest surface point a vertex
+            coeffs_v[3 * i + n_id - 4] = 1.0
+
+    #    if use_normals:
+    #        A = np.vstack((vn[nearest_f])).T
+    #        coeffs_n[3 * i:3 * i + 3] = np.linalg.lstsq(A, dist_vec)[0]
+
+    #coeffs = np.hstack((coeffs_v, coeffs_n))
+    #rows = np.hstack((rows, rows))
+    #cols = np.hstack((cols, source.v.shape[0] + cols))
+    matrix = sp.csc_matrix((coeffs_v, (rows, cols)), shape=(target.v.shape[0], source.v.shape[0]))
+    return matrix
+
+
+def generate_transform_matrices(mesh, factors):
+    """Generates len(factors) meshes, each of them is scaled by factors[i] and
+       computes the transformations between them.
+
+    Returns:
+       M: a set of meshes downsampled from mesh by a factor specified in factors.
+       A: Adjacency matrix for each of the meshes
+       D: Downsampling transforms between each of the meshes
+       U: Upsampling transforms between each of the meshes
+    """
+
+    factors = map(lambda x: 1.0 / x, factors)
+    M, A, D, U = [], [], [], []
+    A.append(get_vert_connectivity(mesh.v, mesh.f).tocoo())
+    M.append(mesh)
+
+    for i,factor in enumerate(factors):
+        ds_f, ds_D = qslim_decimator_transformer(M[-1], factor=factor)
+        D.append(ds_D.tocoo())
+        new_mesh_v = ds_D.dot(M[-1].v)
+        new_mesh = Mesh(v=new_mesh_v, f=ds_f)
+        M.append(new_mesh)
+        A.append(get_vert_connectivity(new_mesh.v, new_mesh.f).tocoo())
+        U.append(setup_deformation_transfer(M[-1], M[-2]).tocoo())
+
+    return M, A, D, U

src/graph_networks/graphcmr/utils_mesh.py

@@ -0,0 +1,138 @@
+# code from https://github.com/nkolot/GraphCMR/blob/master/utils/mesh.py
+
+from __future__ import division
+import torch
+import numpy as np
+import scipy.sparse
+
+# from models import SMPL
+import os
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
+from graph_networks.graphcmr.graph_layers import spmm
+
+def scipy_to_pytorch(A, U, D):
+    """Convert scipy sparse matrices to pytorch sparse matrix."""
+    ptU = []
+    ptD = []
+    
+    for i in range(len(U)):
+        u = scipy.sparse.coo_matrix(U[i])
+        i = torch.LongTensor(np.array([u.row, u.col]))
+        v = torch.FloatTensor(u.data)
+        ptU.append(torch.sparse.FloatTensor(i, v, u.shape))
+    
+    for i in range(len(D)):
+        d = scipy.sparse.coo_matrix(D[i])
+        i = torch.LongTensor(np.array([d.row, d.col]))
+        v = torch.FloatTensor(d.data)
+        ptD.append(torch.sparse.FloatTensor(i, v, d.shape)) 
+
+    return ptU, ptD
+
+
+def adjmat_sparse(adjmat, nsize=1):
+    """Create row-normalized sparse graph adjacency matrix."""
+    adjmat = scipy.sparse.csr_matrix(adjmat)
+    if nsize > 1:
+        orig_adjmat = adjmat.copy()
+        for _ in range(1, nsize):
+            adjmat = adjmat * orig_adjmat
+    adjmat.data = np.ones_like(adjmat.data)
+    for i in range(adjmat.shape[0]):
+        adjmat[i,i] = 1
+    num_neighbors = np.array(1 / adjmat.sum(axis=-1))
+    adjmat = adjmat.multiply(num_neighbors)
+    adjmat = scipy.sparse.coo_matrix(adjmat)
+    row = adjmat.row
+    col = adjmat.col
+    data = adjmat.data
+    i = torch.LongTensor(np.array([row, col]))
+    v = torch.from_numpy(data).float()
+    adjmat = torch.sparse.FloatTensor(i, v, adjmat.shape)
+    return adjmat
+
+def get_graph_params(filename, nsize=1):
+    """Load and process graph adjacency matrix and upsampling/downsampling matrices."""
+    data = np.load(filename, encoding='latin1', allow_pickle=True) # np.load(filename, encoding='latin1')
+    A = data['A']
+    U = data['U']
+    D = data['D']
+    U, D = scipy_to_pytorch(A, U, D)
+    A = [adjmat_sparse(a, nsize=nsize) for a in A]
+    return A, U, D
+
+class Mesh(object):
+    """Mesh object that is used for handling certain graph operations."""
+    def __init__(self, filename='data/mesh_downsampling.npz',
+                 num_downsampling=1, nsize=1, body_model=None, device=torch.device('cuda')):
+        self._A, self._U, self._D = get_graph_params(filename=filename, nsize=nsize)
+        self._A = [a.to(device) for a in self._A]
+        self._U = [u.to(device) for u in self._U]
+        self._D = [d.to(device) for d in self._D]
+        self.num_downsampling = num_downsampling
+
+        # load template vertices from SMPL and normalize them
+        if body_model is None:
+            smpl = SMPL()
+        else:
+            smpl = body_model
+        ref_vertices = smpl.v_template
+        center = 0.5*(ref_vertices.max(dim=0)[0] + ref_vertices.min(dim=0)[0])[None]
+        ref_vertices -= center
+        ref_vertices /= ref_vertices.abs().max().item()
+
+        self._ref_vertices = ref_vertices.to(device)
+        self.faces = smpl.faces.int().to(device)
+
+    @property
+    def adjmat(self):
+        """Return the graph adjacency matrix at the specified subsampling level."""
+        return self._A[self.num_downsampling].float()
+
+    @property
+    def ref_vertices(self):
+        """Return the template vertices at the specified subsampling level."""
+        ref_vertices = self._ref_vertices
+        for i in range(self.num_downsampling):
+            ref_vertices = torch.spmm(self._D[i], ref_vertices)
+        return ref_vertices
+
+    def get_ref_vertices(self, n_downsample):
+        """Return the template vertices at any desired subsampling level."""
+        ref_vertices = self._ref_vertices
+        for i in range(n_downsample):
+            ref_vertices = torch.spmm(self._D[i], ref_vertices)
+        return ref_vertices
+
+    def downsample(self, x, n1=0, n2=None):
+        """Downsample mesh."""
+        if n2 is None:
+            n2 = self.num_downsampling
+        if x.ndimension() < 3:
+            for i in range(n1, n2):
+                x = spmm(self._D[i], x)
+        elif x.ndimension() == 3:
+            out = []
+            for i in range(x.shape[0]):
+                y = x[i]
+                for j in range(n1, n2):
+                    y = spmm(self._D[j], y)
+                out.append(y)
+            x = torch.stack(out, dim=0)
+        return x
+
+    def upsample(self, x, n1=1, n2=0):
+        """Upsample mesh."""
+        if x.ndimension() < 3:
+            for i in reversed(range(n2, n1)):
+                x = spmm(self._U[i], x)
+        elif x.ndimension() == 3:
+            out = []
+            for i in range(x.shape[0]):
+                y = x[i]
+                for j in reversed(range(n2, n1)):
+                    y = spmm(self._U[j], y)
+                out.append(y)
+            x = torch.stack(out, dim=0)
+        return x

src/graph_networks/losses_for_vertex_wise_predictions/calculate_distance_between_points_on_mesh.py

@@ -0,0 +1,245 @@
+
+"""
+code adapted from: https://github.com/mikedh/trimesh/blob/main/examples/shortest.py
+shortest.py
+----------------
+Given a mesh and two vertex indices find the shortest path
+between the two vertices while only traveling along edges
+of the mesh.
+"""
+
+# python src/graph_networks/losses_for_vertex_wise_predictions/calculate_distance_between_points_on_mesh.py
+
+
+import os
+import sys
+import glob
+import csv
+import json
+import shutil
+import tqdm
+import numpy as np
+import pickle as pkl
+import trimesh
+import networkx as nx
+
+
+
+
+
+def read_csv(csv_file):
+    with open(csv_file,'r') as f:
+        reader = csv.reader(f)
+        headers = next(reader)
+        row_list = [{h:x for (h,x) in zip(headers,row)} for row in reader]
+    return row_list
+
+
+def load_all_template_mesh_distances(root_out_path, filename='all_vertex_distances.npy'):
+    vert_dists = np.load(root_out_path + filename)
+    return vert_dists
+
+
+def prepare_graph_from_template_mesh_and_calculate_all_distances(path_mesh, root_out_path, calc_dist_mat=False):
+    # root_out_path = ROOT_OUT_PATH
+    '''
+    from smal_pytorch.smal_model.smal_torch_new import SMAL
+    smal = SMAL()
+    verts = smal.v_template.detach().cpu().numpy()
+    faces = smal.faces.detach().cpu().numpy()
+    '''
+    # path_mesh = ROOT_PATH_MESH + 'mesh_downsampling_meshesmy_smpl_39dogsnorm_Jr_4_dog_template_downsampled0.obj'
+    my_mesh = trimesh.load_mesh(path_mesh, process=False,  maintain_order=True)
+    verts = my_mesh.vertices
+    faces = my_mesh.faces
+    # edges without duplication
+    edges = my_mesh.edges_unique
+    # the actual length of each unique edge
+    length = my_mesh.edges_unique_length
+    # create the graph with edge attributes for length (option A)
+    #   g = nx.Graph()
+    #   for edge, L in zip(edges, length): g.add_edge(*edge, length=L)
+    # you can create the graph with from_edgelist and
+    # a list comprehension (option B)
+    ga = nx.from_edgelist([(e[0], e[1], {'length': L}) for e, L in zip(edges, length)])
+    # calculate the distances between all vertex pairs
+    if calc_dist_mat:
+        # calculate distances between all possible vertex pairs
+        # shortest_path = nx.shortest_path(ga, source=ind_v0, target=ind_v1, weight='length')
+        # shortest_dist = nx.shortest_path_length(ga, source=ind_v0, target=ind_v1, weight='length')
+        dis = dict(nx.shortest_path_length(ga, weight='length', method='dijkstra'))
+        vertex_distances = np.zeros((n_verts_smal, n_verts_smal))
+        for ind_v0 in range(n_verts_smal):
+            print(ind_v0)
+            for ind_v1 in range(ind_v0, n_verts_smal):
+                vertex_distances[ind_v0, ind_v1] = dis[ind_v0][ind_v1]
+                vertex_distances[ind_v1, ind_v0] = dis[ind_v0][ind_v1]
+        # save those distances
+        np.save(root_out_path + 'all_vertex_distances.npy', vertex_distances)
+        vert_dists = vertex_distances
+    else:
+        vert_dists = np.load(root_out_path + 'all_vertex_distances.npy')
+    return ga, vert_dists
+
+
+def calculate_vertex_overview_for_gc_annotation(name, gc_info_raw, vert_dists, root_out_path_vis=None, verts=None, faces=None, img_v12_dir=None):
+    # input:
+    #   root_out_path_vis = ROOT_OUT_PATH
+    #   img_v12_dir = IMG_V12_DIR
+    #   name = images_with_gc_labelled[ind_img]
+    #   gc_info_raw = gc_dict['bite/' + name] 
+    # output:
+    #    vertex_overview: np array of shape (n_verts_smal, 3) with [first: no-contact=0 contact=1     second: index of vertex     third: dist]
+    n_verts_smal = 3889
+    gc_vertices = []
+    gc_info_np = np.zeros((n_verts_smal))
+    for ind_v in gc_info_raw:
+        if ind_v < n_verts_smal:
+            gc_vertices.append(ind_v)
+            gc_info_np[ind_v] = 1
+    # save a visualization of those annotations 
+    if root_out_path_vis is not None:
+        my_mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False,  maintain_order=True)
+    if img_v12_dir is not None and root_out_path_vis is not None:
+        vert_colors = np.repeat(255*gc_info_np[:, None], 3, 1)
+        my_mesh.visual.vertex_colors = vert_colors
+        my_mesh.export(root_out_path_vis + (name).replace('.jpg', '_withgc.obj'))
+        img_path = img_v12_dir + name
+        shutil.copy(img_path, root_out_path_vis + name)
+    # calculate for each vertex the distance to the closest element of the other group
+    non_gc_vertices = list(set(range(n_verts_smal)) - set(gc_vertices))
+    print('vertices in contact: ' + str(len(gc_vertices)))
+    print('vertices without contact: ' + str(len(non_gc_vertices)))
+    vertex_overview = np.zeros((n_verts_smal, 3))   # first: no-contact=0 contact=1     second: index of vertex     third: dist
+    vertex_overview[:, 0] = gc_info_np
+    # loop through all contact vertices
+    for ind_v in gc_vertices:
+        min_length = 100
+        for ind_v_ps in non_gc_vertices:    # possible solution
+            # this_path = nx.shortest_path(ga, source=ind_v, target=ind_v_ps, weight='length')
+            # this_length = nx.shortest_path_length(ga, source=ind_v, target=ind_v_ps, weight='length')
+            this_length = vert_dists[ind_v, ind_v_ps]
+            if this_length < min_length:
+                min_length = this_length
+                vertex_overview[ind_v, 1] = ind_v_ps
+                vertex_overview[ind_v, 2] = this_length
+    # loop through all non-contact vertices
+    for ind_v in non_gc_vertices:
+        min_length = 100
+        for ind_v_ps in gc_vertices:    # possible solution
+            # this_path = nx.shortest_path(ga, source=ind_v, target=ind_v_ps, weight='length')
+            # this_length = nx.shortest_path_length(ga, source=ind_v, target=ind_v_ps, weight='length')
+            this_length = vert_dists[ind_v, ind_v_ps]
+            if this_length < min_length:
+                min_length = this_length
+                vertex_overview[ind_v, 1] = ind_v_ps
+                vertex_overview[ind_v, 2] = this_length
+    if root_out_path_vis is not None:
+        # save a colored mesh
+        my_mesh_dists = my_mesh.copy()
+        scale_0 = (vertex_overview[vertex_overview[:, 0]==0, 2]).max()
+        scale_1 = (vertex_overview[vertex_overview[:, 0]==1, 2]).max()
+        vert_col = np.zeros((n_verts_smal, 3)) 
+        vert_col[vertex_overview[:, 0]==0, 1] = vertex_overview[vertex_overview[:, 0]==0, 2] * 255 / scale_0     # green
+        vert_col[vertex_overview[:, 0]==1, 0] = vertex_overview[vertex_overview[:, 0]==1, 2] * 255 / scale_1     # red
+        my_mesh_dists.visual.vertex_colors = np.uint8(vert_col)
+        my_mesh_dists.export(root_out_path_vis + (name).replace('.jpg', '_withgcdists.obj'))
+    return vertex_overview
+
+
+
+
+
+
+
+
+def main():
+
+    ROOT_PATH_MESH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/graphcmr/data/meshes/'
+    ROOT_PATH_ANNOT = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage3/'
+    IMG_V12_DIR = '/ps/scratch/nrueegg/new_projects/Animals/data/dog_datasets/Stanford_Dogs_Dataset/StanfordExtra_V12/StanExtV12_Images/'		
+    # ROOT_OUT_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/losses_for_vertex_wise_predictions/debugging_results/'
+    ROOT_OUT_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage3/'
+    ROOT_OUT_PATH_VIS = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage3/vis/'
+    ROOT_OUT_PATH_DISTSGCNONGC = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage3/vertex_distances_gc_nongc/'
+    ROOT_PATH_ALL_VERT_DIST_TEMPLATE  = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/'
+
+    # load all vertex distances
+    path_mesh = ROOT_PATH_MESH + 'mesh_downsampling_meshesmy_smpl_39dogsnorm_Jr_4_dog_template_downsampled0.obj'
+    my_mesh = trimesh.load_mesh(path_mesh, process=False,  maintain_order=True)
+    verts = my_mesh.vertices
+    faces = my_mesh.faces
+    # vert_dists, ga = prepare_graph_from_template_mesh_and_calculate_all_distances(path_mesh, ROOT_OUT_PATH, calc_dist_mat=False)
+    vert_dists = load_all_template_mesh_distances(ROOT_PATH_ALL_VERT_DIST_TEMPLATE, filename='all_vertex_distances.npy')
+
+
+
+
+    all_keys = []
+    gc_dict = {}
+    # data/stanext_related_data/ground_contact_annotations/stage3/main_partA1667_20221021_140108.csv
+    # for csv_file in ['main_partA500_20221018_131139.csv', 'pilot_20221017_104201.csv', 'my_gcannotations_qualification.csv']:        
+    # for csv_file in ['main_partA1667_20221021_140108.csv', 'main_partA500_20221018_131139.csv', 'pilot_20221017_104201.csv', 'my_gcannotations_qualification.csv']:
+    for csv_file in ['main_partA1667_20221021_140108.csv', 'main_partA500_20221018_131139.csv', 'main_partB20221023_150926.csv', 'pilot_20221017_104201.csv', 'my_gcannotations_qualification.csv']:
+        # load all ground contact annotations
+        gc_annot_csv = ROOT_PATH_ANNOT + csv_file   # 'my_gcannotations_qualification.csv'
+        gc_row_list = read_csv(gc_annot_csv)
+        for ind_row in range(len(gc_row_list)):
+            json_acceptable_string = (gc_row_list[ind_row]['vertices']).replace("'", "\"")
+            gc_dict_temp = json.loads(json_acceptable_string)
+            all_keys.extend(gc_dict_temp.keys())
+            gc_dict.update(gc_dict_temp)
+        print(len(gc_dict.keys()))
+
+    print('number of labeled images: ' + str(len(gc_dict.keys())))      # WHY IS THIS ONLY 699?
+
+    import pdb; pdb.set_trace()
+
+
+    # prepare and save contact annotations including distances
+    vertex_overview_dict = {}
+    for ind_img, name_ingcdict in enumerate(gc_dict.keys()):     #  range(len(gc_dict.keys())):
+        name = name_ingcdict.split('bite/')[1]
+        # name = images_with_gc_labelled[ind_img]
+        print('work on image ' + str(ind_img) + ': ' + name)
+        # gc_info_raw = gc_dict['bite/' + name]      # a list with all vertex numbers that are in ground contact
+        gc_info_raw = gc_dict[name_ingcdict]      # a list with all vertex numbers that are in ground contact
+
+        if not os.path.exists(ROOT_OUT_PATH_VIS + name.split('/')[0]): os.makedirs(ROOT_OUT_PATH_VIS + name.split('/')[0])
+        if not os.path.exists(ROOT_OUT_PATH_DISTSGCNONGC + name.split('/')[0]): os.makedirs(ROOT_OUT_PATH_DISTSGCNONGC + name.split('/')[0])
+
+        vertex_overview = calculate_vertex_overview_for_gc_annotation(name, gc_info_raw, vert_dists, root_out_path_vis=ROOT_OUT_PATH_VIS, verts=verts, faces=faces, img_v12_dir=None)
+        np.save(ROOT_OUT_PATH_DISTSGCNONGC + name.replace('.jpg', '_gc_vertdists_overview.npy'), vertex_overview)
+
+        vertex_overview_dict[name.split('.')[0]] = {'gc_vertdists_overview': vertex_overview, 'gc_index_list': gc_info_raw}
+
+
+
+        
+
+    # import pdb; pdb.set_trace()
+
+    with open(ROOT_OUT_PATH + 'gc_annots_overview_stage3complete_withtraintestval_xx.pkl', 'wb') as fp:
+        pkl.dump(vertex_overview_dict, fp)
+
+
+
+
+
+
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+
+
+
+

src/graph_networks/losses_for_vertex_wise_predictions/calculate_distance_between_points_on_mesh_forfourpaws.py

@@ -0,0 +1,213 @@
+
+"""
+code adapted from: https://github.com/mikedh/trimesh/blob/main/examples/shortest.py
+shortest.py
+----------------
+Given a mesh and two vertex indices find the shortest path
+between the two vertices while only traveling along edges
+of the mesh.
+"""
+
+# python src/graph_networks/losses_for_vertex_wise_predictions/calculate_distance_between_points_on_mesh_forfourpaws.py
+
+
+import os
+import sys
+import glob
+import csv
+import json
+import shutil
+import tqdm
+import numpy as np
+import pickle as pkl
+import trimesh
+import networkx as nx
+
+
+
+
+
+def read_csv(csv_file):
+    with open(csv_file,'r') as f:
+        reader = csv.reader(f)
+        headers = next(reader)
+        row_list = [{h:x for (h,x) in zip(headers,row)} for row in reader]
+    return row_list
+
+
+def load_all_template_mesh_distances(root_out_path, filename='all_vertex_distances.npy'):
+    vert_dists = np.load(root_out_path + filename)
+    return vert_dists
+
+
+def prepare_graph_from_template_mesh_and_calculate_all_distances(path_mesh, root_out_path, calc_dist_mat=False):
+    # root_out_path = ROOT_OUT_PATH
+    '''
+    from smal_pytorch.smal_model.smal_torch_new import SMAL
+    smal = SMAL()
+    verts = smal.v_template.detach().cpu().numpy()
+    faces = smal.faces.detach().cpu().numpy()
+    '''
+    # path_mesh = ROOT_PATH_MESH + 'mesh_downsampling_meshesmy_smpl_39dogsnorm_Jr_4_dog_template_downsampled0.obj'
+    my_mesh = trimesh.load_mesh(path_mesh, process=False,  maintain_order=True)
+    verts = my_mesh.vertices
+    faces = my_mesh.faces
+    # edges without duplication
+    edges = my_mesh.edges_unique
+    # the actual length of each unique edge
+    length = my_mesh.edges_unique_length
+    # create the graph with edge attributes for length (option A)
+    #   g = nx.Graph()
+    #   for edge, L in zip(edges, length): g.add_edge(*edge, length=L)
+    # you can create the graph with from_edgelist and
+    # a list comprehension (option B)
+    ga = nx.from_edgelist([(e[0], e[1], {'length': L}) for e, L in zip(edges, length)])
+    # calculate the distances between all vertex pairs
+    if calc_dist_mat:
+        # calculate distances between all possible vertex pairs
+        # shortest_path = nx.shortest_path(ga, source=ind_v0, target=ind_v1, weight='length')
+        # shortest_dist = nx.shortest_path_length(ga, source=ind_v0, target=ind_v1, weight='length')
+        dis = dict(nx.shortest_path_length(ga, weight='length', method='dijkstra'))
+        vertex_distances = np.zeros((n_verts_smal, n_verts_smal))
+        for ind_v0 in range(n_verts_smal):
+            print(ind_v0)
+            for ind_v1 in range(ind_v0, n_verts_smal):
+                vertex_distances[ind_v0, ind_v1] = dis[ind_v0][ind_v1]
+                vertex_distances[ind_v1, ind_v0] = dis[ind_v0][ind_v1]
+        # save those distances
+        np.save(root_out_path + 'all_vertex_distances.npy', vertex_distances)
+        vert_dists = vertex_distances
+    else:
+        vert_dists = np.load(root_out_path + 'all_vertex_distances.npy')
+    return ga, vert_dists
+
+
+def calculate_vertex_overview_for_gc_annotation(name, gc_info_raw, vert_dists, root_out_path_vis=None, verts=None, faces=None, img_v12_dir=None):
+    # input:
+    #   root_out_path_vis = ROOT_OUT_PATH
+    #   img_v12_dir = IMG_V12_DIR
+    #   name = images_with_gc_labelled[ind_img]
+    #   gc_info_raw = gc_dict['bite/' + name] 
+    # output:
+    #    vertex_overview: np array of shape (n_verts_smal, 3) with [first: no-contact=0 contact=1     second: index of vertex     third: dist]
+    n_verts_smal = 3889
+    gc_vertices = []
+    gc_info_np = np.zeros((n_verts_smal))
+    for ind_v in gc_info_raw:
+        if ind_v < n_verts_smal:
+            gc_vertices.append(ind_v)
+            gc_info_np[ind_v] = 1
+    # save a visualization of those annotations 
+    if root_out_path_vis is not None:
+        my_mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False,  maintain_order=True)
+    if img_v12_dir is not None and root_out_path_vis is not None:
+        vert_colors = np.repeat(255*gc_info_np[:, None], 3, 1)
+        my_mesh.visual.vertex_colors = vert_colors
+        my_mesh.export(root_out_path_vis + (name).replace('.jpg', '_withgc.obj'))
+        img_path = img_v12_dir + name
+        shutil.copy(img_path, root_out_path_vis + name)
+    # calculate for each vertex the distance to the closest element of the other group
+    non_gc_vertices = list(set(range(n_verts_smal)) - set(gc_vertices))
+    print('vertices in contact: ' + str(len(gc_vertices)))
+    print('vertices without contact: ' + str(len(non_gc_vertices)))
+    vertex_overview = np.zeros((n_verts_smal, 3))   # first: no-contact=0 contact=1     second: index of vertex     third: dist
+    vertex_overview[:, 0] = gc_info_np
+    # loop through all contact vertices
+    for ind_v in gc_vertices:
+        min_length = 100
+        for ind_v_ps in non_gc_vertices:    # possible solution
+            # this_path = nx.shortest_path(ga, source=ind_v, target=ind_v_ps, weight='length')
+            # this_length = nx.shortest_path_length(ga, source=ind_v, target=ind_v_ps, weight='length')
+            this_length = vert_dists[ind_v, ind_v_ps]
+            if this_length < min_length:
+                min_length = this_length
+                vertex_overview[ind_v, 1] = ind_v_ps
+                vertex_overview[ind_v, 2] = this_length
+    # loop through all non-contact vertices
+    for ind_v in non_gc_vertices:
+        min_length = 100
+        for ind_v_ps in gc_vertices:    # possible solution
+            # this_path = nx.shortest_path(ga, source=ind_v, target=ind_v_ps, weight='length')
+            # this_length = nx.shortest_path_length(ga, source=ind_v, target=ind_v_ps, weight='length')
+            this_length = vert_dists[ind_v, ind_v_ps]
+            if this_length < min_length:
+                min_length = this_length
+                vertex_overview[ind_v, 1] = ind_v_ps
+                vertex_overview[ind_v, 2] = this_length
+    if root_out_path_vis is not None:
+        # save a colored mesh
+        my_mesh_dists = my_mesh.copy()
+        scale_0 = (vertex_overview[vertex_overview[:, 0]==0, 2]).max()
+        scale_1 = (vertex_overview[vertex_overview[:, 0]==1, 2]).max()
+        vert_col = np.zeros((n_verts_smal, 3)) 
+        vert_col[vertex_overview[:, 0]==0, 1] = vertex_overview[vertex_overview[:, 0]==0, 2] * 255 / scale_0     # green
+        vert_col[vertex_overview[:, 0]==1, 0] = vertex_overview[vertex_overview[:, 0]==1, 2] * 255 / scale_1     # red
+        my_mesh_dists.visual.vertex_colors = np.uint8(vert_col)
+        my_mesh_dists.export(root_out_path_vis + (name).replace('.jpg', '_withgcdists.obj'))
+    return vertex_overview
+
+
+
+
+
+
+
+
+
+def main():
+
+    ROOT_PATH_MESH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/graphcmr/data/meshes/'
+    IMG_V12_DIR = '/ps/scratch/nrueegg/new_projects/Animals/data/dog_datasets/Stanford_Dogs_Dataset/StanfordExtra_V12/StanExtV12_Images/'		
+    # ROOT_OUT_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/losses_for_vertex_wise_predictions/debugging_results/'
+    ROOT_OUT_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stages12together/'
+    ROOT_PATH_ALL_VERT_DIST_TEMPLATE  = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/'
+
+    # load all vertex distances
+    path_mesh = ROOT_PATH_MESH + 'mesh_downsampling_meshesmy_smpl_39dogsnorm_Jr_4_dog_template_downsampled0.obj'
+    my_mesh = trimesh.load_mesh(path_mesh, process=False,  maintain_order=True)
+    verts = my_mesh.vertices
+    faces = my_mesh.faces
+    # vert_dists, ga = prepare_graph_from_template_mesh_and_calculate_all_distances(path_mesh, ROOT_OUT_PATH, calc_dist_mat=False)
+    vert_dists = load_all_template_mesh_distances(ROOT_PATH_ALL_VERT_DIST_TEMPLATE, filename='all_vertex_distances.npy')
+
+    # paw vertices:
+    # left and right is a bit different, but that is ok (we will anyways mirror data at training time)
+    right_front_paw = [3829,+3827,+3825,+3718,+3722,+3723,+3743,+3831,+3719,+3726,+3716,+3724,+3828,+3717,+3721,+3725,+3832,+3830,+3720,+3288,+3740,+3714,+3826,+3715,+3728,+3712,+3287,+3284,+3727,+3285,+3742,+3291,+3710,+3697,+3711,+3289,+3730,+3713,+3739,+3282,+3738,+3708,+3709,+3741,+3698,+3696,+3308,+3695,+3706,+3700,+3707,+3306,+3305,+3737,+3304,+3303,+3307,+3736,+3735,+3250,+3261,+3732,+3734,+3733,+3731,+3729,+3299,+3297,+3298,+3295,+3293,+3296,+3294,+3292,+3312,+3311,+3314,+3309,+3290,+3313,+3410,+3315,+3411,+3412,+3316,+3421,+3317,+3415,+3445,+3327,+3328,+3283,+3343,+3326,+3325,+3330,+3286,+3399,+3398,+3329,+3446,+3400,+3331,+3401,+3281,+3332,+3279,+3402,+3419,+3407,+3356,+3358,+3357,+3280,+3354,+3277,+3278,+3346,+3347,+3377,+3378,+3345,+3386,+3379,+3348,+3384,+3418,+3372,+3276,+3275,+3374,+3274,+3373,+3375,+3369,+3371,+3376,+3273,+3396,+3397,+3395,+3388,+3360,+3370,+3361,+3394,+3387,+3420,+3359,+3389,+3272,+3391,+3393,+3390,+3392,+3363,+3362,+3367,+3365,+3705,+3271,+3704,+3703,+3270,+3269,+3702,+3268,+3224,+3267,+3701,+3225,+3699,+3265,+3264,+3266,+3263,+3262,+3249,+3228,+3230,+3251,+3301,+3300,+3302,+3252]
+    right_back_paw = [3472,+3627,+3470,+3469,+3471,+3473,+3626,+3625,+3475,+3655,+3519,+3468,+3629,+3466,+3476,+3624,+3521,+3654,+3657,+3838,+3518,+3653,+3839,+3553,+3474,+3516,+3656,+3628,+3834,+3535,+3630,+3658,+3477,+3520,+3517,+3595,+3522,+3597,+3596,+3501,+3534,+3503,+3478,+3500,+3479,+3502,+3607,+3499,+3608,+3496,+3605,+3609,+3504,+3606,+3642,+3614,+3498,+3480,+3631,+3610,+3613,+3506,+3659,+3660,+3632,+3841,+3661,+3836,+3662,+3633,+3663,+3664,+3634,+3635,+3486,+3665,+3636,+3637,+3666,+3490,+3837,+3667,+3493,+3638,+3492,+3495,+3616,+3644,+3494,+3835,+3643,+3833,+3840,+3615,+3650,+3668,+3652,+3651,+3645,+3646,+3647,+3649,+3648,+3622,+3617,+3448,+3621,+3618,+3623,+3462,+3464,+3460,+3620,+3458,+3461,+3463,+3465,+3573,+3571,+3467,+3569,+3557,+3558,+3572,+3570,+3556,+3585,+3593,+3594,+3459,+3566,+3592,+3567,+3568,+3538,+3539,+3555,+3537,+3536,+3554,+3575,+3574,+3583,+3541,+3550,+3576,+3581,+3639,+3577,+3551,+3582,+3580,+3552,+3578,+3542,+3549,+3579,+3523,+3526,+3598,+3525,+3600,+3640,+3599,+3601,+3602,+3603,+3529,+3604,+3530,+3533,+3532,+3611,+3612,+3482,+3481,+3505,+3452,+3455,+3456,+3454,+3457,+3619,+3451,+3450,+3449,+3591,+3589,+3641,+3584,+3561,+3587,+3559,+3488,+3484,+3483]
+    left_front_paw = [1791,+1950,+1948,+1790,+1789,+1746,+1788,+1747,+1949,+1944,+1792,+1945,+1356,+1775,+1759,+1777,+1787,+1946,+1757,+1761,+1745,+1943,+1947,+1744,+1309,+1786,+1771,+1354,+1774,+1765,+1767,+1768,+1772,+1763,+1770,+1773,+1769,+1764,+1766,+1758,+1760,+1762,+1336,+1333,+1330,+1325,+1756,+1323,+1755,+1753,+1749,+1754,+1751,+1321,+1752,+1748,+1750,+1312,+1319,+1315,+1313,+1317,+1318,+1316,+1314,+1311,+1310,+1299,+1276,+1355,+1297,+1353,+1298,+1300,+1352,+1351,+1785,+1784,+1349,+1783,+1782,+1781,+1780,+1779,+1778,+1776,+1343,+1341,+1344,+1339,+1342,+1340,+1360,+1335,+1338,+1362,+1357,+1361,+1363,+1458,+1337,+1459,+1456,+1460,+1493,+1332,+1375,+1376,+1331,+1374,+1378,+1334,+1373,+1494,+1377,+1446,+1448,+1379,+1449,+1329,+1327,+1404,+1406,+1405,+1402,+1328,+1426,+1432,+1434,+1403,+1394,+1395,+1433,+1425,+1286,+1380,+1466,+1431,+1290,+1401,+1381,+1427,+1450,+1393,+1430,+1326,+1396,+1428,+1397,+1429,+1398,+1420,+1324,+1422,+1417,+1419,+1421,+1443,+1418,+1423,+1444,+1442,+1424,+1445,+1495,+1440,+1441,+1468,+1436,+1408,+1322,+1435,+1415,+1439,+1409,+1283,+1438,+1416,+1407,+1437,+1411,+1413,+1414,+1320,+1273,+1272,+1278,+1469,+1463,+1457,+1358,+1464,+1465,+1359,+1372,+1391,+1390,+1455,+1447,+1454,+1467,+1453,+1452,+1451,+1383,+1345,+1347,+1348,+1350,+1364,+1392,+1410,+1412]
+    left_back_paw = [1957,+1958,+1701,+1956,+1951,+1703,+1715,+1702,+1700,+1673,+1705,+1952,+1955,+1674,+1699,+1675,+1953,+1704,+1954,+1698,+1677,+1671,+1672,+1714,+1706,+1676,+1519,+1523,+1686,+1713,+1692,+1685,+1543,+1664,+1712,+1691,+1959,+1541,+1684,+1542,+1496,+1663,+1540,+1497,+1499,+1498,+1500,+1693,+1665,+1694,+1716,+1666,+1695,+1501,+1502,+1696,+1667,+1503,+1697,+1504,+1668,+1669,+1506,+1670,+1508,+1510,+1507,+1509,+1511,+1512,+1621,+1606,+1619,+1605,+1513,+1620,+1618,+1604,+1633,+1641,+1642,+1607,+1617,+1514,+1632,+1614,+1689,+1640,+1515,+1586,+1616,+1516,+1517,+1603,+1615,+1639,+1585,+1521,+1602,+1587,+1584,+1601,+1623,+1622,+1631,+1598,+1624,+1629,+1589,+1687,+1625,+1599,+1630,+1569,+1570,+1628,+1626,+1597,+1627,+1590,+1594,+1571,+1568,+1567,+1574,+1646,+1573,+1645,+1648,+1564,+1688,+1647,+1643,+1649,+1650,+1651,+1577,+1644,+1565,+1652,+1566,+1578,+1518,+1524,+1583,+1582,+1520,+1581,+1522,+1525,+1549,+1551,+1580,+1552,+1550,+1656,+1658,+1554,+1657,+1659,+1548,+1655,+1690,+1660,+1556,+1653,+1558,+1661,+1544,+1662,+1654,+1547,+1545,+1527,+1560,+1526,+1678,+1679,+1528,+1708,+1707,+1680,+1529,+1530,+1709,+1546,+1681,+1710,+1711,+1682,+1532,+1531,+1683,+1534,+1533,+1536,+1538,+1600,+1553]
+
+
+    all_contact_vertices = right_front_paw + right_back_paw + left_front_paw + left_back_paw
+
+    name = 'all4pawsincontact.jpg'
+    print('work on 4paw images')
+    gc_info_raw = all_contact_vertices      # a list with all vertex numbers that are in ground contact
+
+    vertex_overview = calculate_vertex_overview_for_gc_annotation(name, gc_info_raw, vert_dists, root_out_path_vis=ROOT_OUT_PATH, verts=verts, faces=faces, img_v12_dir=None)
+    np.save(ROOT_OUT_PATH  + name.replace('.jpg', '_gc_vertdists_overview.npy'), vertex_overview)
+
+    vertex_overview_dict = {}
+    vertex_overview_dict[name.split('.')[0]] = {'gc_vertdists_overview': vertex_overview, 'gc_index_list': gc_info_raw}
+    with open(ROOT_OUT_PATH + 'gc_annots_overview_all4pawsincontact_xx.pkl', 'wb') as fp:
+        pkl.dump(vertex_overview_dict, fp)
+
+
+
+
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+
+
+
+

src/graph_networks/losses_for_vertex_wise_predictions/calculate_distance_between_points_on_mesh_forpaws.py

@@ -0,0 +1,317 @@
+
+"""
+code adapted from: https://github.com/mikedh/trimesh/blob/main/examples/shortest.py
+shortest.py
+----------------
+Given a mesh and two vertex indices find the shortest path
+between the two vertices while only traveling along edges
+of the mesh.
+"""
+
+# python src/graph_networks/losses_for_vertex_wise_predictions/calculate_distance_between_points_on_mesh_forpaws.py
+
+
+import os
+import sys
+import glob
+import csv
+import json
+import shutil
+import tqdm
+import numpy as np
+import pickle as pkl
+import trimesh
+import networkx as nx
+
+
+
+
+
+def read_csv(csv_file):
+    with open(csv_file,'r') as f:
+        reader = csv.reader(f)
+        headers = next(reader)
+        row_list = [{h:x for (h,x) in zip(headers,row)} for row in reader]
+    return row_list
+
+
+def load_all_template_mesh_distances(root_out_path, filename='all_vertex_distances.npy'):
+    vert_dists = np.load(root_out_path + filename)
+    return vert_dists
+
+
+def prepare_graph_from_template_mesh_and_calculate_all_distances(path_mesh, root_out_path, calc_dist_mat=False):
+    # root_out_path = ROOT_OUT_PATH
+    '''
+    from smal_pytorch.smal_model.smal_torch_new import SMAL
+    smal = SMAL()
+    verts = smal.v_template.detach().cpu().numpy()
+    faces = smal.faces.detach().cpu().numpy()
+    '''
+    # path_mesh = ROOT_PATH_MESH + 'mesh_downsampling_meshesmy_smpl_39dogsnorm_Jr_4_dog_template_downsampled0.obj'
+    my_mesh = trimesh.load_mesh(path_mesh, process=False,  maintain_order=True)
+    verts = my_mesh.vertices
+    faces = my_mesh.faces
+    # edges without duplication
+    edges = my_mesh.edges_unique
+    # the actual length of each unique edge
+    length = my_mesh.edges_unique_length
+    # create the graph with edge attributes for length (option A)
+    #   g = nx.Graph()
+    #   for edge, L in zip(edges, length): g.add_edge(*edge, length=L)
+    # you can create the graph with from_edgelist and
+    # a list comprehension (option B)
+    ga = nx.from_edgelist([(e[0], e[1], {'length': L}) for e, L in zip(edges, length)])
+    # calculate the distances between all vertex pairs
+    if calc_dist_mat:
+        # calculate distances between all possible vertex pairs
+        # shortest_path = nx.shortest_path(ga, source=ind_v0, target=ind_v1, weight='length')
+        # shortest_dist = nx.shortest_path_length(ga, source=ind_v0, target=ind_v1, weight='length')
+        dis = dict(nx.shortest_path_length(ga, weight='length', method='dijkstra'))
+        vertex_distances = np.zeros((n_verts_smal, n_verts_smal))
+        for ind_v0 in range(n_verts_smal):
+            print(ind_v0)
+            for ind_v1 in range(ind_v0, n_verts_smal):
+                vertex_distances[ind_v0, ind_v1] = dis[ind_v0][ind_v1]
+                vertex_distances[ind_v1, ind_v0] = dis[ind_v0][ind_v1]
+        # save those distances
+        np.save(root_out_path + 'all_vertex_distances.npy', vertex_distances)
+        vert_dists = vertex_distances
+    else:
+        vert_dists = np.load(root_out_path + 'all_vertex_distances.npy')
+    return ga, vert_dists
+
+
+def calculate_vertex_overview_for_gc_annotation(name, gc_info_raw, vert_dists, root_out_path_vis=None, verts=None, faces=None, img_v12_dir=None):
+    # input:
+    #   root_out_path_vis = ROOT_OUT_PATH
+    #   img_v12_dir = IMG_V12_DIR
+    #   name = images_with_gc_labelled[ind_img]
+    #   gc_info_raw = gc_dict['bite/' + name] 
+    # output:
+    #    vertex_overview: np array of shape (n_verts_smal, 3) with [first: no-contact=0 contact=1     second: index of vertex     third: dist]
+    n_verts_smal = 3889
+    gc_vertices = []
+    gc_info_np = np.zeros((n_verts_smal))
+    for ind_v in gc_info_raw:
+        if ind_v < n_verts_smal:
+            gc_vertices.append(ind_v)
+            gc_info_np[ind_v] = 1
+    # save a visualization of those annotations 
+    if root_out_path_vis is not None:
+        my_mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False,  maintain_order=True)
+    if img_v12_dir is not None and root_out_path_vis is not None:
+        vert_colors = np.repeat(255*gc_info_np[:, None], 3, 1)
+        my_mesh.visual.vertex_colors = vert_colors
+        my_mesh.export(root_out_path_vis + (name).replace('.jpg', '_withgc.obj'))
+        img_path = img_v12_dir + name
+        shutil.copy(img_path, root_out_path_vis + name)
+    # calculate for each vertex the distance to the closest element of the other group
+    non_gc_vertices = list(set(range(n_verts_smal)) - set(gc_vertices))
+    print('vertices in contact: ' + str(len(gc_vertices)))
+    print('vertices without contact: ' + str(len(non_gc_vertices)))
+    vertex_overview = np.zeros((n_verts_smal, 3))   # first: no-contact=0 contact=1     second: index of vertex     third: dist
+    vertex_overview[:, 0] = gc_info_np
+    # loop through all contact vertices
+    for ind_v in gc_vertices:
+        min_length = 100
+        for ind_v_ps in non_gc_vertices:    # possible solution
+            # this_path = nx.shortest_path(ga, source=ind_v, target=ind_v_ps, weight='length')
+            # this_length = nx.shortest_path_length(ga, source=ind_v, target=ind_v_ps, weight='length')
+            this_length = vert_dists[ind_v, ind_v_ps]
+            if this_length < min_length:
+                min_length = this_length
+                vertex_overview[ind_v, 1] = ind_v_ps
+                vertex_overview[ind_v, 2] = this_length
+    # loop through all non-contact vertices
+    for ind_v in non_gc_vertices:
+        min_length = 100
+        for ind_v_ps in gc_vertices:    # possible solution
+            # this_path = nx.shortest_path(ga, source=ind_v, target=ind_v_ps, weight='length')
+            # this_length = nx.shortest_path_length(ga, source=ind_v, target=ind_v_ps, weight='length')
+            this_length = vert_dists[ind_v, ind_v_ps]
+            if this_length < min_length:
+                min_length = this_length
+                vertex_overview[ind_v, 1] = ind_v_ps
+                vertex_overview[ind_v, 2] = this_length
+    if root_out_path_vis is not None:
+        # save a colored mesh
+        my_mesh_dists = my_mesh.copy()
+        scale_0 = (vertex_overview[vertex_overview[:, 0]==0, 2]).max()
+        scale_1 = (vertex_overview[vertex_overview[:, 0]==1, 2]).max()
+        vert_col = np.zeros((n_verts_smal, 3)) 
+        vert_col[vertex_overview[:, 0]==0, 1] = vertex_overview[vertex_overview[:, 0]==0, 2] * 255 / scale_0     # green
+        vert_col[vertex_overview[:, 0]==1, 0] = vertex_overview[vertex_overview[:, 0]==1, 2] * 255 / scale_1     # red
+        my_mesh_dists.visual.vertex_colors = np.uint8(vert_col)
+        my_mesh_dists.export(root_out_path_vis + (name).replace('.jpg', '_withgcdists.obj'))
+    return vertex_overview
+
+
+def summarize_results_stage2b(row_list, display_worker_performance=False):
+    # four catch trials are included in every batch
+    annot_n02088466_3184 = {'paw_rb': 0, 'paw_rf': 1, 'paw_lb': 1, 'paw_lf': 1, 'additional_part': 0, 'no_contact': 0}
+    annot_n02100583_9922 = {'paw_rb': 1, 'paw_rf': 0, 'paw_lb': 0, 'paw_lf': 0, 'additional_part': 0, 'no_contact': 0}
+    annot_n02105056_2798 = {'paw_rb': 1, 'paw_rf': 1, 'paw_lb': 1, 'paw_lf': 1, 'additional_part': 1, 'no_contact': 0}
+    annot_n02091831_2288 = {'paw_rb': 0, 'paw_rf': 1, 'paw_lb': 1, 'paw_lf': 0, 'additional_part': 0, 'no_contact': 0}
+    all_comments = []
+    all_annotations = {}
+    for row in row_list:
+        all_comments.append(row['Answer.submitComments'])
+        worker_id = row['WorkerId']
+        if display_worker_performance:
+            print('----------------------------------------------------------------------------------------------')
+            print('Worker ID: ' + worker_id)
+        n_wrong = 0
+        n_correct = 0
+        for ind in range(0, len(row['Answer.submitValuesNotSure'].split(';')) - 1):
+            input_image = (row['Input.images'].split(';')[ind]).split('StanExtV12_Images/')[-1]
+            paw_rb = row['Answer.submitValuesRightBack'].split(';')[ind] 
+            paw_rf = row['Answer.submitValuesRightFront'].split(';')[ind] 
+            paw_lb = row['Answer.submitValuesLeftBack'].split(';')[ind] 
+            paw_lf = row['Answer.submitValuesLeftFront'].split(';')[ind]  
+            addpart = row['Answer.submitValuesAdditional'].split(';')[ind] 
+            no_contact = row['Answer.submitValuesNoContact'].split(';')[ind]  
+            unsure = row['Answer.submitValuesNotSure'].split(';')[ind] 
+            annot = {'paw_rb': paw_rb, 'paw_rf': paw_rf, 'paw_lb': paw_lb, 'paw_lf': paw_lf, 
+                    'additional_part': addpart, 'no_contact': no_contact, 'not_sure': unsure, 
+                    'worker_id': worker_id}     # , 'input_image': input_image}
+            if ind == 0:
+                gt = annot_n02088466_3184
+            elif ind == 1:
+                gt = annot_n02105056_2798
+            elif ind == 2:
+                gt = annot_n02091831_2288
+            elif ind == 3:
+                gt = annot_n02100583_9922
+            else:
+                pass
+            if ind < 4:
+                for key in gt.keys():
+                    if str(annot[key]) == str(gt[key]):
+                        n_correct += 1
+                    else:
+                        if display_worker_performance:
+                            print(input_image)
+                            print(key + ':[ expected: ' + str(gt[key]) + '   predicted: ' + str(annot[key]) + ' ]')
+                        n_wrong += 1
+            else:
+                all_annotations[input_image] = annot
+        if display_worker_performance:
+            print('n_correct: ' + str(n_correct))
+            print('n_wrong: ' + str(n_wrong))
+    return all_annotations, all_comments
+
+
+
+
+
+
+def main():
+
+    ROOT_PATH_MESH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/graphcmr/data/meshes/'
+    ROOT_PATH_ANNOT = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage2b/'
+    IMG_V12_DIR = '/ps/scratch/nrueegg/new_projects/Animals/data/dog_datasets/Stanford_Dogs_Dataset/StanfordExtra_V12/StanExtV12_Images/'		
+    # ROOT_OUT_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/src/graph_networks/losses_for_vertex_wise_predictions/debugging_results/'
+    ROOT_OUT_PATH = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage2b/'
+    ROOT_OUT_PATH_VIS = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage2b/vis/'
+    ROOT_OUT_PATH_DISTSGCNONGC = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/stage2b/vertex_distances_gc_nongc/'
+    ROOT_PATH_ALL_VERT_DIST_TEMPLATE  = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/data/stanext_related_data/ground_contact_annotations/'
+
+    # load all vertex distances
+    path_mesh = ROOT_PATH_MESH + 'mesh_downsampling_meshesmy_smpl_39dogsnorm_Jr_4_dog_template_downsampled0.obj'
+    my_mesh = trimesh.load_mesh(path_mesh, process=False,  maintain_order=True)
+    verts = my_mesh.vertices
+    faces = my_mesh.faces
+    # vert_dists, ga = prepare_graph_from_template_mesh_and_calculate_all_distances(path_mesh, ROOT_OUT_PATH, calc_dist_mat=False)
+    vert_dists = load_all_template_mesh_distances(ROOT_PATH_ALL_VERT_DIST_TEMPLATE, filename='all_vertex_distances.npy')
+
+
+
+
+
+    # paw vertices:
+    # left and right is a bit different, but that is ok (we will anyways mirror data at training time)
+    right_front_paw = [3829,+3827,+3825,+3718,+3722,+3723,+3743,+3831,+3719,+3726,+3716,+3724,+3828,+3717,+3721,+3725,+3832,+3830,+3720,+3288,+3740,+3714,+3826,+3715,+3728,+3712,+3287,+3284,+3727,+3285,+3742,+3291,+3710,+3697,+3711,+3289,+3730,+3713,+3739,+3282,+3738,+3708,+3709,+3741,+3698,+3696,+3308,+3695,+3706,+3700,+3707,+3306,+3305,+3737,+3304,+3303,+3307,+3736,+3735,+3250,+3261,+3732,+3734,+3733,+3731,+3729,+3299,+3297,+3298,+3295,+3293,+3296,+3294,+3292,+3312,+3311,+3314,+3309,+3290,+3313,+3410,+3315,+3411,+3412,+3316,+3421,+3317,+3415,+3445,+3327,+3328,+3283,+3343,+3326,+3325,+3330,+3286,+3399,+3398,+3329,+3446,+3400,+3331,+3401,+3281,+3332,+3279,+3402,+3419,+3407,+3356,+3358,+3357,+3280,+3354,+3277,+3278,+3346,+3347,+3377,+3378,+3345,+3386,+3379,+3348,+3384,+3418,+3372,+3276,+3275,+3374,+3274,+3373,+3375,+3369,+3371,+3376,+3273,+3396,+3397,+3395,+3388,+3360,+3370,+3361,+3394,+3387,+3420,+3359,+3389,+3272,+3391,+3393,+3390,+3392,+3363,+3362,+3367,+3365,+3705,+3271,+3704,+3703,+3270,+3269,+3702,+3268,+3224,+3267,+3701,+3225,+3699,+3265,+3264,+3266,+3263,+3262,+3249,+3228,+3230,+3251,+3301,+3300,+3302,+3252]
+    right_back_paw = [3472,+3627,+3470,+3469,+3471,+3473,+3626,+3625,+3475,+3655,+3519,+3468,+3629,+3466,+3476,+3624,+3521,+3654,+3657,+3838,+3518,+3653,+3839,+3553,+3474,+3516,+3656,+3628,+3834,+3535,+3630,+3658,+3477,+3520,+3517,+3595,+3522,+3597,+3596,+3501,+3534,+3503,+3478,+3500,+3479,+3502,+3607,+3499,+3608,+3496,+3605,+3609,+3504,+3606,+3642,+3614,+3498,+3480,+3631,+3610,+3613,+3506,+3659,+3660,+3632,+3841,+3661,+3836,+3662,+3633,+3663,+3664,+3634,+3635,+3486,+3665,+3636,+3637,+3666,+3490,+3837,+3667,+3493,+3638,+3492,+3495,+3616,+3644,+3494,+3835,+3643,+3833,+3840,+3615,+3650,+3668,+3652,+3651,+3645,+3646,+3647,+3649,+3648,+3622,+3617,+3448,+3621,+3618,+3623,+3462,+3464,+3460,+3620,+3458,+3461,+3463,+3465,+3573,+3571,+3467,+3569,+3557,+3558,+3572,+3570,+3556,+3585,+3593,+3594,+3459,+3566,+3592,+3567,+3568,+3538,+3539,+3555,+3537,+3536,+3554,+3575,+3574,+3583,+3541,+3550,+3576,+3581,+3639,+3577,+3551,+3582,+3580,+3552,+3578,+3542,+3549,+3579,+3523,+3526,+3598,+3525,+3600,+3640,+3599,+3601,+3602,+3603,+3529,+3604,+3530,+3533,+3532,+3611,+3612,+3482,+3481,+3505,+3452,+3455,+3456,+3454,+3457,+3619,+3451,+3450,+3449,+3591,+3589,+3641,+3584,+3561,+3587,+3559,+3488,+3484,+3483]
+    left_front_paw = [1791,+1950,+1948,+1790,+1789,+1746,+1788,+1747,+1949,+1944,+1792,+1945,+1356,+1775,+1759,+1777,+1787,+1946,+1757,+1761,+1745,+1943,+1947,+1744,+1309,+1786,+1771,+1354,+1774,+1765,+1767,+1768,+1772,+1763,+1770,+1773,+1769,+1764,+1766,+1758,+1760,+1762,+1336,+1333,+1330,+1325,+1756,+1323,+1755,+1753,+1749,+1754,+1751,+1321,+1752,+1748,+1750,+1312,+1319,+1315,+1313,+1317,+1318,+1316,+1314,+1311,+1310,+1299,+1276,+1355,+1297,+1353,+1298,+1300,+1352,+1351,+1785,+1784,+1349,+1783,+1782,+1781,+1780,+1779,+1778,+1776,+1343,+1341,+1344,+1339,+1342,+1340,+1360,+1335,+1338,+1362,+1357,+1361,+1363,+1458,+1337,+1459,+1456,+1460,+1493,+1332,+1375,+1376,+1331,+1374,+1378,+1334,+1373,+1494,+1377,+1446,+1448,+1379,+1449,+1329,+1327,+1404,+1406,+1405,+1402,+1328,+1426,+1432,+1434,+1403,+1394,+1395,+1433,+1425,+1286,+1380,+1466,+1431,+1290,+1401,+1381,+1427,+1450,+1393,+1430,+1326,+1396,+1428,+1397,+1429,+1398,+1420,+1324,+1422,+1417,+1419,+1421,+1443,+1418,+1423,+1444,+1442,+1424,+1445,+1495,+1440,+1441,+1468,+1436,+1408,+1322,+1435,+1415,+1439,+1409,+1283,+1438,+1416,+1407,+1437,+1411,+1413,+1414,+1320,+1273,+1272,+1278,+1469,+1463,+1457,+1358,+1464,+1465,+1359,+1372,+1391,+1390,+1455,+1447,+1454,+1467,+1453,+1452,+1451,+1383,+1345,+1347,+1348,+1350,+1364,+1392,+1410,+1412]
+    left_back_paw = [1957,+1958,+1701,+1956,+1951,+1703,+1715,+1702,+1700,+1673,+1705,+1952,+1955,+1674,+1699,+1675,+1953,+1704,+1954,+1698,+1677,+1671,+1672,+1714,+1706,+1676,+1519,+1523,+1686,+1713,+1692,+1685,+1543,+1664,+1712,+1691,+1959,+1541,+1684,+1542,+1496,+1663,+1540,+1497,+1499,+1498,+1500,+1693,+1665,+1694,+1716,+1666,+1695,+1501,+1502,+1696,+1667,+1503,+1697,+1504,+1668,+1669,+1506,+1670,+1508,+1510,+1507,+1509,+1511,+1512,+1621,+1606,+1619,+1605,+1513,+1620,+1618,+1604,+1633,+1641,+1642,+1607,+1617,+1514,+1632,+1614,+1689,+1640,+1515,+1586,+1616,+1516,+1517,+1603,+1615,+1639,+1585,+1521,+1602,+1587,+1584,+1601,+1623,+1622,+1631,+1598,+1624,+1629,+1589,+1687,+1625,+1599,+1630,+1569,+1570,+1628,+1626,+1597,+1627,+1590,+1594,+1571,+1568,+1567,+1574,+1646,+1573,+1645,+1648,+1564,+1688,+1647,+1643,+1649,+1650,+1651,+1577,+1644,+1565,+1652,+1566,+1578,+1518,+1524,+1583,+1582,+1520,+1581,+1522,+1525,+1549,+1551,+1580,+1552,+1550,+1656,+1658,+1554,+1657,+1659,+1548,+1655,+1690,+1660,+1556,+1653,+1558,+1661,+1544,+1662,+1654,+1547,+1545,+1527,+1560,+1526,+1678,+1679,+1528,+1708,+1707,+1680,+1529,+1530,+1709,+1546,+1681,+1710,+1711,+1682,+1532,+1531,+1683,+1534,+1533,+1536,+1538,+1600,+1553]
+
+
+
+
+
+    all_keys = []
+    gc_dict = {}
+    vertex_overview_nocontact = {}
+    # data/stanext_related_data/ground_contact_annotations/stage3/main_partA1667_20221021_140108.csv
+    for csv_file in ['Stage2b_finalResults.csv']:
+        # load all ground contact annotations
+        gc_annot_csv = ROOT_PATH_ANNOT + csv_file   # 'my_gcannotations_qualification.csv'
+        gc_row_list = read_csv(gc_annot_csv)
+        all_annotations, all_comments = summarize_results_stage2b(gc_row_list, display_worker_performance=False)
+        for key, value in all_annotations.items():
+            if value['not_sure'] == '0':
+                if value['no_contact'] == '1':
+                    vertex_overview_nocontact[key.split('.')[0]] = {'gc_vertdists_overview': 'no contact', 'gc_index_list': None}
+                else:
+                    all_contact_vertices = []
+                    if value['paw_rf'] == '1':
+                        all_contact_vertices.extend(right_front_paw)
+                    if value['paw_rb'] == '1':
+                        all_contact_vertices.extend(right_back_paw)
+                    if value['paw_lf'] == '1':
+                        all_contact_vertices.extend(left_front_paw)
+                    if value['paw_lb'] == '1':
+                        all_contact_vertices.extend(left_back_paw)
+                    gc_dict[key] = all_contact_vertices
+    print('number of labeled images: ' + str(len(gc_dict.keys())))     
+    print('number of images without contact: ' + str(len(vertex_overview_nocontact.keys())))      
+
+    # prepare and save contact annotations including distances
+    vertex_overview_dict = {}
+    for ind_img, name_ingcdict in enumerate(gc_dict.keys()):     #  range(len(gc_dict.keys())):
+        name = name_ingcdict      # name_ingcdict.split('bite/')[1]
+        # name = images_with_gc_labelled[ind_img]
+        print('work on image ' + str(ind_img) + ': ' + name)
+        # gc_info_raw = gc_dict['bite/' + name]      # a list with all vertex numbers that are in ground contact
+        gc_info_raw = gc_dict[name_ingcdict]      # a list with all vertex numbers that are in ground contact
+
+        if not os.path.exists(ROOT_OUT_PATH_VIS + name.split('/')[0]): os.makedirs(ROOT_OUT_PATH_VIS + name.split('/')[0])
+        if not os.path.exists(ROOT_OUT_PATH_DISTSGCNONGC + name.split('/')[0]): os.makedirs(ROOT_OUT_PATH_DISTSGCNONGC + name.split('/')[0])
+
+        vertex_overview = calculate_vertex_overview_for_gc_annotation(name, gc_info_raw, vert_dists, root_out_path_vis=ROOT_OUT_PATH_VIS, verts=verts, faces=faces, img_v12_dir=None)
+        np.save(ROOT_OUT_PATH_DISTSGCNONGC + name.replace('.jpg', '_gc_vertdists_overview.npy'), vertex_overview)
+
+        vertex_overview_dict[name.split('.')[0]] = {'gc_vertdists_overview': vertex_overview, 'gc_index_list': gc_info_raw}
+
+
+
+        
+
+    # import pdb; pdb.set_trace()
+
+    with open(ROOT_OUT_PATH + 'gc_annots_overview_stage2b_contact_complete_xx.pkl', 'wb') as fp:
+        pkl.dump(vertex_overview_dict, fp)
+
+    with open(ROOT_OUT_PATH + 'gc_annots_overview_stage2b_nocontact_complete_xx.pkl', 'wb') as fp:
+        pkl.dump(vertex_overview_nocontact, fp)
+
+
+
+
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+
+
+
+