rpx_benchmark.tasks.vqa_generation

  1"""RPX VQA ground-truth generation.
  2
  3Generates deterministic Q/A pairs from annotations. No human labeling needed.
  4No LLM. No fuzzy matching. MCQ with 4 options, random baseline = 25%.
  5
  6Design principles (Karpathy/Girshick/Fox/He):
  7  - Every question has ONE correct answer derivable from masks + questionnaire.
  8  - Every question maps to a robot operation.
  9  - Phase variation is the anti-bias mechanism — same question, different answer.
 10  - MCQ distractors come from real attributes of real objects in the dataset.
 11  - Option order is shuffled per question with a fixed seed for reproducibility.
 12
 13Question types:
 14  Q1. Existence   — "Is there a red mug?"        → from visibility mask
 15  Q2. Attribute   — "What color is the mug?"      → from questionnaire
 16  Q3. Count       — "How many cups are visible?"   → from mask instance count
 17  Q4. Spatial     — "What is left of the bowl?"    → from bbox centroids (image plane)
 18
 19Evaluation: MCQ accuracy. Per-phase breakdown → feeds STR.
 20"""
 21
 22from __future__ import annotations
 23
 24import hashlib
 25import random
 26from dataclasses import dataclass, field
 27from typing import Dict, List, Tuple
 28
 29# ── Question templates ──────────────────────────────────────────────────────
 30# Each template maps to a robot operation. Multiple phrasings per type
 31# to reduce template memorization, but all are simple imperative/interrogative.
 32
 33EXISTENCE_TEMPLATES = [
 34    "Is there a {attr} {category} in this scene?",
 35    "Can you see a {attr} {category}?",
 36    "Is a {attr} {category} present?",
 37]
 38
 39ATTRIBUTE_TEMPLATES = {
 40    "color": [
 41        "What color is the {category}?",
 42        "What is the color of the {category}?",
 43    ],
 44    "material": [
 45        "What material is the {category} made of?",
 46        "What is the material of the {category}?",
 47    ],
 48    "shape": [
 49        "What shape is the {category}?",
 50        "What is the shape of the {category}?",
 51    ],
 52}
 53
 54COUNT_TEMPLATES = [
 55    "How many {category} objects are visible?",
 56    "How many {category} can you see?",
 57    "Count the {category} in this scene.",
 58]
 59
 60SPATIAL_TEMPLATES = [
 61    "What object is to the {direction} of the {object}?",
 62    "What is {direction} of the {object}?",
 63]
 64
 65SPATIAL_DIRECTIONS = ["left", "right", "above", "below"]
 66
 67
 68@dataclass
 69class QAPair:
 70    """One question-answer pair with MCQ options."""
 71
 72    scene_id: str
 73    phase: int
 74    frame_id: str
 75    question_type: str  # existence | attribute | count | spatial
 76    question: str
 77    correct_answer: str
 78    options: List[str]  # 4 options, correct answer included, shuffled
 79    correct_index: int  # index of correct answer in options
 80    metadata: Dict = field(default_factory=dict)
 81
 82    def to_dict(self) -> dict:
 83        return {
 84            "scene_id": self.scene_id,
 85            "phase": self.phase,
 86            "frame_id": self.frame_id,
 87            "question_type": self.question_type,
 88            "question": self.question,
 89            "correct_answer": self.correct_answer,
 90            "options": self.options,
 91            "correct_index": self.correct_index,
 92            "metadata": self.metadata,
 93        }
 94
 95
 96def _deterministic_seed(scene_id: str, phase: int, frame_id: str, q: str) -> int:
 97    """Reproducible seed per question for option shuffling."""
 98    h = hashlib.md5(f"{scene_id}/{phase}/{frame_id}/{q}".encode()).hexdigest()
 99    return int(h[:8], 16)
100
101
102def _shuffle_options(correct: str, distractors: List[str], seed: int) -> Tuple[List[str], int]:
103    """Shuffle 4 options with deterministic seed. Return (options, correct_index)."""
104    opts = [correct] + distractors[:3]
105    rng = random.Random(seed)
106    rng.shuffle(opts)
107    return opts, opts.index(correct)
108
109
110def _pick_template(templates: list, seed: int) -> str:
111    """Pick a template deterministically."""
112    return templates[seed % len(templates)]
113
114
115# ── Generators ──────────────────────────────────────────────────────────────
116
117
118def generate_existence_questions(
119    visible_objects: List[Dict],  # [{"category": "mug", "color": "red", ...}, ...]
120    all_objects_in_dataset: List[Dict],  # for generating "no" distractors
121    scene_id: str,
122    phase: int,
123    frame_id: str,
124) -> List[QAPair]:
125    """Generate balanced existence questions: 50% yes, 50% no.
126
127    'Yes' questions ask about objects visible in this frame.
128    'No' questions ask about objects NOT visible (from other scenes in the dataset).
129    This prevents a model from scoring well by always answering 'yes'.
130    """
131    pairs = []
132    visible_set = {(o["category"], o.get("color", "")) for o in visible_objects}
133
134    # Yes questions — pick up to 3 visible objects
135    for obj in visible_objects[:3]:
136        attr = obj.get("color", obj.get("material", ""))
137        if not attr:
138            continue
139        seed = _deterministic_seed(scene_id, phase, frame_id, f"exist_yes_{obj['category']}")
140        q = _pick_template(EXISTENCE_TEMPLATES, seed).format(attr=attr, category=obj["category"])
141        opts, idx = _shuffle_options("yes", ["no", "not sure", "partially visible"], seed)
142        pairs.append(
143            QAPair(
144                scene_id=scene_id,
145                phase=phase,
146                frame_id=frame_id,
147                question_type="existence",
148                question=q,
149                correct_answer="yes",
150                options=opts,
151                correct_index=idx,
152                metadata={"object": obj["category"], "expected": True},
153            )
154        )
155
156    # No questions — ask about objects absent from this frame
157    absent = [
158        o for o in all_objects_in_dataset if (o["category"], o.get("color", "")) not in visible_set
159    ]
160    for obj in absent[: len(pairs)]:  # match count for 50/50 balance
161        attr = obj.get("color", obj.get("material", ""))
162        if not attr:
163            continue
164        seed = _deterministic_seed(scene_id, phase, frame_id, f"exist_no_{obj['category']}")
165        q = _pick_template(EXISTENCE_TEMPLATES, seed).format(attr=attr, category=obj["category"])
166        opts, idx = _shuffle_options("no", ["yes", "not sure", "partially visible"], seed)
167        pairs.append(
168            QAPair(
169                scene_id=scene_id,
170                phase=phase,
171                frame_id=frame_id,
172                question_type="existence",
173                question=q,
174                correct_answer="no",
175                options=opts,
176                correct_index=idx,
177                metadata={"object": obj["category"], "expected": False},
178            )
179        )
180
181    return pairs
182
183
184def generate_attribute_questions(
185    visible_objects: List[Dict],
186    all_objects_in_dataset: List[Dict],  # for distractors
187    scene_id: str,
188    phase: int,
189    frame_id: str,
190) -> List[QAPair]:
191    """Generate attribute questions with distractors from real objects."""
192    pairs = []
193    # Collect all attribute values by type for distractor pool
194    attr_pool: Dict[str, List[str]] = {"color": [], "material": [], "shape": []}
195    for o in all_objects_in_dataset:
196        for attr_type in attr_pool:
197            v = o.get(attr_type, "")
198            if v and v not in attr_pool[attr_type]:
199                attr_pool[attr_type].append(v)
200
201    for obj in visible_objects[:5]:
202        for attr_type, templates in ATTRIBUTE_TEMPLATES.items():
203            gt_val = obj.get(attr_type, "")
204            if not gt_val:
205                continue
206            # Get 3 distractors from real attribute values
207            pool = [v for v in attr_pool.get(attr_type, []) if v != gt_val]
208            if len(pool) < 3:
209                continue
210
211            seed = _deterministic_seed(
212                scene_id, phase, frame_id, f"attr_{attr_type}_{obj['category']}"
213            )
214            rng = random.Random(seed)
215            distractors = rng.sample(pool, 3)
216            q = _pick_template(templates, seed).format(category=obj["category"])
217            opts, idx = _shuffle_options(gt_val, distractors, seed)
218            pairs.append(
219                QAPair(
220                    scene_id=scene_id,
221                    phase=phase,
222                    frame_id=frame_id,
223                    question_type="attribute",
224                    question=q,
225                    correct_answer=gt_val,
226                    options=opts,
227                    correct_index=idx,
228                    metadata={"object": obj["category"], "attribute": attr_type},
229                )
230            )
231
232    return pairs
233
234
235def generate_count_questions(
236    visible_objects: List[Dict],
237    scene_id: str,
238    phase: int,
239    frame_id: str,
240) -> List[QAPair]:
241    """Generate counting questions. GT from mask instance count."""
242    pairs = []
243    # Count by category
244    counts: Dict[str, int] = {}
245    for obj in visible_objects:
246        counts[obj["category"]] = counts.get(obj["category"], 0) + 1
247
248    # Ask about categories with ≥1 instance
249    for cat, count in counts.items():
250        seed = _deterministic_seed(scene_id, phase, frame_id, f"count_{cat}")
251        q = _pick_template(COUNT_TEMPLATES, seed).format(category=cat)
252        gt = str(count)
253        # Distractors: nearby integers, always including 0
254        distractor_pool = [str(i) for i in range(0, max(count + 4, 6)) if str(i) != gt]
255        rng = random.Random(seed)
256        distractors = rng.sample(distractor_pool, min(3, len(distractor_pool)))
257        if len(distractors) < 3:
258            continue
259        opts, idx = _shuffle_options(gt, distractors, seed)
260        pairs.append(
261            QAPair(
262                scene_id=scene_id,
263                phase=phase,
264                frame_id=frame_id,
265                question_type="count",
266                question=q,
267                correct_answer=gt,
268                options=opts,
269                correct_index=idx,
270                metadata={"category": cat, "count": count},
271            )
272        )
273
274    return pairs
275
276
277def generate_spatial_questions(
278    visible_objects: List[Dict],  # must include "bbox_cx", "bbox_cy"
279    scene_id: str,
280    phase: int,
281    frame_id: str,
282) -> List[QAPair]:
283    """Generate spatial questions. Relations defined in image plane.
284
285    "Left" = smaller x-coordinate. "Above" = smaller y-coordinate.
286    This matches how VLMs process images and avoids viewpoint ambiguity.
287    """
288    pairs = []
289    if len(visible_objects) < 2:
290        return pairs
291
292    for ref_obj in visible_objects[:3]:
293        ref_cx = ref_obj.get("bbox_cx")
294        ref_cy = ref_obj.get("bbox_cy")
295        if ref_cx is None or ref_cy is None:
296            continue
297
298        for direction in SPATIAL_DIRECTIONS:
299            # Find the nearest object in that direction
300            best = None
301            best_dist = float("inf")
302            for other in visible_objects:
303                if other["category"] == ref_obj["category"] and other.get(
304                    "instance_id"
305                ) == ref_obj.get("instance_id"):
306                    continue
307                ox, oy = other.get("bbox_cx"), other.get("bbox_cy")
308                if ox is None or oy is None:
309                    continue
310
311                # Check direction
312                if direction == "left" and ox >= ref_cx:
313                    continue
314                if direction == "right" and ox <= ref_cx:
315                    continue
316                if direction == "above" and oy >= ref_cy:
317                    continue
318                if direction == "below" and oy <= ref_cy:
319                    continue
320
321                dist = abs(ox - ref_cx) + abs(oy - ref_cy)
322                if dist < best_dist:
323                    best_dist = dist
324                    best = other
325
326            if best is None:
327                continue
328
329            seed = _deterministic_seed(
330                scene_id, phase, frame_id, f"spatial_{direction}_{ref_obj['category']}"
331            )
332            q = _pick_template(SPATIAL_TEMPLATES, seed).format(
333                direction=direction, object=ref_obj["category"]
334            )
335            gt = best["category"]
336
337            # Distractors: other visible objects
338            others = [
339                o["category"]
340                for o in visible_objects
341                if o["category"] != gt and o["category"] != ref_obj["category"]
342            ]
343            others = list(set(others))
344            if len(others) < 2:
345                others += ["nothing", "unknown"]
346            rng = random.Random(seed)
347            distractors = rng.sample(others, min(3, len(others)))
348            while len(distractors) < 3:
349                distractors.append("nothing")
350            opts, idx = _shuffle_options(gt, distractors[:3], seed)
351
352            pairs.append(
353                QAPair(
354                    scene_id=scene_id,
355                    phase=phase,
356                    frame_id=frame_id,
357                    question_type="spatial",
358                    question=q,
359                    correct_answer=gt,
360                    options=opts,
361                    correct_index=idx,
362                    metadata={
363                        "reference": ref_obj["category"],
364                        "direction": direction,
365                        "target": best["category"],
366                    },
367                )
368            )
369            break  # one spatial question per reference object per direction
370
371    return pairs
372
373
374# ── Evaluation ──────────────────────────────────────────────────────────────
375
376
377def evaluate_mcq(predictions: List[int], ground_truth: List[QAPair]) -> Dict:
378    """Evaluate MCQ predictions. Returns accuracy overall and per type/phase.
379
380    Args:
381        predictions: list of predicted option indices (0-3)
382        ground_truth: list of QAPair objects
383
384    Returns:
385        dict with overall accuracy + per-type + per-phase breakdowns
386    """
387    assert len(predictions) == len(ground_truth)
388
389    correct = 0
390    by_type: Dict[str, List[bool]] = {}
391    by_phase: Dict[int, List[bool]] = {}
392
393    for pred, qa in zip(predictions, ground_truth, strict=False):
394        hit = pred == qa.correct_index
395        correct += hit
396
397        by_type.setdefault(qa.question_type, []).append(hit)
398        by_phase.setdefault(qa.phase, []).append(hit)
399
400    n = len(predictions)
401    result = {
402        "accuracy": correct / n if n > 0 else 0.0,
403        "n_questions": n,
404        "per_type": {t: sum(v) / len(v) for t, v in by_type.items()},
405        "per_phase": {p: sum(v) / len(v) for p, v in by_phase.items()},
406    }
407
408    # STR for VQA = phase accuracy difference
409    phases = sorted(by_phase.keys())
410    if len(phases) >= 2:
411        phase_acc = {p: sum(v) / len(v) for p, v in by_phase.items()}
412        if 0 in phase_acc and 1 in phase_acc:
413            result["str_clu_to_int"] = phase_acc[1] - phase_acc[0]
414        if 1 in phase_acc and 2 in phase_acc:
415            result["str_int_to_cln"] = phase_acc[2] - phase_acc[1]
416
417    return result