diff options
Diffstat (limited to 'worldalign/region_oracle.py')
| -rw-r--r-- | worldalign/region_oracle.py | 216 |
1 files changed, 216 insertions, 0 deletions
diff --git a/worldalign/region_oracle.py b/worldalign/region_oracle.py new file mode 100644 index 0000000..61054ee --- /dev/null +++ b/worldalign/region_oracle.py @@ -0,0 +1,216 @@ +"""Is Visual Genome capped, or is our part correspondence the problem? + +Visual Genome's cross-modal anchor bound is 0.336 against a within-modality +ceiling of 0.805 for text and 0.994 for vision. Both sides identify their own +scenes; they disagree about which scenes resemble each other. That leaves two +very different explanations and they call for opposite responses. + +Either the two encoders would agree if only their parts were put in +correspondence -- our spectral segments and the annotators' phrases carve the +image differently, and the aggregation has to guess which goes with which -- +or they would not, because a DINOv2 patch descriptor and a written phrase +simply do not covary over what makes two regions alike. + +Visual Genome settles it, because each region box is index-aligned with its own +description. Encode the crop inside each box and the phrase describing it, and +the parts are in correspondence **by annotation**. If the cross-modal bound +jumps, part correspondence is the whole problem and it is an engineering target. +If it stays near 0.336, the encoders do not agree at the part level either, and +no aggregation will rescue this corpus. + +DIAGNOSTIC ONLY. Box-to-phrase alignment is cross-modal supervision that a +deployed system would not have. This measures a ceiling, it is not a method. +""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from scipy.optimize import linear_sum_assignment +from tqdm import tqdm +from transformers import AutoImageProcessor, AutoModel + +from .common import seed_everything, write_json +from .natural_families import load_phrases +from .natural_pipeline import content_directions, word_vectors +from .synth_cc_battery import moment_field + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--vg-dir", default="artifacts/vg_5k") + parser.add_argument("--image-cache", default="/tmp/yurenh2-worldalign-vg-images") + parser.add_argument("--model", default="facebook/dinov2-base") + parser.add_argument("--crop-size", type=int, default=126) + parser.add_argument("--samples", type=int, default=256) + parser.add_argument("--fit-scenes", type=int, default=900) + parser.add_argument("--word-vectors", type=int, default=128) + parser.add_argument("--min-count", type=int, default=60) + parser.add_argument("--max-vocabulary", type=int, default=600) + parser.add_argument("--context-window", type=int, default=2) + parser.add_argument("--keep", type=int, default=64) + parser.add_argument("--shrinkage", type=float, default=0.05) + parser.add_argument("--weight-power", type=float, default=0.5) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--device", default="cuda:3") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--output", default="artifacts/vg_5k/region_oracle.json") + return parser.parse_args() + + +def standardise(matrix: np.ndarray) -> np.ndarray: + mask = ~np.eye(len(matrix), dtype=bool) + values = matrix[mask] + out = (matrix - values.mean()) / values.std() + np.fill_diagonal(out, 0.0) + return out + + +def anchor_bound(visual: np.ndarray, text: np.ndarray, repeats: int = 5) -> float: + size = len(visual) + scores = [] + for repeat in range(repeats): + generator = np.random.default_rng(repeat) + shuffle = generator.permutation(size) + half = size // 2 + anchors, probe = shuffle[:half], shuffle[half:] + order = generator.permutation(len(probe)) + + def profile(field, rows): + block = field[np.ix_(rows, anchors)] + block = block - block.mean(1, keepdims=True) + return block / block.std(1, keepdims=True).clip(1e-9) + + similarity = profile(visual, probe) @ profile(text, probe[order]).T / half + _, columns = linear_sum_assignment(-similarity) + scores.append(float((order[columns] == np.arange(len(probe))).mean())) + return float(np.mean(scores)) + + +@torch.inference_mode() +def main() -> None: + args = parse_args() + seed_everything(args.seed) + vg_dir = Path(args.vg_dir) + truth = [json.loads(line) for line in + (vg_dir / "ground_truth.private.jsonl").read_text().splitlines() if line.strip()] + vision = {json.loads(line)["node_id"]: json.loads(line) for line in + (vg_dir / "vision_nodes.private.jsonl").read_text().splitlines() if line.strip()} + records = {json.loads(line)["node_id"]: json.loads(line) for line in + (vg_dir / "text_nodes.jsonl").read_text().splitlines() if line.strip()} + + pairs = [p for p in truth + if p["vision_node_id"] in vision and p["text_node_id"] in records] + needed = pairs[: args.samples + args.fit_scenes] + print(f"{len(needed)} scenes", flush=True) + + processor = AutoImageProcessor.from_pretrained(args.model) + model = AutoModel.from_pretrained(args.model, torch_dtype=torch.float32).to(args.device) + model.eval() + mean = torch.tensor(processor.image_mean).view(3, 1, 1) + std = torch.tensor(processor.image_std).view(3, 1, 1) + + # one crop per annotated region, in the order the phrases are given + crops, owners = [], [] + for index, pair in enumerate(tqdm(needed, desc="crop")): + record = vision[pair["vision_node_id"]] + phrases = records[pair["text_node_id"]]["region_closed"] + path = Path(args.image_cache, f"{record['source_image_id']}.jpg") + if not path.exists(): + continue + with Image.open(path) as handle: + picture = handle.convert("RGB") + width, height = picture.size + for position, region in enumerate(record["regions"][: len(phrases)]): + x0 = max(0, min(int(region["x"]), width - 2)) + y0 = max(0, min(int(region["y"]), height - 2)) + x1 = max(x0 + 2, min(int(region["x"] + region["width"]), width)) + y1 = max(y0 + 2, min(int(region["y"] + region["height"]), height)) + patch = picture.crop((x0, y0, x1, y1)).resize( + (args.crop_size, args.crop_size), Image.BILINEAR + ) + crops.append(torch.from_numpy(np.asarray(patch).copy()) + .permute(2, 0, 1).float() / 255.0) + owners.append((index, position)) + + features = [] + for start in tqdm(range(0, len(crops), args.batch_size), desc="encode"): + batch = torch.stack(crops[start : start + args.batch_size]) + normalised = ((batch - mean) / std).to(args.device) + hidden = model(pixel_values=normalised, return_dict=True).last_hidden_state + features.append(hidden[:, 0].float().cpu().numpy()) + features = np.concatenate(features) + print(f"{len(features)} region crops encoded", flush=True) + + vectors = word_vectors(load_phrases(vg_dir / "text_nodes.jsonl", "region_closed"), args) + + def phrase_vector(phrase: str) -> np.ndarray | None: + hits = [vectors[t] for t in re.findall(r"[a-z]+", phrase.lower()) if t in vectors] + return np.mean(hits, axis=0) if hits else None + + # collect per scene, keeping only regions whose phrase encodes -- so the + # two sides carry exactly the same parts in exactly the same order + vision_parts: dict[int, list] = {} + text_parts: dict[int, list] = {} + for (index, position), feature in zip(owners, features): + phrase = records[needed[index]["text_node_id"]]["region_closed"][position] + encoded = phrase_vector(phrase) + if encoded is None: + continue + vision_parts.setdefault(index, []).append(feature.astype(np.float64)) + text_parts.setdefault(index, []).append(encoded) + + usable = [i for i in sorted(vision_parts) if len(vision_parts[i]) >= 4] + evaluated = [i for i in usable if i < args.samples][: args.samples] + fit = [i for i in usable if i >= args.samples] + print(f"{len(evaluated)} evaluated, {len(fit)} for fitting", flush=True) + + results = {} + for name, parts in (("vision", vision_parts), ("text", text_parts)): + fitted = content_directions([np.stack(parts[i]) for i in fit], args.shrinkage) + centre, basis, values = fitted + width = min(args.keep, basis.shape[1]) + scale = values[:width] ** args.weight_power + results[name] = [ + F.normalize(torch.tensor( + (np.stack(parts[i]) - centre) @ basis[:, :width] * scale, + dtype=torch.float32), dim=-1) + for i in evaluated + ] + + visual_field = standardise(moment_field(results["vision"]).double().numpy()) + text_field = standardise(moment_field(results["text"]).double().numpy()) + mask = ~np.eye(len(evaluated), dtype=bool) + summary = { + "protocol": ( + "DIAGNOSTIC, NOT A METHOD. Each annotated region box is encoded as a " + "crop and paired with its own description, so the two modalities " + "carry the same parts in the same order. That box-to-phrase link is " + "cross-modal supervision a deployed system would not have." + ), + "scenes": len(evaluated), + "mean_parts": float(np.mean([len(vision_parts[i]) for i in evaluated])), + "correlation": float(np.corrcoef(visual_field[mask], text_field[mask])[0, 1]), + "anchor_bound_with_part_oracle": anchor_bound(visual_field, text_field), + "anchor_bound_unsupervised_reference": 0.336, + "within_modality_ceiling_text": 0.805, + "reading": ( + "A large jump means part correspondence is the whole problem and is " + "an engineering target. Staying near 0.336 means the two encoders " + "do not covary at the part level either, and no aggregation " + "rescues this corpus." + ), + } + print(json.dumps(summary, indent=2)) + write_json(args.output, summary) + + +if __name__ == "__main__": + main() |