rpx_benchmark.vqa.hub_rgb

  1"""Fetch one or two members from RPX's uncompressed RGB tar shards via HTTP
  2ranges. Normal (single-image) samples fetch only the target image, exactly
  3as before. In-context (two-image) samples independently fetch the target
  4image AND the reference image, then reconstruct Image 1 by cropping the
  5fetched reference frame with the sample's own reference_crop_bbox and
  6verifying the result against reference_crop_sha256 -- the same deterministic
  7crop+PNG-encode convention data/mask_annotation/visual_grounding_gt/vqa_gt/
  8sos_reference.py used when the hash was originally recorded."""
  9
 10from __future__ import annotations
 11
 12import hashlib
 13import io
 14import json
 15import tarfile
 16import time
 17import urllib.error
 18import urllib.request
 19from collections import defaultdict
 20from pathlib import Path
 21from typing import Any, Iterable
 22
 23from PIL import Image
 24
 25from ..exceptions import DownloadError
 26from .contract import VQASample, normalize_shard
 27
 28
 29def image_cache_name(sample: VQASample) -> str:
 30    identity = json.dumps(sample.image, sort_keys=True, separators=(",", ":"))
 31    digest = hashlib.sha256(identity.encode()).hexdigest()[:20]
 32    return f"rgb_{digest}{Path(sample.image['member']).suffix}"
 33
 34
 35def reference_cache_name(sample: VQASample) -> str:
 36    """Cache key for the reconstructed (cropped) reference image -- distinct
 37    from image_cache_name because the crop_bbox, not just the raw locator,
 38    determines the resulting bytes."""
 39    assert sample.reference_image is not None and sample.reference_crop_bbox is not None
 40    identity = json.dumps(
 41        {"locator": sample.reference_image, "crop_bbox": list(sample.reference_crop_bbox)},
 42        sort_keys=True,
 43        separators=(",", ":"),
 44    )
 45    digest = hashlib.sha256(identity.encode()).hexdigest()[:20]
 46    return f"ref_{digest}.png"
 47
 48
 49class HTTPRangeReader(io.RawIOBase):
 50    """Small seekable reader with block caching; avoids full multi-GB shard downloads."""
 51
 52    def __init__(
 53        self,
 54        url: str,
 55        block_size: int = 8 * 1024 * 1024,
 56        timeout: float = 60.0,
 57        retries: int = 3,
 58    ) -> None:
 59        self.url = url
 60        self.block_size = block_size
 61        self.timeout = timeout
 62        self.retries = retries
 63        self.position = 0
 64        self.length = self._content_length()
 65        self._block_start = -1
 66        self._block = b""
 67
 68    def _content_length(self) -> int:
 69        request = urllib.request.Request(
 70            self.url, headers={"Range": "bytes=0-0", "Accept-Encoding": "identity"}
 71        )
 72        with self._open(request) as response:
 73            content_range = response.headers.get("Content-Range")
 74            if not content_range:
 75                raise OSError("Hub endpoint does not support byte ranges")
 76            return int(content_range.rsplit("/", 1)[1])
 77
 78    def _open(self, request: urllib.request.Request):
 79        """Open one range request with bounded waits and transient retries."""
 80        for attempt in range(self.retries + 1):
 81            try:
 82                return urllib.request.urlopen(request, timeout=self.timeout)
 83            except urllib.error.HTTPError as exc:
 84                # Authentication/not-found/range errors are deterministic. Server
 85                # errors are commonly transient and safe to retry.
 86                if exc.code < 500 or attempt == self.retries:
 87                    raise
 88            except (TimeoutError, urllib.error.URLError):
 89                if attempt == self.retries:
 90                    raise
 91            time.sleep(2**attempt)
 92        raise DownloadError(
 93            f"exhausted HTTP retries for {self.url}",
 94            hint="check Hub connectivity and retry; completed cache files are preserved",
 95        )
 96
 97    def readable(self) -> bool:
 98        return True
 99
100    def seekable(self) -> bool:
101        return True
102
103    def tell(self) -> int:
104        return self.position
105
106    def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
107        if whence == io.SEEK_CUR:
108            offset += self.position
109        elif whence == io.SEEK_END:
110            offset += self.length
111        if offset < 0:
112            raise OSError("negative seek")
113        self.position = offset
114        return offset
115
116    def readinto(self, buffer: Any) -> int:
117        data = self.read(len(buffer))
118        buffer[: len(data)] = data
119        return len(data)
120
121    def read(self, size: int = -1) -> bytes:
122        if self.position >= self.length:
123            return b""
124        if size < 0:
125            size = self.length - self.position
126        output = bytearray()
127        while size and self.position < self.length:
128            block_start = (self.position // self.block_size) * self.block_size
129            if block_start != self._block_start:
130                block_end = min(self.length - 1, block_start + self.block_size - 1)
131                request = urllib.request.Request(
132                    self.url,
133                    headers={
134                        "Range": f"bytes={block_start}-{block_end}",
135                        "Accept-Encoding": "identity",
136                    },
137                )
138                with self._open(request) as response:
139                    self._block = response.read()
140                self._block_start = block_start
141            offset = self.position - self._block_start
142            chunk = self._block[offset : offset + size]
143            if not chunk:
144                break
145            output.extend(chunk)
146            self.position += len(chunk)
147            size -= len(chunk)
148        return bytes(output)
149
150
151def hub_url(image: VQASample | dict[str, str]) -> str:
152    locator = image.image if isinstance(image, VQASample) else image
153    shard = normalize_shard(str(locator["shard"]))
154    return (
155        f"https://huggingface.co/datasets/{locator['repo_id']}/resolve/"
156        f"{locator['revision']}/{shard}"
157    )
158
159
160def _fetch_tar_member(locator: dict[str, str]) -> bytes:
161    reader = HTTPRangeReader(hub_url(locator))
162    with tarfile.open(fileobj=reader, mode="r:") as archive:
163        member = next((entry for entry in archive if entry.name == locator["member"]), None)
164        if member is None:
165            raise FileNotFoundError(locator["member"])
166        extracted = archive.extractfile(member)
167        if extracted is None:
168            raise FileNotFoundError(locator["member"])
169        return extracted.read()
170
171
172def _write_rgb_payload(sample: VQASample, target: Path, payload: bytes) -> None:
173    with Image.open(io.BytesIO(payload)) as image:
174        if image.size != (sample.img_w, sample.img_h):
175            raise DownloadError(
176                f"{sample.sample_id}: image is {image.size}, expected {(sample.img_w, sample.img_h)}",
177                hint="check that the manifest and pinned Hub revision were generated together",
178            )
179        image.verify()
180    target.parent.mkdir(parents=True, exist_ok=True)
181    target.write_bytes(payload)
182
183
184def _write_reference_payload(sample: VQASample, target: Path, payload: bytes) -> None:
185    assert sample.reference_crop_bbox is not None
186    x0, y0, x1, y1 = sample.reference_crop_bbox
187    with Image.open(io.BytesIO(payload)) as verified:
188        verified.verify()
189    with Image.open(io.BytesIO(payload)) as source:
190        raw = source.convert("RGB")
191        if not (0 <= x0 <= x1 < raw.width and 0 <= y0 <= y1 < raw.height):
192            raise DownloadError(
193                f"{sample.sample_id}: reference_crop_bbox {sample.reference_crop_bbox} "
194                f"is outside the fetched reference frame {raw.size}",
195                hint="check that the manifest and pinned Hub revision were generated together",
196            )
197        cropped = raw.crop((x0, y0, x1 + 1, y1 + 1))
198        buffer = io.BytesIO()
199        cropped.save(buffer, format="PNG", compress_level=6)
200        crop_bytes = buffer.getvalue()
201    digest = hashlib.sha256(crop_bytes).hexdigest()
202    if digest != sample.reference_crop_sha256:
203        raise DownloadError(
204            f"{sample.sample_id}: reference crop sha256 mismatch: got {digest}, "
205            f"expected {sample.reference_crop_sha256}",
206            hint=(
207                "the reconstructed Image 1 does not match the hash recorded at generation "
208                "time -- check Pillow version compatibility with sos_reference.py's PNG "
209                "encoder before assuming data corruption"
210            ),
211        )
212    target.parent.mkdir(parents=True, exist_ok=True)
213    target.write_bytes(crop_bytes)
214
215
216def fetch_rgb(sample: VQASample, output_root: str | Path) -> Path:
217    """Fetch, dimension-check, and cache the exact TARGET image used by a VQA
218    row (Image 2 for in-context rows, the only image for normal rows)."""
219    target = Path(output_root) / image_cache_name(sample)
220    if target.is_file():
221        return target
222    target.parent.mkdir(parents=True, exist_ok=True)
223    payload = _fetch_tar_member(sample.image)
224    _write_rgb_payload(sample, target, payload)
225    return target
226
227
228def fetch_reference_crop(sample: VQASample, output_root: str | Path) -> Path:
229    """Fetch the RAW reference (SOS) frame, crop it with the sample's own
230    reference_crop_bbox (inclusive-pixel xyxy, same convention as
231    answer_bbox), and verify the PNG-encoded result against
232    reference_crop_sha256. Raises DownloadError on any hash mismatch --
233    never silently serves an unverified Image 1."""
234    if sample.reference_image is None or sample.reference_crop_bbox is None:
235        raise DownloadError(
236            f"{sample.sample_id}: not an in-context sample", hint="only call this for in-context rows"
237        )
238    target = Path(output_root) / reference_cache_name(sample)
239    if target.is_file():
240        return target
241    target.parent.mkdir(parents=True, exist_ok=True)
242    payload = _fetch_tar_member(sample.reference_image)
243    _write_reference_payload(sample, target, payload)
244    return target
245
246
247def fetch_images(sample: VQASample, output_root: str | Path) -> tuple[Path, ...]:
248    """Fetch every image a model call needs, in the same [reference, target]
249    order as sample.images: independently fetches and verifies each locator
250    (item 3) rather than assuming one implies the other."""
251    if sample.is_in_context:
252        return (fetch_reference_crop(sample, output_root), fetch_rgb(sample, output_root))
253    return (fetch_rgb(sample, output_root),)
254
255
256def fetch_images_many(
257    samples: Iterable[VQASample],
258    output_root: str | Path,
259    *,
260    missing_report: list[dict[str, str]] | None = None,
261) -> dict[str, tuple[Path, ...]]:
262    """Fetch a manifest efficiently by scanning each remote tar shard once.
263
264    The old per-row path reopened a shard and scanned from byte zero for every
265    requested member. That is tolerable for one row but becomes quadratic in
266    shard position for acceptance/benchmark manifests. Here, missing members
267    are grouped by shard and extracted during one forward scan. Existing cache
268    files remain authoritative and are never downloaded again. Strict mode is
269    the default. Supplying ``missing_report`` inventories absent tar members,
270    continues through every remaining shard, and returns only fully cached
271    samples; callers can then create an explicit audited subset manifest.
272    """
273    sample_list = list(samples)
274    root = Path(output_root)
275    # url -> member -> [(kind, sample, destination)]
276    pending: dict[
277        str, dict[str, list[tuple[str, VQASample, Path]]]
278    ] = defaultdict(lambda: defaultdict(list))
279
280    for sample in sample_list:
281        rgb_target = root / image_cache_name(sample)
282        if not rgb_target.is_file():
283            pending[hub_url(sample.image)][str(sample.image["member"])].append(
284                ("rgb", sample, rgb_target)
285            )
286        if sample.is_in_context:
287            assert sample.reference_image is not None
288            ref_target = root / reference_cache_name(sample)
289            if not ref_target.is_file():
290                pending[hub_url(sample.reference_image)][
291                    str(sample.reference_image["member"])
292                ].append(("reference", sample, ref_target))
293
294    total_shards = len(pending)
295    total_members = sum(len(members) for members in pending.values())
296    print(
297        f"image prefetch: {total_members} uncached members across {total_shards} shards",
298        flush=True,
299    )
300    for shard_index, (url, members) in enumerate(pending.items(), 1):
301        print(
302            f"image prefetch shard [{shard_index}/{total_shards}]: "
303            f"{len(members)} members {url}",
304            flush=True,
305        )
306        remaining = dict(members)
307        reader = HTTPRangeReader(url)
308        with tarfile.open(fileobj=reader, mode="r:") as archive:
309            for entry in archive:
310                consumers = remaining.pop(entry.name, None)
311                if consumers is None:
312                    continue
313                extracted = archive.extractfile(entry)
314                if extracted is None:
315                    raise FileNotFoundError(entry.name)
316                payload = extracted.read()
317                for kind, sample, target in consumers:
318                    if kind == "rgb":
319                        _write_rgb_payload(sample, target, payload)
320                    else:
321                        _write_reference_payload(sample, target, payload)
322                if not remaining:
323                    break
324        if remaining:
325            if missing_report is None:
326                raise FileNotFoundError(
327                    f"members missing from {url}: {sorted(remaining)[:10]}"
328                )
329            for member, consumers in sorted(remaining.items()):
330                for kind, sample, _target in consumers:
331                    missing_report.append(
332                        {
333                            "sample_id": sample.sample_id,
334                            "role": kind,
335                            "url": url,
336                            "member": member,
337                        }
338                    )
339
340    blocked = {
341        entry["sample_id"] for entry in (missing_report or [])
342    }
343    return {
344        sample.sample_id: fetch_images(sample, root)
345        for sample in sample_list
346        if sample.sample_id not in blocked
347    }