src/hist_cub.py

import itertools
import functools
import math
import multiprocessing
from pathlib import Path

import matplotlib
matplotlib.rcParams.update({'font.size': 24})
matplotlib.rcParams.update({
  "text.usetex": True,
  "text.latex.preamble": r"\usepackage{biolinum} \usepackage{libertineRoman} \usepackage{libertineMono} \usepackage{biolinum} \usepackage[libertine]{newtxmath}",
  'ps.usedistiller': "xpdf",
})

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import tqdm
from scipy.stats import wasserstein_distance

import pose_estimation


def cub(x, a, b, c):
    x2 = x * x
    x3 = x2 * x

    y = a * x3 + b * x2 + c * x

    return y


def subsample(a, p=0.0005, seed=0):
    np.random.seed(seed)
    N = len(a)
    inds = np.random.choice(range(N), size=int(p * N))
    a = a[inds].copy()

    return a


def read_cos_opt(path, fname="cos_hist.npy"):
    cos_opt = []
    for p in Path(path).rglob(fname):
        d = np.load(p)
        cos_opt.append(d)

    cos_opt = np.array(cos_opt)

    return cos_opt


def plot_hist(cos_opt_dir, hist_smpl_fpath, params, out_dir, bins=10, xy=None):
    cos_opt = read_cos_opt(cos_opt_dir)
    angle_opt = np.arccos(cos_opt)
    angle_opt2 = cub(angle_opt, *params)

    cos_opt2 = np.cos(angle_opt2)
    cos_smpl = np.load(hist_smpl_fpath)
    # cos_smpl = subsample(cos_smpl)
    print(cos_smpl.shape)

    cos_smpl = np.clip(cos_smpl, -1, 1)

    cos_opt = angle_opt
    cos_opt2 = angle_opt2
    cos_smpl = np.arccos(cos_smpl)

    cos_opt = 180 / math.pi * cos_opt
    cos_opt2 = 180 / math.pi * cos_opt2
    cos_smpl = 180 / math.pi * cos_smpl
    max_range = 90  # math.pi / 2

    xticks = [0, 15, 30, 45, 60, 75, 90]
    for idx, bone in enumerate(pose_estimation.SKELETON):
        i, j = bone
        i_name = pose_estimation.KPS[i]
        j_name = pose_estimation.KPS[j]
        if i_name != "Left Upper Leg":
            continue

        name = f"{i_name}_{j_name}"

        gs = gridspec.GridSpec(2, 4)
        fig = plt.figure(tight_layout=True, figsize=(16, 8), dpi=300)

        ax0 = fig.add_subplot(gs[0, 0])
        ax0.hist(cos_smpl[:, idx], bins=bins, range=(0, max_range), density=True)
        ax0.set_xticks(xticks)
        ax0.tick_params(labelbottom=False, labelleft=True)

        ax1 = fig.add_subplot(gs[1, 0], sharex=ax0)
        ax1.hist(cos_opt[:, idx], bins=bins, range=(0, max_range), density=True)
        ax1.set_xticks(xticks)

        if xy is not None:
            ax2 = fig.add_subplot(gs[:, 1:3])
            ax2.plot(xy[0], xy[1], linewidth=8)
            ax2.plot(xy[0], xy[0], linewidth=4, linestyle="dashed")
            ax2.set_xticks(xticks)
            ax2.set_yticks(xticks)

        ax3 = fig.add_subplot(gs[0, 3], sharey=ax0)
        ax3.hist(cos_opt2[:, idx], bins=bins, range=(0, max_range), density=True)
        ax3.set_xticks(xticks)
        ax3.tick_params(labelbottom=False, labelleft=False)

        ax4 = fig.add_subplot(gs[1, 3], sharex=ax3, sharey=ax1)
        alpha = 0.5
        ax4.hist(cos_opt[:, idx], bins=bins, range=(0, max_range), density=True, label=r"$\mathcal{B}_i$", alpha=alpha)
        ax4.hist(cos_opt2[:, idx], bins=bins, range=(0, max_range), density=True, label=r"$f(\mathcal{B}_i)$", alpha=alpha)
        ax4.hist(cos_smpl[:, idx], bins=bins, range=(0, max_range), density=True, label=r"$\mathcal{A}_i$", alpha=alpha)
        ax4.set_xticks(xticks)
        ax4.tick_params(labelbottom=True, labelleft=False)
        ax4.legend()

        fig.savefig(out_dir / f"hist_{name}.png")
        plt.close()


def kldiv(p_hist, q_hist):
    wd = wasserstein_distance(p_hist, q_hist)

    return wd


def calc_histogram(x, bins=10, range=(0, 1)):
    h, _ = np.histogram(x, bins=bins, range=range, density=True)

    return h

def step(params, angles_opt, p_hist, bone_idx=None):
    if sum(params) > 1:
        return math.inf, params

    kl = 0
    for i, _ in enumerate(pose_estimation.SKELETON):
        if bone_idx is not None and i != bone_idx:
            continue

        angles_opt2 = cub(angles_opt[:, i], *params)
        if angles_opt2.max() > 1 or angles_opt2.min() < 0:
            kl = math.inf

            break

        q_hist = calc_histogram(angles_opt2)

        kl += kldiv(p_hist[i], q_hist)

    return kl, params


def optimize(cos_opt_dir, hist_smpl_fpath, bone_idx=None):
    cos_opt = read_cos_opt(cos_opt_dir)
    angles_opt = np.arccos(cos_opt) / (math.pi / 2)
    cos_smpl = np.load(hist_smpl_fpath)
    # cos_smpl = subsample(cos_smpl)
    print(cos_smpl.shape)
    cos_smpl = np.clip(cos_smpl, -1, 1)
    mask = cos_smpl <= 1
    assert np.all(mask), (~mask).mean()
    mask = cos_smpl >= 0
    assert np.all(mask), (~mask).mean()
    angles_smpl = np.arccos(cos_smpl) / (math.pi / 2)
    p_hist = [
        calc_histogram(angles_smpl[:, i])
        for i, _ in enumerate(pose_estimation.SKELETON)
    ]

    with multiprocessing.Pool(8) as p:
        results = list(
            tqdm.tqdm(
                p.imap_unordered(
                    functools.partial(step, angles_opt=angles_opt, p_hist=p_hist, bone_idx=bone_idx),
                    itertools.product(
                        np.linspace(0, 20, 100),
                        np.linspace(-20, 20, 200),
                        np.linspace(-20, 1, 100),
                    ),
                ),
                total=(100 * 200 * 100),
            )
        )

    kls, params = zip(*results)
    ind = np.argmin(kls)
    best_params = params[ind]

    print(kls[ind], best_params)

    inds = np.argsort(kls)
    for i in inds[:10]:
        print(kls[i])
        print(params[i])
        print()

    return best_params


def main():
    cos_opt_dir = "paper_single2_150mse"
    hist_smpl_fpath = "./data/hist_smpl.npy"
    # hist_smpl_fpath = "./testtest.npy"
    params = optimize(cos_opt_dir, hist_smpl_fpath)
    # params = (1.2121212121212122, -1.105527638190953, 0.787878787878789)
    # params = (0.20202020202020202, 0.30150753768844396, 0.3636363636363633)
    print(params)

    x = np.linspace(0, math.pi / 2, 100)
    y = cub(x / (math.pi / 2), *params) * (math.pi / 2)
    x = x * 180 / math.pi
    y = y * 180 / math.pi

    out_dir = Path("hists")
    out_dir.mkdir(parents=True, exist_ok=True)
    plot_hist(cos_opt_dir, hist_smpl_fpath, params, out_dir, xy=(x, y))

    plt.figure(figsize=(4, 4), dpi=300)
    plt.plot(x, y, linewidth=6)
    plt.plot(x, x, linewidth=2, linestyle="dashed")
    xticks = [0, 15, 30, 45, 60, 75, 90]
    plt.xticks(xticks)
    plt.yticks(xticks)
    plt.axis("equal")
    plt.tight_layout()
    plt.savefig(out_dir / "new_out.png")


if __name__ == "__main__":
    main()