features/extract_handcrafted.py

"""
Hand-crafted features for the gold 4x1000 subset — corrected pipeline.

Mirrors notebooks/handcrafted.ipynb (2026-07-07):
  1. crop_padding             — strip near-black letterbox borders
  2. guarded rembg mask       — PRE-COMPUTED: data/masks/
                                (verdicts.csv; mask applied only when it removes a
                                solid, convex, edge-touching, uniform background)
  3. applied images           — spatial views cropped to mask bbox; mask limits
                                color/light stats. rejected images — full image.

Differences vs extract_handcrafted.py: masking is guarded (not unconditional),
spatial features get bbox-crop instead of nothing, no inline rembg (cache only).

Output: data/features/handcrafted.parquet   (key: id, str)
Usage:  python features/extract_handcrafted_gold.py [--workers 4]
"""

import argparse
import os
import sys
from multiprocessing.pool import ThreadPool
from pathlib import Path

import cv2
import numpy as np
import pandas as pd
from PIL import Image
from skimage.feature import local_binary_pattern
from skimage.segmentation import slic
from tqdm import tqdm

ROOT = Path(__file__).resolve().parent.parent
IMAGES = ROOT / "data/images"
MASK_DIR = ROOT / "data/masks"
VERDICT_CSV = MASK_DIR / "verdicts.csv"
SELECTED = ROOT / "data/artwork_metadata.csv"
OUTPUT = ROOT / "data/features/handcrafted.parquet"

HIST_BINS = (8, 4, 4)
LBP_P = 8
LBP_BINS = LBP_P * (LBP_P - 1) + 3
MAX_SIDE = 1024  # cap resolution: texture/edge features stay comparable across
                 # museum scan sizes; unbounded FFT/SLIC on 33 MP scans OOMs

Image.MAX_IMAGE_PIXELS = None


def crop_padding(img_rgb, threshold=5):
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    rows = np.where(gray.max(axis=1) > threshold)[0]
    cols = np.where(gray.max(axis=0) > threshold)[0]
    if len(rows) == 0 or len(cols) == 0:
        return img_rgb
    return img_rgb[rows[0]:rows[-1] + 1, cols[0]:cols[-1] + 1]


def grey_world_normalize(img_rgb):
    f = img_rgb.astype(np.float32)
    means = f.reshape(-1, 3).mean(0) + 1e-6
    f *= means.mean() / means
    return np.clip(f, 0, 255).astype(np.uint8)


VER = pd.read_csv(VERDICT_CSV, dtype=str).set_index("filename")


def load_final(filename):
    """(rgb, hsv, gray, edges, mask) under the corrected pipeline."""
    img = crop_padding(np.array(Image.open(IMAGES / filename).convert("RGB")))
    mask = None
    if filename in VER.index:
        v = VER.loc[filename]
        if v["verdict"] == "applied":
            m = cv2.imread(str(MASK_DIR / (Path(filename).stem + ".png")),
                           cv2.IMREAD_GRAYSCALE)
            if m is not None and m.shape == img.shape[:2]:
                y0, y1 = int(v["y0"]), int(v["y1"])
                x0, x1 = int(v["x0"]), int(v["x1"])
                img, mask = img[y0:y1 + 1, x0:x1 + 1], m[y0:y1 + 1, x0:x1 + 1]
    scale = MAX_SIDE / max(img.shape[:2])
    if scale < 1.0:
        size = (round(img.shape[1] * scale), round(img.shape[0] * scale))
        img = cv2.resize(img, size, interpolation=cv2.INTER_AREA)
        if mask is not None:
            mask = cv2.resize(mask, size, interpolation=cv2.INTER_NEAREST)
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    return img, hsv, gray, cv2.Canny(gray, 100, 200), mask


def extract_one(filename):
    try:
        img, hsv, gray, edges, mask = load_final(filename)
        row = {"filename": filename, "hc_fg_applied": int(mask is not None)}

        # color
        hist = cv2.calcHist([hsv], [0, 1, 2], mask, list(HIST_BINS),
                            [0, 180, 0, 256, 0, 256])
        cv2.normalize(hist, hist)
        for i, v in enumerate(hist.flatten()):
            row[f"hc_hist_{i}"] = float(v)
        hsv_n = cv2.cvtColor(grey_world_normalize(img), cv2.COLOR_RGB2HSV)
        hist_n = cv2.calcHist([hsv_n], [0, 1, 2], mask, list(HIST_BINS),
                              [0, 180, 0, 256, 0, 256])
        cv2.normalize(hist_n, hist_n)
        for i, v in enumerate(hist_n.flatten()):
            row[f"hc_norm_hist_{i}"] = float(v)
        h, s, v_ = cv2.split(hsv)
        for pref, ch, bins, rng in [("h", h, 16, [0, 180]),
                                    ("s", s, 8, [0, 256]),
                                    ("v", v_, 8, [0, 256])]:
            c = cv2.calcHist([ch], [0], mask, [bins], rng).flatten()
            c /= c.sum() + 1e-10
            for i, x in enumerate(c):
                row[f"hc_{pref}_hist_{i}"] = float(x)

        # light + color scalars (masked), edge density (full frame)
        sel_g = gray if mask is None else gray[mask > 0]
        sel_hsv = hsv.reshape(-1, 3) if mask is None else hsv[mask > 0]
        row["hc_avg_hue"] = float(sel_hsv[:, 0].mean())
        row["hc_avg_sat"] = float(sel_hsv[:, 1].mean())
        row["hc_brightness"] = float(sel_hsv[:, 2].mean())
        row["hc_contrast"] = float(sel_g.std())
        row["hc_darkness"] = float((sel_g < 64).mean())
        row["hc_edge_density"] = float((edges > 0).mean())

        # symmetry
        hh, ww = gray.shape
        row["hc_sym_lr"] = 1.0 - float(np.abs(
            gray[:, :ww // 2].astype(np.float32)
            - np.fliplr(gray[:, ww - ww // 2:]).astype(np.float32)).mean()) / 255.0
        row["hc_sym_tb"] = 1.0 - float(np.abs(
            gray[:hh // 2, :].astype(np.float32)
            - np.flipud(gray[hh - hh // 2:, :]).astype(np.float32)).mean()) / 255.0

        # flatness
        lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB).astype(np.float32)
        segments = slic(lab / 255.0, n_segments=200, compactness=10,
                        start_label=0, channel_axis=2)
        stds = [[], [], []]
        for sid in np.unique(segments):
            region = lab[segments == sid]
            if len(region) > 1:
                for c in range(3):
                    stds[c].append(float(region[:, c].std()))
        sl, sa, sb = (float(np.mean(s_)) if s_ else 0.0 for s_ in stds)
        row["hc_flatness_l"], row["hc_flatness_a"] = sl, sa
        row["hc_flatness_b"], row["hc_flatness_mean"] = sb, (sl + sa + sb) / 3.0

        # geometry
        power = np.abs(np.fft.fftshift(np.fft.fft2(gray.astype(np.float32)))) ** 2
        ph, pw = power.shape
        cy, cx = ph // 2, pw // 2
        r = np.sqrt((np.arange(pw) - cx) ** 2 + (np.arange(ph)[:, None] - cy) ** 2)
        edges_r = np.logspace(0, np.log10(min(cx, cy)), 9)
        bands = []
        for lo, hi in zip(edges_r[:-1], edges_r[1:]):
            sel = (r >= lo) & (r < hi)
            bands.append(float(power[sel].mean()) if sel.any() else 0.0)
        total = sum(bands) + 1e-10
        for i, v in enumerate(bands):
            row[f"hc_fft_band_{i}"] = v / total
        lbp = local_binary_pattern(gray, P=LBP_P, R=1, method="nri_uniform")
        lh, _ = np.histogram(lbp.ravel(), bins=LBP_BINS, range=(0, LBP_BINS))
        lh = lh / (lh.sum() + 1e-10)
        for i, v in enumerate(lh):
            row[f"hc_lbp_{i}"] = float(v)

        # lines
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=50,
                                minLineLength=20, maxLineGap=5)
        if lines is not None:
            total_len = float(sum(np.hypot(x2 - x1, y2 - y1)
                                  for x1, y1, x2, y2 in lines[:, 0]))
            hough_mask = np.zeros_like(edges)
            for x1, y1, x2, y2 in lines[:, 0]:
                cv2.line(hough_mask, (x1, y1), (x2, y2), 255, 1)
            row["hc_hough_count"] = len(lines)
            row["hc_hough_density"] = total_len / (hh * ww)
            row["hc_straight_ratio"] = (float((hough_mask > 0).sum())
                                        / (float((edges > 0).sum()) + 1e-10))
        else:
            row["hc_hough_count"] = 0
            row["hc_hough_density"] = 0.0
            row["hc_straight_ratio"] = 0.0
        gx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        gy = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        mag = np.hypot(gx, gy)
        angle = np.degrees(np.arctan2(gy, gx)) % 180
        ah, _ = np.histogram(angle.ravel(), bins=8, range=(0, 180),
                             weights=mag.ravel())
        ah = ah / (ah.sum() + 1e-10)
        for i, v in enumerate(ah):
            row[f"hc_angle_hist_{i}"] = float(v)

        return row
    except Exception as e:
        sys.stderr.write(f"FAIL {filename}: {e}\n")
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--chunk", type=int, default=100)
    args = ap.parse_args()

    gold = (pd.read_csv(SELECTED, dtype=str)
            .drop_duplicates("filename")[["filename"]])
    existing = (pd.read_parquet(OUTPUT) if OUTPUT.exists()
                else pd.DataFrame(columns=["filename"]))
    have = set(existing["filename"])
    todo = gold[~gold["filename"].isin(have)]
    print(f"gold={len(gold)}  done={len(have)}  todo={len(todo)}")

    tasks = todo["filename"].tolist()
    new_rows, failed = [], 0
    with ThreadPool(args.workers) as pool:
        for row in tqdm(pool.imap_unordered(extract_one, tasks, chunksize=4),
                        total=len(tasks)):
            if row is None:
                failed += 1
                continue
            new_rows.append(row)
            if len(new_rows) >= args.chunk:
                existing = pd.concat([existing, pd.DataFrame(new_rows)],
                                     ignore_index=True)
                existing.to_parquet(OUTPUT, index=False)
                new_rows = []
    if new_rows:
        existing = pd.concat([existing, pd.DataFrame(new_rows)],
                             ignore_index=True)
    existing.to_parquet(OUTPUT, index=False)
    print(f"Wrote {OUTPUT}: {len(existing)} rows, "
          f"{len(existing.columns)} cols. Failures: {failed}")


if __name__ == "__main__":
    main()