File size: 4,322 Bytes
ebeed67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """COCO replay: public images mixed into a detector fine-tune.
A fine-tune on a few thousand frames of ~10 cameras forgets what the COCO head
knew about everything those cameras never show. Mixing a fixed slice of COCO
val2017 back in keeps the 80 classes alive. The slice is pinned: the images of
val2017 that hold at least one person / vehicle / animal, in file-name order,
the first ``count`` of them — the same images on every run.
Nothing here is uploaded by the hub: the container fetches the public files
once into the provider's cache volume (``/cache/coco``) and reuses them. All
80 COCO classes are labelled (the head is kept whole).
``fetch`` is injectable: tests never touch the network.
"""
from __future__ import annotations
import io
import json
import urllib.request
import zipfile
from pathlib import Path
from typing import Callable
from ._vendor.camstack_eval.classes import DETECTOR_CLASSES, macro_of
ANNOTATIONS_URL = "http://images.cocodataset.org/annotations/annotations_trainval2017.zip"
ANNOTATIONS_MEMBER = "annotations/instances_val2017.json"
IMAGE_URL = "http://images.cocodataset.org/val2017/{file}"
Fetch = Callable[[str], bytes]
def http_fetch(url: str) -> bytes:
with urllib.request.urlopen(url, timeout=120) as response: # noqa: S310 (fixed public URLs)
return response.read()
def load_annotations(cache: Path, fetch: Fetch) -> dict:
"""``instances_val2017.json``, fetched once into ``cache``."""
target = cache / "instances_val2017.json"
if not target.is_file():
cache.mkdir(parents=True, exist_ok=True)
archive = zipfile.ZipFile(io.BytesIO(fetch(ANNOTATIONS_URL)))
target.write_bytes(archive.read(ANNOTATIONS_MEMBER))
return json.loads(target.read_text("utf-8"))
def select_images(annotations: dict, count: int) -> list[dict]:
"""The pinned slice: images with a detector-class object, file-name order."""
names = {c["id"]: c["name"] for c in annotations["categories"]}
wanted: set[int] = set()
for a in annotations["annotations"]:
if not a.get("iscrowd") and macro_of(names[a["category_id"]]) in DETECTOR_CLASSES:
wanted.add(a["image_id"])
images = sorted((im for im in annotations["images"] if im["id"] in wanted), key=lambda im: im["file_name"])
return images[:count]
def yolo_lines(annotations: dict, image: dict, coco80: list[str]) -> list[str]:
names = {c["id"]: c["name"] for c in annotations["categories"]}
w, h = float(image["width"]), float(image["height"])
lines = []
for a in annotations["annotations"]:
if a["image_id"] != image["id"] or a.get("iscrowd"):
continue
name = names[a["category_id"]]
if name not in coco80:
continue
x, y, bw, bh = a["bbox"]
if bw <= 0 or bh <= 0:
continue
lines.append(
f"{coco80.index(name)} {(x + bw / 2) / w:.6f} {(y + bh / 2) / h:.6f} {bw / w:.6f} {bh / h:.6f}"
)
return lines
def prepare_replay(count: int, cache: Path, out_root: Path, coco80: list[str], fetch: Fetch = http_fetch) -> int:
"""Write ``count`` replay images + labels under ``out_root/images|labels/replay``."""
if count <= 0:
return 0
annotations = load_annotations(cache, fetch)
images_dir = out_root / "images" / "replay"
labels_dir = out_root / "labels" / "replay"
images_dir.mkdir(parents=True, exist_ok=True)
labels_dir.mkdir(parents=True, exist_ok=True)
by_image: dict[int, list[dict]] = {}
for a in annotations["annotations"]:
by_image.setdefault(a["image_id"], []).append(a)
indexed = {**annotations, "annotations": []}
written = 0
for image in select_images(annotations, count):
cached = cache / "val2017" / image["file_name"]
if not cached.is_file():
cached.parent.mkdir(parents=True, exist_ok=True)
cached.write_bytes(fetch(IMAGE_URL.format(file=image["file_name"])))
stem = Path(image["file_name"]).stem
(images_dir / f"{stem}.jpg").write_bytes(cached.read_bytes())
lines = yolo_lines({**indexed, "annotations": by_image.get(image["id"], [])}, image, coco80)
(labels_dir / f"{stem}.txt").write_text("".join(line + "\n" for line in lines), "utf-8")
written += 1
return written
|