Image Linting - 生成画像の検品を重ねると正しい絵も減る

"""features.json（lint_features.py）と labels.json（目で付けた合否）から、
検査ごとの誤った棄却・誤った合格と、検査を直列に重ねた時の通過率を数える。
最後に、顔の切り抜きで「同じ服装か」を CLIP とタグの cos で比べる。

閾値と必須のタグは、合否のラベルを付ける前に決めて、この下の定数に書いた値のまま動かしていない。

使い方:
    python3 evaluate.py <features.json> <labels.json> [結果の JSON]
"""
import itertools
import json
import math
import sys

import numpy as np

TAG = 0.35                       # タグが「出ている」とみなす確率
CLIP_MIN = 0.80                  # 基準の絵との CLIP の cos の下限
BORDER_MAX = 30.0                # 左右の端の平均色と白の距離（RGB、0〜441）の上限
NEON_MAX = 0.02                  # 彩度も明るさも高い画素の割合の上限

OUTFIT = {                       # 服装ごとに必須のタグ（| はどれか 1 つ）
    "A": ["red_jacket", "black_hair", "denim|jeans"],
    "B": ["green_sweater", "blonde_hair", "black_skirt"],
    "C": ["brown_hair", "suit|blue_suit", "red_necktie"],
    "D": ["grey_hair|white_hair", "raincoat|yellow_raincoat", "grey_pants"],
}
SUBJECT = {"A": "1girl", "B": "1girl", "C": "1boy", "D": "1girl"}
CROWD = ["multiple_girls", "2girls", "3girls", "multiple_boys", "2boys", "3boys",
         "multiple_views", "1girl", "1boy"]          # 自分の SUBJECT は後で外す
BROKEN = ["bad_hands", "bad_anatomy", "extra_arms", "extra_digits", "bad_feet",
          "text", "english_text", "watermark", "signature"]
RATINGS = ["general", "sensitive", "questionable", "explicit"]


def make_checks(names):
    ix = {n: i for i, n in enumerate(names)}

    def p(rec, tag):
        return rec["tags"][ix[tag]] if tag in ix else 0.0

    def any_of(rec, alt):
        return max(p(rec, t) for t in alt.split("|")) >= TAG

    def crowd_tags(rec, d):
        bad = [t for t in CROWD if t != SUBJECT[d] and p(rec, t) >= TAG]
        return not bad and p(rec, SUBJECT[d]) >= TAG

    return [
        ("人数（タグ）", crowd_tags),
        ("人数（検出器）", lambda r, d: len(r["people"]) == 1 and len(r["faces"]) == 1),
        ("崩れ・文字（タグ）", lambda r, d: all(p(r, t) < TAG for t in BROKEN)),
        ("露出（レーティング）",
         lambda r, d: max(RATINGS, key=lambda t: p(r, t)) == "general"),
        ("服（タグ）", lambda r, d: all(any_of(r, a) for a in OUTFIT[d])),
        ("色（統計）", lambda r, d: r["color"]["border_dist_from_white"] <= BORDER_MAX
         and r["color"]["neon_ratio"] <= NEON_MAX),
        ("基準の絵との近さ（CLIP）", lambda r, d: r["clip_to_ref"] >= CLIP_MIN),
    ]


def cos(a, b):
    a, b = np.asarray(a), np.asarray(b)
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))


# 顔の切り抜きのタグの cos から外す語（表情・視線・構図・背景・人数・レーティング）
STOP = {"solo", "1girl", "1boy", "portrait", "looking_at_viewer", "closed_mouth",
        "open_mouth", "parted_lips", "smile", "expressionless", "blush", "upper_body",
        "white_background", "simple_background", "transparent_background",
        "grey_background", "close-up", "light_smile", "male_focus", "general",
        "sensitive", "questionable", "explicit"}


def main():
    feats = json.load(open(sys.argv[1]))
    labels = json.load(open(sys.argv[2]))
    out_path = sys.argv[3] if len(sys.argv) > 3 else None
    names, imgs = feats["tag_names"], feats["images"]
    for name, rec in imgs.items():
        ref = imgs[f"{name[0]}_ref_20261100"]
        rec["clip_to_ref"] = cos(rec["clip"], ref["clip"])
    checks = make_checks(names)
    vocab = set(names)
    used = set(CROWD + BROKEN + RATINGS) | {t for v in OUTFIT.values() for a in v
                                            for t in a.split("|")}
    print("タグ付けのモデルの語彙に無いタグ（確率は常に 0 として扱われる）:",
          sorted(used - vocab))

    targets = sorted(n for n in imgs if "_ref_" not in n)
    good = [n for n in targets if labels[n]["ok"]]
    bad = [n for n in targets if not labels[n]["ok"]]
    print(f"評価する画像 {len(targets)} 枚（目で見て合格 {len(good)}、不合格 {len(bad)}）")
    print("基準の絵 4 枚の合否:",
          {n: all(f(imgs[n], n[0]) for _, f in checks) for n in sorted(imgs) if "_ref_" in n})

    passed = {c: {n: f(imgs[n], n[0]) for n in targets} for c, f in checks}
    rows = []
    print(f"\n{'検査':<22}{'誤った棄却':>10}{'誤った合格':>10}{'単独で止めた不合格':>12}")
    for c, _ in checks:
        fr = [n for n in good if not passed[c][n]]
        fa = [n for n in bad if passed[c][n]]
        only = [n for n in bad if not passed[c][n]
                and all(passed[o][n] for o, _ in checks if o != c)]
        rows.append({"check": c, "false_reject": len(fr), "false_accept": len(fa),
                     "caught_bad": len(bad) - len(fa), "only_this": len(only),
                     "false_reject_names": fr})
        print(f"{c:<22}{len(fr):>4} / {len(good):<4}{len(fa):>5} / {len(bad):<4}{len(only):>8}")

    # 直列に重ねた時、合格の絵が全部を通る割合（実測）と、独立を仮定した積
    print("\n段数  実測（合格の絵が全部を通る）  各段の通過率の積  (1-p)^n（p は各段の平均）")
    stages = []
    for k in range(1, len(checks) + 1):
        cs = [c for c, _ in checks[:k]]
        all_pass = sum(all(passed[c][n] for c in cs) for n in good) / len(good)
        rates = [sum(passed[c][n] for n in good) / len(good) for c in cs]
        prod = math.prod(rates)
        pbar = 1 - sum(rates) / len(rates)
        stages.append({"n": k, "measured": all_pass, "product": prod,
                       "mean_p": pbar, "formula": (1 - pbar) ** k})
        print(f"{k:>3}   {all_pass:>8.3f}                     {prod:>8.3f}        "
              f"{(1 - pbar) ** k:>8.3f}  (p={pbar:.3f})")
    caught = sum(any(not passed[c][n] for c, _ in checks) for n in bad)
    print(f"\n7 段全部で止めた不合格: {caught} / {len(bad)}")
    print(f"5% の誤りを 5 段: (1-0.05)^5 = {0.95 ** 5:.3f}")

    # 顔の切り抜きで、同じ服装か別の服装かを見分けられるか（neon 以外で合格の絵と基準の絵）
    keep = [i for i, n in enumerate(names) if n not in STOP and feats.get("tag_categories",
            [0] * len(names))[i] == 0]
    faces = sorted(n for n in imgs if "_neon_" not in n and "face_clip" in imgs[n]
                   and labels.get(n, {"ok": True})["ok"])
    same, diff = {"clip": [], "tag": []}, {"clip": [], "tag": []}
    for a, b in itertools.combinations(faces, 2):
        ta = [imgs[a]["face_tags"][i] for i in keep]
        tb = [imgs[b]["face_tags"][i] for i in keep]
        bucket = same if a[0] == b[0] else diff
        bucket["clip"].append(cos(imgs[a]["face_clip"], imgs[b]["face_clip"]))
        bucket["tag"].append(cos(ta, tb))
    print(f"\n顔の切り抜き {len(faces)} 枚（neon 以外で目で見て合格の絵と、基準の絵）")
    for k in ("clip", "tag"):
        print(f"  {k:<4} 同じ服装 {len(same[k])} 組 {min(same[k]):.2f}〜{max(same[k]):.2f}"
              f"（中央値 {np.median(same[k]):.2f}） / 別の服装 {len(diff[k])} 組 "
              f"{min(diff[k]):.2f}〜{max(diff[k]):.2f}（中央値 {np.median(diff[k]):.2f}）")

    if out_path:
        json.dump({"n_good": len(good), "n_bad": len(bad), "checks": rows,
                   "stages": stages, "identity": {"same": same, "diff": diff}},
                  open(out_path, "w"), ensure_ascii=False, indent=1)


if __name__ == "__main__":
    main()

"""生成した PNG から、検品に使う値を計算して features.json に書く。

使うモデル（全て ONNX、CPU で動かす）:
  WD ViT Tagger v3          : タグの確率（人数・露出・服・崩れ）
  CLIP ViT-B/32（画像側）   : 画像の埋め込み（基準の絵との近さ）
  anime_face_detection v1.4_s / anime_person_detection v1.3_s : 顔と人物の数
色の統計は NumPy だけで計算する。

使い方:
    python3 lint_features.py <PNG のディレクトリ> <モデルのディレクトリ> <出力 JSON>
"""
import json
import os
import sys

import numpy as np
import onnxruntime as ort
from PIL import Image

OPTS = ort.SessionOptions()
OPTS.intra_op_num_threads = 8


def session(path):
    return ort.InferenceSession(path, OPTS, providers=["CPUExecutionProvider"])


# ---- WD tagger -------------------------------------------------------------
def load_tagger(d):
    import csv
    with open(os.path.join(d, "selected_tags.csv")) as f:
        rows = list(csv.DictReader(f))
    return session(os.path.join(d, "model.onnx")), [r["name"] for r in rows], \
        [int(r["category"]) for r in rows]


def tag_probs(sess, img):
    """白で正方形に埋めて 448 に縮め、BGR・0〜255 の float で渡す（モデルの前処理）。"""
    w, h = img.size
    s = max(w, h)
    canvas = Image.new("RGB", (s, s), (255, 255, 255))
    canvas.paste(img, ((s - w) // 2, (s - h) // 2))
    x = np.asarray(canvas.resize((448, 448), Image.BICUBIC), dtype=np.float32)[:, :, ::-1]
    return sess.run(None, {"input": x[None].copy()})[0][0]


# ---- CLIP ------------------------------------------------------------------
MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)


def clip_embed(sess, img):
    """短い辺を 224 に縮めて中央を切り抜く（preprocessor_config.json と同じ）。"""
    w, h = img.size
    r = 224 / min(w, h)
    img = img.resize((round(w * r), round(h * r)), Image.BICUBIC)
    w, h = img.size
    l, t = (w - 224) // 2, (h - 224) // 2
    x = np.asarray(img.crop((l, t, l + 224, t + 224)), dtype=np.float32) / 255.0
    x = ((x - MEAN) / STD).transpose(2, 0, 1)[None]
    e = sess.run(None, {"pixel_values": x})[0][0]
    return e / np.linalg.norm(e)


# ---- YOLO 系の検出器（顔・人物） -------------------------------------------
def detect(sess, img, thr, size=640, iou=0.7):
    w, h = img.size
    r = size / max(w, h)
    nw, nh = round(w * r), round(h * r)
    canvas = Image.new("RGB", (size, size), (114, 114, 114))
    canvas.paste(img.resize((nw, nh), Image.BILINEAR), (0, 0))
    x = (np.asarray(canvas, dtype=np.float32) / 255.0).transpose(2, 0, 1)[None]
    out = sess.run(None, {"images": x})[0][0]           # [4 + クラス数, 候補数]
    if out.shape[0] > out.shape[1]:
        out = out.T
    boxes, scores = out[:4].T, out[4:].max(axis=0)
    keep = scores >= thr
    boxes, scores = boxes[keep], scores[keep]
    xyxy = np.stack([boxes[:, 0] - boxes[:, 2] / 2, boxes[:, 1] - boxes[:, 3] / 2,
                     boxes[:, 0] + boxes[:, 2] / 2, boxes[:, 1] + boxes[:, 3] / 2], 1) / r
    order, picked = scores.argsort()[::-1], []
    while order.size:                                    # 重なった候補を 1 つにまとめる（NMS）
        i = order[0]
        picked.append(i)
        a = xyxy[i]
        b = xyxy[order[1:]]
        inter = (np.clip(np.minimum(a[2], b[:, 2]) - np.maximum(a[0], b[:, 0]), 0, None) *
                 np.clip(np.minimum(a[3], b[:, 3]) - np.maximum(a[1], b[:, 1]), 0, None))
        area = lambda z: (z[..., 2] - z[..., 0]) * (z[..., 3] - z[..., 1])
        order = order[1:][inter / (area(a) + area(b) - inter + 1e-9) < iou]
    return [[float(v) for v in xyxy[i]] + [float(scores[i])] for i in picked]


# ---- 色の統計 ---------------------------------------------------------------
def color_stats(img):
    x = np.asarray(img, dtype=np.float32) / 255.0
    h, w, _ = x.shape
    m = max(4, w // 20)                                  # 左右の端の帯（背景のはずの所）
    border = np.concatenate([x[:, :m].reshape(-1, 3), x[:, -m:].reshape(-1, 3)])
    mx, mn = x.max(axis=2), x.min(axis=2)
    sat = np.where(mx > 0, (mx - mn) / np.maximum(mx, 1e-6), 0)
    neon = float(((sat > 0.85) & (mx > 0.85)).mean())   # 彩度も明るさも高い画素の割合
    return {
        "border_rgb": [float(v) for v in border.mean(axis=0) * 255],
        "border_dist_from_white": float(np.linalg.norm((1 - border.mean(axis=0)) * 255)),
        "neon_ratio": neon,
        "mean_saturation": float(sat.mean()),
    }


def main():
    png_dir, model_dir, out_path = sys.argv[1:4]
    tagger, names, cats = load_tagger(os.path.join(model_dir, "wd-vit-tagger-v3"))
    clip = session(os.path.join(model_dir, "clip-vit-base-patch32-onnx", "vision_model.onnx"))
    dets = {}
    for key, d in [("face", "face_detect_v1.4_s"), ("person", "person_detect_v1.3_s")]:
        with open(os.path.join(model_dir, d, "threshold.json")) as f:
            dets[key] = (session(os.path.join(model_dir, d, "model.onnx")),
                         json.load(f)["threshold"])

    result = {"tag_names": names, "tag_categories": cats, "images": {}}
    for fn in sorted(os.listdir(png_dir)):
        if not fn.endswith(".png"):
            continue
        img = Image.open(os.path.join(png_dir, fn)).convert("RGB")
        p = tag_probs(tagger, img)
        faces = detect(dets["face"][0], img, dets["face"][1])
        people = detect(dets["person"][0], img, dets["person"][1])
        rec = {
            "tags": [round(float(v), 4) for v in p],
            "clip": [round(float(v), 6) for v in clip_embed(clip, img)],
            "faces": faces, "people": people, "color": color_stats(img),
        }
        if faces:                                         # 一番確からしい顔の切り抜き
            x0, y0, x1, y1, _ = max(faces, key=lambda b: b[4])
            pad = 0.25 * (x1 - x0)
            crop = img.crop((max(0, x0 - pad), max(0, y0 - pad),
                             min(img.width, x1 + pad), min(img.height, y1 + pad)))
            rec["face_tags"] = [round(float(v), 4) for v in tag_probs(tagger, crop)]
            rec["face_clip"] = [round(float(v), 6) for v in clip_embed(clip, crop)]
        result["images"][fn[:-4]] = rec
        print(f"{fn}: faces={len(faces)} people={len(people)} "
              f"border={rec['color']['border_dist_from_white']:.1f} "
              f"neon={rec['color']['neon_ratio']:.3f}", flush=True)
    with open(out_path, "w") as f:
        json.dump(result, f)


if __name__ == "__main__":
    main()