298 lines
10 KiB
Python
298 lines
10 KiB
Python
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
|
||
|