Maps an image to grayscale where pixel brightness is determined by how close each pixel's color is to a chosen target RGB color.

color_to_grayscale.py

"""
color_to_grayscale.py
---------------------
Maps an image to grayscale where pixel brightness is determined by how
close each pixel's color is to a chosen target RGB color.

Uses HSV-aware distance so that white, grey, and saturated colors are
correctly distinguished — e.g. targeting red will NOT make white bright,
even though white contains high R, G, and B values.

- Pixels that MATCH the target color  → WHITE  (255)
- Pixels that are FAR from the target → BLACK  (0)

Key fix vs plain RGB distance
------------------------------
RGB space treats white (255,255,255) as close to red (255,0,0), green (0,255,0)
and blue (0,0,255) equally — all score a moderate distance.
HSV space separates HUE (color identity) from SATURATION (colorfulness) and
VALUE (brightness).  White has saturation ≈ 0, so its hue is meaningless.
By weighting the hue term by min(pixel_saturation, target_saturation), hue
comparison is suppressed whenever either side is achromatic, eliminating
false matches between vivid colours and white/grey/black.

Threshold
---------
--threshold controls how strictly a pixel must match the target color.
  1.0 (default) = soft, full gradient across all distances
  0.5           = only pixels within 50% distance are lit; rest → black
  0.1           = very strict; only near-exact matches remain bright
Pixels inside the threshold are rescaled to fill the full 0–255 range,
preserving smooth gradients within the accepted band.

Usage:
    python color_to_grayscale.py --image input.jpg --color 255 0 0
    python color_to_grayscale.py --image input.jpg --color 255 0 0 --threshold 0.3
    python color_to_grayscale.py --image input.jpg --color 255 255 255 --threshold 0.2
    python color_to_grayscale.py --image input.jpg --color 0 128 255 --mode inverse --threshold 0.4
"""

import cv2
import numpy as np
import argparse
import sys
from pathlib import Path
import matplotlib.pyplot as plt


# ──────────────────────────────────────────────
#  HSV-aware color distance
# ──────────────────────────────────────────────

def hsv_distance(img_bgr: np.ndarray, target_rgb: tuple) -> np.ndarray:
    """
    Compute a perceptually meaningful per-pixel distance to a target color
    using HSV space with saturation-weighted hue comparison.

    Returns
    -------
    np.ndarray shape (H, W), float32, values in [0, 1].
    0 = identical to target, 1 = maximally far.
    """
    # ── Image: BGR → float32 HSV, normalised to [0,1] range ─────────────────
    img_hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
    img_hsv[:, :, 0] /= 180.0   # OpenCV H: 0–180 → 0–1
    img_hsv[:, :, 1] /= 255.0   # S: 0–255 → 0–1
    img_hsv[:, :, 2] /= 255.0   # V: 0–255 → 0–1

    # ── Target: RGB → float32 HSV, normalised ───────────────────────────────
    t_bgr_px = np.uint8([[[target_rgb[2], target_rgb[1], target_rgb[0]]]])
    t_hsv    = cv2.cvtColor(t_bgr_px, cv2.COLOR_BGR2HSV).astype(np.float32)
    t_h = float(t_hsv[0, 0, 0]) / 180.0
    t_s = float(t_hsv[0, 0, 1]) / 255.0
    t_v = float(t_hsv[0, 0, 2]) / 255.0

    p_h = img_hsv[:, :, 0]
    p_s = img_hsv[:, :, 1]
    p_v = img_hsv[:, :, 2]

    # ── Circular hue distance, range [0, 0.5] ───────────────────────────────
    hue_diff = np.abs(p_h - t_h)
    hue_dist = np.minimum(hue_diff, 1.0 - hue_diff)  # wrap-around

    # ── Saturation-based hue weight ──────────────────────────────────────────
    # Weight = 0 when either colour is achromatic → hue term vanishes.
    # This is the core fix: white/grey have s≈0, so their hue is ignored.
    hue_weight = np.minimum(p_s, t_s)

    # ── Weighted combination ─────────────────────────────────────────────────
    d_hue = hue_weight * (hue_dist * 2.0)   # [0, 1] after weighting
    d_sat = np.abs(p_s - t_s)               # [0, 1]
    d_val = np.abs(p_v - t_v)               # [0, 1]

    dist = 0.60 * d_hue + 0.25 * d_sat + 0.15 * d_val

    return np.clip(dist, 0.0, 1.0).astype(np.float32)


# ──────────────────────────────────────────────
#  Core function
# ──────────────────────────────────────────────

def map_by_color(
    image_bgr: np.ndarray,
    target_rgb: tuple,
    mode: str = "normal",
    gamma: float = 1.0,
    threshold: float = 1.0,
) -> np.ndarray:
    """
    Convert a BGR image to grayscale based on proximity to a target RGB color.

    Parameters
    ----------
    image_bgr   : np.ndarray  — Input image in BGR format (as loaded by cv2).
    target_rgb  : tuple       — Target color as (R, G, B), each in [0, 255].
    mode        : str         — 'normal'  → close to color = bright
                                'inverse' → close to color = dark
    gamma       : float       — Gamma correction on output (1.0 = none).
    threshold   : float       — Strictness of the match in [0.0, 1.0].
                                1.0 = accept all distances (soft, full gradient).
                                0.0 = accept only a perfect exact match.
                                Pixels with distance > threshold → black (or white
                                in inverse mode). Pixels within threshold are
                                rescaled to fill the full 0–255 brightness range,
                                so contrast is preserved regardless of strictness.

    Returns
    -------
    np.ndarray — Single-channel grayscale image (uint8).
    """
    dist = hsv_distance(image_bgr, target_rgb)   # 0 = match, 1 = far

    # ── Apply threshold ───────────────────────────────────────────────────────
    # Pixels beyond the threshold are clamped to dist=1 (fully unmatched).
    # Pixels within the threshold are rescaled so the accepted band [0, threshold]
    # maps to [0, 1] — preserving a smooth gradient inside the match region.
    if threshold < 1.0:
        within  = dist <= threshold
        rescaled = np.where(within, dist / max(threshold, 1e-6), 1.0)
        dist = rescaled.astype(np.float32)

    brightness = (1.0 - dist) if mode == "normal" else dist

    if gamma != 1.0:
        brightness = np.power(np.clip(brightness, 0.0, 1.0), 1.0 / gamma)

    return (np.clip(brightness, 0.0, 1.0) * 255).astype(np.uint8)


# ──────────────────────────────────────────────
#  Optional: side-by-side comparison
# ──────────────────────────────────────────────

def make_comparison(
    original_bgr: np.ndarray,
    gray: np.ndarray,
    label_left: str = "Original",
    label_right: str = "Processed",
) -> np.ndarray:
    """
    Stack original and result side-by-side with a 2-px divider and text labels.
    """
    gray_bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    h = max(original_bgr.shape[0], gray_bgr.shape[0])

    def pad(img, target_h):
        ph = target_h - img.shape[0]
        return np.pad(img, ((0, ph), (0, 0), (0, 0)), constant_values=0) if ph > 0 else img

    left  = pad(original_bgr, h).copy()
    right = pad(gray_bgr, h).copy()

    # ── Text label helper ────────────────────────────────────────────────────
    def draw_label(img: np.ndarray, text: str) -> None:
        font       = cv2.FONT_HERSHEY_SIMPLEX
        scale      = max(0.5, img.shape[1] / 1000)
        thickness  = max(1, int(scale * 2))
        (tw, th), baseline = cv2.getTextSize(text, font, scale, thickness)
        x, y = 12, th + 12
        # Shadow for readability on any background
        cv2.putText(img, text, (x + 1, y + 1), font, scale, (0,   0,   0  ), thickness + 1, cv2.LINE_AA)
        cv2.putText(img, text, (x,     y    ), font, scale, (255, 255, 255), thickness,     cv2.LINE_AA)

    draw_label(left,  label_left)
    draw_label(right, label_right)

    # ── 2-pixel white divider ────────────────────────────────────────────────
    divider = np.full((h, 2, 3), 200, dtype=np.uint8)

    return np.hstack([left, divider, right])


# ──────────────────────────────────────────────
#  CLI
# ──────────────────────────────────────────────

def parse_args():
    parser = argparse.ArgumentParser(
        description="Map image to grayscale based on closeness to a chosen RGB color (HSV-aware).",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
threshold guide:
  1.0   soft — full gradient, every pixel contributes (default)
  0.5   medium — only pixels within 50%% distance are lit
  0.3   strict — noticeable falloff, clear isolation of target color
  0.1   very strict — near-exact matches only; most pixels go black
        """,
    )
    parser.add_argument("--image",     required=True,          help="Path to input image")
    parser.add_argument("--color",     required=True, nargs=3, type=int,
                        metavar=("R", "G", "B"),               help="Target RGB color, e.g. 255 0 0")
    parser.add_argument("--threshold", default=1.0, type=float,
                        help="Match strictness [0.0–1.0]. Lower = stricter. (default: 1.0)")
    parser.add_argument("--output",    default="",             help="Output file path (default: auto-named)")
    parser.add_argument("--mode",      default="normal",       choices=["normal", "inverse"],
                        help="normal=close→bright | inverse=close→dark  (default: normal)")
    parser.add_argument("--gamma",     default=1.0, type=float,
                        help="Gamma for contrast shaping (default: 1.0)")
    parser.add_argument("--compare",   action="store_true",
                        help="Save a side-by-side comparison image to disk")
    parser.add_argument("--preview",   action="store_true",
                        help="Show side-by-side original vs processed in a window (no file saved)")
    parser.add_argument("--show",      action="store_true",
                        help="Display result in a window (requires a GUI environment)")
    return parser.parse_args()


def main():
    args = parse_args()

    for ch, val in zip("RGB", args.color):
        if not (0 <= val <= 255):
            sys.exit(f"Error: {ch} value {val} is out of range [0, 255].")

    if not (0.0 <= args.threshold <= 1.0):
        sys.exit(f"Error: --threshold must be between 0.0 and 1.0, got {args.threshold}.")

    img_path = Path(args.image)
    if not img_path.exists():
        sys.exit(f"Error: file not found → {img_path}")

    image_bgr = cv2.imread(str(img_path))
    if image_bgr is None:
        sys.exit(f"Error: cv2 could not read the image at {img_path}")

    print(f"Loaded    : {img_path}  ({image_bgr.shape[1]}×{image_bgr.shape[0]} px)")
    print(f"Target    : RGB{tuple(args.color)}")
    print(f"Threshold : {args.threshold}  |  Mode: {args.mode}  |  Gamma: {args.gamma}")

    gray = map_by_color(
        image_bgr,
        target_rgb=tuple(args.color),
        mode=args.mode,
        gamma=args.gamma,
        threshold=args.threshold,
    )

    r, g, b = args.color
    out_path = Path(args.output) if args.output else \
        img_path.parent / f"{img_path.stem}_gray_r{r}g{g}b{b}_t{args.threshold}{img_path.suffix}"

    # ── Preview mode: show side-by-side with matplotlib, skip saving ─────────
    if args.preview:
        r, g, b = args.color
        # cv2 loads BGR → convert to RGB for matplotlib
        original_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

        # 3 columns: original | processed | color swatch
        # The swatch column is kept narrow via gridspec width_ratios
        fig, axes = plt.subplots(
            1, 3, figsize=(13, 5),
            gridspec_kw={"width_ratios": [10, 10, 1.4]}
        )
        fig.suptitle(
            f"Target: RGB({r}, {g}, {b})   threshold={args.threshold}   mode={args.mode}",
            fontsize=12, fontweight="bold"
        )

        axes[0].imshow(original_rgb)
        axes[0].set_title("Original")
        axes[0].axis("off")

        axes[1].imshow(gray, cmap="gray", vmin=0, vmax=255)
        axes[1].set_title(f"Processed  (t={args.threshold})")
        axes[1].axis("off")

        # ── Color swatch ─────────────────────────────────────────────────────
        # Fill the third axis with a solid rectangle of the target color
        swatch = axes[2]
        swatch.set_facecolor([r / 255, g / 255, b / 255])
        swatch.set_title("Target\ncolor", fontsize=9)
        swatch.set_xticks([])
        swatch.set_yticks([])
        # Label the hex code below the swatch
        fig.text(
            (axes[2].get_position().x0 + axes[2].get_position().x1) / 2,
            axes[2].get_position().y0 - 0.04,
            f"#{r:02X}{g:02X}{b:02X}\n({r}, {g}, {b})",
            ha="center", va="top", fontsize=8, family="monospace"
        )

        plt.tight_layout()
        plt.show()
        return   # skip saving

    # ── Save output ──────────────────────────────────────────────────────────
    cv2.imwrite(str(out_path), gray)
    print(f"Saved     : {out_path}")

    if args.compare:
        comp = make_comparison(image_bgr, gray)
        comp_path = out_path.parent / f"{out_path.stem}_comparison{out_path.suffix}"
        cv2.imwrite(str(comp_path), comp)
        print(f"Compare   : {comp_path}")

    if args.show:
        cv2.imshow("Original", image_bgr)
        cv2.imshow("Color-mapped Grayscale", gray)
        print("Press any key to close windows.")
        cv2.waitKey(0)
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()