ControlNet-v1-1-nightly-main/annotator/normalbae/utils/utils.py

import os
import numpy as np
from PIL import Image

import torch
import torch.nn as nn
import torch.nn.functional as F

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt



# convert arg line to args
def convert_arg_line_to_args(arg_line):
    for arg in arg_line.split():
        if not arg.strip():
            continue
        yield str(arg)


# save args
def save_args(args, filename):
    with open(filename, 'w') as f:
        for arg in vars(args):
            f.write('{}: {}\n'.format(arg, getattr(args, arg)))


# concatenate images
def concat_image(image_path_list, concat_image_path):
    imgs = [Image.open(i).convert("RGB").resize((640, 480), resample=Image.BILINEAR) for i in image_path_list]
    imgs_list = []
    for i in range(len(imgs)):
        img = imgs[i]
        imgs_list.append(np.asarray(img))

        H, W, _ = np.asarray(img).shape
        imgs_list.append(255 * np.ones((H, 20, 3)).astype('uint8'))

    imgs_comb = np.hstack(imgs_list[:-1])
    imgs_comb = Image.fromarray(imgs_comb)
    imgs_comb.save(concat_image_path)


# load model
def load_checkpoint(fpath, model):
    ckpt = torch.load(fpath, map_location='cpu')['model']

    load_dict = {}
    for k, v in ckpt.items():
        if k.startswith('module.'):
            k_ = k.replace('module.', '')
            load_dict[k_] = v
        else:
            load_dict[k] = v

    model.load_state_dict(load_dict)
    return model


# compute normal errors
def compute_normal_errors(total_normal_errors):
    metrics = {
        'mean': np.average(total_normal_errors),
        'median': np.median(total_normal_errors),
        'rmse': np.sqrt(np.sum(total_normal_errors * total_normal_errors) / total_normal_errors.shape),
        'a1': 100.0 * (np.sum(total_normal_errors < 5) / total_normal_errors.shape[0]),
        'a2': 100.0 * (np.sum(total_normal_errors < 7.5) / total_normal_errors.shape[0]),
        'a3': 100.0 * (np.sum(total_normal_errors < 11.25) / total_normal_errors.shape[0]),
        'a4': 100.0 * (np.sum(total_normal_errors < 22.5) / total_normal_errors.shape[0]),
        'a5': 100.0 * (np.sum(total_normal_errors < 30) / total_normal_errors.shape[0])
    }
    return metrics


# log normal errors
def log_normal_errors(metrics, where_to_write, first_line):
    print(first_line)
    print("mean median rmse 5 7.5 11.25 22.5 30")
    print("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f" % (
        metrics['mean'], metrics['median'], metrics['rmse'],
        metrics['a1'], metrics['a2'], metrics['a3'], metrics['a4'], metrics['a5']))

    with open(where_to_write, 'a') as f:
        f.write('%s\n' % first_line)
        f.write("mean median rmse 5 7.5 11.25 22.5 30\n")
        f.write("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n\n" % (
            metrics['mean'], metrics['median'], metrics['rmse'],
            metrics['a1'], metrics['a2'], metrics['a3'], metrics['a4'], metrics['a5']))


# makedir
def makedir(dirpath):
    if not os.path.exists(dirpath):
        os.makedirs(dirpath)


# makedir from list
def make_dir_from_list(dirpath_list):
    for dirpath in dirpath_list:
        makedir(dirpath)



########################################################################################################################
# Visualization
########################################################################################################################


# unnormalize image
__imagenet_stats = {'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}
def unnormalize(img_in):
    img_out = np.zeros(img_in.shape)
    for ich in range(3):
        img_out[:, :, ich] = img_in[:, :, ich] * __imagenet_stats['std'][ich]
        img_out[:, :, ich] += __imagenet_stats['mean'][ich]
    img_out = (img_out * 255).astype(np.uint8)
    return img_out


# kappa to exp error (only applicable to AngMF distribution)
def kappa_to_alpha(pred_kappa):
    alpha = ((2 * pred_kappa) / ((pred_kappa ** 2.0) + 1)) \
            + ((np.exp(- pred_kappa * np.pi) * np.pi) / (1 + np.exp(- pred_kappa * np.pi)))
    alpha = np.degrees(alpha)
    return alpha


# normal vector to rgb values
def norm_to_rgb(norm):
    # norm: (B, H, W, 3)
    norm_rgb = ((norm[0, ...] + 1) * 0.5) * 255
    norm_rgb = np.clip(norm_rgb, a_min=0, a_max=255)
    norm_rgb = norm_rgb.astype(np.uint8)
    return norm_rgb


# visualize during training
def visualize(args, img, gt_norm, gt_norm_mask, norm_out_list, total_iter):
    B, _, H, W = gt_norm.shape

    pred_norm_list = []
    pred_kappa_list = []
    for norm_out in norm_out_list:
        norm_out = F.interpolate(norm_out, size=[gt_norm.size(2), gt_norm.size(3)], mode='nearest')
        pred_norm = norm_out[:, :3, :, :]  # (B, 3, H, W)
        pred_norm = pred_norm.detach().cpu().permute(0, 2, 3, 1).numpy()  # (B, H, W, 3)
        pred_norm_list.append(pred_norm)

        pred_kappa = norm_out[:, 3:, :, :]  # (B, 1, H, W)
        pred_kappa = pred_kappa.detach().cpu().permute(0, 2, 3, 1).numpy()  # (B, H, W, 1)
        pred_kappa_list.append(pred_kappa)

    # to numpy arrays
    img = img.detach().cpu().permute(0, 2, 3, 1).numpy()                     # (B, H, W, 3)
    gt_norm = gt_norm.detach().cpu().permute(0, 2, 3, 1).numpy()             # (B, H, W, 3)
    gt_norm_mask = gt_norm_mask.detach().cpu().permute(0, 2, 3, 1).numpy()   # (B, H, W, 1)

    # input image
    target_path = '%s/%08d_img.jpg' % (args.exp_vis_dir, total_iter)
    img = unnormalize(img[0, ...])
    plt.imsave(target_path, img)

    # gt norm
    gt_norm_rgb = ((gt_norm[0, ...] + 1) * 0.5) * 255
    gt_norm_rgb = np.clip(gt_norm_rgb, a_min=0, a_max=255)
    gt_norm_rgb = gt_norm_rgb.astype(np.uint8)

    target_path = '%s/%08d_gt_norm.jpg' % (args.exp_vis_dir, total_iter)
    plt.imsave(target_path, gt_norm_rgb * gt_norm_mask[0, ...])

    # pred_norm
    for i in range(len(pred_norm_list)):
        pred_norm = pred_norm_list[i]
        pred_norm_rgb = norm_to_rgb(pred_norm)
        target_path = '%s/%08d_pred_norm_%d.jpg' % (args.exp_vis_dir, total_iter, i)
        plt.imsave(target_path, pred_norm_rgb)

        pred_kappa = pred_kappa_list[i]
        pred_alpha = kappa_to_alpha(pred_kappa)
        target_path = '%s/%08d_pred_alpha_%d.jpg' % (args.exp_vis_dir, total_iter, i)
        plt.imsave(target_path, pred_alpha[0, :, :, 0], vmin=0, vmax=60, cmap='jet')

        # error in angles
        DP = np.sum(gt_norm * pred_norm, axis=3, keepdims=True)  # (B, H, W, 1)
        DP = np.clip(DP, -1, 1)
        E = np.degrees(np.arccos(DP))  # (B, H, W, 1)
        E = E * gt_norm_mask
        target_path = '%s/%08d_pred_error_%d.jpg' % (args.exp_vis_dir, total_iter, i)
        plt.imsave(target_path, E[0, :, :, 0], vmin=0, vmax=60, cmap='jet')