1import cv2
2import torch
3import torch.nn.functional as F
4from transformers import AutoTokenizer, CLIPImageProcessor
5from model.ToothXpert_MOE import ToothXpertForCausalLMMOE
6from model.llava import conversation as conversation_lib
7from model.llava.mm_utils import tokenizer_image_token
8from model.segment_anything.utils.transforms import ResizeLongestSide
9from utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
10 DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX)
11
12# Preprocessing function
13def preprocess(x, pixel_mean=torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1),
14 pixel_std=torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1), img_size=1024):
15 x = (x - pixel_mean) / pixel_std
16 h, w = x.shape[-2:]
17 padh = img_size - h
18 padw = img_size - w
19 x = F.pad(x, (0, padw, 0, padh))
20 return x
21
22# Load model
23model_path = "./ToothXpert_pretrained"
24device = "cuda:0"
25
26tokenizer = AutoTokenizer.from_pretrained(
27 model_path,
28 model_max_length=512,
29 padding_side="right",
30 use_fast=False,
31)
32tokenizer.pad_token = tokenizer.unk_token
33tokenizer.add_tokens("[SEG]")
34seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
35tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
36
37moe_lora_args = {
38 "lora_r": 8,
39 "lora_alpha": 16,
40 "lora_dropout": 0.05,
41 "lora_target_modules": "q_proj,v_proj",
42 "moe_lora": False,
43 "expert_num": 3,
44 "guide": True,
45 "guide_mode": "smmulsm",
46 "vocab_size": len(tokenizer),
47}
48
49model = ToothXpertForCausalLMMOE.from_pretrained(
50 model_path,
51 low_cpu_mem_usage=True,
52 vision_tower="openai/clip-vit-large-patch14",
53 seg_token_idx=seg_token_idx,
54 torch_dtype=torch.bfloat16,
55 train_mask_decoder=True,
56 out_dim=256,
57 moe_lora_args=moe_lora_args,
58)
59
60model.config.eos_token_id = tokenizer.eos_token_id
61model.config.bos_token_id = tokenizer.bos_token_id
62model.config.pad_token_id = tokenizer.pad_token_id
63
64model.get_model().initialize_vision_modules(model.get_model().config)
65vision_tower = model.get_model().get_vision_tower()
66vision_tower.to(dtype=torch.bfloat16, device=device)
67
68model = model.bfloat16().to(device)
69model.eval()
70
71# Load and process image
72image_path = "your_dental_xray.png"
73image_np = cv2.imread(image_path)
74image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
75original_size_list = [image_np.shape[:2]]
76
77clip_image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14")
78transform = ResizeLongestSide(1024)
79
80image_clip = (
81 clip_image_processor.preprocess(image_np, return_tensors="pt")["pixel_values"][0]
82 .unsqueeze(0).to(device).bfloat16()
83)
84
85image = transform.apply_image(image_np)
86resize_list = [image.shape[:2]]
87image = (
88 preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())
89 .unsqueeze(0).to(device).bfloat16()
90)
91
92# Prepare prompt
93question = "Can you describe the image for me?"
94conv = conversation_lib.conv_templates["llava_v1"].copy()
95conv.messages = []
96prompt = DEFAULT_IMAGE_TOKEN + "\n" + question
97prompt = prompt.replace(DEFAULT_IMAGE_TOKEN,
98 DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN)
99
100conv.append_message(conv.roles[0], prompt)
101conv.append_message(conv.roles[1], "")
102prompt = conv.get_prompt()
103
104input_ids = tokenizer_image_token(prompt, tokenizer, return_tensors="pt")
105input_ids = input_ids.unsqueeze(0).to(device)
106
107# Run inference
108with torch.no_grad():
109 output_ids, pred_masks = model.evaluate(
110 image_clip,
111 image,
112 input_ids,
113 resize_list,
114 original_size_list,
115 max_new_tokens=512,
116 tokenizer=tokenizer,
117 )
118
119output_ids = output_ids[0][output_ids[0] != IMAGE_TOKEN_INDEX]
120text_output = tokenizer.decode(output_ids, skip_special_tokens=False)
121text_output = text_output.split('ASSISTANT:')[-1].replace('</s>', '').strip()
122
123print(f"Question: {question}")
124print(f"Answer: {text_output}")
