Upload hist_cub.py

src/hist_cub.py

@@ -0,0 +1,231 @@
+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()