369 lines
13 KiB
Python
369 lines
13 KiB
Python
"""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
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GenerationRoundExportedFile:
|
|
job_id: int
|
|
asset_id: int
|
|
job_type: str
|
|
source_path: str
|
|
output_path: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GenerationRoundExportResult:
|
|
target_dir: str
|
|
files: tuple[GenerationRoundExportedFile, ...]
|
|
skipped_count: int
|
|
cancelled: bool
|
|
|
|
|
|
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 export_generation_round(
|
|
project_id,
|
|
generation_round_key,
|
|
parent_dir,
|
|
*,
|
|
path=None,
|
|
should_stop=None,
|
|
):
|
|
"""Copy usable successful output images from one persisted generation round."""
|
|
|
|
project = image_studio.get_project(project_id, path=path)
|
|
if project is None:
|
|
raise ImageStudioExportError("商品套图项目不存在，无法导出历史记录")
|
|
try:
|
|
parent = _existing_parent_dir(parent_dir)
|
|
except ImageStudioExportError as exc:
|
|
raise ImageStudioExportError("导出父目录不可用") from exc
|
|
jobs = image_studio.list_generation_round_current_jobs(
|
|
project.id,
|
|
generation_round_key,
|
|
path=path,
|
|
)
|
|
planned, skipped_count = _planned_generation_round_files(jobs, path)
|
|
if not planned:
|
|
raise ImageStudioExportError("本轮没有可导出的成功生成图片")
|
|
|
|
target = _create_generation_round_target_dir(project, parent, jobs)
|
|
copied = []
|
|
cancelled = False
|
|
try:
|
|
for item in planned:
|
|
if callable(should_stop) and should_stop():
|
|
cancelled = True
|
|
break
|
|
output_path = os.path.join(target, item["filename"])
|
|
try:
|
|
shutil.copy2(item["source_path"], output_path)
|
|
except OSError:
|
|
skipped_count += 1
|
|
if os.path.isfile(output_path):
|
|
try:
|
|
os.remove(output_path)
|
|
except OSError:
|
|
pass
|
|
continue
|
|
copied.append(
|
|
GenerationRoundExportedFile(
|
|
job_id=int(item["job_id"]),
|
|
asset_id=int(item["asset_id"]),
|
|
job_type=str(item["job_type"]),
|
|
source_path=item["source_path"],
|
|
output_path=output_path,
|
|
)
|
|
)
|
|
except Exception as exc:
|
|
raise ImageStudioExportError("导出本轮生成图片失败") from exc
|
|
|
|
if not copied:
|
|
try:
|
|
os.rmdir(target)
|
|
except OSError:
|
|
pass
|
|
if cancelled:
|
|
return GenerationRoundExportResult("", (), skipped_count, True)
|
|
raise ImageStudioExportError("本轮生成图片导出失败")
|
|
return GenerationRoundExportResult(
|
|
target_dir=target,
|
|
files=tuple(copied),
|
|
skipped_count=skipped_count,
|
|
cancelled=cancelled,
|
|
)
|
|
|
|
|
|
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 _planned_generation_round_files(jobs, db_path):
|
|
planned = []
|
|
skipped_count = 0
|
|
for index, job in enumerate(jobs, 1):
|
|
if str(getattr(job, "status", "") or "") != "succeeded":
|
|
continue
|
|
asset_id = getattr(job, "output_asset_id", None)
|
|
asset = image_studio.get_asset(asset_id, path=db_path) if asset_id else None
|
|
source_path = str(getattr(asset, "local_path", "") or "")
|
|
if (
|
|
asset is None
|
|
or str(getattr(asset, "status", "") or "")
|
|
== image_studio.ASSET_STATUS_MISSING
|
|
or not source_path
|
|
or not os.path.isfile(source_path)
|
|
):
|
|
skipped_count += 1
|
|
continue
|
|
suffix = os.path.splitext(source_path)[1].lower()
|
|
if not suffix or len(suffix) > 8:
|
|
suffix = ".jpg"
|
|
job_type = _safe_export_component(
|
|
getattr(job, "job_type", ""),
|
|
fallback="套图",
|
|
)
|
|
planned.append(
|
|
{
|
|
"job_id": int(job.id),
|
|
"asset_id": int(asset.id),
|
|
"job_type": str(getattr(job, "job_type", "") or "套图"),
|
|
"source_path": os.path.abspath(source_path),
|
|
"filename": "%02d_%s%s" % (index, job_type, suffix),
|
|
}
|
|
)
|
|
return planned, skipped_count
|
|
|
|
|
|
def _create_generation_round_target_dir(project, parent_dir, jobs):
|
|
account = _safe_export_component(
|
|
getattr(project, "account_name", "") or getattr(project, "account_alias", ""),
|
|
fallback="未命名店铺",
|
|
)
|
|
item = _safe_export_component(getattr(project, "item_id", ""), fallback="临时草稿")
|
|
created_at = str(getattr(jobs[0], "created_at", "") or "") if jobs else ""
|
|
digits = "".join(character for character in created_at if character.isdigit())[:14]
|
|
stamp = digits or datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
base = os.path.join(parent_dir, "%s_%s_%s" % (account, item, stamp))
|
|
candidate = base
|
|
index = 2
|
|
while True:
|
|
try:
|
|
os.makedirs(candidate, exist_ok=False)
|
|
return candidate
|
|
except FileExistsError:
|
|
candidate = "%s_%d" % (base, index)
|
|
index += 1
|
|
except OSError as exc:
|
|
raise ImageStudioExportError("无法创建导出目录") from exc
|
|
|
|
|
|
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
|
|
|
|
|
|
def _safe_export_component(value, *, fallback):
|
|
text = str(value or "").strip()
|
|
safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in text).strip("_")
|
|
if safe:
|
|
return safe
|
|
return str(fallback)
|