| """VLM helpers + Qwen2.5-VL engine.""" |
| import re |
| import threading |
|
|
| import cv2 |
| import numpy as np |
| from PIL import Image as PILImage |
|
|
| BOX_PATTERN = re.compile(r"\[(\d+),\s*(\d+),\s*(\d+),\s*(\d+)\]") |
|
|
|
|
| def parse_boxes(text, width, height): |
| """Parse box '[ymin,xmin,ymax,xmax]' (0-1000) -> [(x1,y1,x2,y2)] theo pixel.""" |
| boxes = [] |
| for ymin, xmin, ymax, xmax in BOX_PATTERN.findall(text): |
| x1 = int(int(xmin) * width / 1000.0) |
| y1 = int(int(ymin) * height / 1000.0) |
| x2 = int(int(xmax) * width / 1000.0) |
| y2 = int(int(ymax) * height / 1000.0) |
| x1 = max(0, min(x1, width)) |
| x2 = max(0, min(x2, width)) |
| y1 = max(0, min(y1, height)) |
| y2 = max(0, min(y2, height)) |
| boxes.append((x1, y1, x2, y2)) |
| return boxes |
|
|
|
|
| def draw_boxes(frame_bgr, boxes, label): |
| """Vẽ box + label lên bản copy của frame, trả frame mới.""" |
| out = frame_bgr.copy() |
| for (x1, y1, x2, y2) in boxes: |
| cv2.rectangle(out, (x1, y1), (x2, y2), (0, 255, 0), 2) |
| cv2.putText( |
| out, label, (x1, max(0, y1 - 10)), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2, |
| ) |
| return out |
|
|
|
|
| def encode_frame_jpeg(frame_bgr, quality=80): |
| """Encode BGR frame thành JPEG bytes.""" |
| ok, buf = cv2.imencode(".jpg", frame_bgr, [cv2.IMWRITE_JPEG_QUALITY, quality]) |
| if not ok: |
| raise RuntimeError("JPEG encode failed") |
| return buf.tobytes() |
|
|
|
|
| def build_messages(pil_image, prompt): |
| """Tạo messages cho Qwen chat template.""" |
| return [ |
| { |
| "role": "user", |
| "content": [ |
| {"type": "image", "image": pil_image}, |
| {"type": "text", "text": prompt}, |
| ], |
| } |
| ] |
|
|
|
|
| class VLMEngine: |
| """Wrapper Qwen2.5-VL với streaming token.""" |
|
|
| def __init__(self, model_name="Qwen/Qwen2.5-VL-3B-Instruct", device=None): |
| self.model_name = model_name |
| self.device = device |
| self.model = None |
| self.processor = None |
|
|
| @property |
| def loaded(self): |
| return self.model is not None and self.processor is not None |
|
|
| def load(self): |
| import torch |
| from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor |
|
|
| if self.device is None: |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" |
| if self.device == "cpu": |
| print("[VLMEngine] WARNING: chạy trên CPU, inference sẽ rất chậm.") |
|
|
| self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained( |
| self.model_name, |
| torch_dtype=torch.float16 if self.device == "cuda" else torch.float32, |
| device_map="auto" if self.device == "cuda" else None, |
| ) |
| self.processor = AutoProcessor.from_pretrained(self.model_name) |
| print(f"[VLMEngine] Model loaded on {self.device}.") |
|
|
| def stream_infer(self, frame_bgr, prompt): |
| if not self.loaded: |
| raise RuntimeError("VLMEngine chưa load(). Gọi load() trước.") |
|
|
| from transformers import TextIteratorStreamer |
| from qwen_vl_utils import process_vision_info |
|
|
| rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) |
| pil_image = PILImage.fromarray(rgb) |
| messages = build_messages(pil_image, prompt) |
|
|
| text = self.processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True |
| ) |
| image_inputs, video_inputs = process_vision_info(messages) |
| |
| |
| self.last_image_size = image_inputs[0].size if image_inputs else None |
| inputs = self.processor( |
| text=[text], images=image_inputs, videos=video_inputs, |
| padding=True, return_tensors="pt", |
| ).to(self.device) |
|
|
| streamer = TextIteratorStreamer( |
| self.processor.tokenizer, |
| skip_prompt=True, |
| skip_special_tokens=True, |
| ) |
| gen_kwargs = dict(**inputs, max_new_tokens=256, streamer=streamer) |
| thread = threading.Thread(target=self.model.generate, kwargs=gen_kwargs) |
| thread.start() |
| for chunk in streamer: |
| if chunk: |
| yield chunk |
| thread.join() |
|
|
