init
This commit is contained in:
293
image/image_tools.py
Normal file
293
image/image_tools.py
Normal file
@@ -0,0 +1,293 @@
|
||||
"""
|
||||
图片处理独立工具 - 可单独调用，也可被主流程复用。
|
||||
|
||||
主流程（付款触发）不变，这些工具供 AI 按需组合使用。
|
||||
"""
|
||||
import os
|
||||
import uuid
|
||||
import tempfile
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
_OUTPUT_DIR = os.getenv("RESULT_IMAGE_DIR", "results")
|
||||
os.makedirs(_OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
|
||||
async def _download(url: str) -> str:
|
||||
"""下载图片到临时文件"""
|
||||
import aiohttp
|
||||
tmp = os.path.join(tempfile.gettempdir(), f"img_{uuid.uuid4().hex}.jpg")
|
||||
headers = {
|
||||
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
|
||||
"Referer": "https://www.taobao.com/",
|
||||
}
|
||||
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 remove_background(image_url: str, save_path: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】去背景 → 纯白/纯色背景。
|
||||
输入 URL 或本地路径，输出白底产品图。
|
||||
"""
|
||||
from image.perspective_fix import _gemini_call, PROMPT_WHITE_BG
|
||||
tmp = None
|
||||
try:
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
tmp = await _download(image_url)
|
||||
src = tmp
|
||||
else:
|
||||
src = image_url
|
||||
if not os.path.exists(src):
|
||||
return {"success": False, "result_path": "", "message": f"文件不存在: {src}"}
|
||||
|
||||
out = save_path or os.path.join(_OUTPUT_DIR, f"bg_{uuid.uuid4().hex}.jpg")
|
||||
ok = await _gemini_call(src, out, PROMPT_WHITE_BG, aspect_ratio="auto", label="去背景")
|
||||
if ok:
|
||||
return {"success": True, "result_path": out, "message": "去背景完成"}
|
||||
return {"success": False, "result_path": "", "message": "去背景失败"}
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
if tmp and os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
|
||||
|
||||
async def perspective_correct(image_url: str, save_path: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】透视矫正。
|
||||
输入需为白底图（可先调 remove_background），输出展平后的图。
|
||||
"""
|
||||
import cv2
|
||||
from image.perspective_fix import find_quad, four_point_transform
|
||||
tmp = None
|
||||
try:
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
tmp = await _download(image_url)
|
||||
src = tmp
|
||||
else:
|
||||
src = image_url
|
||||
if not os.path.exists(src):
|
||||
return {"success": False, "result_path": "", "message": f"文件不存在: {src}"}
|
||||
|
||||
img = cv2.imread(src)
|
||||
if img is None:
|
||||
return {"success": False, "result_path": "", "message": "无法读取图片"}
|
||||
pts = find_quad(img)
|
||||
if pts is None:
|
||||
return {"success": False, "result_path": "", "message": "未检测到四边形，无法透视矫正"}
|
||||
warped = four_point_transform(img, pts)
|
||||
out = save_path or os.path.join(_OUTPUT_DIR, f"persp_{uuid.uuid4().hex}.jpg")
|
||||
cv2.imwrite(out, warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
return {"success": True, "result_path": out, "message": "透视矫正完成"}
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
if tmp and os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
|
||||
|
||||
async def extract_pattern(image_url: str, prompt: str = "", aspect_ratio: str = "1:1",
|
||||
save_path: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】印花提取/主处理。
|
||||
按提示词和比例输出处理后的图。
|
||||
"""
|
||||
from services.service_gemini import GeminiExtractV2Service
|
||||
tmp = None
|
||||
try:
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
tmp = await _download(image_url)
|
||||
src = tmp
|
||||
else:
|
||||
src = image_url
|
||||
if not os.path.exists(src):
|
||||
return {"success": False, "result_path": "", "message": f"文件不存在: {src}"}
|
||||
|
||||
out = save_path or os.path.join(_OUTPUT_DIR, f"extract_{uuid.uuid4().hex}.jpg")
|
||||
service = GeminiExtractV2Service()
|
||||
try:
|
||||
ok, msg, _ = await service.extract_pattern(
|
||||
input_path=src, output_path=out,
|
||||
custom_prompt=prompt or None, aspect_ratio=aspect_ratio,
|
||||
)
|
||||
if ok and os.path.exists(out):
|
||||
return {"success": True, "result_path": out, "message": "提取完成"}
|
||||
return {"success": False, "result_path": "", "message": msg or "提取失败"}
|
||||
finally:
|
||||
await service.cleanup()
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
if tmp and os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
|
||||
|
||||
async def enhance_image(image_url: str, save_path: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】高清增强。
|
||||
使用 Qwen RunningHub，失败时降级 Gemini。
|
||||
"""
|
||||
from services.service_qwen import 清晰化_api
|
||||
from image.perspective_fix import _gemini_call, PROMPT_ENHANCE_SIMPLE
|
||||
tmp = None
|
||||
try:
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
tmp = await _download(image_url)
|
||||
src = tmp
|
||||
else:
|
||||
src = image_url
|
||||
if not os.path.exists(src):
|
||||
return {"success": False, "result_path": "", "message": f"文件不存在: {src}"}
|
||||
|
||||
out = save_path or os.path.join(_OUTPUT_DIR, f"enh_{uuid.uuid4().hex}.jpg")
|
||||
ok = await 清晰化_api(img_path=src, save_path=out)
|
||||
if not ok:
|
||||
ok = await _gemini_call(src, out, PROMPT_ENHANCE_SIMPLE, aspect_ratio="auto", label="增强")
|
||||
if ok:
|
||||
return {"success": True, "result_path": out, "message": "高清增强完成"}
|
||||
return {"success": False, "result_path": "", "message": "高清增强失败"}
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
if tmp and os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
|
||||
|
||||
async def color_match_images(orig_url: str, result_url: str, save_path: str = "",
|
||||
strength: float = 0.75) -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】颜色匹配。将 result 的色调匹配到 orig。
|
||||
"""
|
||||
import cv2
|
||||
from image.perspective_fix import _color_match
|
||||
tmp_orig = tmp_result = None
|
||||
try:
|
||||
if orig_url.startswith(("http://", "https://")):
|
||||
tmp_orig = await _download(orig_url)
|
||||
orig_path = tmp_orig
|
||||
else:
|
||||
orig_path = orig_url
|
||||
if result_url.startswith(("http://", "https://")):
|
||||
tmp_result = await _download(result_url)
|
||||
result_path = tmp_result
|
||||
else:
|
||||
result_path = result_url
|
||||
|
||||
orig_img = cv2.imread(orig_path)
|
||||
result_img = cv2.imread(result_path)
|
||||
if orig_img is None or result_img is None:
|
||||
return {"success": False, "result_path": "", "message": "图片读取失败"}
|
||||
matched = _color_match(orig_img, result_img, strength=strength)
|
||||
out = save_path or os.path.join(_OUTPUT_DIR, f"color_{uuid.uuid4().hex}.jpg")
|
||||
cv2.imwrite(out, matched, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
return {"success": True, "result_path": out, "message": f"颜色匹配完成(强度{strength:.0%})"}
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
for t in (tmp_orig, tmp_result):
|
||||
if t and os.path.exists(t):
|
||||
os.remove(t)
|
||||
|
||||
|
||||
async def trim_border(image_url: str, save_path: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】裁切四周背景边（支持任意颜色：白/黄/米等）。
|
||||
"""
|
||||
import cv2
|
||||
from image.perspective_fix import tool_trim_white_border
|
||||
tmp = None
|
||||
try:
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
tmp = await _download(image_url)
|
||||
src = tmp
|
||||
else:
|
||||
src = image_url
|
||||
if not os.path.exists(src):
|
||||
return {"success": False, "result_path": "", "message": f"文件不存在: {src}"}
|
||||
|
||||
img = cv2.imread(src)
|
||||
if img is None:
|
||||
return {"success": False, "result_path": "", "message": "无法读取图片"}
|
||||
trimmed, did_trim, info = tool_trim_white_border(img)
|
||||
out = save_path or os.path.join(_OUTPUT_DIR, f"trim_{uuid.uuid4().hex}.jpg")
|
||||
cv2.imwrite(out, trimmed, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||
return {"success": True, "result_path": out, "message": "裁边完成" if did_trim else "无需裁边"}
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
if tmp and os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
|
||||
|
||||
async def vectorize_to_eps(image_url: str, save_path: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】矢量化 - 将图片转为 EPS 矢量文件。
|
||||
客户要做矢量图、转 EPS、转 AI 格式时调用。
|
||||
"""
|
||||
tmp = None
|
||||
try:
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
tmp = await _download(image_url)
|
||||
src = tmp
|
||||
else:
|
||||
src = image_url
|
||||
if not os.path.exists(src):
|
||||
return {"success": False, "result_path": "", "message": f"文件不存在: {src}"}
|
||||
|
||||
from services.service_vectorizer import VectorizerService
|
||||
svc = VectorizerService()
|
||||
out = save_path or os.path.join(_OUTPUT_DIR, f"vec_{uuid.uuid4().hex}.eps")
|
||||
result_path = await svc.image_to_eps(src, save_eps_path=out)
|
||||
if result_path and os.path.exists(result_path):
|
||||
return {"success": True, "result_path": result_path, "message": "矢量化完成，已生成 EPS 文件"}
|
||||
return {"success": False, "result_path": "", "message": "矢量化失败"}
|
||||
except ImportError as e:
|
||||
return {"success": False, "result_path": "", "message": f"矢量化服务不可用: {e}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
if tmp and os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
|
||||
|
||||
async def meitu_enhance(image_url: str, mode: str = "standard", save_path: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
【独立工具】美图画质增强。
|
||||
模式: crystal(极速重绘) standard(标准) enhance(增强) hdr(HDR) portrait(人像优化)
|
||||
客户要画质增强、清晰化、美图处理时调用。
|
||||
"""
|
||||
tmp = None
|
||||
try:
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
tmp = await _download(image_url)
|
||||
src = tmp
|
||||
else:
|
||||
src = image_url
|
||||
if not os.path.exists(src):
|
||||
return {"success": False, "result_path": "", "message": f"文件不存在: {src}"}
|
||||
|
||||
from pathlib import Path
|
||||
from services.service_meitu import MeituAPIService
|
||||
svc = MeituAPIService()
|
||||
output_dir = Path(_OUTPUT_DIR)
|
||||
result = await svc.process_image(src, mode=mode, output_dir=output_dir)
|
||||
out = result.get("processed_path")
|
||||
if out and os.path.exists(str(out)):
|
||||
if save_path:
|
||||
import shutil
|
||||
shutil.copy(str(out), save_path)
|
||||
out = save_path
|
||||
return {"success": True, "result_path": str(out), "message": f"画质增强完成（{result.get('mode_name', mode)}）"}
|
||||
return {"success": False, "result_path": "", "message": "美图处理失败"}
|
||||
except ImportError as e:
|
||||
return {"success": False, "result_path": "", "message": f"美图服务不可用: {e}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "result_path": "", "message": str(e)}
|
||||
finally:
|
||||
if tmp and os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
Reference in New Issue
Block a user