feat(ai-studio): export final selections
This commit is contained in:
chengma
@@ -0,0 +1,210 @@
|
||||
"""Safe local export helpers for AI image studio final selections."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
||||
from . import appconfig, image_studio
|
||||
|
||||
|
||||
EXISTING_FAIL = "fail"
|
||||
EXISTING_OVERWRITE_MANAGED = "overwrite_managed"
|
||||
EXISTING_TIMESTAMP = "timestamp"
|
||||
EXISTING_MODES = {EXISTING_FAIL, EXISTING_OVERWRITE_MANAGED, EXISTING_TIMESTAMP}
|
||||
|
||||
|
||||
class ImageStudioExportError(RuntimeError):
|
||||
"""Raised when AI studio selections cannot be exported safely."""
|
||||
|
||||
|
||||
class ExportTargetExistsError(ImageStudioExportError):
|
||||
"""Raised when the default target directory already exists."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExportedFile:
|
||||
selection_type: str
|
||||
asset_id: int
|
||||
source_path: str
|
||||
output_path: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExportResult:
|
||||
target_dir: str
|
||||
files: tuple[ExportedFile, ...]
|
||||
main_count: int
|
||||
detail_count: int
|
||||
existing_mode: str
|
||||
|
||||
|
||||
def target_dir_for_project(project, parent_dir, suffix=None):
|
||||
parent = _existing_parent_dir(parent_dir)
|
||||
item = _safe_item_id(getattr(project, "item_id", ""))
|
||||
dirname = item if not suffix else f"{item}_{suffix}"
|
||||
return os.path.abspath(os.path.join(parent, dirname))
|
||||
|
||||
|
||||
def export_project_selection(
|
||||
project_id,
|
||||
parent_dir,
|
||||
*,
|
||||
existing_mode=EXISTING_FAIL,
|
||||
path=None,
|
||||
config=None,
|
||||
timestamp=None,
|
||||
):
|
||||
"""Export current main/detail selections as ordered JPEG files."""
|
||||
|
||||
mode = str(existing_mode or EXISTING_FAIL)
|
||||
if mode not in EXISTING_MODES:
|
||||
raise ImageStudioExportError("导出目录处理方式无效")
|
||||
cfg = appconfig.load_config() if config is None else config
|
||||
database_path = path or appconfig.db_path(cfg)
|
||||
project = image_studio.get_project(project_id, path=database_path)
|
||||
if project is None:
|
||||
raise ImageStudioExportError("AI工场项目不存在")
|
||||
parent = _existing_parent_dir(parent_dir)
|
||||
planned = _planned_files(project, database_path)
|
||||
if not planned:
|
||||
raise ImageStudioExportError("主图和详情图终选都为空，不能导出")
|
||||
target = _choose_target_dir(project, parent, mode, timestamp=timestamp)
|
||||
quality = int(appconfig.ai_config(cfg).get("jpg_quality", 90) or 90)
|
||||
staging = os.path.join(parent, f".{_safe_item_id(project.item_id)}_staging_{uuid.uuid4().hex}")
|
||||
try:
|
||||
os.makedirs(staging, exist_ok=False)
|
||||
staged = []
|
||||
for item in planned:
|
||||
output_path = os.path.join(staging, item["filename"])
|
||||
_save_jpeg(item["source_path"], output_path, quality)
|
||||
staged.append(ExportedFile(item["selection_type"], item["asset_id"], item["source_path"], output_path))
|
||||
os.makedirs(target, exist_ok=True)
|
||||
if os.path.exists(target) and mode == EXISTING_OVERWRITE_MANAGED:
|
||||
_remove_managed_exports(target, project.item_id)
|
||||
final_files = []
|
||||
for staged_file in staged:
|
||||
final_path = os.path.join(target, os.path.basename(staged_file.output_path))
|
||||
os.replace(staged_file.output_path, final_path)
|
||||
final_files.append(
|
||||
ExportedFile(
|
||||
staged_file.selection_type,
|
||||
staged_file.asset_id,
|
||||
staged_file.source_path,
|
||||
final_path,
|
||||
)
|
||||
)
|
||||
return ExportResult(
|
||||
target_dir=target,
|
||||
files=tuple(final_files),
|
||||
main_count=sum(1 for item in final_files if item.selection_type == "main"),
|
||||
detail_count=sum(1 for item in final_files if item.selection_type == "detail"),
|
||||
existing_mode=mode,
|
||||
)
|
||||
except Exception as exc:
|
||||
if isinstance(exc, ImageStudioExportError):
|
||||
raise
|
||||
raise ImageStudioExportError(f"导出终选图片失败：{exc}") from exc
|
||||
finally:
|
||||
if os.path.isdir(staging):
|
||||
shutil.rmtree(staging, ignore_errors=True)
|
||||
|
||||
|
||||
def _choose_target_dir(project, parent_dir, existing_mode, timestamp=None):
|
||||
target = target_dir_for_project(project, parent_dir)
|
||||
if not os.path.exists(target):
|
||||
return target
|
||||
if not os.path.isdir(target):
|
||||
raise ImageStudioExportError(f"导出目标已存在但不是目录：{target}")
|
||||
if existing_mode == EXISTING_FAIL:
|
||||
raise ExportTargetExistsError(f"商品目录已存在：{target}")
|
||||
if existing_mode == EXISTING_OVERWRITE_MANAGED:
|
||||
return target
|
||||
stamp = timestamp or datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
base = target_dir_for_project(project, parent_dir, suffix=stamp)
|
||||
candidate = base
|
||||
index = 2
|
||||
while os.path.exists(candidate):
|
||||
candidate = f"{base}_{index}"
|
||||
index += 1
|
||||
return candidate
|
||||
|
||||
|
||||
def _planned_files(project, db_path):
|
||||
item = _safe_item_id(project.item_id)
|
||||
planned = []
|
||||
for selection_type, label in (("main", "主图"), ("detail", "详情图")):
|
||||
selections = image_studio.list_selections(project.id, selection_type, path=db_path)
|
||||
for index, selection in enumerate(selections, start=1):
|
||||
asset = image_studio.get_asset(selection.asset_id, path=db_path)
|
||||
if asset is None:
|
||||
raise ImageStudioExportError(f"终选照片不存在：#{selection.asset_id}")
|
||||
source_path = str(asset.local_path or "")
|
||||
if not source_path or not os.path.isfile(source_path):
|
||||
raise ImageStudioExportError(f"终选照片本地文件缺失：#{asset.id}")
|
||||
planned.append(
|
||||
{
|
||||
"selection_type": selection_type,
|
||||
"asset_id": int(asset.id),
|
||||
"source_path": os.path.abspath(source_path),
|
||||
"filename": f"{item}_{label}_{index}.jpg",
|
||||
}
|
||||
)
|
||||
return planned
|
||||
|
||||
|
||||
def _save_jpeg(source_path, output_path, quality):
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(source_path) as image:
|
||||
if image.mode in {"RGBA", "LA"} or (
|
||||
image.mode == "P" and "transparency" in getattr(image, "info", {})
|
||||
):
|
||||
rgba = image.convert("RGBA")
|
||||
background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
|
||||
background.alpha_composite(rgba)
|
||||
final_image = background.convert("RGB")
|
||||
else:
|
||||
final_image = image.convert("RGB")
|
||||
final_image.save(output_path, format="JPEG", quality=max(1, min(95, int(quality or 90))))
|
||||
with Image.open(output_path) as check:
|
||||
check.verify()
|
||||
except Exception as exc:
|
||||
raise ImageStudioExportError(f"JPEG 转码失败：{exc}") from exc
|
||||
|
||||
|
||||
def _remove_managed_exports(target_dir, item_id):
|
||||
item = re.escape(_safe_item_id(item_id))
|
||||
pattern = re.compile(rf"^{item}_(主图|详情图)_\d+\.jpg$", re.IGNORECASE)
|
||||
for filename in os.listdir(target_dir):
|
||||
if not pattern.match(filename):
|
||||
continue
|
||||
path = os.path.join(target_dir, filename)
|
||||
if os.path.isfile(path):
|
||||
os.remove(path)
|
||||
|
||||
|
||||
def _existing_parent_dir(parent_dir):
|
||||
parent = os.path.abspath(str(parent_dir or ""))
|
||||
if not os.path.isdir(parent):
|
||||
raise ImageStudioExportError(f"导出父目录不存在：{parent}")
|
||||
return parent
|
||||
|
||||
|
||||
def _safe_item_id(item_id):
|
||||
value = str(item_id or "").strip()
|
||||
if not value:
|
||||
raise ImageStudioExportError("商品ID不能为空")
|
||||
if any(char in value for char in ('/', '\\', os.sep, os.altsep or "\0")):
|
||||
raise ImageStudioExportError("商品ID包含非法路径字符")
|
||||
if value in {".", ".."}:
|
||||
raise ImageStudioExportError("商品ID不能是路径符号")
|
||||
safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in value).strip("_")
|
||||
if not safe:
|
||||
raise ImageStudioExportError("商品ID不能作为目录名")
|
||||
return safe
|
||||
Reference in New Issue
Block a user