614 lines
22 KiB
Python
614 lines
22 KiB
Python
import torch
|
||
from transformers import AutoProcessor, AutoModelForCausalLM, MarianMTModel, MarianTokenizer
|
||
from PIL import Image, ImageDraw
|
||
import requests
|
||
from io import BytesIO
|
||
import asyncio
|
||
import re
|
||
|
||
class FlorenceService:
|
||
def __init__(self):
|
||
self.model = None
|
||
self.processor = None
|
||
self.translator_model = None
|
||
self.translator_tokenizer = 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("📦 加载 Florence-2 模型...")
|
||
self.model = AutoModelForCausalLM.from_pretrained(
|
||
"microsoft/Florence-2-base",
|
||
trust_remote_code=True
|
||
).to(self.device)
|
||
|
||
self.processor = AutoProcessor.from_pretrained(
|
||
"microsoft/Florence-2-base",
|
||
trust_remote_code=True
|
||
)
|
||
|
||
print("✅ Florence-2 模型加载完成")
|
||
|
||
# 加载英中翻译模型
|
||
print("📦 加载英中翻译模型...")
|
||
translator_name = "Helsinki-NLP/opus-mt-en-zh"
|
||
self.translator_tokenizer = MarianTokenizer.from_pretrained(translator_name)
|
||
self.translator_model = MarianMTModel.from_pretrained(translator_name).to(self.device)
|
||
|
||
self.is_loaded = True
|
||
print("✅ 翻译模型加载完成")
|
||
|
||
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:
|
||
"""检测点击位置的物体（使用大区域裁剪 + 详细描述）"""
|
||
|
||
# 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"📊 图片信息:")
|
||
print(f" - 原始尺寸: {actual_width}x{actual_height}")
|
||
print(f" - 传入尺寸: {image_width}x{image_height}")
|
||
print(f" - 点击坐标: ({click_x}, {click_y})")
|
||
|
||
# 2. 在图片上绘制标注点，让模型能"看到"用户的标注
|
||
marked_image = image.copy()
|
||
draw = ImageDraw.Draw(marked_image)
|
||
|
||
# 绘制醒目的标注标记（红色圆圈 + 十字）
|
||
marker_size = 30
|
||
# 外圈（红色）
|
||
draw.ellipse(
|
||
[click_x - marker_size, click_y - marker_size,
|
||
click_x + marker_size, click_y + marker_size],
|
||
outline=(255, 0, 0),
|
||
width=5
|
||
)
|
||
# 十字线（红色）
|
||
draw.line(
|
||
[click_x - marker_size - 10, click_y, click_x + marker_size + 10, click_y],
|
||
fill=(255, 0, 0),
|
||
width=5
|
||
)
|
||
draw.line(
|
||
[click_x, click_y - marker_size - 10, click_x, click_y + marker_size + 10],
|
||
fill=(255, 0, 0),
|
||
width=5
|
||
)
|
||
|
||
print(f"✏️ 已在图片上绘制标注点: ({click_x}, {click_y})")
|
||
|
||
# 3. 使用标记后的完整图片进行物体检测
|
||
print("🔍 正在检测标记后的完整图片...")
|
||
loop = asyncio.get_event_loop()
|
||
detected_object = await loop.run_in_executor(
|
||
None,
|
||
self._detect_object_at_point_sync,
|
||
marked_image,
|
||
click_x,
|
||
click_y
|
||
)
|
||
|
||
# 3. 翻译成中文（如果是英文）
|
||
description_cn = await self._translate_to_chinese(detected_object)
|
||
|
||
# 4. 格式化返回结果（无图标，纯文本）
|
||
formatted_text = f"图{image_index}: {description_cn}"
|
||
|
||
return {
|
||
"description": description_cn,
|
||
"descriptionEn": detected_object,
|
||
"confidence": 0.95, # Florence-2 不返回置信度，固定值
|
||
"imageIndex": image_index,
|
||
"formattedText": formatted_text
|
||
}
|
||
|
||
def _detect_object_at_point_sync(self, image: Image.Image, click_x: int, click_y: int) -> str:
|
||
"""智能裁剪策略：聚焦在点击位置的具体物品，而不是整个人物"""
|
||
|
||
# 步骤1：使用物体检测找到点击位置的物体边界框
|
||
print("🔍 步骤1: 检测所有物体...")
|
||
bbox = self._find_object_bbox(image, click_x, click_y)
|
||
|
||
if bbox:
|
||
x1, y1, x2, y2 = bbox
|
||
bbox_width = x2 - x1
|
||
bbox_height = y2 - y1
|
||
|
||
print(f"📦 检测到物体边界框: ({x1:.0f}, {y1:.0f}) -> ({x2:.0f}, {y2:.0f})")
|
||
print(f" 尺寸: {bbox_width:.0f}x{bbox_height:.0f}")
|
||
|
||
# 判断：如果是大物体（人物），使用点击点周围的固定区域
|
||
# 如果是小物体，使用完整边界框
|
||
if bbox_width > 400 or bbox_height > 600:
|
||
# 大物体（可能是人物）：以点击点为中心裁剪固定区域
|
||
print(" → 检测到大物体（人物），聚焦点击位置")
|
||
crop_size = 300
|
||
x1 = max(0, click_x - crop_size // 2)
|
||
y1 = max(0, click_y - crop_size // 2)
|
||
x2 = min(image.width, click_x + crop_size // 2)
|
||
y2 = min(image.height, click_y + crop_size // 2)
|
||
|
||
print(f"✂️ 聚焦裁剪: ({x1}, {y1}) -> ({x2}, {y2}), 尺寸: {x2-x1}x{y2-y1}")
|
||
else:
|
||
# 小物体：使用完整边界框+扩展
|
||
print(" → 检测到小物体，使用完整边界框")
|
||
padding = 0.15
|
||
pad_x = int(bbox_width * padding)
|
||
pad_y = int(bbox_height * padding)
|
||
|
||
x1 = max(0, int(x1 - pad_x))
|
||
y1 = max(0, int(y1 - pad_y))
|
||
x2 = min(image.width, int(x2 + pad_x))
|
||
y2 = min(image.height, int(y2 + pad_y))
|
||
|
||
print(f"✂️ 扩展边界框: ({x1}, {y1}) -> ({x2}, {y2})")
|
||
|
||
cropped_region = image.crop((x1, y1, x2, y2))
|
||
else:
|
||
# 如果没有检测到物体，使用标注点周围的中等区域
|
||
print("⚠️ 未检测到物体边界框，使用标注点周围区域")
|
||
crop_size = 300
|
||
x1 = max(0, click_x - crop_size // 2)
|
||
y1 = max(0, click_y - crop_size // 2)
|
||
x2 = min(image.width, click_x + crop_size // 2)
|
||
y2 = min(image.height, click_y + crop_size // 2)
|
||
|
||
print(f"✂️ 裁剪区域: ({x1}, {y1}) -> ({x2}, {y2})")
|
||
cropped_region = image.crop((x1, y1, x2, y2))
|
||
|
||
# 步骤2：对裁剪区域进行详细描述
|
||
print("🔍 步骤2: 生成详细描述...")
|
||
description = self._describe_region(cropped_region)
|
||
|
||
print(f"✅ 最终识别: {description}")
|
||
return description
|
||
|
||
def _find_object_bbox(self, image: Image.Image, click_x: int, click_y: int):
|
||
"""使用物体检测找到点击位置的物体边界框"""
|
||
import time
|
||
start_time = time.time()
|
||
|
||
prompt = "<OD>"
|
||
|
||
inputs = self.processor(
|
||
text=prompt,
|
||
images=image,
|
||
return_tensors="pt"
|
||
)
|
||
|
||
# 显式移动所有tensor到GPU
|
||
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||
|
||
print(f" → 推理设备: {next(self.model.parameters()).device}")
|
||
|
||
with torch.no_grad():
|
||
generated_ids = self.model.generate(
|
||
input_ids=inputs["input_ids"],
|
||
pixel_values=inputs["pixel_values"],
|
||
max_new_tokens=1024,
|
||
num_beams=3,
|
||
do_sample=False
|
||
)
|
||
|
||
elapsed = time.time() - start_time
|
||
print(f" → 物体检测耗时: {elapsed:.2f}秒")
|
||
|
||
generated_text = self.processor.batch_decode(
|
||
generated_ids,
|
||
skip_special_tokens=False
|
||
)[0]
|
||
|
||
parsed_results = self.processor.post_process_generation(
|
||
generated_text,
|
||
task="<OD>",
|
||
image_size=(image.width, image.height)
|
||
)
|
||
|
||
if parsed_results and '<OD>' in parsed_results:
|
||
od_result = parsed_results['<OD>']
|
||
bboxes = od_result.get('bboxes', [])
|
||
labels = od_result.get('labels', [])
|
||
|
||
print(f" 检测到 {len(bboxes)} 个物体")
|
||
|
||
# 找到包含点击点的边界框
|
||
for i, (bbox, label) in enumerate(zip(bboxes, labels)):
|
||
x1, y1, x2, y2 = bbox
|
||
if x1 <= click_x <= x2 and y1 <= click_y <= y2:
|
||
print(f" ✓ 找到匹配物体 {i+1}: {label}")
|
||
return bbox
|
||
|
||
# 如果没有包含的，找最近的
|
||
print(" 未找到包含点击点的物体，寻找最近的...")
|
||
min_distance = float('inf')
|
||
closest_bbox = None
|
||
|
||
for bbox in bboxes:
|
||
x1, y1, x2, y2 = bbox
|
||
center_x = (x1 + x2) / 2
|
||
center_y = (y1 + y2) / 2
|
||
distance = ((center_x - click_x) ** 2 + (center_y - click_y) ** 2) ** 0.5
|
||
|
||
if distance < min_distance:
|
||
min_distance = distance
|
||
closest_bbox = bbox
|
||
|
||
if closest_bbox:
|
||
print(f" ✓ 最近物体距离: {min_distance:.0f}px")
|
||
return closest_bbox
|
||
|
||
return None
|
||
|
||
def _describe_region(self, region_image: Image.Image) -> str:
|
||
"""对区域图像进行详细描述"""
|
||
import time
|
||
start_time = time.time()
|
||
|
||
prompt = "<MORE_DETAILED_CAPTION>"
|
||
|
||
inputs = self.processor(
|
||
text=prompt,
|
||
images=region_image,
|
||
return_tensors="pt"
|
||
)
|
||
|
||
# 显式移动所有tensor到GPU
|
||
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||
|
||
with torch.no_grad():
|
||
generated_ids = self.model.generate(
|
||
input_ids=inputs["input_ids"],
|
||
pixel_values=inputs["pixel_values"],
|
||
max_new_tokens=50,
|
||
num_beams=3,
|
||
do_sample=False
|
||
)
|
||
|
||
elapsed = time.time() - start_time
|
||
print(f" → 描述生成耗时: {elapsed:.2f}秒")
|
||
|
||
generated_text = self.processor.batch_decode(
|
||
generated_ids,
|
||
skip_special_tokens=True
|
||
)[0]
|
||
|
||
print(f" 原始描述: {generated_text}")
|
||
|
||
# 提取核心物品名称
|
||
description = self._extract_item_name(generated_text)
|
||
|
||
return description
|
||
|
||
def _extract_item_name(self, description: str) -> str:
|
||
"""从详细描述中提取核心物品名称"""
|
||
# 移除常见的前缀短语（按长度排序，先匹配长的）
|
||
remove_phrases = [
|
||
"the image is a promotional poster ",
|
||
"the image is a portrait ",
|
||
"this image shows ",
|
||
"the image shows ",
|
||
"the image is a ",
|
||
"the image is ",
|
||
"a close-up of ",
|
||
"a close up of ",
|
||
"close-up of ",
|
||
"close up of ",
|
||
"an image of ",
|
||
"a picture of ",
|
||
"a photo of ",
|
||
"there is a ",
|
||
"there is ",
|
||
"this is a ",
|
||
"this is ",
|
||
"it is a ",
|
||
"it is ",
|
||
]
|
||
|
||
description_lower = description.lower().strip()
|
||
|
||
# 持续删除前缀，直到没有匹配的
|
||
changed = True
|
||
while changed:
|
||
changed = False
|
||
for phrase in remove_phrases:
|
||
if description_lower.startswith(phrase):
|
||
description_lower = description_lower[len(phrase):].strip()
|
||
changed = True
|
||
break
|
||
|
||
# 取前4-5个词作为物品名称
|
||
words = description_lower.split()
|
||
|
||
if len(words) >= 4:
|
||
item_name = " ".join(words[:4])
|
||
elif len(words) >= 2:
|
||
item_name = " ".join(words[:2])
|
||
else:
|
||
item_name = description_lower
|
||
|
||
return item_name.strip()
|
||
|
||
def _detect_with_context(self, image: Image.Image) -> str:
|
||
"""使用大区域上下文识别主要物体（在线程池中执行）"""
|
||
# 使用 <MORE_DETAILED_CAPTION> 生成详细描述，包含颜色、材质等
|
||
prompt = "<MORE_DETAILED_CAPTION>"
|
||
|
||
print(f"🎯 使用 prompt: {prompt}")
|
||
|
||
inputs = self.processor(
|
||
text=prompt,
|
||
images=image,
|
||
return_tensors="pt"
|
||
).to(self.device)
|
||
|
||
# 生成描述
|
||
with torch.no_grad():
|
||
generated_ids = self.model.generate(
|
||
input_ids=inputs["input_ids"],
|
||
pixel_values=inputs["pixel_values"],
|
||
max_new_tokens=40, # 增加长度以获取更多细节
|
||
num_beams=3,
|
||
do_sample=False
|
||
)
|
||
|
||
# 解码结果
|
||
generated_text = self.processor.batch_decode(
|
||
generated_ids,
|
||
skip_special_tokens=False
|
||
)[0]
|
||
|
||
# 提取描述（去除提示词和特殊标记）
|
||
description = generated_text.replace(prompt, "").strip()
|
||
description = description.replace("</s>", "").replace("<s>", "").strip()
|
||
description = " ".join(description.split())
|
||
|
||
# 🎯 提取核心信息（颜色+物体）
|
||
description = self._extract_core_description(description)
|
||
|
||
print(f"✅ 识别结果（带上下文）: {description}")
|
||
return description
|
||
|
||
def _detect_region_sync(self, image: Image.Image, bbox: tuple) -> str:
|
||
"""使用区域标注识别完整图片中的指定区域（在线程池中执行）"""
|
||
# 使用 <REGION_TO_CAPTION> 任务，识别完整图片中的指定区域
|
||
# bbox 格式：(x1, y1, x2, y2)，归一化到 [0, 1000]
|
||
x1, y1, x2, y2 = bbox
|
||
prompt = f"<REGION_TO_CAPTION><loc_{x1}><loc_{y1}><loc_{x2}><loc_{y2}>"
|
||
|
||
print(f"🎯 使用 prompt: {prompt}")
|
||
|
||
inputs = self.processor(
|
||
text=prompt,
|
||
images=image,
|
||
return_tensors="pt"
|
||
).to(self.device)
|
||
|
||
# 生成描述
|
||
with torch.no_grad():
|
||
generated_ids = self.model.generate(
|
||
input_ids=inputs["input_ids"],
|
||
pixel_values=inputs["pixel_values"],
|
||
max_new_tokens=30,
|
||
num_beams=3,
|
||
do_sample=False
|
||
)
|
||
|
||
# 解码结果
|
||
generated_text = self.processor.batch_decode(
|
||
generated_ids,
|
||
skip_special_tokens=False
|
||
)[0]
|
||
|
||
# 提取描述（去除提示词和特殊标记）
|
||
# Florence-2 返回格式: <REGION_TO_CAPTION>description</s>
|
||
description = generated_text.replace("<REGION_TO_CAPTION>", "").strip()
|
||
|
||
# 清理特殊标记和位置标记
|
||
description = re.sub(r'<loc_\d+>', '', description)
|
||
description = description.replace("</s>", "").replace("<s>", "").strip()
|
||
description = " ".join(description.split())
|
||
|
||
# 🎯 提取核心信息（颜色+物体）
|
||
description = self._extract_core_description(description)
|
||
|
||
print(f"✅ 区域识别结果: {description}")
|
||
return description
|
||
|
||
def _detect_sync(self, image: Image.Image) -> str:
|
||
"""同步推理（在线程池中执行）- 已弃用，使用 _detect_region_sync"""
|
||
# 使用 <MORE_DETAILED_CAPTION> 生成更详细的描述，包含颜色和细节
|
||
prompt = "<MORE_DETAILED_CAPTION>"
|
||
|
||
inputs = self.processor(
|
||
text=prompt,
|
||
images=image,
|
||
return_tensors="pt"
|
||
).to(self.device)
|
||
|
||
# 生成描述
|
||
with torch.no_grad():
|
||
generated_ids = self.model.generate(
|
||
input_ids=inputs["input_ids"],
|
||
pixel_values=inputs["pixel_values"],
|
||
max_new_tokens=30, # 增加长度以获取更多细节
|
||
num_beams=3,
|
||
do_sample=False
|
||
)
|
||
|
||
# 解码结果
|
||
generated_text = self.processor.batch_decode(
|
||
generated_ids,
|
||
skip_special_tokens=False
|
||
)[0]
|
||
|
||
# 提取描述（去除提示词和特殊标记）
|
||
description = generated_text.replace(prompt, "").strip()
|
||
|
||
# 清理特殊标记
|
||
description = description.replace("</s>", "").replace("<s>", "").strip()
|
||
# 清理多余的空格
|
||
description = " ".join(description.split())
|
||
|
||
# 🎯 提取核心信息（颜色+物体）
|
||
description = self._extract_core_description(description)
|
||
|
||
print(f"✅ 识别结果: {description}")
|
||
return description
|
||
|
||
def _extract_core_description(self, text: str) -> str:
|
||
"""提取核心描述（只保留主要物体名称）"""
|
||
# 移除所有介绍性和描述性短语（包括摄影术语）
|
||
remove_phrases = [
|
||
"a close up of ", # 特写镜头（优先匹配）
|
||
"a close-up of ",
|
||
"close up of ",
|
||
"close-up of ",
|
||
"the image shows ",
|
||
"the photo shows ",
|
||
"this is a ",
|
||
"this is an ",
|
||
"this is ",
|
||
"there is a ",
|
||
"there is an ",
|
||
"there is ",
|
||
"there are ",
|
||
"a picture of ",
|
||
"an image of ",
|
||
"it shows ",
|
||
"shows ",
|
||
"depicts ",
|
||
"featuring ",
|
||
"with ",
|
||
" on top of ",
|
||
" on a ",
|
||
" on the ",
|
||
" in a ",
|
||
" in the ",
|
||
" sitting on ",
|
||
" lying on ",
|
||
" placed on ",
|
||
" next to ",
|
||
" beside ",
|
||
]
|
||
|
||
result = text.lower()
|
||
|
||
# 先移除介绍性短语（按长度降序，优先匹配长短语）
|
||
for phrase in sorted(remove_phrases, key=len, reverse=True):
|
||
result = result.replace(phrase, " ")
|
||
|
||
# 只保留第一个逗号或句号之前的内容
|
||
for sep in [". ", ", ", " and ", " with "]:
|
||
if sep in result:
|
||
result = result.split(sep)[0]
|
||
break
|
||
|
||
# 移除冠词
|
||
result = result.replace(" a ", " ").replace(" an ", " ").replace(" the ", " ")
|
||
|
||
# 清理多余空格
|
||
result = " ".join(result.split()).strip()
|
||
|
||
# 限制长度（最多3-4个词）
|
||
words = result.split()
|
||
if len(words) > 4:
|
||
result = " ".join(words[:4])
|
||
|
||
return result
|
||
|
||
async def _translate_to_chinese(self, text: str) -> str:
|
||
"""使用关键词映射进行准确翻译（替代MarianMT）"""
|
||
if not text or not text.strip():
|
||
return text
|
||
|
||
# 如果已经是中文，直接返回
|
||
if any('\u4e00' <= char <= '\u9fff' for char in text):
|
||
return text
|
||
|
||
try:
|
||
# 使用关键词映射
|
||
translated = self._translate_with_keywords(text)
|
||
print(f"🌐 翻译: {text} -> {translated}")
|
||
return translated
|
||
except Exception as e:
|
||
print(f"⚠️ 翻译失败: {e}，返回原文")
|
||
return text
|
||
|
||
def _translate_with_keywords(self, text: str) -> str:
|
||
"""使用关键词映射进行翻译"""
|
||
keyword_map = {
|
||
# 家具
|
||
"chair": "椅子", "sofa": "沙发", "couch": "沙发",
|
||
"bean bag": "懒人沙发", "beanbag": "懒人沙发",
|
||
"table": "桌子", "desk": "书桌", "bed": "床",
|
||
|
||
# 织物
|
||
"towel": "毛巾", "pillow": "枕头", "cushion": "靠垫",
|
||
"blanket": "毯子", "throw": "抱枕",
|
||
"rug": "地毯", "carpet": "地毯", "mat": "垫子",
|
||
"fabric": "织物", "knit": "针织", "knitted": "针织",
|
||
"fluffy": "毛绒", "fuzzy": "毛茸茸",
|
||
|
||
# 颜色
|
||
"green": "绿色", "red": "红色", "blue": "蓝色",
|
||
"yellow": "黄色", "white": "白色", "black": "黑色",
|
||
"brown": "棕色", "pink": "粉色", "beige": "米色",
|
||
"gray": "灰色", "grey": "灰色",
|
||
|
||
# 物品
|
||
"book": "书", "books": "书籍", "magazine": "杂志",
|
||
"stack": "堆叠", "pile": "堆", "bottle": "瓶子",
|
||
"plant": "植物", "vase": "花瓶",
|
||
|
||
# 衣物
|
||
"sweater": "毛衣", "shirt": "衬衫",
|
||
"jacket": "夹克", "coat": "外套",
|
||
"dress": "连衣裙", "scarf": "围巾",
|
||
|
||
# 形状
|
||
"heart": "心形", "round": "圆形",
|
||
"square": "方形",
|
||
|
||
# 状态
|
||
"lying": "放置", "sitting": "放置",
|
||
"resting": "靠着", "open": "打开",
|
||
"with": "带有", "and": "和",
|
||
}
|
||
|
||
words = text.lower().split()
|
||
translated_words = []
|
||
|
||
for word in words:
|
||
clean_word = word.strip('.,!?;:')
|
||
translated_words.append(keyword_map.get(clean_word, word))
|
||
|
||
return " ".join(translated_words)
|
||
|