import reimport torchfrom transformers import AutoProcessor, AutoModelForVision2Seqfrom transformers.image_utils import load_image DEVICE = "cuda" if torch.cuda.is_available() else "cpu"MODEL_PATH = "docling-project/ScreenVLM"NORM_SIZE = 500 image = load_image("https://example.com/screenshot.png") processor = AutoProcessor.from_pretrained(MODEL_PATH)model = AutoModelForVision2Seq.from_pretrained( MODEL_PATH, torch_dtype=torch.bfloat16, _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "sdpa",).to(DEVICE) messages = [ { "role": "user", "content": [ {"type": "image"}, {"type": "text", "text": "Generate the screen representation for this UI:"}, ], },] prompt = processor.apply_chat_template(messages, add_generation_prompt=True)inputs = processor(text=prompt, images=[image], return_tensors="pt").to(DEVICE) generated_ids = model.generate(**inputs, max_new_tokens=6192)prompt_length = inputs.input_ids.shape[1]output = processor.batch_decode( generated_ids[:, prompt_length:], skip_special_tokens=False,)[0].lstrip() def parse_screentag(text, width, height): pattern = re.compile( r"<(?P<tag>[a-zA-Z][a-zA-Z0-9_]*)>" r"\s*<loc_(?P<l>\d+)><loc_(?P<t>\d+)><loc_(?P<r>\d+)><loc_(?P<b>\d+)>" r"(?P<text>[^<]*)" ) elements = [] for m in pattern.finditer(text): l, t, r, b = [max(0, min(int(m.group(k)), NORM_SIZE)) for k in ("l", "t", "r", "b")] if r < l: l, r = r, l if b < t: t, b = b, t x = l / NORM_SIZE * width y = t / NORM_SIZE * height w = (r - l) / NORM_SIZE * width h = (b - t) / NORM_SIZE * height elements.append({ "label": m.group("tag"), "bbox": (x, y, w, h), "text": m.group("text").strip() or None, }) return elements elements = parse_screentag(output, *image.size)for el in elements: print(f"{el['label']:20s} bbox=({int(el['bbox'][0]):4d},{int(el['bbox'][1]):4d},{int(el['bbox'][2]):4d},{int(el['bbox'][3]):4d}) text={el['text']!r}")