File size: 2,373 Bytes
9fe3f84 | 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 | from __future__ import annotations
import json
import random
from dataclasses import dataclass
from pathlib import Path
WORDS = (
"alpha",
"bravo",
"cedar",
"delta",
"ember",
"falcon",
"green",
"harbor",
"indigo",
"jasper",
)
@dataclass(frozen=True)
class SpatialOcrExample:
id: str
image: str
question: str
answer: str
labels: dict[str, str]
def generate_spatial_ocr_manifest(output_dir: Path, count: int, seed: int = 7) -> tuple[SpatialOcrExample, ...]:
if count <= 0:
raise ValueError("count must be positive")
output_dir.mkdir(parents=True, exist_ok=True)
images_dir = output_dir / "images"
images_dir.mkdir(exist_ok=True)
rng = random.Random(seed)
examples = tuple(_example(index, rng) for index in range(count))
for example in examples:
_write_image(images_dir / Path(example.image).name, example)
manifest = output_dir / "manifest.jsonl"
with manifest.open("w", encoding="utf-8") as handle:
for item in examples:
handle.write(json.dumps(item.__dict__, sort_keys=True) + "\n")
return examples
def _example(index: int, rng: random.Random) -> SpatialOcrExample:
regions = ("top left", "top right", "bottom left", "bottom right")
words = rng.sample(WORDS, 4)
selected = rng.randrange(4)
labels = dict(zip(regions, words))
return SpatialOcrExample(
id=f"spatial_ocr_{index:04d}",
image=f"images/spatial_ocr_{index:04d}.png",
question=f"What word is in the {regions[selected]}?",
answer=labels[regions[selected]],
labels=labels,
)
def _write_image(path: Path, example: SpatialOcrExample) -> None:
try:
from PIL import Image, ImageDraw
except ImportError as exc:
raise RuntimeError("Install data dependencies with `python -m pip install -e '.[data]'`.") from exc
image = Image.new("RGB", (640, 480), "white")
draw = ImageDraw.Draw(image)
boxes = {
"top left": (24, 24, 296, 216),
"top right": (344, 24, 616, 216),
"bottom left": (24, 264, 296, 456),
"bottom right": (344, 264, 616, 456),
}
for region, box in boxes.items():
draw.rectangle(box, outline="black", width=3)
draw.text((box[0] + 24, box[1] + 72), example.labels[region], fill="black")
image.save(path)
|