| """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 |
| 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) |
| 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: |
| 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() |
|
|
