import torch from transformers import Qwen2VLForConditionalGeneration, AutoProcessor from qwen_vl_utils import process_vision_info from PIL import Image, ImageDraw import requests from io import BytesIO import asyncio import time import os class QwenVLService: def __init__(self): self.model = None self.processor = None self.is_loaded = False # 自动检测并使用 GPU self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"🎯 使用设备: {self.device}") if self.device == "cuda": print(f"🚀 GPU 加速已启用: {torch.cuda.get_device_name(0)}") print(f"💾 GPU 显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB") print(f"📊 GPU 数量: {torch.cuda.device_count()}") async def load_model(self): """异步加载模型""" if self.is_loaded: return loop = asyncio.get_event_loop() await loop.run_in_executor(None, self._load_model_sync) def _load_model_sync(self): """同步加载模型（在线程池中执行）""" print("🚀 正在加载 Qwen2-VL-2B 模型...") # 直接使用 HuggingFace（通过 HF_ENDPOINT 环境变量使用镜像） model_path = "Qwen/Qwen2-VL-2B-Instruct" print(f"📦 从 HuggingFace 下载模型: {model_path}") print(f"🌐 使用镜像: {os.environ.get('HF_ENDPOINT', 'https://huggingface.co')}") print(f"💾 缓存目录: {os.environ.get('HF_HOME', './models')}") # 加载模型（会自动从 HF_ENDPOINT 下载） print(f"🔄 正在下载并加载模型到 {self.device}...") self.model = Qwen2VLForConditionalGeneration.from_pretrained( model_path, torch_dtype=torch.float16 if self.device == "cuda" else torch.float32, device_map="auto" if self.device == "cuda" else None, trust_remote_code=True ) if self.device == "cpu": self.model = self.model.to(self.device) # 加载处理器 print("🔄 正在加载处理器...") self.processor = AutoProcessor.from_pretrained( model_path, trust_remote_code=True ) self.is_loaded = True print("✅ Qwen2-VL-2B 模型加载完成！") async def detect_object_at_point( self, image_url: str, click_x: int, click_y: int, image_index: int, image_width: int, image_height: int ) -> dict: """使用Qwen2-VL检测点击位置的物体""" start_time = time.time() # 1. 下载图片 print(f"📥 下载图片: {image_url}") response = requests.get(image_url, timeout=10) image = Image.open(BytesIO(response.content)).convert("RGB") actual_width, actual_height = image.size print(f"📊 原始图片: {actual_width}x{actual_height}") # 缩放大图以加快推理速度（保持宽高比） max_size = 768 # 降低到768以加快推理（从1024） if max(actual_width, actual_height) > max_size: scale = max_size / max(actual_width, actual_height) new_width = int(actual_width * scale) new_height = int(actual_height * scale) # 同步缩放坐标 click_x = int(click_x * scale) click_y = int(click_y * scale) image = image.resize((new_width, new_height), Image.LANCZOS) print(f"⚡ 已缩放至: {new_width}x{new_height} (加速推理)") print(f"📍 标注坐标: ({click_x}, {click_y})") # 2. 在图片上绘制标注点 marked_image = image.copy() draw = ImageDraw.Draw(marked_image) # 绘制超级明显的红色标记（放大3倍） marker_size = 50 # 从30增加到50 line_width = 10 # 从6增加到10 # 绘制3层嵌套圆圈，更明显 for i in range(3): size = marker_size - i * 10 draw.ellipse( [click_x - size, click_y - size, click_x + size, click_y + size], outline=(255, 0, 0), width=line_width ) # 绘制十字线 draw.line( [click_x - marker_size - 20, click_y, click_x + marker_size + 20, click_y], fill=(255, 0, 0), width=line_width ) draw.line( [click_x, click_y - marker_size - 20, click_x, click_y + marker_size + 20], fill=(255, 0, 0), width=line_width ) # 在圆心绘制实心点 center_size = 8 draw.ellipse( [click_x - center_size, click_y - center_size, click_x + center_size, click_y + center_size], fill=(255, 0, 0) ) print(f"✏️ 已在图片上绘制标注点") # 3. 构建对话式问答prompt messages = [ { "role": "user", "content": [ { "type": "image", "image": marked_image, }, { "type": "text", "text": "看图中红色圆圈的正中心位置，那里有一个物体。只用2-4个词描述这个被圆圈中心点标记的物体（格式：颜色+材质/物品，如：白色沙发、红色毛衣）。只说这一个物体，不要其他。" } ], } ] # 4. 准备输入 text = self.processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) image_inputs, video_inputs = process_vision_info(messages) inputs = self.processor( text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt", ) # 移动到设备 inputs = inputs.to(self.device) # 5. 生成回答 print("🔍 正在识别物体...") inference_start = time.time() loop = asyncio.get_event_loop() description = await loop.run_in_executor( None, self._generate_sync, inputs ) inference_time = time.time() - inference_start print(f"⚡ 推理耗时: {inference_time:.2f}秒") # 6. 清理结果 description = self._clean_description(description) # 7. 格式化返回结果 formatted_text = f"图{image_index}: {description}" total_time = time.time() - start_time print(f"✅ 识别完成: {description} (总耗时: {total_time:.2f}秒)") return { "description": description, "descriptionEn": "", # Qwen2-VL直接输出中文 "confidence": 0.95, "imageIndex": image_index, "formattedText": formatted_text } def _generate_sync(self, inputs): """同步生成（在线程池中执行）""" with torch.no_grad(): generated_ids = self.model.generate( **inputs, max_new_tokens=20, # 进一步减少生成长度加速（30→20） do_sample=False, temperature=0.0, use_cache=True # 启用KV缓存加速 ) # 只解码生成的部分 generated_ids_trimmed = [ out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) ] output_text = self.processor.batch_decode( generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False )[0] return output_text def _clean_description(self, text: str) -> str: """清理描述文本，移除不需要的前缀和格式""" import re text = text.strip() # 移除常见的前缀（Qwen2-VL可能生成的） prefixes_to_remove = [ r"红色圆圈标记中心位置的物体是\s*", r"红色圆圈标记中心位置的\s*", r"标记中心位置的物体是\s*", r"红色圆圈标记的物体是\s*", r"红色圆圈标记的是\s*", r"红色圆圈标记[：:]\s*", r"红色标记的物体是\s*", r"红色标记[：:]\s*", r"标记的物体是\s*", r"标记位置的物体是\s*", r"标记位置[：:]\s*", r"标记的位置[：:]\s*", r"标记处[：:]\s*", r"标记中心位置的\s*", r"中心位置的物体是\s*", r"中心位置的\s*", r"物体是\s*", r"图片中[，,、]?\s*", r"这是\s*", r"这个是\s*", r"标记的是\s*", ] for prefix in prefixes_to_remove: text = re.sub(f"^{prefix}", "", text, flags=re.IGNORECASE) # 移除开头的标点符号 text = re.sub(r"^[：:，,、。\s]+", "", text) # 如果包含多个物品（用顿号、逗号分隔），只取第一个 if '、' in text: text = text.split('、')[0] elif '，' in text: text = text.split('，')[0] elif ',' in text: text = text.split(',')[0] # 限制长度（最多20个字，避免截断） if len(text) > 20: text = text[:20] return text.strip() # 全局单例 _qwen_service = None def get_qwen_service() -> QwenVLService: """获取全局Qwen服务实例""" global _qwen_service if _qwen_service is None: _qwen_service = QwenVLService() return _qwen_service