cmbot/src/core/models.py

from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import List, Optional


# ---------------------------------------------------------------------------
# 素材
# ---------------------------------------------------------------------------

@dataclass
class ImageAsset:
    """一张图片素材（衣服图或印花图）。"""
    path: Path
    selected: bool = True        # 是否参与批量合成


# ---------------------------------------------------------------------------
# 印花变换状态
# ---------------------------------------------------------------------------

@dataclass
class TransformState:
    """印花在衣服原图像素坐标系中的变换参数。

    x/y: 未旋转包围盒左上角坐标（原图像素）。
    width/height: 缩放后尺寸（原图像素）。
    rotation: 围绕印花中心旋转角度（度）。
    """
    x: float = 0.0
    y: float = 0.0
    width: float = 0.0
    height: float = 0.0
    rotation: float = 0.0
    keep_aspect_ratio: bool = True

    def center(self):
        """返回印花中心坐标 (cx, cy)。"""
        return self.x + self.width / 2.0, self.y + self.height / 2.0


# ---------------------------------------------------------------------------
# 模板
# ---------------------------------------------------------------------------

@dataclass
class Template:
    """可复用的合成参数模板，使用比例坐标适配任意尺寸衣服图。

    x_ratio / y_ratio: 印花左上角相对于衣服宽高的比例。
    width_ratio / height_ratio: 印花宽高相对于衣服宽高的比例。
    type: "builtin" 为内置模板，"custom" 为用户自定义模板。
    """
    name: str
    x_ratio: float = 0.0
    y_ratio: float = 0.0
    width_ratio: float = 0.3
    height_ratio: float = 0.3
    rotation: float = 0.0
    type: str = "custom"         # "builtin" | "custom"

    def to_transform_state(
        self,
        garment_width: float,
        garment_height: float,
        print_width: float = None,
        print_height: float = None,
    ) -> TransformState:
        """将比例参数转换为针对指定衣服尺寸的像素坐标 TransformState。

        width_ratio / height_ratio 定义衣服上的目标框。若提供印花原始尺寸
        (print_width / print_height)，印花按原始宽高比缩放后 contain 进目标框
        并居中，避免被拉伸变形；未提供时退化为直接铺满目标框（旧行为）。
        """
        box_x = self.x_ratio * garment_width
        box_y = self.y_ratio * garment_height
        box_w = self.width_ratio * garment_width
        box_h = self.height_ratio * garment_height

        if print_width and print_height and print_width > 0 and print_height > 0:
            scale = min(box_w / print_width, box_h / print_height)
            draw_w = print_width * scale
            draw_h = print_height * scale
            cx = box_x + box_w / 2.0
            cy = box_y + box_h / 2.0
            return TransformState(
                x=cx - draw_w / 2.0,
                y=cy - draw_h / 2.0,
                width=draw_w,
                height=draw_h,
                rotation=self.rotation,
            )

        return TransformState(
            x=box_x,
            y=box_y,
            width=box_w,
            height=box_h,
            rotation=self.rotation,
        )


# ---------------------------------------------------------------------------
# 导出选项
# ---------------------------------------------------------------------------

@dataclass
class ExportOptions:
    """单张或批量导出时的输出配置。"""
    output_dir: str = ""          # 空字符串 = 使用程序默认 output/ 目录
    output_format: str = "PNG"    # "PNG" | "JPG"
    quality: int = 95             # JPG 质量 1-95，PNG 时忽略


# ---------------------------------------------------------------------------
# 批量选项
# ---------------------------------------------------------------------------

class BatchMode(str, Enum):
    """批量合成的配对模式。"""
    ONE_TO_ONE = "one_to_one"         # 一一匹配
    MANY_GARMENTS = "many_garments"   # 多衣服 × 单印花
    MANY_PRINTS = "many_prints"       # 单衣服 × 多印花
    FULL_COMBO = "full_combo"         # 全组合（矩阵）


class VisibilityStatus(str, Enum):
    """衣服目标区域与印花的可见度状态。"""
    NORMAL = "正常"
    LOW = "偏低"
    UNCLEAR = "不明显"
    UNKNOWN = "无法判断"


@dataclass
class BatchOptions:
    """批量任务配置。"""
    mode: BatchMode = BatchMode.FULL_COMBO
    export_options: ExportOptions = field(default_factory=ExportOptions)


@dataclass
class VisibilityResult:
    """衣服目标区域与印花颜色对比分析结果。"""
    status: VisibilityStatus
    score: float = 0.0
    rgb_distance: float = 0.0
    brightness_difference: float = 0.0
    error: str = ""


# ---------------------------------------------------------------------------
# AI 穿搭
# ---------------------------------------------------------------------------

@dataclass
class OutfitTask:
    """AI 穿搭单行任务。

    row_index: Excel 行号（从 1 开始，便于写回 D/E/F 列）。
    status: 待处理 / 生成中 / 完成 / 失败 / 跳过。
    """
    row_index: int
    title: str
    product_id: str
    garment_path: str
    status: str = "待处理"


@dataclass
class OutfitResult:
    """AI 穿搭单行生成结果。

    output_path: 单文件行 = 结果图路径；目录行 = 输出子目录路径（写回 Excel D 列）。
    output_paths: 目录行下每张结果图的路径（供缩略图逐张展示）；单文件行留空。
    """
    task: OutfitTask
    success: bool
    output_path: str = ""
    error: str = ""
    attempts: int = 0
    output_paths: List[str] = field(default_factory=list)


@dataclass
class TitleResult:
    """AI 标题生成单行结果（docs/11 §17）。

    generated_title: 成功时为 AI 生成、清洗后的单行标题（写回 Excel A 列）。
    """
    task: OutfitTask
    success: bool
    generated_title: str = ""
    error: str = ""
    attempts: int = 0


# ---------------------------------------------------------------------------
# 合成结果
# ---------------------------------------------------------------------------

@dataclass
class ComposeResult:
    """单张合成任务的结果。"""
    success: bool
    garment_path: Optional[Path] = None
    print_path: Optional[Path] = None
    output_path: Optional[Path] = None   # 成功时有效
    error: str = ""                       # 失败时的错误说明
    visibility: Optional[VisibilityResult] = None


@dataclass
class BatchResult:
    """批量合成任务的汇总结果。"""
    results: List[ComposeResult] = field(default_factory=list)
    total: int = 0
    success_count: int = 0
    failure_count: int = 0