File size: 3,699 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Zero-shot localization of the runic inscription with Qwen3.5 (bounding box in 0-1000 coordinates).
Writes one JSON line per image: {"image", "raw", "box"} where box is [x1, y1, x2, y2] in pixels of
the stored image, or null when no box could be parsed. Resumable.
"""
import argparse
import json
import os
import re
import sys
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from qwen_common import clean_generation, load_model, stop_ids # noqa: E402
from PIL import Image
PROMPT = ("Locate the runic inscription in this image: the area that contains the carved or painted runes, "
"including all lines of runes. Output only its bounding box as JSON: "
"{\"bbox_2d\": [x1, y1, x2, y2]} with coordinates between 0 and 1000.")
def parse_box(text, w, h):
m = re.search(r"\[([^\[\]]*)\]", text) # first coordinate list; the "2" in the key "bbox_2d" is not a coordinate
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]]
if max(x1, y1, x2, y2) <= 1.0: # relative 0-1
scale_x, scale_y = w, h
else:
scale_x, scale_y = w / 1000.0, h / 1000.0
x1, x2 = sorted((x1 * scale_x, x2 * scale_x))
y1, y2 = sorted((y1 * scale_y, y2 * scale_y))
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(w, x2), min(h, y2)
if x2 - x1 < 0.05 * w or y2 - y1 < 0.05 * h:
return None
return [round(x1), round(y1), round(x2), round(y2)]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--inputs", required=True)
ap.add_argument("--image_root", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--bs", type=int, default=16)
args = ap.parse_args()
images = []
for p in args.inputs.split(","):
for l in open(p):
im = json.loads(l)["image"]
if im not in images:
images.append(im)
done = set()
if os.path.exists(args.out):
done = {json.loads(l)["image"] for l in open(args.out)}
todo = [i for i in images if i not in done]
print("images", len(images), "todo", len(todo), flush=True)
processor, model = load_model(args.model)
model.to("cuda").eval()
tok = processor.tokenizer
tok.padding_side = "left"
msgs = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": PROMPT}]}]
prompt = processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)
fout = open(args.out, "a")
for b in range(0, len(todo), args.bs):
chunk = todo[b:b + args.bs]
ims = []
for name in chunk:
im = Image.open(os.path.join(args.image_root, name)).convert("RGB")
ims.append(im)
enc = processor(text=[prompt] * len(ims), images=ims, return_tensors="pt", padding=True)
enc = {k: v.to("cuda") for k, v in enc.items()}
with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
out = model.generate(**enc, max_new_tokens=40, do_sample=False, eos_token_id=stop_ids(tok),
pad_token_id=tok.pad_token_id)
gen = out[:, enc["input_ids"].shape[1]:]
for name, im, g in zip(chunk, ims, gen):
raw = clean_generation(tok.decode(g, skip_special_tokens=False))
fout.write(json.dumps({"image": name, "raw": raw, "box": parse_box(raw, *im.size)}) + "\n")
fout.flush()
print(f"{min(b + args.bs, len(todo))}/{len(todo)}", flush=True)
print("DONE", flush=True)
if __name__ == "__main__":
main()
|