312 lines
11 KiB
Python
312 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
|
|
def load_json(path: Path) -> dict[str, Any]:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def crop_sprite(atlas: Image.Image, meta: dict[str, Any]) -> Image.Image:
|
|
crop = atlas.crop((
|
|
int(meta["x"]), int(meta["y"]),
|
|
int(meta["x"]) + int(meta["width"]),
|
|
int(meta["y"]) + int(meta["height"]),
|
|
)).convert("RGBA")
|
|
alpha = np.asarray(crop.getchannel("A"))
|
|
ys, xs = np.where(alpha > 8)
|
|
if len(xs):
|
|
crop = crop.crop((int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1))
|
|
return crop
|
|
|
|
|
|
def resize_center(crop: Image.Image, size: int = 64, pad: int = 5) -> Image.Image:
|
|
w, h = crop.size
|
|
target = size - pad * 2
|
|
scale = min(target / max(w, 1), target / max(h, 1))
|
|
nw, nh = max(1, round(w * scale)), max(1, round(h * scale))
|
|
crop = crop.resize((nw, nh), Image.Resampling.LANCZOS)
|
|
canvas = Image.new("RGBA", (size, size), (255, 255, 255, 0))
|
|
canvas.alpha_composite(crop, ((size - nw) // 2, (size - nh) // 2))
|
|
return canvas
|
|
|
|
|
|
def dct_matrix(n: int) -> np.ndarray:
|
|
mat = np.zeros((n, n), dtype=np.float32)
|
|
factor = math.pi / n
|
|
for k in range(n):
|
|
scale = math.sqrt(1 / n) if k == 0 else math.sqrt(2 / n)
|
|
for i in range(n):
|
|
mat[k, i] = scale * math.cos((i + 0.5) * k * factor)
|
|
return mat
|
|
|
|
|
|
DCT32 = dct_matrix(32)
|
|
|
|
|
|
def phash(gray32: np.ndarray) -> np.ndarray:
|
|
coeff = DCT32 @ gray32 @ DCT32.T
|
|
low = coeff[:8, :8].copy()
|
|
vals = low.ravel()[1:]
|
|
med = float(np.median(vals))
|
|
return (low.ravel()[1:] > med).astype(np.uint8)
|
|
|
|
|
|
def color_hist(rgb: np.ndarray, alpha: np.ndarray, bins: int = 4) -> np.ndarray:
|
|
mask = alpha > 0.08
|
|
if not np.any(mask):
|
|
return np.zeros(bins ** 3, dtype=np.float32)
|
|
pix = np.clip(rgb[mask], 0, 0.999999)
|
|
q = (pix * bins).astype(np.int32)
|
|
idx = q[:, 0] * bins * bins + q[:, 1] * bins + q[:, 2]
|
|
hist = np.bincount(idx, minlength=bins ** 3).astype(np.float32)
|
|
s = hist.sum()
|
|
return hist / s if s else hist
|
|
|
|
|
|
def descriptor(crop: Image.Image) -> dict[str, Any]:
|
|
original_w, original_h = crop.size
|
|
canvas = resize_center(crop, 64, 5)
|
|
arr = np.asarray(canvas, dtype=np.float32) / 255.0
|
|
alpha = arr[:, :, 3]
|
|
rgb = arr[:, :, :3]
|
|
composited = rgb * alpha[:, :, None] + (1.0 - alpha[:, :, None])
|
|
gray_img = Image.fromarray(np.uint8(np.clip(
|
|
(0.299 * composited[:, :, 0] + 0.587 * composited[:, :, 1] + 0.114 * composited[:, :, 2]) * 255,
|
|
0, 255
|
|
)))
|
|
gray32 = np.asarray(gray_img.resize((32, 32), Image.Resampling.LANCZOS), dtype=np.float32) / 255.0
|
|
edge = np.asarray(gray_img.filter(ImageFilter.FIND_EDGES), dtype=np.float32) / 255.0
|
|
|
|
mask = alpha > 0.08
|
|
xproj = mask.sum(axis=0).astype(np.float32)
|
|
yproj = mask.sum(axis=1).astype(np.float32)
|
|
if xproj.sum():
|
|
xproj /= xproj.sum()
|
|
if yproj.sum():
|
|
yproj /= yproj.sum()
|
|
|
|
return {
|
|
"alpha": alpha,
|
|
"gray": gray32,
|
|
"edge": edge,
|
|
"phash": phash(gray32),
|
|
"hist": color_hist(rgb, alpha),
|
|
"xproj": xproj,
|
|
"yproj": yproj,
|
|
"aspect": original_w / max(1, original_h),
|
|
"occupancy": float(mask.mean()),
|
|
}
|
|
|
|
|
|
def cosine(a: np.ndarray, b: np.ndarray) -> float:
|
|
denom = float(np.linalg.norm(a) * np.linalg.norm(b))
|
|
if denom == 0:
|
|
return 0.0
|
|
return float(np.dot(a.ravel(), b.ravel()) / denom)
|
|
|
|
|
|
def similarity(a: dict[str, Any], b: dict[str, Any]) -> float:
|
|
alpha = cosine(a["alpha"], b["alpha"])
|
|
edge = cosine(a["edge"], b["edge"])
|
|
gray = 1.0 - float(np.mean((a["gray"] - b["gray"]) ** 2))
|
|
hamming = float(np.mean(a["phash"] != b["phash"]))
|
|
ph = 1.0 - hamming
|
|
hist = cosine(a["hist"], b["hist"])
|
|
proj = 0.5 * (cosine(a["xproj"], b["xproj"]) + cosine(a["yproj"], b["yproj"]))
|
|
ratio = min(a["aspect"], b["aspect"]) / max(a["aspect"], b["aspect"])
|
|
occ = 1.0 - min(1.0, abs(a["occupancy"] - b["occupancy"]) / max(a["occupancy"], b["occupancy"], 0.02))
|
|
return (
|
|
0.24 * alpha +
|
|
0.17 * edge +
|
|
0.16 * gray +
|
|
0.14 * ph +
|
|
0.12 * hist +
|
|
0.08 * proj +
|
|
0.05 * ratio +
|
|
0.04 * occ
|
|
)
|
|
|
|
|
|
def resource_records(catalog: dict[str, Any], standard: bool) -> list[dict[str, Any]]:
|
|
seen: set[tuple[Any, ...]] = set()
|
|
out = []
|
|
for obj in catalog.get("visual_objects", []):
|
|
if not obj.get("resource_key") or not obj.get("resource_meta"):
|
|
continue
|
|
layer = obj.get("standard_layer_name") if standard else None
|
|
key = (
|
|
obj.get("component_role"),
|
|
obj.get("resource_key"),
|
|
layer,
|
|
)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
out.append({
|
|
"resource_key": obj["resource_key"],
|
|
"resource_meta": obj["resource_meta"],
|
|
"component_role": obj.get("component_role"),
|
|
"standard_layer_name": layer,
|
|
"visual_object_id": obj.get("visual_object_id"),
|
|
})
|
|
return out
|
|
|
|
|
|
def build_desc(records: list[dict[str, Any]], atlas: Image.Image) -> list[dict[str, Any]]:
|
|
return [descriptor(crop_sprite(atlas, r["resource_meta"])) for r in records]
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--standard", required=True)
|
|
ap.add_argument("--cm95", required=True)
|
|
ap.add_argument("--output", required=True)
|
|
ap.add_argument("--min-score", type=float, default=0.84)
|
|
ap.add_argument("--min-margin", type=float, default=0.025)
|
|
args = ap.parse_args()
|
|
|
|
std_path = Path(args.standard)
|
|
cm_path = Path(args.cm95)
|
|
out = Path(args.output)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
|
|
std_cat = load_json(std_path)
|
|
cm_cat = load_json(cm_path)
|
|
std_dir = std_path.parent.parent / "navsea-anonymous"
|
|
cm_dir = cm_path.parent
|
|
std_atlas = Image.open(std_dir / "_sprite" / "sprite.png").convert("RGBA")
|
|
cm_atlas = Image.open(cm_dir / "_sprite" / "sprite.png").convert("RGBA")
|
|
|
|
std = resource_records(std_cat, True)
|
|
cm = resource_records(cm_cat, False)
|
|
std_d = build_desc(std, std_atlas)
|
|
cm_d = build_desc(cm, cm_atlas)
|
|
|
|
std_by_role: dict[str, list[int]] = defaultdict(list)
|
|
cm_by_role: dict[str, list[int]] = defaultdict(list)
|
|
for i, r in enumerate(std):
|
|
std_by_role[r["component_role"]].append(i)
|
|
for i, r in enumerate(cm):
|
|
cm_by_role[r["component_role"]].append(i)
|
|
|
|
cm_rankings: dict[int, list[tuple[float, int]]] = {}
|
|
std_best: dict[int, tuple[float, int]] = {}
|
|
|
|
for role, cm_indices in cm_by_role.items():
|
|
std_indices = std_by_role.get(role, [])
|
|
if not std_indices:
|
|
continue
|
|
for ci in cm_indices:
|
|
vals = sorted(
|
|
((similarity(cm_d[ci], std_d[si]), si) for si in std_indices),
|
|
reverse=True,
|
|
)
|
|
cm_rankings[ci] = vals[:5]
|
|
for si in std_indices:
|
|
vals = max(
|
|
((similarity(std_d[si], cm_d[ci]), ci) for ci in cm_indices),
|
|
default=(0.0, -1),
|
|
)
|
|
std_best[si] = vals
|
|
|
|
anchors = []
|
|
candidates = []
|
|
for ci, ranked in cm_rankings.items():
|
|
if not ranked:
|
|
continue
|
|
best_score, si = ranked[0]
|
|
second_score = ranked[1][0] if len(ranked) > 1 else 0.0
|
|
mutual = std_best.get(si, (0.0, -1))[1] == ci
|
|
margin = best_score - second_score
|
|
row = {
|
|
"cm95_resource_key": cm[ci]["resource_key"],
|
|
"component_role": cm[ci]["component_role"],
|
|
"best_standard_resource_key": std[si]["resource_key"],
|
|
"best_standard_layer": std[si]["standard_layer_name"],
|
|
"score": best_score,
|
|
"margin": margin,
|
|
"mutual_nearest": mutual,
|
|
"standard_visual_object_id": std[si]["visual_object_id"],
|
|
"alternatives": [
|
|
{
|
|
"standard_resource_key": std[sj]["resource_key"],
|
|
"standard_layer": std[sj]["standard_layer_name"],
|
|
"score": score,
|
|
}
|
|
for score, sj in ranked[1:4]
|
|
],
|
|
}
|
|
candidates.append(row)
|
|
if mutual and best_score >= args.min_score and margin >= args.min_margin:
|
|
row = dict(row)
|
|
row["status"] = "MATCH"
|
|
anchors.append(row)
|
|
|
|
# Map accepted visual anchors back to CM95 anonymous Visual Object occurrences.
|
|
anchor_by_key_role = {
|
|
(x["cm95_resource_key"], x["component_role"]): x for x in anchors
|
|
}
|
|
occurrences = []
|
|
for obj in cm_cat.get("visual_objects", []):
|
|
k = (obj.get("resource_key"), obj.get("component_role"))
|
|
anchor = anchor_by_key_role.get(k)
|
|
if not anchor:
|
|
continue
|
|
occurrences.append({
|
|
"visual_object_id": obj["visual_object_id"],
|
|
"group_id": obj["group_id"],
|
|
"cm95_resource_key": obj["resource_key"],
|
|
"assigned_standard_layer": anchor["best_standard_layer"],
|
|
"score": anchor["score"],
|
|
"basis": "VISUAL_ANCHOR_MUTUAL_NEAREST",
|
|
})
|
|
|
|
summary = {
|
|
"standard_visual_resources": len(std),
|
|
"cm95_visual_resources": len(cm),
|
|
"candidate_pairs": len(candidates),
|
|
"accepted_visual_anchors": len(anchors),
|
|
"cm95_visual_object_occurrences_named": len(occurrences),
|
|
"standard_layers_reached": sorted(set(x["best_standard_layer"] for x in anchors if x.get("best_standard_layer"))),
|
|
}
|
|
result = {
|
|
"schema": "style-visual-anchor-match/v1",
|
|
"rules": {
|
|
"cm95_source_layer_used": False,
|
|
"provider_name_used": False,
|
|
"resource_key_text_semantics_used": False,
|
|
"standard_layer_is_output_only": True,
|
|
"min_score": args.min_score,
|
|
"min_margin": args.min_margin,
|
|
"requires_mutual_nearest": True,
|
|
},
|
|
"summary": summary,
|
|
"anchors": sorted(anchors, key=lambda x: x["score"], reverse=True),
|
|
"cm95_named_occurrences": occurrences,
|
|
"top_candidates": sorted(candidates, key=lambda x: x["score"], reverse=True)[:200],
|
|
}
|
|
(out / "visual-anchor-matches.json").write_text(
|
|
json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
for x in result["anchors"][:80]:
|
|
print(
|
|
f'{x["score"]:.4f} margin={x["margin"]:.4f} '
|
|
f'{x["cm95_resource_key"]} -> {x["best_standard_resource_key"]} '
|
|
f'[{x["best_standard_layer"]}]'
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|