299 lines
10 KiB
Python
299 lines
10 KiB
Python
"""
|
||
VLM图像处理工具模块
|
||
包含图像识别、图像压缩、图像保存等功能
|
||
"""
|
||
import os
|
||
import re
|
||
import json
|
||
import math
|
||
import base64
|
||
import uuid
|
||
from datetime import datetime
|
||
from typing import Optional, Dict, Any
|
||
from PIL import Image
|
||
from io import BytesIO
|
||
import tiktoken
|
||
from langchain_core.tools import tool
|
||
|
||
from config import VLM_CONFIG, SERVER_CONFIG
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
from utils.function_tracker import track_function_calls
|
||
|
||
|
||
VLM_API_BASE = VLM_CONFIG["api_base"]
|
||
VLM_MODEL_NAME = VLM_CONFIG["model"]
|
||
MAX_MODEL_LEN = 32768
|
||
INPUT_TOKEN_SAFETY_RATIO = 0.8
|
||
MIN_IMAGE_SIZE = 224
|
||
JPEG_QUALITY = 85
|
||
|
||
SYSTEM_PROMPT_VLM = """你是一个专业的工业设备巡检专家助手。
|
||
你的任务是协助工程师分析现场情况。
|
||
|
||
当你遇到以下情况时：
|
||
1. 用户提供了图片路径（如 .png, .jpg, /app/data/... 等）。
|
||
2. 用户询问有关设备外观、颜色、仪表读数或物理损坏的问题。
|
||
|
||
你应该：
|
||
- 调用 `vlm_image_to_text` 工具。
|
||
- 在 `prompt` 参数中，根据用户的具体意图（如：检查漏油、读仪表、看指示灯）量身定制分析指令。
|
||
- 获取 VLM 结果后，结合专业知识给出建议。
|
||
"""
|
||
|
||
|
||
def encode_image(image_path: str) -> str:
|
||
"""将本地图片转换为 Base64 编码"""
|
||
with open(image_path, "rb") as f:
|
||
return base64.b64encode(f.read()).decode("utf-8")
|
||
|
||
|
||
def estimate_qwen_vl_visual_tokens(width: int, height: int) -> int:
|
||
"""
|
||
根据 Qwen-VL 的图像分块策略估算视觉 tokens。
|
||
公式来源：Qwen-VL 官方技术文档（Naive resizing + 192 tokens per unit）
|
||
"""
|
||
units_h = math.ceil(height / 336)
|
||
units_w = math.ceil(width / 336)
|
||
total_units = units_h * units_w + 1
|
||
return total_units * 192
|
||
|
||
|
||
def count_text_tokens(text: str) -> int:
|
||
"""使用 tiktoken 估算文本 token 数（兼容中英文）"""
|
||
enc = tiktoken.get_encoding("cl100k_base")
|
||
return len(enc.encode(text))
|
||
|
||
|
||
def smart_resize_base64_image(
|
||
base64_image: str,
|
||
prompt_text: str,
|
||
max_input_tokens: int = int(MAX_MODEL_LEN * INPUT_TOKEN_SAFETY_RATIO),
|
||
min_size: int = MIN_IMAGE_SIZE,
|
||
jpeg_quality: int = JPEG_QUALITY
|
||
) -> str:
|
||
"""
|
||
智能压缩 base64 图像，确保 (文本 tokens + 视觉 tokens) <= max_input_tokens
|
||
"""
|
||
if not base64_image.startswith("data:image"):
|
||
raise ValueError("Invalid base64 image format. Must start with 'data:image'.")
|
||
|
||
try:
|
||
header, b64_data = base64_image.split(",", 1)
|
||
image_data = base64.b64decode(b64_data)
|
||
img = Image.open(BytesIO(image_data)).convert("RGB")
|
||
orig_w, orig_h = img.size
|
||
except Exception as e:
|
||
raise ValueError(f"Failed to decode base64 image: {e}")
|
||
|
||
text_tokens = count_text_tokens(prompt_text)
|
||
visual_tokens = estimate_qwen_vl_visual_tokens(orig_w, orig_h)
|
||
total_tokens = text_tokens + visual_tokens
|
||
|
||
print(f"[VLM Preprocess] 原始尺寸: {orig_w}x{orig_h} | 文本 tokens: {text_tokens} | 视觉 tokens: {visual_tokens} | 总计: {total_tokens}")
|
||
|
||
if total_tokens <= max_input_tokens:
|
||
print("[VLM Preprocess] ✅ 无需压缩")
|
||
return base64_image
|
||
|
||
scale = 1.0
|
||
target_tokens = total_tokens
|
||
while target_tokens > max_input_tokens and scale > 0.1:
|
||
scale *= 0.8
|
||
new_w = max(min_size, int(orig_w * scale))
|
||
new_h = max(min_size, int(orig_h * scale))
|
||
target_tokens = text_tokens + estimate_qwen_vl_visual_tokens(new_w, new_h)
|
||
|
||
final_w = max(min_size, int(orig_w * scale))
|
||
final_h = max(min_size, int(orig_h * scale))
|
||
|
||
print(f"[VLM Preprocess] ⚠️ 需压缩 → 新尺寸: {final_w}x{final_h} | 预估总 tokens: {target_tokens}")
|
||
|
||
resized_img = img.resize((final_w, final_h), Image.Resampling.LANCZOS)
|
||
buffer = BytesIO()
|
||
resized_img.save(buffer, format="JPEG", quality=jpeg_quality)
|
||
compressed_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||
return f"data:image/jpeg;base64,{compressed_b64}"
|
||
|
||
|
||
def save_base64_image_and_record(
|
||
base64_image: str,
|
||
description: str,
|
||
output_dir: str = "image_input",
|
||
record_file: str = "image_input/image_records.json"
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
保存 base64 图片到指定目录，并记录图片名称和描述到 JSON 文件。
|
||
|
||
Args:
|
||
base64_image: base64 编码的图片数据（data:image/xxx;base64,... 格式）
|
||
description: 图片的描述文本
|
||
output_dir: 图片保存目录
|
||
record_file: 记录 JSON 文件路径
|
||
|
||
Returns:
|
||
包含保存结果的字典：{"success": bool, "image_name": str, "image_path": str, "error": str}
|
||
"""
|
||
try:
|
||
if not base64_image.startswith("data:image"):
|
||
return {"success": False, "error": "Invalid base64 image format"}
|
||
|
||
header, b64_data = base64_image.split(",", 1)
|
||
|
||
if "jpeg" in header or "jpg" in header:
|
||
ext = "jpg"
|
||
elif "png" in header:
|
||
ext = "png"
|
||
elif "gif" in header:
|
||
ext = "gif"
|
||
elif "webp" in header:
|
||
ext = "webp"
|
||
else:
|
||
ext = "jpg"
|
||
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
image_name = f"{uuid.uuid4().hex[:12]}_{datetime.now().strftime('%Y%m%d%H%M%S')}.{ext}"
|
||
image_path = os.path.join(output_dir, image_name)
|
||
|
||
image_data = base64.b64decode(b64_data)
|
||
with open(image_path, "wb") as f:
|
||
f.write(image_data)
|
||
|
||
record_entry = {
|
||
"image_name": image_name,
|
||
"image_path": image_path,
|
||
"description": description,
|
||
"timestamp": datetime.now().isoformat()
|
||
}
|
||
|
||
os.makedirs(os.path.dirname(record_file), exist_ok=True)
|
||
|
||
existing_records = []
|
||
if os.path.exists(record_file):
|
||
try:
|
||
with open(record_file, "r", encoding="utf-8") as f:
|
||
existing_records = json.load(f)
|
||
if not isinstance(existing_records, list):
|
||
existing_records = []
|
||
except (json.JSONDecodeError, IOError):
|
||
existing_records = []
|
||
|
||
existing_records.append(record_entry)
|
||
|
||
with open(record_file, "w", encoding="utf-8") as f:
|
||
json.dump(existing_records, f, ensure_ascii=False, indent=2)
|
||
|
||
print(f"[ImageSaved] 图片已保存: {image_path}")
|
||
|
||
return {
|
||
"success": True,
|
||
"image_name": image_name,
|
||
"image_path": image_path,
|
||
"error": None
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
"success": False,
|
||
"image_name": None,
|
||
"image_path": None,
|
||
"error": str(e)
|
||
}
|
||
|
||
|
||
@tool
|
||
async def vlm_image_to_text(
|
||
image_path: str,
|
||
prompt: Optional[str] = None
|
||
) -> str:
|
||
"""
|
||
调用 Qwen3.5-35B 进行工业巡检分析。
|
||
- 自动移除所有 <think...</think 包裹的内容（包括多行）
|
||
- 仅返回最终结果文本
|
||
- 支持 40K token 上下文，自动压缩大图
|
||
"""
|
||
|
||
default_prompt = (
|
||
"【重要指令】禁止输出任何形式的 <think...</think、<reasoning></reasoning> 或其他内部思考过程。"
|
||
"仅输出最终的专业分析结果。\n\n"
|
||
"你是一位资深的工业巡检专家。请仔细观察这张图片并按以下结构报告：\n"
|
||
"1. 设备识别：图中是什么设备？型号或标签是否可见？\n"
|
||
"2. 状态评估：仪表盘数值是否在正常区间？指示灯颜色是什么？\n"
|
||
"3. 异常检测：是否存在漏油、生锈、电缆破损、螺丝松动或烟雾？\n"
|
||
"4. 环境风险：周围是否有杂物堆放或安全隐患？\n"
|
||
"请用专业、客观、简短的语言描述。"
|
||
)
|
||
|
||
final_prompt = f"{default_prompt}\n特别注意：{prompt}" if prompt else default_prompt
|
||
|
||
safe_image = smart_resize_base64_image(
|
||
base64_image=image_path,
|
||
prompt_text=final_prompt,
|
||
max_input_tokens=int(MAX_MODEL_LEN * INPUT_TOKEN_SAFETY_RATIO)
|
||
)
|
||
|
||
try:
|
||
result = await OpenaiAPI.open_api_vl_without_thinking(safe_image)
|
||
|
||
print(f"✅ 最终输出:\n{result}")
|
||
|
||
save_result = save_base64_image_and_record(
|
||
base64_image=image_path,
|
||
description=result,
|
||
output_dir="picture",
|
||
record_file="picture/image_records.json"
|
||
)
|
||
if save_result["success"]:
|
||
print(f"📷 图片已保存: {save_result['image_name']}")
|
||
image_url = f"{SERVER_CONFIG['base_url']}/picture/{save_result['image_name']}"
|
||
result = f"{result}\n\n"
|
||
else:
|
||
print(f"⚠️ 图片保存失败: {save_result['error']}")
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
error_msg = f"VLM 调用失败: {str(e)}\n{traceback.format_exc()}"
|
||
print(error_msg)
|
||
return error_msg
|
||
|
||
|
||
@track_function_calls
|
||
async def image_rag_describe(
|
||
image_path: str,
|
||
query: str,
|
||
top_k: int = 5
|
||
) -> str:
|
||
"""
|
||
图像RAG描述工具
|
||
先用VLM分析图像，然后用RAG检索相关知识
|
||
|
||
Args:
|
||
image_path: 图像路径（base64格式）
|
||
query: 查询文本
|
||
top_k: RAG检索数量
|
||
|
||
Returns:
|
||
分析结果文本
|
||
"""
|
||
try:
|
||
vlm_result = await vlm_image_to_text(image_path, prompt=query)
|
||
|
||
from .rag_tools import rag_search
|
||
rag_result = await rag_search(query=f"{query} {vlm_result}", top_k=top_k)
|
||
|
||
if rag_result.get("success", False):
|
||
source_citation = rag_result.get("sourceCitation", {})
|
||
rag_text = ""
|
||
for resource, items in source_citation.items():
|
||
for item in items:
|
||
rag_text += f"【{resource}】{item.get('text', '')}\n\n"
|
||
|
||
return f"图像分析结果：\n{vlm_result}\n\n相关知识：\n{rag_text}"
|
||
else:
|
||
return vlm_result
|
||
|
||
except Exception as e:
|
||
return f"图像RAG分析失败: {str(e)}"
|