"""Per-baseline sample loading for the baseline-intervention sweep. Both datasets answer the same question at every baseline; only the image changes. A `Sample` therefore carries one question/answer plus one image path per baseline, and `Task` bundles the per-dataset prompt suffix, the candidate answer set (used by likelihood scoring) and the answer normaliser. VSR Questions come from V7_testing.json. Baseline 0.00 is the untouched original image; the other baselines are the GenStereo views recorded in each sample's "stereo" block. Youtube_self_depth_QA Levels 1-3, each a jsonl of QA over a rendered visualisation. Every sample already stores one visualisation per baseline ("visualizations"), rendered with the annotation points warped into that view, so baseline 0.00 is the original view with the same annotation. """ import itertools import json import os import re from .baselines import BASELINES, by_baseline, canonical REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) VSR_TEST_JSON = os.path.join(REPO_ROOT, "02_data/VSR/V7_testing.json") VSR_ORIGINAL_DIR = os.path.join(REPO_ROOT, "00_data/gen_output_VSR/original_images") VSR_GENERATED_ROOT = os.path.join(REPO_ROOT, "00_data/gen_output_VSR/generate_image") YOUTUBE_ROOT = os.path.join(REPO_ROOT, "02_data/Youtube_self_depth_QA") IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".webp", ".bmp") LETTERS = ["A", "B", "C", "D"] ORDERINGS = [" > ".join(p) for p in itertools.permutations(LETTERS)] class Sample: """One question, one answer, one image per baseline.""" __slots__ = ("sample_id", "question", "answer", "images", "meta") def __init__(self, sample_id, question, answer, images, meta=None): self.sample_id = sample_id self.question = question self.answer = answer self.images = images # canonical baseline -> absolute image path self.meta = meta or {} def image(self, baseline): try: return self.images[canonical(baseline)] except KeyError: raise KeyError( f"Sample {self.sample_id} has no view for baseline {baseline}; " f"available: {sorted(self.images)}" ) from None class Task: """A dataset split plus everything needed to prompt and score it.""" def __init__(self, name, samples, answer_instruction, choices, normalize): self.name = name self.samples = samples self.answer_instruction = answer_instruction self.choices = choices self._normalize = normalize def normalize(self, text, sample=None): """Parse a model answer (or a reference) into a canonical form. `sample` is optional: the level-2 multiple-choice normaliser uses it to recognise a model that answered with the ordering instead of the option letter, and every other normaliser ignores it. """ return self._normalize(text, sample) def prompt(self, sample): if self.answer_instruction: return f"{sample.question}\n{self.answer_instruction}" return sample.question # -------------------------------------------------------------------------- # answer normalisation # -------------------------------------------------------------------------- def normalize_true_false(text, sample=None): """First standalone true/false in the text, lowercased.""" if not text: return "" match = re.search(r"\b(true|false)\b", text.lower()) return match.group(1) if match else text.strip().lower() def normalize_letter(text, sample=None): """First standalone A-D in the text, uppercased.""" if not text: return "" match = re.search(r"\b([ABCD])\b", text.upper()) return match.group(1) if match else text.strip().upper() def normalize_ordering(text, sample=None): """The first four distinct A-D letters, rendered as "B > C > A > D". Accepts any separator the model happens to use (">", ",", "then", ...) and tolerates a preamble, so a correct ranking is not scored wrong purely on formatting. """ if not text: return "" letters = [] for letter in re.findall(r"\b([ABCD])\b", text.upper()): if letter not in letters: letters.append(letter) if len(letters) == 4: break if len(letters) != 4: return text.strip().upper() return " > ".join(letters) ORDERING_AFTER_LETTER = re.compile(r"\b([ABCD])\b\s*>") def normalize_choice_letter(text, sample=None): """Option letter for a multiple-choice answer. Level 2 labels its four OPTIONS A-D while the image labels its four POINTS A-D too, so a model may answer either with the option letter ("C") or by writing out the ordering that option holds ("B > C > A > D"). Only the former is the expected format, but the latter is unambiguous and is mapped back through the sample's own option table rather than scored wrong. The two are told apart by what follows the first standalone letter: a ">" means the letter opens an ordering, anything else ("C", "(C)", "C) B > C > A > D", "Answer: C") means it is the option label. """ if not text: return "" match = re.search(r"\b([ABCD])\b", text.upper()) options = (sample.meta.get("options") if sample is not None else None) or {} if match and options and ORDERING_AFTER_LETTER.match(text.upper(), match.start()): ordering = normalize_ordering(text) for label, option in options.items(): if option == ordering: return label return normalize_letter(text) # -------------------------------------------------------------------------- # VSR # -------------------------------------------------------------------------- def _load_records(path): """Read either a JSON list or a jsonl file.""" with open(path, "r", encoding="utf-8") as f: if path.endswith(".jsonl"): return [json.loads(line) for line in f if line.strip()] return json.load(f) def _resolve_image(directory, stem): for ext in IMAGE_EXTENSIONS: candidate = os.path.join(directory, f"{stem}{ext}") if os.path.isfile(candidate): return candidate return os.path.join(directory, f"{stem}.jpg") def load_vsr(test_json=VSR_TEST_JSON, original_dir=VSR_ORIGINAL_DIR, generated_root=VSR_GENERATED_ROOT, answer_instruction=None): samples = [] for record in _load_records(test_json): sample_id = str(record["id"]) question = record["conversations"][0]["value"].replace("", "").strip() answer = next( (turn.get("value", "") for turn in record["conversations"] if turn.get("from") == "gpt"), "", ) images = {"0.00": _resolve_image(original_dir, sample_id)} for baseline, entry in by_baseline(record.get("stereo", {})).items(): images[baseline] = os.path.join(generated_root, entry["image"]) samples.append(Sample(sample_id, question, answer, images)) if answer_instruction is None: # The VSR question text already ends with its own instruction. answer_instruction = "" return Task("vsr", samples, answer_instruction, ["True", "False"], normalize_true_false) # -------------------------------------------------------------------------- # Youtube_self_depth_QA # -------------------------------------------------------------------------- LETTER_INSTRUCTION = "Answer with a single letter: A, B, C, or D." def already_instructed(records): """True when the annotations write their own answer instruction into the question. Which levels do this has changed with every regeneration of the data (level 2 started doing it when it became multiple choice, level 3 when it was rebuilt), so it is detected per file instead of hard-coded per level. Appending a second copy is not harmless: the duplicated sentence is exactly the kind of prompt noise that shifts a borderline answer, and it would differ between splits. """ return all("Answer with" in record["question"] for record in records) def load_youtube(level, root=YOUTUBE_ROOT, split="test", answer_instruction=None): """One Youtube level. Levels 1 and 3 ask for one of the four labelled points/trajectories, so the answer space is A-D and the prompt gets a format instruction appended. Level 2 comes in two formats. In the current one each record carries an `options` table and the question already embeds the four candidate orderings as (A)-(D) plus its own answer instruction, which makes it an A-D single choice like the other two levels. The older format asked the model to write the ordering out, so the answer space was the 24 permutations. Which one is in the file is detected from the records themselves rather than assumed, so the older annotations (and the `.ordering_bak` copies) still load correctly. """ level = int(level) if level not in (1, 2, 3): raise ValueError(f"Youtube_self_depth_QA has levels 1-3, got {level}") path = os.path.join(root, f"level{level}", f"{split}.jsonl") records = _load_records(path) multiple_choice = level == 2 and all("options" in record for record in records) if level == 2 and not multiple_choice: # Legacy level 2: the model wrote the ordering out, so the answer space # is the 24 permutations rather than A-D. choices, normalize = ORDERINGS, normalize_ordering default_instruction = "Answer with the ranking only." else: choices = LETTERS normalize = normalize_choice_letter if multiple_choice else normalize_letter default_instruction = "" if already_instructed(records) else LETTER_INSTRUCTION samples = [] for record in records: images = { baseline: os.path.join(root, relative) for baseline, relative in by_baseline(record["visualizations"]).items() } meta = {"task_type": record.get("task_type"), "level": record.get("level")} if record.get("options"): meta["options"] = record["options"] meta["answer_type"] = record.get("answer_type") samples.append( Sample(record["sample_id"], record["question"], record["answer"], images, meta) ) if answer_instruction is None: answer_instruction = default_instruction return Task(f"youtube_level{level}", samples, answer_instruction, choices, normalize) def load_task(dataset, level=None, **kwargs): """dataset in {vsr, youtube}; `level` selects the Youtube split.""" if dataset == "vsr": return load_vsr(**kwargs) if dataset == "youtube": if level is None: raise ValueError("--level is required for the youtube dataset") return load_youtube(level, **kwargs) raise ValueError(f"Unknown dataset {dataset!r}; expected 'vsr' or 'youtube'") def missing_images(task, baseline): return [s.sample_id for s in task.samples if not os.path.isfile(s.image(baseline))] __all__ = [ "BASELINES", "ORDERINGS", "Sample", "Task", "canonical", "load_task", "load_vsr", "load_youtube", "missing_images", "normalize_choice_letter", "normalize_letter", "normalize_ordering", "normalize_true_false", ]