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) -> TransformState:
        """将比例参数转换为针对指定衣服尺寸的像素坐标 TransformState。"""
        return TransformState(
            x=self.x_ratio * garment_width,
            y=self.y_ratio * garment_height,
            width=self.width_ratio * garment_width,
            height=self.height_ratio * garment_height,
            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 = ""


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

@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