File size: 2,934 Bytes
ed46c5d | 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 | """Crop photographs to the zero-shot inscription box (data/ground_4b.jsonl) with a 20% margin and
rescale each crop to about 0.64 MP, so that runes occupy more of the recognizer's input.
Boxes are re-parsed from the raw model output (the "box" field of runs before the parse_box fix is wrong).
Images without a usable box are kept whole. Writes crops/ and splits/ocr_crop_{train,val,test}.jsonl.
Runs as the first step of the ocr4b_crop job (exp_v4.sh)."""
import json
import os
import re
from PIL import Image
R = "/home/jovyan/humonen/runic_access"
TARGET = 640_000
MARGIN = 0.2
def parse_box(text, w, h): # same as ground.parse_box, without importing torch
m = re.search(r"\[([^\[\]]*)\]", text)
nums = re.findall(r"-?\d+(?:\.\d+)?", m.group(1) if m else "")
if len(nums) < 4:
return None
x1, y1, x2, y2 = [float(v) for v in nums[:4]]
sx, sy = (w, h) if max(x1, y1, x2, y2) <= 1.0 else (w / 1000.0, h / 1000.0)
x1, x2 = sorted((x1 * sx, x2 * sx))
y1, y2 = sorted((y1 * sy, y2 * sy))
x1, y1, x2, y2 = max(0, x1), max(0, y1), min(w, x2), min(h, y2)
if x2 - x1 < 0.05 * w or y2 - y1 < 0.05 * h:
return None
return [x1, y1, x2, y2]
def crop_one(src, raw_text, dst):
"""Crop src to the parsed box with a margin, rescale to ~TARGET pixels, save to dst; return the crop box or None."""
im = Image.open(src).convert("RGB")
box = parse_box(raw_text, im.width, im.height)
if box:
x1, y1, x2, y2 = box
mx, my = MARGIN * (x2 - x1), MARGIN * (y2 - y1)
box = [max(0, x1 - mx), max(0, y1 - my), min(im.width, x2 + mx), min(im.height, y2 + my)]
im = im.crop(tuple(round(v) for v in box))
s = (TARGET / (im.width * im.height)) ** 0.5
im = im.resize((max(28, round(im.width * s)), max(28, round(im.height * s))), Image.BICUBIC)
im.save(dst + ".tmp.jpg", quality=92)
os.replace(dst + ".tmp.jpg", dst)
return box
if __name__ == "__main__":
raw = {}
for l in open(f"{R}/data/ground_4b.jsonl"):
r = json.loads(l)
raw[r["image"]] = r["raw"]
os.makedirs(f"{R}/crops", exist_ok=True)
stats = {"cropped": 0, "whole": 0}
boxes_out = open(f"{R}/data/crop_boxes.jsonl", "w")
for split in ("train", "val", "test"):
with open(f"{R}/splits/ocr_crop_{split}.jsonl", "w") as f:
for l in open(f"{R}/splits/ocr_real_{split}.jsonl"):
r = json.loads(l)
dst_rel = "crops/" + os.path.basename(r["image"])
box = crop_one(f"{R}/{r['image']}", raw.get(r["image"], ""), f"{R}/{dst_rel}")
stats["cropped" if box else "whole"] += 1
boxes_out.write(json.dumps({"image": r["image"], "crop_box": box}) + "\n")
f.write(json.dumps(dict(r, image=dst_rel, image_full=r["image"]), ensure_ascii=False) + "\n")
print(split, stats, flush=True)
boxes_out.close()
print("DONE", stats)
|