Python script to plot a 3-dimensional HSV histogram from a color image.

hsv_histogram.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-


"""
HSV image histogram script.
"""


import os
#
from omegaconf import OmegaConf, MISSING
import numpy as np
import cv2
from matplotlib.colors import hsv_to_rgb
import matplotlib.pyplot as plt


# ##############################################################################
# # HISTOGRAM FUNCTION
# ##############################################################################
def hist3d_hsv(rgb_img, bins=20,
               scale_dots=1.0,
               figsize=(10, 5),
               range=((0, 255), (0, 255), (0, 255)),
               bg_color=(0, 0.2, 0)):
    """
    rgb_img must be in range 0, 255!
    HSV will be in range 0..179, 0..255, 0..255
    https://docs.opencv.org/3.0-beta/doc/py_tutorials/py_imgproc/
    py_colorspaces/py_colorspaces.html
    """

    h, w, ch = rgb_img.shape
    assert ch == 3, "Only 3-channel images allowed!"
    hsv_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2HSV)
    hist, edges = np.histogramdd(hsv_img.reshape(h * w, ch),
                                 range=range,
                                 bins=bins)
    hidxs = hist.nonzero()
    n_hidxs = hidxs[0].shape[0]
    sizevals = np.zeros(n_hidxs, dtype=np.float64)
    hsv_vals = np.zeros((n_hidxs, ch), dtype=np.float64)
    for i, hid in enumerate(zip(*hidxs)):
        sizevals[i] = hist[hid]
        hsv_vals[i] = [edg[j] for j, edg in zip(hid, edges)]
    # scale the dots
    try:
        sizevals /= sizevals.max()
    except ZeroDivisionError:
        raise RuntimeError("This should never happen")
    #
    # max dot diameter to avoid collision is this
    binspreads = [(float(rmax)-rmin) / bins for rmin, rmax in range]
    # max dot surface to avoid collision is this (pi*r^2)
    max_surf = 3.141592 * (min(binspreads) / 2)**2
    sizevals *= (max_surf * scale_dots)
    #
    rgb_vals = hsv_vals.copy()
    rgb_vals[:, 0] /= 179
    rgb_vals[:, 1] /= 255
    rgb_vals[:, 2] /= 255
    rgb_vals = hsv_to_rgb(rgb_vals)
    #
    fig = plt.figure(figsize=figsize)
    ax_im = fig.add_subplot(1, 2, 1)
    ax_hist = fig.add_subplot(1, 2, 2, projection='3d')
    #
    ax_im.imshow(rgb_img)
    #
    ax_hist.set_xlabel("Hue")
    ax_hist.set_ylabel("Saturation")
    ax_hist.set_zlabel("Value")
    ax_hist.set_facecolor(bg_color)
    ax_hist.scatter(*hsv_vals.T,
                    depthshade=0,
                    c=rgb_vals,  # !!!
                    s=sizevals)
    fig.tight_layout()
    return fig


# ##############################################################################
# # GLOBALS
# ##############################################################################
CONF = OmegaConf.create()
CONF.IMG_PATH = MISSING
CONF.BINS = 20
CONF.SCALE_DOTS = 7.0
CONF.DPI = 200
CONF.SUFFIX_OUT = "_plot.png"
CONF.FIGSIZE = (11, 5)  # h, w
CONF.BG_RGB = (0.8, 1, 0.8)
CONF.PLOT = False
#
cli_conf = OmegaConf.from_cli()
CONF = OmegaConf.merge(CONF, cli_conf)
print(OmegaConf.to_yaml(CONF))


# ##############################################################################
# # MAIN ROUTINE
# ##############################################################################
img_rgb = cv2.cvtColor(cv2.imread(CONF.IMG_PATH), cv2.COLOR_BGR2RGB)
fig = hist3d_hsv(img_rgb,
                 bins=CONF.BINS,
                 scale_dots=CONF.SCALE_DOTS,
                 figsize=CONF.FIGSIZE,
                 bg_color=CONF.BG_RGB)

if CONF.PLOT:
    fig.show()
    breakpoint()
else:
    out_path = os.path.splitext(CONF.IMG_PATH)[0] + CONF.SUFFIX_OUT
    fig.savefig(out_path, dpi=CONF.DPI)