main.py

import cv2 as cv
import numpy as np
import argparse, os, json
from collections import Counter

# ---------------- Conversions ----------------
def bgr_to_rgb(bgr): return cv.cvtColor(bgr, cv.COLOR_BGR2RGB)
def bgr_to_hsv(bgr): return cv.cvtColor(bgr, cv.COLOR_BGR2HSV)
def bgr_to_lab(bgr): return cv.cvtColor(bgr, cv.COLOR_BGR2LAB)

# ---------------- Stats ----------------
def img_stats(img, space_name):
    # img is uint8, shape HxWxC
    means = img.reshape(-1, img.shape[2]).mean(axis=0)
    stds  = img.reshape(-1, img.shape[2]).std(axis=0)
    return {
        "space": space_name,
        "mean": [float(x) for x in means],
        "std":  [float(x) for x in stds]
    }

# ---------------- Dominant colors ----------------
def dominant_colors_kmeans(bgr, k=3, max_iter=10):
    data = bgr.reshape((-1, 3)).astype(np.float32)
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, max_iter, 1.0)
    flags = cv.KMEANS_PP_CENTERS
    compactness, labels, centers = cv.kmeans(data, k, None, criteria, 3, flags)
    centers_u8 = np.clip(centers, 0, 255).astype(np.uint8)
    counts = Counter(labels.flatten())
    total = float(len(labels))

    idx_sorted = [i for i,_ in counts.most_common()]
    palette = []
    for idx in idx_sorted:
        bgr_c = centers_u8[idx].tolist()
        rgb_c = bgr_to_rgb(np.array([[bgr_c]], dtype=np.uint8)).reshape(-1).tolist()
        hsv_c = bgr_to_hsv(np.array([[bgr_c]], dtype=np.uint8)).reshape(-1).tolist()
        lab_c = bgr_to_lab(np.array([[bgr_c]], dtype=np.uint8)).reshape(-1).tolist()
        share = counts[idx] / total
        palette.append({
            "share": float(share),
            "BGR": [int(x) for x in bgr_c],
            "RGB": [int(x) for x in rgb_c],
            "HSV": [int(x) for x in hsv_c],
            "Lab": [int(x) for x in lab_c],
        })
    return palette

def make_palette_image(palette, width=600, height=80):
    img = np.zeros((height, width, 3), dtype=np.uint8)
    x = 0
    for p in palette:
        w = max(1, int(p["share"] * width))
        r,g,b = p["RGB"]  # stored as RGB
        cv.rectangle(img, (x, 0), (min(width-1, x+w-1), height-1), (b,g,r), -1)  # convert to BGR for draw
        x += w
    for i in range(1, len(palette)):
        x_sep = int(sum([pp["share"] for pp in palette[:i]]) * width)
        cv.line(img, (x_sep, 0), (x_sep, height-1), (30,30,30), 1)
    return img

# ---------------- Heuristics ----------------
def rust_zinc_indicators(bgr, delta=6):
    """
    Heuristic only. Uses Lab:
    - rustish_ratio: fraction of pixels with a* > median(a*) + delta
    - zincish_ratio: fraction of pixels with b* > median(b*) + delta
    """
    lab = bgr_to_lab(bgr)
    L, a, b = cv.split(lab)
    a_med, b_med = np.median(a), np.median(b)
    a_thr = a_med + delta
    b_thr = b_med + delta

    rustish = (a.astype(np.float32) > a_thr).mean()
    zincish = (b.astype(np.float32) > b_thr).mean()
    return {
        "rustish_ratio": float(rustish),
        "zincish_ratio": float(zincish),
        "a_median": float(a_med),
        "b_median": float(b_med),
        "a_thresh": float(a_thr),
        "b_thresh": float(b_thr),
        "delta": float(delta)
    }

# ---------------- Classification ----------------
def classify_from_ratios(rustish_ratio, zincish_ratio, rust_thr=0.002, zinc_thr=0.01):
    """
    Your rule:
    - zinc if zincish_ratio > 0.01
    - else rust if rustish_ratio > 0.002
    - else normal
    """
    if zincish_ratio > zinc_thr:
        return "zinc"
    elif rustish_ratio > rust_thr:
        return "rust"
    else:
        return "normal"

# ---------------- Main ----------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--img", required=True, help="path to image")
    ap.add_argument("--k", type=int, default=3, help="number of dominant colors")
    ap.add_argument("--resize_max", type=int, default=1200, help="resize longer side to this (0=off)")
    ap.add_argument("--outdir", default="color_out")
    # thresholds you defined:
    ap.add_argument("--rust_thr", type=float, default=0.01)
    ap.add_argument("--zinc_thr", type=float, default=0.02)
    # indicator sensitivity (Lab delta)
    ap.add_argument("--lab_delta", type=float, default=6.0)
    args = ap.parse_args()

    os.makedirs(args.outdir, exist_ok=True)

    bgr = cv.imread(args.img, cv.IMREAD_COLOR)
    if bgr is None:
        raise RuntimeError(f"Cannot read image: {args.img}")

    # optional resize
    h, w = bgr.shape[:2]
    if args.resize_max > 0:
        s = max(h, w)
        if s > args.resize_max:
            scale = args.resize_max / float(s)
            bgr = cv.resize(bgr, (int(w*scale), int(h*scale)), interpolation=cv.INTER_AREA)

    # color stats
    rgb = bgr_to_rgb(bgr)
    hsv = bgr_to_hsv(bgr)
    lab = bgr_to_lab(bgr)
    stats = [
        img_stats(rgb, "RGB"),
        img_stats(hsv, "HSV"),
        img_stats(lab, "Lab"),
    ]

    # dominant colors
    palette = dominant_colors_kmeans(bgr, k=max(1, args.k))

    # heuristics
    indicators = rust_zinc_indicators(bgr, delta=args.lab_delta)

    # classification using your thresholds
    cls = classify_from_ratios(
        rustish_ratio=indicators["rustish_ratio"],
        zincish_ratio=indicators["zincish_ratio"],
        rust_thr=args.rust_thr,
        zinc_thr=args.zinc_thr
    )

    # save palette image
    base = os.path.splitext(os.path.basename(args.img))[0]
    pal_img = make_palette_image(palette)
    pal_path = os.path.join(args.outdir, f"{base}_palette.png")
    cv.imwrite(pal_path, pal_img)

    # build + save JSON
    report = {
        "input": os.path.basename(args.img),
        "size_hw": [int(bgr.shape[0]), int(bgr.shape[1])],
        "color_stats": stats,
        "dominant_colors": palette,    # ordered by share desc
        "heuristics": indicators,
        "classification": cls,
        "thresholds": {"rust_thr": args.rust_thr, "zinc_thr": args.zinc_thr},
        "palette_image": pal_path
    }
    rep_path = os.path.join(args.outdir, f"{base}_color_report.json")
    with open(rep_path, "w") as f:
        json.dump(report, f, indent=2)

    # console summary
    print(json.dumps({
        "input": report["input"],
        "classification": cls,
        "rustish_ratio": round(indicators["rustish_ratio"], 4),
        "zincish_ratio": round(indicators["zincish_ratio"], 4),
        "top_colors_rgb": [p["RGB"] for p in palette],
        "top_colors_share": [round(p["share"], 4) for p in palette],
        "report_path": rep_path,
        "palette_image": pal_path
    }, indent=2))



if __name__ == "__main__":
    main()