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
36from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
37from projects.samtok.models import DirectResize, VQ_SAM2, VQ_SAM2Config, SAM2Config
38
39# build VLM
40model = Qwen3VLForConditionalGeneration.from_pretrained(
41 "zhouyik/Qwen3-VL-8B-SAMTok", torch_dtype="auto"
42).cuda().eval()
43processor = AutoProcessor.from_pretrained("zhouyik/Qwen3-VL-4B-SAMTok")
44
45# build SAMTok
46CODEBOOK_SIZE = 256
47CODEBOOK_DEPTH = 2
48sam2_config = SAM2Config(
49 ckpt_path="zhouyik/Qwen3-VL-4B-SAMTok/sam2.1_hiera_large.pt",
50)
51vq_sam2_config = VQ_SAM2Config(
52 sam2_config=sam2_config,
53 codebook_size=CODEBOOK_SIZE,
54 codebook_depth=CODEBOOK_DEPTH,
55 shared_codebook=False,
56 latent_dim=256,
57)
58vq_sam2 = VQ_SAM2(vq_sam2_config).cuda().eval()
59state = torch.load("zhouyik/Qwen3-VL-4B-SAMTok/mask_tokenizer_256x2.pth", map_location="cpu")
60vq_sam2.load_state_dict(state)
61sam2_image_processor = DirectResize(1024)
62
63# message
64image_path = "figs/totoro.jpg"
65question = "Could you please give me a detail description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer."
66image = Image.open(image_path).convert('RGB')
67ori_width, ori_height = image.size
68messages = [
69 {
70 "role": "user",
71 "content": [
72 {
73 "type": "image",
74 "image": image_path,
75 },
76 {"type": "text", "text": question},
77 ],
78 }
79]
80
81# VLM inferece
82inputs = processor.apply_chat_template(
83 messages,
84 tokenize=True,
85 add_generation_prompt=True,
86 return_dict=True,
87 return_tensors="pt"
88)
89inputs = inputs.to(model.device)
90
91generated_ids = model.generate(
92 **inputs,
93 max_new_tokens=512,
94 do_sample=False,
95 top_p=1.0,
96)
97generated_ids_trimmed = [
98 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
99]
100output_text = processor.batch_decode(
101 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
102)
103
104# decode mask
105quant_ids = extract_mt_token_ids_v1(output_text[0])
106if len(quant_ids) % CODEBOOK_DEPTH != 0:
107 output_text = [fix_mt_format_comprehensive(output_text[0])]
108 quant_ids = extract_mt_token_ids_v2(output_text[0])
109
110batch_size = len(quant_ids) // CODEBOOK_DEPTH
111remap_quant_ids = []
112tags = []
113for bs_id in range(batch_size):
114 chunk_quant_ids = quant_ids[bs_id*CODEBOOK_DEPTH:(bs_id+1)*CODEBOOK_DEPTH]
115 tags.append(f"{chunk_quant_ids[0]}-{chunk_quant_ids[1]}")
116 remap_chunk_quant_ids = [quant_id - book_id*CODEBOOK_SIZE for book_id, quant_id in enumerate(chunk_quant_ids)]
117 code1 = remap_chunk_quant_ids[0]
118 code2 = remap_chunk_quant_ids[1]
119 if not (code2 >= 0 and code2 < CODEBOOK_SIZE):
120 code2 = -1
121 remap_chunk_quant_ids_error_handle = [code1, code2]
122 remap_quant_ids.append(remap_chunk_quant_ids_error_handle)
123
124batch_size = len(remap_quant_ids)
125sam2_image = np.array(image)
126sam2_image = sam2_image_processor.apply_image(sam2_image)
127sam2_pixel_values = torch.from_numpy(sam2_image).permute(2, 0, 1).contiguous()
128sam2_pixel_values = sam2_pixel_values.unsqueeze(0).to(vq_sam2.dtype).to(vq_sam2.device)
129sam2_pixel_values = sam2_pixel_values.repeat(batch_size, 1, 1, 1)
130
131quant_ids = torch.LongTensor(remap_quant_ids).to(vq_sam2.device)
132
133with torch.no_grad():
134 _pred_masks = vq_sam2.forward_with_codes(sam2_pixel_values, quant_ids)
135_pred_masks = torch.nn.functional.interpolate(_pred_masks, size=(ori_height, ori_width), mode='bilinear')
136_pred_masks = _pred_masks > 0.5
137_pred_masks = _pred_masks[:, 0, :, :].cpu().numpy().astype(np.uint8)
138text_token_2d_mask_mapping = {tag: _pred_mask for tag, _pred_mask in zip(tags, _pred_masks)}