ZuoWei
a6c42d505a
新增功能: - 天网协作系统 (HTTP API 端口 6060) - 三种工作流 (查找图片/处理图片/转人工派单) - 图片任务数据库 (支持客户后续增加需求) - 图绘派单系统集成 (API: 8005) - 文字检测与加价 (60-80 元高价值订单) - 风险评估与接单判断 - 作图失败自动转人工 新增文档: - 项目功能汇总.md - 三种工作流功能说明.md - 文字加价功能说明.md - 风险评估功能说明.md - 图片任务数据库功能说明.md - 图绘派单系统集成说明.md - 作图失败转接人工说明.md - DEPLOYMENT.md - TIANWANG_INTEGRATION.md 核心修改: - core/pydantic_ai_agent.py - core/workflow.py - core/websocket_client.py - image/image_analyzer.py - services/service_tuhui_dispatch.py - db/image_tasks_db.py 版本：v1.0 日期：2026-02-28
329 lines
14 KiB
Python
Executable File
329 lines
14 KiB
Python
Executable File
"""图片处理模块 - 调用 Gemini 作图API，含质检与自动重试"""
|
||
import os
|
||
import uuid
|
||
import tempfile
|
||
from typing import Optional, Dict, Any
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
_OUTPUT_DIR = os.getenv("RESULT_IMAGE_DIR", "results")
|
||
_MAX_RETRIES = int(os.getenv("PROCESS_MAX_RETRIES", "2")) # 含首次共最多处理几次
|
||
|
||
|
||
class ImageProcessor:
|
||
"""图片处理 - 对接 GeminiExtractV2Service，含质检与重试"""
|
||
|
||
def __init__(self):
|
||
os.makedirs(_OUTPUT_DIR, exist_ok=True)
|
||
|
||
# ─── 内部工具 ────────────────────────────────────────────
|
||
|
||
async def _download(self, url: str) -> str:
|
||
"""下载图片到临时文件，返回本地路径"""
|
||
import aiohttp
|
||
tmp = os.path.join(tempfile.gettempdir(), f"gemini_in_{uuid.uuid4().hex}.jpg")
|
||
headers = {
|
||
"User-Agent": (
|
||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||
"Chrome/122.0.0.0 Safari/537.36"
|
||
),
|
||
"Referer": "https://www.taobao.com/",
|
||
"Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
|
||
}
|
||
async with aiohttp.ClientSession(headers=headers) as session:
|
||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=30)) as resp:
|
||
if resp.status != 200:
|
||
raise RuntimeError(f"下载图片失败: HTTP {resp.status}")
|
||
with open(tmp, "wb") as f:
|
||
f.write(await resp.read())
|
||
return tmp
|
||
|
||
async def _do_perspective(self, service, src: str, level: str) -> str:
|
||
"""透视矫正，返回矫正后文件路径（失败则返回原路径）"""
|
||
out = os.path.join(tempfile.gettempdir(), f"gemini_persp_{uuid.uuid4().hex}.jpg")
|
||
ok, msg, _ = await service.correct_perspective(src, out, level=level)
|
||
if ok:
|
||
print(f"[ImageProcessor] 透视矫正完成")
|
||
return out
|
||
else:
|
||
print(f"[ImageProcessor] 透视矫正失败 ({msg})，跳过")
|
||
if os.path.exists(out):
|
||
os.remove(out)
|
||
return src
|
||
|
||
@staticmethod
|
||
def _build_retry_prompt(gemini_prompt: str, qa_issue: str, qa_suggestion: str) -> str:
|
||
"""
|
||
根据 QA 质检问题类型，智能调整重试提示词。
|
||
比简单追加建议更有针对性，让 Gemini 知道上次哪里出了问题。
|
||
"""
|
||
base = gemini_prompt or ""
|
||
issue = (qa_issue or "").lower()
|
||
suggestion = qa_suggestion if qa_suggestion and qa_suggestion != "无" else ""
|
||
|
||
# 背景不干净
|
||
if any(kw in issue for kw in ["背景", "杂物", "多余", "白色不纯"]):
|
||
prefix = "【重要：背景必须是纯白色 #FFFFFF，去掉所有杂物和阴影】"
|
||
return prefix + ("\n" + base if base else "")
|
||
|
||
# 清晰度/细节不足
|
||
if any(kw in issue for kw in ["模糊", "清晰", "细节", "锐化", "分辨率"]):
|
||
prefix = "【重要：提升整体清晰度和细节，输出高分辨率版本，不要模糊】"
|
||
return prefix + ("\n" + base if base else "")
|
||
|
||
# 内容缺失/截断
|
||
if any(kw in issue for kw in ["缺失", "截断", "不完整", "边缘", "裁剪"]):
|
||
prefix = "【重要：保留主体完整内容，不要截断边缘，确保四角全部保留】"
|
||
return prefix + ("\n" + base if base else "")
|
||
|
||
# 颜色偏差
|
||
if any(kw in issue for kw in ["颜色", "色彩", "偏色", "色调"]):
|
||
prefix = "【重要：忠实还原原图颜色，不要改变色调或过度饱和】"
|
||
return prefix + ("\n" + base if base else "")
|
||
|
||
# AI幻觉/变形
|
||
if any(kw in issue for kw in ["幻觉", "变形", "失真", "扭曲", "ai生成"]):
|
||
prefix = "【重要：严格按原图内容处理，不要添加或改变任何图案细节】"
|
||
return prefix + ("\n" + base if base else "")
|
||
|
||
# 没有匹配到具体类型，直接用质检建议
|
||
if suggestion:
|
||
return (base + f"\n【上次问题：{qa_issue}。本次改进方向：{suggestion}】").strip()
|
||
|
||
return base
|
||
|
||
async def _do_main(self, service, src: str, gemini_prompt: str, aspect_ratio: str,
|
||
attempt: int, qa_issue: str = "", qa_suggestion: str = "") -> tuple[bool, str, str]:
|
||
"""
|
||
执行一次主处理。
|
||
重试时根据 QA 问题类型智能调整提示词。
|
||
|
||
Returns:
|
||
(success, output_path, message)
|
||
"""
|
||
out_name = f"result_{uuid.uuid4().hex}.jpg"
|
||
output_path = os.path.join(_OUTPUT_DIR, out_name)
|
||
|
||
if attempt == 1:
|
||
prompt = gemini_prompt or None
|
||
else:
|
||
prompt = self._build_retry_prompt(gemini_prompt, qa_issue, qa_suggestion)
|
||
print(f"[ImageProcessor] 重试策略 | 问题: {qa_issue} | 提示词: {(prompt or '')[:80]}...")
|
||
|
||
print(f"[ImageProcessor] 主处理第 {attempt} 次 (比例={aspect_ratio})...")
|
||
success, message, _ = await service.extract_pattern(
|
||
input_path=src,
|
||
output_path=output_path,
|
||
custom_prompt=prompt,
|
||
aspect_ratio=aspect_ratio,
|
||
)
|
||
return success, output_path, message
|
||
|
||
# ─── 主入口 ──────────────────────────────────────────────
|
||
|
||
async def process_image(
|
||
self,
|
||
image_url: str,
|
||
operation: str,
|
||
requirements: str = "",
|
||
gemini_prompt: str = "",
|
||
aspect_ratio: str = "1:1",
|
||
perspective: str = "no",
|
||
proc_type: str = "",
|
||
subject: str = "",
|
||
quality: str = "",
|
||
params: Optional[Dict[str, Any]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
完整处理流程：下载 → 透视矫正（可选）→ Gemini主处理 → 质检 → 重试（可选）
|
||
|
||
Returns:
|
||
{
|
||
"success": bool,
|
||
"result_path": str,
|
||
"message": str,
|
||
"qa_score": int, # 质检得分 0-100
|
||
"qa_pass": bool, # 是否通过质检
|
||
"qa_issue": str, # 质检发现的问题
|
||
"attempts": int, # 共处理了几次
|
||
}
|
||
"""
|
||
from services.service_gemini import GeminiExtractV2Service
|
||
from image.image_qa import image_qa
|
||
|
||
# Step 1: 下载原图
|
||
try:
|
||
tmp_input = await self._download(image_url)
|
||
except Exception as e:
|
||
return {
|
||
"success": False, "result_path": "", "message": str(e),
|
||
"qa_score": 0, "qa_pass": False, "qa_issue": "下载失败", "attempts": 0,
|
||
}
|
||
|
||
# Step 1.5: 敏感图片检测
|
||
try:
|
||
from utils.content_filter import is_sensitive_image
|
||
sensitive, reason = await is_sensitive_image(tmp_input)
|
||
if sensitive:
|
||
if os.path.exists(tmp_input):
|
||
os.remove(tmp_input)
|
||
return {
|
||
"success": False, "result_path": "", "message": reason,
|
||
"qa_score": 0, "qa_pass": False, "qa_issue": "敏感图片", "attempts": 0,
|
||
}
|
||
except Exception as e:
|
||
print(f"[ImageProcessor] 敏感图片检测异常: {e}，继续处理")
|
||
|
||
# Step 1.6: 预检（尺寸/格式/损坏）
|
||
try:
|
||
from image.image_precheck import precheck
|
||
ok, msg = precheck(tmp_input)
|
||
if not ok:
|
||
if os.path.exists(tmp_input):
|
||
os.remove(tmp_input)
|
||
return {
|
||
"success": False, "result_path": "", "message": msg,
|
||
"qa_score": 0, "qa_pass": False, "qa_issue": "预检不通过", "attempts": 0,
|
||
}
|
||
except Exception as e:
|
||
print(f"[ImageProcessor] 预检异常: {e}，继续处理")
|
||
|
||
service = GeminiExtractV2Service()
|
||
tmp_files = [tmp_input]
|
||
try:
|
||
# Step 2: 透视矫正
|
||
current_input = tmp_input
|
||
if perspective in ("mild", "strong"):
|
||
print(f"[ImageProcessor] 透视矫正中 (level={perspective})...")
|
||
corrected = await self._do_perspective(service, tmp_input, perspective)
|
||
if corrected != tmp_input:
|
||
tmp_files.append(corrected)
|
||
current_input = corrected
|
||
|
||
# Step 3: 主处理 + 质检，最多 _MAX_RETRIES 次
|
||
qa_result = {"score": 0, "pass": False, "issue": "未质检", "suggestion": "无"}
|
||
output_path = ""
|
||
last_message = ""
|
||
qa_issue = ""
|
||
qa_suggestion = ""
|
||
|
||
for attempt in range(1, _MAX_RETRIES + 1):
|
||
ok, output_path, last_message = await self._do_main(
|
||
service, current_input, gemini_prompt, aspect_ratio,
|
||
attempt=attempt, qa_issue=qa_issue, qa_suggestion=qa_suggestion,
|
||
)
|
||
|
||
if not ok:
|
||
print(f"[ImageProcessor] 第 {attempt} 次处理失败: {last_message}")
|
||
if attempt < _MAX_RETRIES:
|
||
continue
|
||
return {
|
||
"success": False, "result_path": "", "message": last_message,
|
||
"qa_score": 0, "qa_pass": False, "qa_issue": "Gemini处理失败", "attempts": attempt,
|
||
}
|
||
|
||
# Step 4: 质检
|
||
print(f"[ImageProcessor] 质检中 (第 {attempt} 次结果)...")
|
||
qa_result = await image_qa.check(
|
||
original_path=current_input,
|
||
result_path=output_path,
|
||
proc_type=proc_type,
|
||
subject=subject,
|
||
quality=quality,
|
||
gemini_prompt=gemini_prompt,
|
||
)
|
||
qa_issue = qa_result.get("issue", "")
|
||
qa_suggestion = qa_result.get("suggestion", "无")
|
||
|
||
if qa_result["pass"]:
|
||
print(f"[ImageProcessor] 质检通过 ({qa_result['score']}分)，共处理 {attempt} 次")
|
||
break
|
||
else:
|
||
print(f"[ImageProcessor] 质检不合格 ({qa_result['score']}分)，问题: {qa_result['issue']}")
|
||
if attempt < _MAX_RETRIES:
|
||
# 清理这次不合格的结果
|
||
if os.path.exists(output_path):
|
||
os.remove(output_path)
|
||
print(f"[ImageProcessor] 准备第 {attempt + 1} 次重试...")
|
||
else:
|
||
print(f"[ImageProcessor] 已达最大重试次数 {_MAX_RETRIES}，保留最后结果，人工跟进")
|
||
|
||
return {
|
||
"success": True,
|
||
"result_path": output_path,
|
||
"message": last_message,
|
||
"qa_score": qa_result.get("score", 0),
|
||
"qa_pass": qa_result.get("pass", False),
|
||
"qa_issue": qa_result.get("issue", ""),
|
||
"attempts": attempt,
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
"success": False, "result_path": "", "message": f"处理异常: {e}",
|
||
"qa_score": 0, "qa_pass": False, "qa_issue": str(e), "attempts": 0,
|
||
}
|
||
finally:
|
||
await service.cleanup()
|
||
for f in tmp_files:
|
||
if os.path.exists(f):
|
||
os.remove(f)
|
||
|
||
async def enhance(self, image_url: str) -> Dict[str, Any]:
|
||
return await self.process_image(image_url, "enhance")
|
||
|
||
async def remove_bg(self, image_url: str) -> Dict[str, Any]:
|
||
return await self.process_image(image_url, "remove_bg")
|
||
|
||
async def resize(self, image_url: str, width: int, height: int = 0) -> Dict[str, Any]:
|
||
"""
|
||
改尺寸：下载图片（或读取本地路径），按指定宽高缩放，保存到 results/。
|
||
|
||
Args:
|
||
image_url: 图片 URL 或本地路径
|
||
width: 目标宽度（像素）
|
||
height: 目标高度（0=按宽度等比缩放）
|
||
|
||
Returns:
|
||
{"success": bool, "result_path": str, "message": str}
|
||
"""
|
||
from PIL import Image
|
||
is_temp = image_url.startswith(("http://", "https://"))
|
||
try:
|
||
if is_temp:
|
||
tmp = await self._download(image_url)
|
||
else:
|
||
tmp = image_url
|
||
if not os.path.exists(tmp):
|
||
return {"success": False, "result_path": "", "message": f"文件不存在: {tmp}"}
|
||
except Exception as e:
|
||
return {"success": False, "result_path": "", "message": str(e)}
|
||
|
||
try:
|
||
img = Image.open(tmp).convert("RGB")
|
||
w_orig, h_orig = img.size
|
||
if width <= 0 or width > 10000:
|
||
return {"success": False, "result_path": "", "message": f"宽度无效: {width}"}
|
||
if height == 0:
|
||
ratio = width / w_orig
|
||
height = int(h_orig * ratio)
|
||
elif height <= 0 or height > 10000:
|
||
return {"success": False, "result_path": "", "message": f"高度无效: {height}"}
|
||
resized = img.resize((width, height), Image.Resampling.LANCZOS)
|
||
out_name = f"resize_{uuid.uuid4().hex}.jpg"
|
||
out_path = os.path.join(_OUTPUT_DIR, out_name)
|
||
resized.save(out_path, "JPEG", quality=95)
|
||
print(f"[ImageProcessor] 改尺寸完成: {w_orig}x{h_orig} → {width}x{height}")
|
||
return {"success": True, "result_path": out_path, "message": f"已改为 {width}x{height}"}
|
||
except Exception as e:
|
||
return {"success": False, "result_path": "", "message": str(e)}
|
||
finally:
|
||
if is_temp and os.path.exists(tmp):
|
||
os.remove(tmp)
|
||
|
||
|
||
# 全局实例
|
||
image_processor = ImageProcessor()
|