Newer
Older
DetectColorChecker24 / colorchecker.py
"""ColorChecker Classic 24 detection and RGB measurement."""

from enum import IntEnum
from os import PathLike
from pathlib import Path

import cv2
import numpy as np


class ResultCode(IntEnum):
    """画像ファイルを処理した結果。整数としても比較できます。"""

    SUCCESS = 0
    IMAGE_READ_ERROR = 1
    DETECTION_FAILED = 2
    PROCESSING_ERROR = 3


def read_image(filename: str | PathLike[str]) -> np.ndarray | None:
    """日本語を含むパスからBGR画像を読み込みます。"""
    data = np.frombuffer(Path(filename).read_bytes(), dtype=np.uint8)
    if data.size == 0:
        return None
    return cv2.imdecode(data, cv2.IMREAD_COLOR)


def detect_colorchecker(
        filename: str | PathLike[str] | np.ndarray
) -> tuple[ResultCode, np.ndarray, dict[str, np.ndarray]]:
    """画像ファイルまたは画像配列から24色票のRGB値を取得します。

    Args:
        filename: ファイルパス、またはuint8型・shape=(高さ, 幅, 3)のBGR画像。
            入力配列は変更しません。引数名は既存のキーワード呼び出しと互換です。

    Returns:
        (結果コード, RGB配列, 処理画像の辞書)。RGB配列は常にshape=(3, 24)、dtype=float64。
        行はR, G, B、列はPATCH_NAMES順(向き補正後の左上から行優先)。
        成功時の値は0~255の平均値、失敗時は全要素NaNです。
        読み込み失敗・不正な入力=1、検出失敗=2、計測などの処理失敗=3。

        処理画像はwarped(向き補正後のBGR画像)、mask(入力座標の二値画像)、
        debug(計測領域と色票番号を描画したBGR画像)。すべてuint8型です。
        失敗時の処理画像の辞書は空です。

    この関数は標準出力やCSV・デバッグ画像の保存を行いません。
    """
    empty = np.full((3, 24), np.nan, dtype=np.float64)
    if isinstance(filename, np.ndarray):
        if (filename.dtype != np.uint8 or filename.ndim != 3
                or filename.shape[2] != 3 or filename.size == 0):
            return ResultCode.IMAGE_READ_ERROR, empty, {}
        image = filename
    elif isinstance(filename, (str, PathLike)):
        try:
            image = read_image(filename)
        except (OSError, ValueError, TypeError, cv2.error):
            return ResultCode.IMAGE_READ_ERROR, empty, {}
        if image is None:
            return ResultCode.IMAGE_READ_ERROR, empty, {}
    else:
        return ResultCode.IMAGE_READ_ERROR, empty, {}

    try:
        warped, _, detection_info = detect_colorchecker_image(image)
    except RuntimeError:
        return ResultCode.DETECTION_FAILED, empty, {}
    except (cv2.error, ValueError, FloatingPointError):
        return ResultCode.PROCESSING_ERROR, empty, {}

    try:
        results, debug = measure_patches(warped)
        rgb = np.array(
            [[patch[channel] for patch in results] for channel in ("R", "G", "B")],
            dtype=np.float64,
        )
        if rgb.shape != (3, 24) or not np.isfinite(rgb).all():
            return ResultCode.PROCESSING_ERROR, empty, {}
    except (cv2.error, ValueError, FloatingPointError):
        return ResultCode.PROCESSING_ERROR, empty, {}

    return ResultCode.SUCCESS, rgb, {
        "warped": warped,
        "mask": detection_info["mask"],
        "debug": debug,
    }


PATCH_NAMES = [
    "Dark Skin", "Light Skin", "Blue Sky",
    "Foliage", "Blue Flower", "Bluish Green",

    "Orange", "Purplish Blue", "Moderate Red",
    "Purple", "Yellow Green", "Orange Yellow",

    "Blue", "Green", "Red",
    "Yellow", "Magenta", "Cyan",

    "White", "Neutral 8", "Neutral 6.5",
    "Neutral 5", "Neutral 3.5", "Black"
]


# ============================================================
# 1次元クラスタリング
# ============================================================

def cluster_1d(values, k):

    data = np.asarray(
        values,
        dtype=np.float32
    ).reshape(-1, 1)

    if len(data) < k:
        raise RuntimeError(
            "候補数が不足しています"
        )

    criteria = (
        cv2.TERM_CRITERIA_EPS
        + cv2.TERM_CRITERIA_MAX_ITER,
        100,
        0.1
    )

    _, labels, centers = cv2.kmeans(
        data,
        k,
        None,
        criteria,
        20,
        cv2.KMEANS_PP_CENTERS
    )

    # 座標の小さい順に
    # 0,1,2,... となるようラベルを並べ直す
    order = np.argsort(centers[:, 0])

    remap = np.empty(k, dtype=np.int32)

    remap[order] = np.arange(k)

    labels = remap[labels.ravel()]

    centers = centers[order, 0]

    return labels, centers


# ============================================================
# 色票候補の検出
# ============================================================

def find_patch_candidates(image, threshold):

    hsv = cv2.cvtColor(
        image,
        cv2.COLOR_BGR2HSV
    )

    value = hsv[:, :, 2]

    mask = np.uint8(
        value > threshold
    ) * 255

    # 小さなノイズ除去
    mask = cv2.morphologyEx(
        mask,
        cv2.MORPH_OPEN,
        np.ones((3, 3), np.uint8)
    )

    # 色票内部の小さな穴を埋める
    mask = cv2.morphologyEx(
        mask,
        cv2.MORPH_CLOSE,
        np.ones((5, 5), np.uint8)
    )

    contours, hierarchy = cv2.findContours(
        mask,
        cv2.RETR_TREE,
        cv2.CHAIN_APPROX_SIMPLE
    )

    height, width = image.shape[:2]

    image_area = height * width

    candidates = []

    for index, contour in enumerate(contours):

        # 明るい背景の内側にある色票も探索する。
        # 奇数階層は暗い領域(穴)の輪郭なので候補に含めない。
        depth = 0
        parent = hierarchy[0, index, 3]
        while parent != -1:
            depth += 1
            parent = hierarchy[0, parent, 3]
        if depth % 2:
            continue

        x, y, w, h = cv2.boundingRect(contour)

        area = cv2.contourArea(contour)

        # 画像全体に対する面積で絞る
        if not (
            0.0015 * image_area
            < area
            < 0.035 * image_area
        ):
            continue

        aspect = w / max(h, 1)

        # 色票はおおむね正方形
        if not (
            0.50
            < aspect
            < 2.0
        ):
            continue

        rectangularity = (
            area / max(w * h, 1)
        )

        # 長方形として十分埋まっているか
        if rectangularity < 0.68:
            continue

        candidates.append({
            "contour": contour,
            "x": x,
            "y": y,
            "w": w,
            "h": h,
            "cx": x + w / 2.0,
            "cy": y + h / 2.0,
            "area": area
        })

    return candidates, mask


# ============================================================
# 6 x 4 格子への割り当て
# ============================================================

def fit_grid(candidates,
             rows=4,
             cols=6):

    # 多少のパッチが検出できなくてもよい
    if len(candidates) < 12:
        return None

    row_labels, row_centers = cluster_1d(
        [p["cy"] for p in candidates],
        rows
    )

    col_labels, col_centers = cluster_1d(
        [p["cx"] for p in candidates],
        cols
    )

    cells = {}

    for i, patch in enumerate(candidates):

        row = int(row_labels[i])
        col = int(col_labels[i])

        key = (row, col)

        # 同一セルに複数候補が入った場合は
        # 大きいものを採用
        if (
            key not in cells
            or patch["area"]
            > cells[key]["area"]
        ):
            cells[key] = patch

    if len(cells) < 12:
        return None

    # --------------------------------------------------------
    # 格子間隔の均一性
    # --------------------------------------------------------

    dx = np.diff(col_centers)
    dy = np.diff(row_centers)

    if (
        np.min(dx) <= 0
        or np.min(dy) <= 0
    ):
        return None

    regularity = (
        np.std(dx) / np.mean(dx)
        +
        np.std(dy) / np.mean(dy)
    )

    # --------------------------------------------------------
    # 色票サイズの均一性
    # --------------------------------------------------------

    areas = np.array(
        [
            p["area"]
            for p in cells.values()
        ],
        dtype=np.float64
    )

    area_cv = (
        np.std(areas)
        /
        np.mean(areas)
    )

    # パッチ数が多いほど高得点
    # 格子が乱れているほど減点
    score = (
        len(cells)
        - 2.0 * regularity
        - area_cv
    )

    return score, cells


# ============================================================
# ColorCheckerを正面向きへ変換
# ============================================================

def warp_from_cells(image,
                    cells,
                    cell_size=140,
                    rows=4,
                    cols=6):

    src_points = []
    dst_points = []

    for (row, col), patch in cells.items():

        # 実画像上の色票中心
        src_points.append([
            patch["cx"],
            patch["cy"]
        ])

        # 正規化画像上の色票中心
        if (rows, cols) == (6, 4):
            # 縦向きの格子は90度回転した座標へ直接射影する。
            dst_points.append([
                (rows - row - 0.5) * cell_size,
                (col + 0.5) * cell_size
            ])
        else:
            dst_points.append([
                (col + 0.5) * cell_size,
                (row + 0.5) * cell_size
            ])

    src_points = np.asarray(
        src_points,
        dtype=np.float32
    )

    dst_points = np.asarray(
        dst_points,
        dtype=np.float32
    )

    # 多数の色票中心を使って
    # RANSACでHomographyを推定
    H, inliers = cv2.findHomography(
        src_points,
        dst_points,
        cv2.RANSAC,
        5.0
    )

    if H is None:
        raise RuntimeError(
            "Homographyを推定できませんでした"
        )

    warped = cv2.warpPerspective(
        image,
        H,
        (
            6 * cell_size,
            4 * cell_size
        )
    )

    return warped, H, inliers


# ============================================================
# ColorCheckerの向きを判定
# ============================================================

def patch_hsv_statistics(
        warped,
        fraction=0.50):

    height, width = warped.shape[:2]

    cell_w = width / 6.0
    cell_h = height / 4.0

    hsv = cv2.cvtColor(
        warped,
        cv2.COLOR_BGR2HSV
    )

    saturation = np.zeros(
        (4, 6),
        dtype=np.float64
    )

    value = np.zeros(
        (4, 6),
        dtype=np.float64
    )

    for row in range(4):

        for col in range(6):

            cx = (
                col + 0.5
            ) * cell_w

            cy = (
                row + 0.5
            ) * cell_h

            half_w = (
                cell_w
                * fraction
                / 2.0
            )

            half_h = (
                cell_h
                * fraction
                / 2.0
            )

            x0 = int(cx - half_w)
            x1 = int(cx + half_w)

            y0 = int(cy - half_h)
            y1 = int(cy + half_h)

            roi = hsv[
                y0:y1,
                x0:x1
            ]

            saturation[
                row,
                col
            ] = roi[:, :, 1].mean()

            value[
                row,
                col
            ] = roi[:, :, 2].mean()

    return saturation, value


def normalize_orientation(warped):

    # 通常画像だけでなく
    # 180度回転や左右反転にも対応
    variants = [
        warped,
        cv2.rotate(
            warped,
            cv2.ROTATE_180
        ),
        cv2.flip(
            warped,
            1
        ),
        cv2.flip(
            warped,
            0
        )
    ]

    best_score = -1e100

    best_image = None

    for img in variants:

        sat, val = (
            patch_hsv_statistics(img)
        )

        # ----------------------------------------------------
        # 最下段は無彩色列のはず
        # ----------------------------------------------------

        row_saturation = (
            sat.mean(axis=1)
        )

        gray_advantage = (
            np.mean(
                row_saturation[:3]
            )
            -
            row_saturation[3]
        )

        # ----------------------------------------------------
        # 白 → 黒
        # 左から右へ暗くなるはず
        # ----------------------------------------------------

        x = np.arange(
            6,
            dtype=np.float64
        )

        slope = np.polyfit(
            x,
            val[3],
            1
        )[0]

        score = gray_advantage

        # 負の傾きなら加点
        if slope < 0:
            score += (
                -slope * 0.6
            )
        else:
            score -= (
                slope * 2.0
            )

        if score > best_score:

            best_score = score

            best_image = img

    return best_image


# ============================================================
# RGB測定
# ============================================================

def measure_patches(
        warped,
        cell_size=140,
        roi_fraction=0.55):

    results = []

    debug = warped.copy()

    for row in range(4):

        for col in range(6):

            cx = (
                col + 0.5
            ) * cell_size

            cy = (
                row + 0.5
            ) * cell_size

            half_w = (
                cell_size
                * roi_fraction
                / 2
            )

            half_h = (
                cell_size
                * roi_fraction
                / 2
            )

            x0 = int(cx - half_w)
            x1 = int(cx + half_w)

            y0 = int(cy - half_h)
            y1 = int(cy + half_h)

            roi = warped[
                y0:y1,
                x0:x1
            ]

            pixels = (
                roi
                .reshape(-1, 3)
                .astype(np.float32)
            )

            # ------------------------------------------------
            # ハイライト・ゴミ等への対策
            #
            # RGB各チャンネルを独立に切ると
            # 色自体が変化する可能性があるため、
            # 輝度で上下5%を除外する
            # ------------------------------------------------

            luminance = (
                0.114 * pixels[:, 0]
                +
                0.587 * pixels[:, 1]
                +
                0.299 * pixels[:, 2]
            )

            low, high = np.percentile(
                luminance,
                [5, 95]
            )

            valid = (
                (luminance >= low)
                &
                (luminance <= high)
            )

            filtered = pixels[valid]

            if len(filtered) < 10:
                filtered = pixels

            mean_bgr = (
                filtered.mean(axis=0)
            )

            B, G, R = mean_bgr

            index = (
                row * 6
                + col
            )

            results.append({
                "index": index + 1,
                "name": PATCH_NAMES[index],
                "R": float(R),
                "G": float(G),
                "B": float(B)
            })

            # デバッグ描画
            cv2.rectangle(
                debug,
                (x0, y0),
                (x1, y1),
                (0, 255, 0),
                2
            )

            cv2.putText(
                debug,
                str(index + 1),
                (x0 + 5, y0 + 22),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (0, 255, 0),
                2,
                cv2.LINE_AA
            )

    return results, debug


# ============================================================
# ColorChecker検出
# ============================================================

def detect_colorchecker_image(
        image,
        cell_size=140):

    hsv = cv2.cvtColor(
        image,
        cv2.COLOR_BGR2HSV
    )

    value = hsv[:, :, 2]

    # --------------------------------------------------------
    # Otsu法から基準閾値を求める
    # --------------------------------------------------------

    otsu_threshold, _ = cv2.threshold(
        value,
        0,
        255,
        cv2.THRESH_BINARY
        + cv2.THRESH_OTSU
    )

    # --------------------------------------------------------
    # 複数の閾値を試す
    #
    # 黒票・濃い灰色票は検出されなくてもよい
    # --------------------------------------------------------

    thresholds = sorted(
        set(
            int(
                np.clip(
                    otsu_threshold * scale,
                    25,
                    180
                )
            )
            for scale in [
                0.45,
                0.55,
                0.65,
                0.80,
                1.00,
                1.15
            ]
        )
    )

    best = None

    for threshold in thresholds:

        candidates, mask = (
            find_patch_candidates(
                image,
                threshold
            )
        )

        for rows, cols in ((4, 6), (6, 4)):
            grid = fit_grid(candidates, rows=rows, cols=cols)
            if grid is None:
                continue

            score, cells = grid
            if best is None or score > best["score"]:
                best = {
                    "score": score,
                    "threshold": threshold,
                    "candidates": candidates,
                    "cells": cells,
                    "mask": mask,
                    "rows": rows,
                    "cols": cols,
                }

    if best is None:

        raise RuntimeError(
            "ColorCheckerの"
            "6x4色票格子を"
            "検出できませんでした"
        )

    warped, H, inliers = (
        warp_from_cells(
            image,
            best["cells"],
            cell_size,
            rows=best["rows"],
            cols=best["cols"],
        )
    )

    warped = normalize_orientation(
        warped
    )

    return (
        warped,
        H,
        best
    )