Views
No views yet
1import re
2import torch
3import requests
4from PIL import Image, ImageDraw
5from transformers import AutoProcessor, Kosmos2_5ForConditionalGeneration, infer_device
6
7repo = "microsoft/kosmos-2.5"
8device = "cuda:0"
9dtype = torch.bfloat16
10model = Kosmos2_5ForConditionalGeneration.from_pretrained(repo, device_map=device, dtype=dtype)
11processor = AutoProcessor.from_pretrained(repo)
12
13# sample image
14url = "https://huggingface.co/microsoft/kosmos-2.5/resolve/main/receipt_00008.png"
15image = Image.open(requests.get(url, stream=True).raw)
16
17prompt = "<md>"
18inputs = processor(text=prompt, images=image, return_tensors="pt")
19
20height, width = inputs.pop("height"), inputs.pop("width")
21raw_width, raw_height = image.size
22scale_height = raw_height / height
23scale_width = raw_width / width
24
25inputs = {k: v.to(device) if v is not None else None for k, v in inputs.items()}
26inputs["flattened_patches"] = inputs["flattened_patches"].to(dtype)
27generated_ids = model.generate(
28 **inputs,
29 max_new_tokens=1024,
30)
31
32generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
33print(generated_text[0])1import re
2import torch
3import requests
4from PIL import Image, ImageDraw
5from transformers import AutoProcessor, Kosmos2_5ForConditionalGeneration, infer_device
6
7repo = "microsoft/kosmos-2.5"
8device = "cuda:0"
9dtype = torch.bfloat16
10model = Kosmos2_5ForConditionalGeneration.from_pretrained(repo, device_map=device, dtype=dtype)
11processor = AutoProcessor.from_pretrained(repo)
12
13# sample image
14url = "https://huggingface.co/microsoft/kosmos-2.5/resolve/main/receipt_00008.png"
15image = Image.open(requests.get(url, stream=True).raw)
16
17# bs = 1
18prompt = "<ocr>"
19inputs = processor(text=prompt, images=image, return_tensors="pt")
20height, width = inputs.pop("height"), inputs.pop("width")
21raw_width, raw_height = image.size
22scale_height = raw_height / height
23scale_width = raw_width / width
24
25# bs > 1, batch generation
26# inputs = processor(text=[prompt, prompt], images=[image,image], return_tensors="pt")
27# height, width = inputs.pop("height"), inputs.pop("width")
28# raw_width, raw_height = image.size
29# scale_height = raw_height / height[0]
30# scale_width = raw_width / width[0]
31
32inputs = {k: v.to(device) if v is not None else None for k, v in inputs.items()}
33inputs["flattened_patches"] = inputs["flattened_patches"].to(dtype)
34generated_ids = model.generate(
35 **inputs,
36 max_new_tokens=1024,
37)
38
39generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
40def post_process(y, scale_height, scale_width):
41 y = y.replace(prompt, "")
42 if "<md>" in prompt:
43 return y
44 pattern = r"<bbox><x_\d+><y_\d+><x_\d+><y_\d+></bbox>"
45 bboxs_raw = re.findall(pattern, y)
46 lines = re.split(pattern, y)[1:]
47 bboxs = [re.findall(r"\d+", i) for i in bboxs_raw]
48 bboxs = [[int(j) for j in i] for i in bboxs]
49 info = ""
50 for i in range(len(lines)):
51 box = bboxs[i]
52 x0, y0, x1, y1 = box
53 if not (x0 >= x1 or y0 >= y1):
54 x0 = int(x0 * scale_width)
55 y0 = int(y0 * scale_height)
56 x1 = int(x1 * scale_width)
57 y1 = int(y1 * scale_height)
58 info += f"{x0},{y0},{x1},{y0},{x1},{y1},{x0},{y1},{lines[i]}"
59 return info
60
61output_text = post_process(generated_text[0], scale_height, scale_width)
62print(output_text)
63
64draw = ImageDraw.Draw(image)
65lines = output_text.split("\n")
66for line in lines:
67 # draw the bounding box
68 line = list(line.split(","))
69 if len(line) < 8:
70 continue
71 line = list(map(int, line[:8]))
72 draw.polygon(line, outline="red")
73image.save("output.png")@article{lv2023kosmos,
title={Kosmos-2.5: A multimodal literate model},
author={Lv, Tengchao and Huang, Yupan and Chen, Jingye and Cui, Lei and Ma, Shuming and Chang, Yaoyao and Huang, Shaohan and Wang, Wenhui and Dong, Li and Luo, Weiyao and others},
journal={arXiv preprint arXiv:2309.11419},
year={2023}
}