Views
No views yet

transformers implementation of the original Kosmos-2 model from Microsoft.1import requests
2
3from PIL import Image
4from transformers import AutoProcessor, AutoModelForVision2Seq
5
6
7model = AutoModelForVision2Seq.from_pretrained("ydshieh/kosmos-2-patch14-224", trust_remote_code=True)
8processor = AutoProcessor.from_pretrained("ydshieh/kosmos-2-patch14-224", trust_remote_code=True)
9
10prompt = "<grounding>An image of"
11
12url = "https://huggingface.co/ydshieh/kosmos-2-patch14-224/resolve/main/snowman.jpg"
13image = Image.open(requests.get(url, stream=True).raw)
14
15# The original Kosmos-2 demo saves the image first then reload it. For some images, this will give slightly different image input and change the generation outputs.
16# Uncomment the following 2 lines if you want to match the original demo's outputs.
17# (One example is the `two_dogs.jpg` from the demo)
18# image.save("new_image.jpg")
19# image = Image.open("new_image.jpg")
20
21inputs = processor(text=prompt, images=image, return_tensors="pt")
22
23generated_ids = model.generate(
24 pixel_values=inputs["pixel_values"],
25 input_ids=inputs["input_ids"][:, :-1],
26 attention_mask=inputs["attention_mask"][:, :-1],
27 img_features=None,
28 img_attn_mask=inputs["img_attn_mask"][:, :-1],
29 use_cache=True,
30 max_new_tokens=64,
31)
32generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
33
34# Specify `cleanup_and_extract=False` in order to see the raw model generation.
35processed_text = processor.post_processor_generation(generated_text, cleanup_and_extract=False)
36
37print(processed_text)
38# `<grounding> An image of<phrase> a snowman</phrase><object><patch_index_0044><patch_index_0863></object> warming himself by<phrase> a fire</phrase><object><patch_index_0005><patch_index_0911></object>.`
39
40# By default, the generated text is cleanup and the entities are extracted.
41processed_text, entities = processor.post_processor_generation(generated_text)
42
43print(processed_text)
44# `An image of a snowman warming himself by a fire.`
45
46print(entities)
47# `[('a snowman', (12, 21), [(0.390625, 0.046875, 0.984375, 0.828125)]), ('a fire', (41, 47), [(0.171875, 0.015625, 0.484375, 0.890625)])]`entities, you can use the following helper function to draw their bounding bboxes on the image:1import cv2
2import numpy as np
3import os
4import requests
5import torch
6import torchvision.transforms as T
7
8from PIL import Image
9
10
11def is_overlapping(rect1, rect2):
12 x1, y1, x2, y2 = rect1
13 x3, y3, x4, y4 = rect2
14 return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)
15
16
17def draw_entity_boxes_on_image(image, entities, show=False, save_path=None):
18 """_summary_
19 Args:
20 image (_type_): image or image path
21 collect_entity_location (_type_): _description_
22 """
23 if isinstance(image, Image.Image):
24 image_h = image.height
25 image_w = image.width
26 image = np.array(image)[:, :, [2, 1, 0]]
27 elif isinstance(image, str):
28 if os.path.exists(image):
29 pil_img = Image.open(image).convert("RGB")
30 image = np.array(pil_img)[:, :, [2, 1, 0]]
31 image_h = pil_img.height
32 image_w = pil_img.width
33 else:
34 raise ValueError(f"invaild image path, {image}")
35 elif isinstance(image, torch.Tensor):
36 # pdb.set_trace()
37 image_tensor = image.cpu()
38 reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[:, None, None]
39 reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[:, None, None]
40 image_tensor = image_tensor * reverse_norm_std + reverse_norm_mean
41 pil_img = T.ToPILImage()(image_tensor)
42 image_h = pil_img.height
43 image_w = pil_img.width
44 image = np.array(pil_img)[:, :, [2, 1, 0]]
45 else:
46 raise ValueError(f"invaild image format, {type(image)} for {image}")
47
48 if len(entities) == 0:
49 return image
50
51 new_image = image.copy()
52 previous_bboxes = []
53 # size of text
54 text_size = 1
55 # thickness of text
56 text_line = 1 # int(max(1 * min(image_h, image_w) / 512, 1))
57 box_line = 3
58 (c_width, text_height), _ = cv2.getTextSize("F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
59 base_height = int(text_height * 0.675)
60 text_offset_original = text_height - base_height
61 text_spaces = 3
62
63 for entity_name, (start, end), bboxes in entities:
64 for (x1_norm, y1_norm, x2_norm, y2_norm) in bboxes:
65 orig_x1, orig_y1, orig_x2, orig_y2 = int(x1_norm * image_w), int(y1_norm * image_h), int(x2_norm * image_w), int(y2_norm * image_h)
66 # draw bbox
67 # random color
68 color = tuple(np.random.randint(0, 255, size=3).tolist())
69 new_image = cv2.rectangle(new_image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, box_line)
70
71 l_o, r_o = box_line // 2 + box_line % 2, box_line // 2 + box_line % 2 + 1
72
73 x1 = orig_x1 - l_o
74 y1 = orig_y1 - l_o
75
76 if y1 < text_height + text_offset_original + 2 * text_spaces:
77 y1 = orig_y1 + r_o + text_height + text_offset_original + 2 * text_spaces
78 x1 = orig_x1 + r_o
79
80 # add text background
81 (text_width, text_height), _ = cv2.getTextSize(f" {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
82 text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2 = x1, y1 - (text_height + text_offset_original + 2 * text_spaces), x1 + text_width, y1
83
84 for prev_bbox in previous_bboxes:
85 while is_overlapping((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox):
86 text_bg_y1 += (text_height + text_offset_original + 2 * text_spaces)
87 text_bg_y2 += (text_height + text_offset_original + 2 * text_spaces)
88 y1 += (text_height + text_offset_original + 2 * text_spaces)
89
90 if text_bg_y2 >= image_h:
91 text_bg_y1 = max(0, image_h - (text_height + text_offset_original + 2 * text_spaces))
92 text_bg_y2 = image_h
93 y1 = image_h
94 break
95
96 alpha = 0.5
97 for i in range(text_bg_y1, text_bg_y2):
98 for j in range(text_bg_x1, text_bg_x2):
99 if i < image_h and j < image_w:
100 if j < text_bg_x1 + 1.35 * c_width:
101 # original color
102 bg_color = color
103 else:
104 # white
105 bg_color = [255, 255, 255]
106 new_image[i, j] = (alpha * new_image[i, j] + (1 - alpha) * np.array(bg_color)).astype(np.uint8)
107
108 cv2.putText(
109 new_image, f" {entity_name}", (x1, y1 - text_offset_original - 1 * text_spaces), cv2.FONT_HERSHEY_COMPLEX, text_size, (0, 0, 0), text_line, cv2.LINE_AA
110 )
111 # previous_locations.append((x1, y1))
112 previous_bboxes.append((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2))
113
114 pil_image = Image.fromarray(new_image[:, :, [2, 1, 0]])
115 if save_path:
116 pil_image.save(save_path)
117 if show:
118 pil_image.show()
119
120 return new_image
121
122
123# (The same image from the previous code example)
124url = "https://huggingface.co/ydshieh/kosmos-2-patch14-224/resolve/main/snowman.jpg"
125image = Image.open(requests.get(url, stream=True).raw)
126
127# From the previous code example
128entities = [('a snowman', (12, 21), [(0.390625, 0.046875, 0.984375, 0.828125)]), ('a fire', (41, 47), [(0.171875, 0.015625, 0.484375, 0.890625)])]
129
130# Draw the bounding bboxes
131draw_entity_boxes_on_image(image, entities, show=True)