fourplots.py

from contextlib import contextmanager
from collections.abc import Iterable
import matplotlib.pyplot as plt
import cv2
import numpy as np


@contextmanager
def fourplots(cols=4, figsize=(11, 2.5), dpi=100, set_axis=False, grid_props=None, **kwargs):
    
    if grid_props is True:
        grid_props = dict(color='black', linestyle='--')
        
    print(f"fourplots: canvas={(figsize[0]*dpi,figsize[1]*dpi)}")
    
    fig, axs = plt.subplots(1, cols, figsize=figsize, dpi=dpi, **kwargs)
    if not isinstance(axs, Iterable):
        axs = np.asarray([axs])
    for ax in axs.flatten():
        ax.set_axis_off()
    try:
        yield axs
    finally:

        for ax in axs.flatten():
            if ax.has_data():
                if set_axis:
                    ax.set_axis_on()
                if grid_props is not None:
                    ax.grid(True, **grid_props)
        
        fig.tight_layout()
        plt.show()


def fourimages(images, *, titles=None, overlay_fns=None, **kwargs):
    with fourplots(**kwargs) as axs:

        cols = axs.size
        
        if len(images) < cols:
            images = images + ([None] * (cols - len(images)))
        if titles is None:
            titles = [None] * cols
        if overlay_fns is None:
            overlay_fns = [None] * cols

        for i in range(cols):
            im = images[i]
            if im is None:
                continue
            ax = axs[i]
            title = titles[i]
            fn = overlay_fns[i]
            ax.imshow(im, vmin=0, vmax=255, cmap="gray")
            if title:
                ax.set_title(title)
            if fn:
                fn(ax=ax)

def fourhists(dfs, *, titles=None, **kwargs):
    figsize = kwargs.get("figsize", (11, 2.5))

    kwargs.pop("figsize", None)
    
    with fourplots(figsize=figsize, **kwargs) as axs:

        cols = axs.size
        titles = titles or ([None] * cols)
        
        for i in range(cols):
            df = dfs[i]
            if df is None:
                continue
            title = titles[i] or df.columns[0]
            ax = axs[i]
            df.hist(ax=ax)
            ax.grid(True)
            ax.set_axis_on()
            ax.set_title(title)

    figsize2 = (figsize[0], figsize[1] * 1.75)
    with fourplots(figsize=figsize2, **kwargs) as axs:

        cols = axs.size
        titles = titles or ([None] * cols)
        
        kws = dict(showfliers=1, flierprops=dict(alpha=0.3))
        for i in range(cols):
            df = dfs[i]
            if df is None:
                continue
            title = titles[i] or df.columns[0]
            ax = axs[i]
            df.boxplot(**kws, ax=ax)
            ax.grid(True)
            ax.set_axis_on()
            ax.set_title(title)


def canvas_circles(image, circle_ss, *, colors=None, alpha=0.3, thickness=None):
    if image.ndim < 3:
        canvas = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    else:
        canvas = image.copy()
    if colors is None:
        colors = [(255, 0, 0)] * len(circle_ss)
    thickness = thickness or -(-min(*image.shape[:2])//48) # 48x48 이미지에 대해 1픽셀

    half_thick = -(-thickness // 2)

    overlay = canvas.copy()

    for idx, results in enumerate(circle_ss):
        for result in results:
            if isinstance(result, np.ndarray) and result.shape[-1] == 2:
                cv2.drawContours(overlay, [result], 0, colors[idx], thickness)
                (cx, cy), radius = cv2.minEnclosingCircle(result)
                cx = int(round(cx))
                cy = int(round(cy))
                cv2.circle(overlay, (cx, cy), half_thick, colors[idx], half_thick)
            else:
                for cx, cy, cr in result:
                    cv2.circle(overlay, (cx, cy), cr, colors[idx], thickness)
                    cv2.circle(overlay, (cx, cy), half_thick, colors[idx], thickness)

    canvas = cv2.addWeighted(canvas, 1.0 - alpha, overlay, alpha, 0.0)

    return canvas