refine_fitting.py

import torch
import numpy as np
from tqdm import tqdm
import os
from termcolor import colored
import argparse


# from losses import ChamferDistance, MaxMixturePrior, summed_L2, LossTracker
from visualization import viz_error_curves, viz_iteration, set_init_plot, viz_final_fit
from utils import (check_scan_prequisites_fit_bm, cleanup, get_normals, 
                   load_config, save_configs,
                   load_landmarks,load_scan,
                   get_already_fitted_scan_names, get_skipped_scan_names, 
                   initialize_fit_bm_loss_weights, load_loss_weights_config, 
                   print_loss_weights, print_losses, print_params, 
                   process_body_model_path, process_default_dtype, 
                   process_landmarks, process_visualize_steps, 
                   create_results_directory, to_txt,
                   setup_socket, send_to_socket,
                   )
from body_models import BodyModel
from body_parameters import BodyParameters
from datasets import FAUST, CAESAR, FourDHumanOutfit
from dash_app import run_dash_app_as_subprocess
import losses


def fit_body_model(input_dict: dict, cfg: dict):
    """
    Fit a body body model (SMPL/SMPLX) to the input scan using
    data, landmark and prior losses

    :param: input_dict (dict): with keys:
            "name": name of the scan
            "vertices": numpy array (N,3)
            "faces": numpy array (N,3) or None if no faces
            "landmarks": dictionary with keys as landmark names and 
                        values as list [x,y,z] or np.ndarray (3,)
    :param: cfg (dict): config file defined in configs/config.yaml
    """

    DEFAULT_DTYPE = cfg['default_dtype']
    VERBOSE = cfg['verbose']
    VISUALIZE = cfg['visualize']
    VISUALIZE_STEPS = cfg['visualize_steps']
    VISUALIZE_LOGSCALE = cfg["error_curves_logscale"]
    SAVE_PATH = cfg['save_path']
    SOCKET_TYPE = cfg["socket_type"]
    
    if VISUALIZE:
        socket = cfg["socket"]


    # inputs
    input_name = input_dict["name"]
    input_sequence_name = input_dict["sequence_name"] if "sequence_name" in input_dict else ""
    input_vertices = input_dict["vertices"]
    input_faces = input_dict["faces"]
    input_landmarks = input_dict["landmarks"]
    USE_LANDMARKS = True if ("landmarks" in cfg["use_losses"]) else False
    input_index = input_dict["scan_index"]
    input_pose = input_dict["pose"]
    input_shape = input_dict["shape"]
    input_shape_dim = input_shape.shape[0] if input_shape.shape[0] != 1 else input_shape.shape[1]
    input_trans = input_dict["trans"]
    input_gender = input_dict["gender"] if not isinstance(input_dict["gender"],type(None)) else "NEUTRAL"
    USE_NORMALS = True if ("normal" in cfg["use_losses"]) else False 

    # process save path
    SAVE_PATH = os.path.join(SAVE_PATH, input_sequence_name)
    if not os.path.exists(SAVE_PATH):
        os.makedirs(SAVE_PATH)

    # process inputs
    input_vertices = torch.from_numpy(input_vertices).type(DEFAULT_DTYPE).unsqueeze(0).cuda()
    input_faces = torch.from_numpy(input_faces).type(DEFAULT_DTYPE) if \
                            (not isinstance(input_faces,type(None))) else None

    if USE_LANDMARKS:
        landmarks_order = sorted(list(input_landmarks.keys()))
        input_landmarks = np.array([input_landmarks[k] for k in landmarks_order])
        input_landmarks = torch.from_numpy(input_landmarks)
        input_landmarks = input_landmarks.type(DEFAULT_DTYPE).cuda()
    else:
        input_landmarks = []
        landmarks_order = []

    if USE_NORMALS:
        input_normals = get_normals(input_vertices[0].detach().cpu())
        input_normals = input_normals.cuda() # (N,3)
    else:
        input_normals = None
        template_normals = None

    # setup body model
    cfg["body_model_gender"] = input_gender
    cfg["body_model_num_betas"] = input_shape_dim
    body_model = BodyModel(cfg)
    body_model.cuda()
    cfg["init_params"] = {"pose":input_pose, "shape":input_shape, "trans":input_trans}
    body_model_params = BodyParameters(cfg).cuda()
    body_model_landmark_inds = body_model.landmark_indices(landmarks_order)
    print(f"Using {len(input_landmarks)}/{len(body_model.all_landmark_indices)} landmarks.")

    # configure optimization
    ITERATIONS = cfg['iterations']
    LR = cfg['lr']
    START_LR_DECAY = cfg['start_lr_decay_iteration']
    loss_weights = cfg['loss_weights']
    # loss_tracker = LossTracker(loss_weights[0].keys())



    body_optimizer = torch.optim.Adam(body_model_params.parameters(), lr=LR)
    # chamfer_distance = ChamferDistance()
    # prior = MaxMixturePrior(prior_folder=cfg["prior_path"], num_gaussians=8)
    # prior = prior.cuda()
    loss_func = losses.Losses(cfg, cfg["loss_weights"])



    if VISUALIZE:
        starting_verts = body_model.deform_verts(input_pose.cuda(),
                                                input_shape.cuda(),
                                                input_trans.cuda(),
                                                torch.ones(1).cuda()*1).detach().cpu()
        fig = set_init_plot(input_vertices[0].detach().cpu(), 
                            starting_verts,  
                            title=f"Fitting ({input_name}) - initial setup",
                            subsample_input_scan_verts=20)
        send_to_socket(fig, socket, SOCKET_TYPE)

    iterator = tqdm(range(ITERATIONS))
    for i in iterator:

        if VERBOSE: print(colored(f"iteration {i}","red"))

        # if i in loss_weights.keys():
        #     if VERBOSE: print(colored(f"\tChanging loss weights","red"))
        #     data_loss_weight = loss_weights[i]['data']
        #     landmark_loss_weight = loss_weights[i]['landmark']
        #     prior_loss_weight = loss_weights[i]['prior']
        #     beta_loss_weight = loss_weights[i]['beta']

        # if VERBOSE: print_loss_weights(data_loss_weight,landmark_loss_weight,
        #                                 prior_loss_weight,beta_loss_weight,
        #                                 "loss weights:")

        # forward
        pose, beta, trans, scale = body_model_params.forward()
        if VERBOSE: print_params(pose,beta,trans,scale)
        body_model_verts = body_model.deform_verts(pose,
                                                   beta,
                                                   trans,
                                                   scale)     

        if USE_NORMALS:
            template_normals = get_normals(body_model_verts.detach().cpu())
            template_normals = template_normals.cuda() # (6890,3)

        # compute loss
        loss_dict = dict(scan_vertices=input_vertices,
                         scan_normals=input_normals,
                         template_vertices = body_model_verts.unsqueeze(0),
                         template_normals=template_normals,
                         pose=pose[:, 3:],
                         betas=beta
                        )
        
        if USE_LANDMARKS:
            lm_dict = {"scan_landmarks": input_landmarks,
                      "landmarks": body_model_verts[body_model_landmark_inds,:]}
            loss_dict.update(lm_dict)
        loss = loss_func(**loss_dict)

        
        # if VERBOSE: 
        #     print_losses(data_loss,landmark_loss,prior_loss,beta_loss,"losses")
        #     print_losses(data_loss_weighted,landmark_loss_weighted,
        #                  prior_loss_weighted,beta_loss_weighted,"losses weighted")
        iterator.set_description(f"Loss {loss.item():.4f}")

        # optimize
        body_optimizer.zero_grad()
        loss.backward()
        body_optimizer.step()

        if i >= START_LR_DECAY:
            for param_group in body_optimizer.param_groups:
                param_group['lr'] = LR*(ITERATIONS-i)/ITERATIONS
            if VERBOSE: print(colored(f"\tlr: {param_group['lr']}","yellow"))

        if VISUALIZE and (i in VISUALIZE_STEPS):
            new_title = f"Fitting {input_name} - iteration {i}"
            fig = viz_iteration(fig, body_model_verts.detach().cpu(), i , new_title)
            send_to_socket(fig, socket, SOCKET_TYPE)

            new_title = f"Fitting {input_name} losses - iteration {i}"
            fig_losses = viz_error_curves(loss_func.loss_tracker.losses, 
                                          loss_func.loss_weights, 
                                          new_title, VISUALIZE_LOGSCALE)
            send_to_socket(fig_losses, socket, SOCKET_TYPE)

        loss_func.update_loss_weights(i)


    if VISUALIZE:
        fig = viz_final_fit(input_vertices[0].detach().cpu(), 
                          body_model_verts.detach().cpu(),
                          input_faces,
                          title=f"Fitting {input_name} - final fit")
        send_to_socket(fig, socket, SOCKET_TYPE)

    with torch.no_grad():
        pose, beta, trans, scale = body_model_params.forward()
        body_model_verts = body_model.deform_verts(pose, beta, trans, scale)
        fitted_body_model_verts = body_model_verts.detach().cpu().data.numpy()
        fitted_pose = pose.detach().cpu().numpy()
        fitted_shape = beta.detach().cpu().numpy()
        trans = trans.detach().cpu().numpy()
        scale = scale.detach().cpu().numpy()

        save_to = os.path.join(SAVE_PATH,f"{input_name}.npz")
        np.savez(save_to, 
                 body_model = body_model.body_model_name,
                 vertices=fitted_body_model_verts, 
                 pose=fitted_pose, 
                 shape=fitted_shape, 
                 trans=trans, 
                 scale=scale, 
                 name=input_name,
                 scan_index=input_index)


def fit_body_model_onto_dataset(cfg: dict):
    
    # get dataset
    dataset_name = cfg["dataset_name"]
    cfg_dataset = cfg[dataset_name]
    cfg_dataset["use_landmarks"] = cfg["use_landmarks"]
    dataset = eval(cfg["dataset_name"])(**cfg_dataset)


    wait_after_fit_func = input if cfg["pause_script_after_fitting"] else print
    wait_after_fit_func_text = "Fitting completed - press any key to continue!" \
                        if cfg["pause_script_after_fitting"] else "Fitting completed!"

    # if continuing fitting process, get fitted and skipped scans
    fitted_scans = get_already_fitted_scan_names(cfg)
    skipped_scans = get_skipped_scan_names(cfg)

    for i in range(len(dataset)):
        input_example = dataset[i]
        scan_name = input_example["name"]
        sequence_name = input_example["sequence_name"] if "sequence_name" in input_example else ""
        print(f"Fitting scan {sequence_name}/{scan_name} -----------------")

        if (scan_name in fitted_scans) or \
            (scan_name in skipped_scans):
            continue
        
        process_scan = True #check_scan_prequisites_fit_bm(input_example)
        if process_scan:
            input_example["scan_index"] = i
            fit_body_model(input_example, cfg)
        else:
            skipped_scans.append(input_example["name"])
            to_txt(skipped_scans, cfg["save_path"], "skipped_scans.txt")
        print(wait_after_fit_func_text)
        wait_after_fit_func("-----------------------------------")
    print(f"Fitting for {dataset_name} dataset completed!")


def fit_body_model_onto_scan(cfg: dict):
    
    wait_after_fit_func = input if cfg["pause_script_after_fitting"] else print

    scan_name = cfg["scan_path"].split("/")[-1].split(".")[0]
    scan_vertices, scan_faces = load_scan(cfg["scan_path"])
    scan_vertices = scan_vertices / cfg["scale_scan"]

    landmarks = load_landmarks(cfg["landmark_path"],
                              cfg["use_landmarks"],
                              scan_vertices)
    landmarks = {lm_name: (np.array(lm_coord) / cfg["scale_landmarks"]).tolist()
                 for lm_name, lm_coord in landmarks.items()}

    input_example = {"name": scan_name,
                    "vertices": scan_vertices,
                    "faces": scan_faces,
                    "landmarks": landmarks,
                    }

    process_scan = check_scan_prequisites_fit_bm(input_example)
    if process_scan:
        input_example["scan_index"] = 0
        fit_body_model(input_example, cfg)
   
        print(f"Fitting completed - press any key to continue!")
        wait_after_fit_func("-----------------------------------")


if __name__ == "__main__":

    # parse arguments
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(
        help="Subparsers determine the fitting mode: onto_scan or onto_dataset.")
    
    parser_scan = subparsers.add_parser('onto_scan')
    parser_scan.add_argument("--scan_path", type=str, required=True)
    parser_scan.add_argument("--scale_scan", type=float, default=1.0, 
                             help="Scale (divide) the scan vertices by this factor.")
    parser_scan.add_argument("--landmark_path", type=str, required=True)
    parser_scan.add_argument("--scale_landmarks", type=float, default=1.0,
                             help="Scale (divide) the scan landmarks by this factor.")
    parser_scan.set_defaults(func=fit_body_model_onto_scan)

    parser_dataset = subparsers.add_parser('onto_dataset')
    parser_dataset.add_argument("-D","--dataset_name", type=str, required=True)
    parser_dataset.add_argument("-C", "--continue_run", type=str, default=None,
        help="Path to results folder of YYYY_MM_DD_HH_MM_SS format to continue fitting.")
    parser_dataset.set_defaults(func=fit_body_model_onto_dataset)
    
    args = parser.parse_args()


    # load configs
    cfg = load_config()
    cfg_optimization = cfg["refine_bm_fitting"]
    cfg_datasets = cfg["datasets"]
    cfg_paths = cfg["paths"]
    cfg_general = cfg["general"]
    cfg_web_visualization = cfg["web_visualization"]
    cfg_loss_weights = load_loss_weights_config(
            which_strategy="fit_bm_loss_weight_strategy",
            which_option=cfg_optimization["loss_weight_option"])
    cfg_loss_weights = initialize_fit_bm_loss_weights(cfg_loss_weights)

    # merge configs
    cfg = {}
    cfg.update(cfg_optimization)
    cfg.update(cfg_datasets)
    cfg.update(cfg_paths)
    cfg.update(cfg_general)
    cfg.update(cfg_web_visualization)
    cfg.update(vars(args))
    cfg["loss_weights"] = cfg_loss_weights
    cfg["continue_run"] = cfg["continue_run"] if "continue_run" in cfg.keys() else None

    # process configs
    cfg["save_path"] = create_results_directory(save_path=cfg["save_path"], 
                                                continue_run=cfg["continue_run"])
    cfg = process_default_dtype(cfg)
    cfg = process_visualize_steps(cfg)
    cfg = process_landmarks(cfg)
    cfg = process_body_model_path(cfg)

    # check if landmarks to use are defined
    # assert cfg["use_landmarks"] != [], "Please define landmarks to use in config file!"

    # save configs into results dir
    save_configs(cfg)

    # create web visualization
    if cfg["visualize"]:
        cfg["socket"] = setup_socket(cfg["socket_type"])
        dash_app_process, dash_app_pid = run_dash_app_as_subprocess(cfg['socket_port'])
        print(f"Fitting visualization on http://localhost:{cfg['socket_port']}/")


    # wrapped in a try-except to make sure that the 
    # web visualization socket is closed properly
    try:
        args.func(cfg)
    except (Exception,KeyboardInterrupt) as e:
        print(e)

    if cfg["visualize"]:
        cleanup(cfg["visualize"], cfg["socket"], dash_app_process, dash_app_pid)