Views
No views yet
1import re
2import numpy as np
3
4def extract_mt_token_ids_v1(text):
5 pattern = r"<\|mt_(\d{4})\|>"
6 return [int(x) for x in re.findall(pattern, text)]
7
8def extract_mt_token_ids_v2(text):
9 pattern = re.compile(r'<\|mt_start\|><\|mt_(\d{4})\|><\|mt_(\d{4})\|><\|mt_end\|>')
10 matches = pattern.findall(text)
11 ret_list = []
12 for num1, num2 in matches:
13 ret_list.append(int(num1))
14 ret_list.append(int(num2))
15 return ret_list
16
17def find_first_index(arr, value):
18 indices = np.where(arr == value)[0]
19
20 return indices[0] if len(indices) > 0 else -1
21
22def fix_mt_format_comprehensive(text):
23 pattern_too_many = r'(<\|mt_start\|>)(<\|mt_\d+\|>)(<\|mt_\d+\|>)(?:<\|mt_\d+\|>)+<\|mt_end\|>'
24 replacement_too_many = r'\1\2\3<|mt_end|>'
25 text = re.sub(pattern_too_many, replacement_too_many, text)
26
27 pattern_too_few_with_end = r'(<\|mt_start\|>)(<\|mt_\d+\|>)(<\|mt_end\|>)'
28 replacement_too_few = r'\1\2<|mt_9999|><|mt_end|>'
29 text = re.sub(pattern_too_few_with_end, replacement_too_few, text)
30
31 pattern_too_few_no_end = r'(<\|mt_start\|>)(<\|mt_\d+\|>)(?!<\|mt_)'
32 replacement_too_few_no_end = r'\1\2<|mt_9999|><|mt_end|>'
33 text = re.sub(pattern_too_few_no_end, replacement_too_few_no_end, text)
34 return text
35
36def extract_think_and_answer_robust(response: str) -> Tuple[Optional[str], Optional[str]]:
37 think_content = None
38 answer_content = None
39 think_pattern = re.compile(r"<think>(.*?)</think>", re.DOTALL)
40 answer_pattern = re.compile(r"<answer>(.*?)</answer>", re.DOTALL)
41 think_match = think_pattern.search(response)
42 if think_match:
43 think_content = think_match.group(1)
44 answer_match = answer_pattern.search(response)
45 if answer_match:
46 answer_content = answer_match.group(1)
47
48 if answer_content is None or think_content is None:
49 if '<answer>' in response:
50 head, tail = response.split('<answer>', 1)
51 if think_content is None:
52 think_content = head
53 if answer_content is None:
54 answer_content = tail
55 elif '</think>' in response:
56 head, tail = response.split('</think>', 1)
57 if think_content is None:
58 think_content = head
59 if answer_content is None:
60 answer_content = tail
61
62 return think_content, answer_content
63
64from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
65from projects.samtok.models import DirectResize, VQ_SAM2, VQ_SAM2Config, SAM2Config
66
67# build VLM
68model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
69 "zhouyik/Qwen2.5-VL-3B-SAMTok-gres-rl", torch_dtype="auto"
70).cuda().eval()
71processor = AutoProcessor.from_pretrained("zhouyik/Qwen2.5-VL-3B-SAMTok-gres-rl")
72
73# build SAMTok
74CODEBOOK_SIZE = 256
75CODEBOOK_DEPTH = 2
76sam2_config = SAM2Config(
77 ckpt_path="zhouyik/Qwen2.5-VL-3B-SAMTok-gres-rl/sam2.1_hiera_large.pt",
78)
79vq_sam2_config = VQ_SAM2Config(
80 sam2_config=sam2_config,
81 codebook_size=CODEBOOK_SIZE,
82 codebook_depth=CODEBOOK_DEPTH,
83 shared_codebook=False,
84 latent_dim=256,
85)
86vq_sam2 = VQ_SAM2(vq_sam2_config).cuda().eval()
87state = torch.load("zhouyik/Qwen2.5-VL-3B-SAMTok-gres-rl/mask_tokenizer_256x2.pth", map_location="cpu")
88vq_sam2.load_state_dict(state)
89sam2_image_processor = DirectResize(1024)
90
91# message
92image_path = "figs/totoro.jpg"
93phrase = "the biggest totoro"
94question = f"Please segment {phrase} in this image. A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>"
95image = Image.open(image_path).convert('RGB')
96ori_width, ori_height = image.size
97messages = [
98 {
99 "role": "user",
100 "content": [
101 {
102 "type": "image",
103 "image": image_path,
104 },
105 {"type": "text", "text": question},
106 ],
107 }
108]
109
110# VLM inferece
111text = processor.apply_chat_template(
112 messages, tokenize=False, add_generation_prompt=True
113)
114
115image_inputs, video_inputs = process_vision_info(messages)
116inputs = processor(
117 text=[text],
118 images=image_inputs,
119 videos=video_inputs,
120 padding=True,
121 return_tensors="pt",
122)
123inputs = inputs.to("cuda")
124
125generated_ids = model.generate(
126 **inputs,
127 max_new_tokens=512,
128 do_sample=False,
129 top_p=1.0,
130)
131generated_ids_trimmed = [
132 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
133]
134output_text = processor.batch_decode(
135 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
136)
137
138# decode mask
139thinking_content, answer_content = extract_think_and_answer_robust(output_text[0])
140quant_ids = extract_mt_token_ids_v1(answer_content)
141if len(quant_ids) % CODEBOOK_DEPTH != 0:
142 output_text = [fix_mt_format_comprehensive(answer_content)]
143 quant_ids = extract_mt_token_ids_v2(answer_content)
144
145batch_size = len(quant_ids) // CODEBOOK_DEPTH
146remap_quant_ids = []
147tags = []
148for bs_id in range(batch_size):
149 chunk_quant_ids = quant_ids[bs_id*CODEBOOK_DEPTH:(bs_id+1)*CODEBOOK_DEPTH]
150 tags.append(f"{chunk_quant_ids[0]}-{chunk_quant_ids[1]}")
151 remap_chunk_quant_ids = [quant_id - book_id*CODEBOOK_SIZE for book_id, quant_id in enumerate(chunk_quant_ids)]
152 code1 = remap_chunk_quant_ids[0]
153 code2 = remap_chunk_quant_ids[1]
154 if not (code2 >= 0 and code2 < CODEBOOK_SIZE):
155 code2 = -1
156 remap_chunk_quant_ids_error_handle = [code1, code2]
157 remap_quant_ids.append(remap_chunk_quant_ids_error_handle)
158
159batch_size = len(remap_quant_ids)
160sam2_image = np.array(image)
161sam2_image = sam2_image_processor.apply_image(sam2_image)
162sam2_pixel_values = torch.from_numpy(sam2_image).permute(2, 0, 1).contiguous()
163sam2_pixel_values = sam2_pixel_values.unsqueeze(0).to(vq_sam2.dtype).to(vq_sam2.device)
164sam2_pixel_values = sam2_pixel_values.repeat(batch_size, 1, 1, 1)
165
166quant_ids = torch.LongTensor(remap_quant_ids).to(vq_sam2.device)
167
168with torch.no_grad():
169 _pred_masks = vq_sam2.forward_with_codes(sam2_pixel_values, quant_ids)
170_pred_masks = torch.nn.functional.interpolate(_pred_masks, size=(ori_height, ori_width), mode='bilinear')
171_pred_masks = _pred_masks > 0.5
172_pred_masks = _pred_masks[:, 0, :, :].cpu().numpy().astype(np.uint8)
173text_token_2d_mask_mapping = {tag: _pred_mask for tag, _pred_mask in zip(tags, _pred_masks)}