feat(ai-studio): export final selections
This commit is contained in:
chengma
@@ -25,6 +25,7 @@ if QT_IMPORT_ERROR is None:
|
||||
CollectWorker,
|
||||
GenerateWorker,
|
||||
ImageStudioDownloadOriginalWorker,
|
||||
ImageStudioExportWorker,
|
||||
ImageStudioGenerateJobsWorker,
|
||||
ImageStudioPullImagesWorker,
|
||||
WriteBackWorker,
|
||||
|
||||
@@ -7,12 +7,13 @@ import os
|
||||
from PySide6.QtCore import QMimeData
|
||||
from PySide6.QtWidgets import QListWidget, QListWidgetItem
|
||||
|
||||
from ... import accounts, appconfig, db, image_studio, prompts
|
||||
from ... import accounts, appconfig, db, image_studio, image_studio_export, prompts
|
||||
from .. import file_manager
|
||||
from ..widgets import *
|
||||
from ..workers import (
|
||||
ImageStudioDownloadOriginalWorker as _RealImageStudioDownloadOriginalWorker,
|
||||
)
|
||||
from ..workers import ImageStudioExportWorker as _RealImageStudioExportWorker
|
||||
from ..workers import ImageStudioGenerateJobsWorker as _RealImageStudioGenerateJobsWorker
|
||||
from ..workers import ImageStudioPullImagesWorker as _RealImageStudioPullImagesWorker
|
||||
|
||||
@@ -47,6 +48,15 @@ def ImageStudioGenerateJobsWorker(*args, **kwargs):
|
||||
)
|
||||
|
||||
|
||||
def ImageStudioExportWorker(*args, **kwargs):
|
||||
return _call_package_attr(
|
||||
"ImageStudioExportWorker",
|
||||
_RealImageStudioExportWorker,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _drop_event_position(event):
|
||||
if hasattr(event, "position"):
|
||||
return event.position().toPoint()
|
||||
@@ -469,8 +479,20 @@ class ImageStudioTab(QWidget):
|
||||
self.detail_selection_list.setMinimumHeight(96)
|
||||
detail_layout.addWidget(self.detail_selection_label)
|
||||
detail_layout.addWidget(self.detail_selection_list)
|
||||
action_panel = QWidget()
|
||||
action_layout = QVBoxLayout(action_panel)
|
||||
action_layout.setContentsMargins(0, 0, 0, 0)
|
||||
self.export_button = QPushButton("导出终选")
|
||||
self.export_button.setObjectName("imageStudioExportButton")
|
||||
self.export_hint_label = QLabel("可部分导出，不要求主图/详情图满额")
|
||||
self.export_hint_label.setObjectName("imageStudioExportHintLabel")
|
||||
self.export_hint_label.setWordWrap(True)
|
||||
action_layout.addWidget(self.export_button)
|
||||
action_layout.addWidget(self.export_hint_label)
|
||||
action_layout.addStretch(1)
|
||||
layout.addWidget(main_panel, 1)
|
||||
layout.addWidget(detail_panel, 1)
|
||||
layout.addWidget(action_panel, 0)
|
||||
return panel
|
||||
|
||||
def _connect_signals(self):
|
||||
@@ -491,6 +513,7 @@ class ImageStudioTab(QWidget):
|
||||
self.prompt_edit.textChanged.connect(self._save_project_prompt)
|
||||
self.start_button.clicked.connect(self.start_generation)
|
||||
self.stop_button.clicked.connect(self.stop_generation)
|
||||
self.export_button.clicked.connect(self.export_selections)
|
||||
|
||||
def refresh_accounts(self):
|
||||
self.account_combo.clear()
|
||||
@@ -1032,6 +1055,91 @@ class ImageStudioTab(QWidget):
|
||||
self._append_log("[AI工场] 已请求停止，正在等待安全边界")
|
||||
self._status("AI工场生成已请求停止", "warning")
|
||||
|
||||
def export_selections(self, checked=False):
|
||||
if self.current_project is None:
|
||||
self._message("项目未打开", "请先打开一个AI工场项目。")
|
||||
return
|
||||
main_count = len(self._selection_asset_ids("main"))
|
||||
detail_count = len(self._selection_asset_ids("detail"))
|
||||
if main_count + detail_count <= 0:
|
||||
self._message("没有可导出图片", "主图和详情图终选都为空。")
|
||||
return
|
||||
parent_dir = QFileDialog.getExistingDirectory(self, "选择导出父目录")
|
||||
if not parent_dir:
|
||||
return
|
||||
mode = self._export_existing_mode(parent_dir, main_count, detail_count)
|
||||
if mode is None:
|
||||
return
|
||||
worker = ImageStudioExportWorker(
|
||||
self.current_project.id,
|
||||
parent_dir,
|
||||
existing_mode=mode,
|
||||
db_path=self.db_path,
|
||||
config=self.config,
|
||||
)
|
||||
worker.log.connect(self._append_log)
|
||||
worker.finished.connect(self._on_export_finished)
|
||||
worker.failed.connect(self._on_worker_failed)
|
||||
self._start_worker(worker, "AI工场导出终选")
|
||||
|
||||
def _export_existing_mode(self, parent_dir, main_count, detail_count):
|
||||
try:
|
||||
target = image_studio_export.target_dir_for_project(self.current_project, parent_dir)
|
||||
except Exception as exc:
|
||||
self._message("导出目录不可用", str(exc))
|
||||
return None
|
||||
if not os.path.exists(target):
|
||||
return image_studio_export.EXISTING_FAIL
|
||||
box = QMessageBox(self)
|
||||
box.setWindowTitle("商品目录已存在")
|
||||
box.setText(
|
||||
f"目标目录已存在：{target}\n"
|
||||
f"本次将导出主图 {main_count} 张、详情图 {detail_count} 张。请选择处理方式。"
|
||||
)
|
||||
overwrite_button = box.addButton("覆盖本软件导出的图片", QMessageBox.AcceptRole)
|
||||
timestamp_button = box.addButton("新建带时间目录", QMessageBox.ActionRole)
|
||||
box.addButton("取消", QMessageBox.RejectRole)
|
||||
box.setDefaultButton(timestamp_button)
|
||||
box.exec()
|
||||
clicked = box.clickedButton()
|
||||
if clicked is overwrite_button:
|
||||
return image_studio_export.EXISTING_OVERWRITE_MANAGED
|
||||
if clicked is timestamp_button:
|
||||
return image_studio_export.EXISTING_TIMESTAMP
|
||||
return None
|
||||
|
||||
def _on_export_finished(self, summary):
|
||||
if summary.get("ok") is False:
|
||||
if self._running_worker is not None:
|
||||
self._on_worker_failed(-1, summary.get("error") or "导出终选失败")
|
||||
return
|
||||
self._finish_worker()
|
||||
target_dir = summary.get("target_dir") or ""
|
||||
main_count = int(summary.get("main_count") or 0)
|
||||
detail_count = int(summary.get("detail_count") or 0)
|
||||
file_count = int(summary.get("file_count") or 0)
|
||||
self._append_log(f"[AI工场] 导出完成：主图{main_count}，详情图{detail_count}，文件{file_count}")
|
||||
self._status(f"AI工场终选已导出：{file_count} 个文件", "success")
|
||||
self._show_export_success(target_dir, main_count, detail_count, file_count)
|
||||
|
||||
def _show_export_success(self, target_dir, main_count, detail_count, file_count):
|
||||
box = QMessageBox(self)
|
||||
box.setWindowTitle("导出完成")
|
||||
box.setText(
|
||||
f"已导出 {file_count} 个 JPEG 文件。\n"
|
||||
f"主图 {main_count} 张，详情图 {detail_count} 张。\n"
|
||||
f"目录：{target_dir}"
|
||||
)
|
||||
open_button = box.addButton("打开目录", QMessageBox.AcceptRole)
|
||||
box.addButton("确定", QMessageBox.RejectRole)
|
||||
box.setDefaultButton(open_button)
|
||||
box.exec()
|
||||
if box.clickedButton() is open_button:
|
||||
try:
|
||||
file_manager.open_in_file_manager(target_dir)
|
||||
except Exception as exc:
|
||||
self._message("打开目录失败", str(exc))
|
||||
|
||||
def _on_generate_progress(self, payload):
|
||||
total = max(1, int(payload.get("total") or self.progress_bar.maximum() or 1))
|
||||
done = min(total, int(payload.get("done") or 0))
|
||||
@@ -1096,6 +1204,7 @@ class ImageStudioTab(QWidget):
|
||||
self.aspect_combo.setEnabled(not running)
|
||||
self.start_button.setEnabled(not running)
|
||||
self.stop_button.setEnabled(running)
|
||||
self.export_button.setEnabled(not running)
|
||||
|
||||
def _set_account_combo(self, alias):
|
||||
index = self.account_combo.findData(alias)
|
||||
|
||||
+39
-1
@@ -6,7 +6,7 @@ import datetime as _dt
|
||||
import re
|
||||
import threading
|
||||
|
||||
from .. import image_studio, image_studio_generation, image_studio_images
|
||||
from .. import image_studio, image_studio_export, image_studio_generation, image_studio_images
|
||||
from .widgets import *
|
||||
|
||||
|
||||
@@ -192,6 +192,44 @@ class ImageStudioGenerateJobsWorker(BaseWorker):
|
||||
return summary
|
||||
|
||||
|
||||
class ImageStudioExportWorker(BaseWorker):
|
||||
"""Export AI studio final selections to local JPEG files."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
project_id,
|
||||
parent_dir,
|
||||
*,
|
||||
existing_mode=image_studio_export.EXISTING_FAIL,
|
||||
db_path=None,
|
||||
config=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.project_id = int(project_id)
|
||||
self.parent_dir = parent_dir
|
||||
self.existing_mode = existing_mode
|
||||
self.db_path = db_path
|
||||
self.config = config
|
||||
|
||||
def execute(self):
|
||||
self.log.emit("[AI工场] 导出终选：开始")
|
||||
result = image_studio_export.export_project_selection(
|
||||
self.project_id,
|
||||
self.parent_dir,
|
||||
existing_mode=self.existing_mode,
|
||||
path=self.db_path,
|
||||
config=self.config,
|
||||
)
|
||||
self.log.emit("[AI工场] 导出终选：成功")
|
||||
return {
|
||||
"target_dir": result.target_dir,
|
||||
"main_count": result.main_count,
|
||||
"detail_count": result.detail_count,
|
||||
"file_count": len(result.files),
|
||||
"existing_mode": result.existing_mode,
|
||||
}
|
||||
|
||||
|
||||
def _generation_mode_label(mode):
|
||||
mode = appconfig.normalize_generate_mode(mode)
|
||||
return {
|
||||
|
||||
@@ -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