"""SPU/SKU 数据库查询 + 底图/模特图资源查找(产品图生成流水线的数据层)。 数据关系(已核实 spu_sku.db): - SPU.code = 款号(如 DG004) - SKU.code = "款号-颜色编码"(如 DG004-BL01),SKU.color = 中文色名(黑/灰/...) - basemap 目录 = basemap/<款号>//xxx.jpg - material_library/<品类>/ 存放模特图 - SKU.img_url_2~5 = CDN 图 URL(底图/细节/模特图,仅作参考字段) """ import sqlite3 from pathlib import Path from typing import Any, Dict, List, Optional from PIL import Image # 支持的图片格式:模特图/底图均按此识别(png/jpg 等常见格式全覆盖;AVIF/GIF/TIFF 亦支持) IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".avif", ".gif", ".tiff", ".tif"} def _connect(db_path) -> sqlite3.Connection: conn = sqlite3.connect(str(db_path)) conn.row_factory = sqlite3.Row return conn def list_spus(db_path, country: Optional[str] = None) -> List[Dict[str, Any]]: """全部 SPU(可选按国家过滤)。""" conn = _connect(db_path) sql = "SELECT id, code, style, material, printing_type, target_audience, pattern, country, mark FROM SPU" params: list = [] if country: sql += " WHERE country = ?" params.append(country) sql += " ORDER BY code" rows = conn.execute(sql, params).fetchall() conn.close() return [dict(r) for r in rows] def list_colors(db_path, spu_code: str) -> List[Dict[str, Any]]: """款号 → 颜色列表(SKU.code 去重,附中文色名、CDN 底图 URL、最低价)。""" conn = _connect(db_path) rows = conn.execute( """SELECT s.code AS sku_code, s.color, s.img_url_2 AS img_url, MIN(s.price) AS price FROM SKU s JOIN SPU p ON s.spu_id = p.id WHERE p.code = ? GROUP BY s.code, s.color ORDER BY s.code""", (spu_code,)).fetchall() conn.close() return [dict(r) for r in rows] def _norm_name(s: str) -> str: """去空格(半角+全角)+ 小写,用于 SKU code 与文件夹名比较。""" return str(s or "").replace(" ", "").replace(" ", "").strip().lower() def find_basemap(basemap_root, spu_code: str, sku_code: str) -> Optional[Path]: """basemap/<款号>// 下第一张图片;无返回 None。 SKU code 与文件夹名比较时两边都去空格(兼容 db code 或文件夹名带空格)。 """ root = Path(basemap_root) / spu_code if not root.exists(): return None target = _norm_name(sku_code) if not target: return None d = root / sku_code if not (d.exists() and d.is_dir()): d = None for cand in sorted(root.iterdir()): if cand.is_dir() and _norm_name(cand.name) == target: d = cand break if d is None: return None for f in sorted(d.iterdir()): if f.is_file() and f.suffix.lower() in IMG_EXTS: return f return None def list_model_images(material_root, category: str = "T-shirt") -> List[Path]: """material_library/<品类>/ 下所有图片;无返回空列表。""" d = Path(material_root) / category if not d.exists(): return [] return [f for f in sorted(d.iterdir()) if f.is_file() and f.suffix.lower() in IMG_EXTS] def find_first_model_folder(material_root, preferred: Optional[str] = None): """material_library 下「第一个有图片的子目录」及其图片列表。 - preferred(如 config 的 model_category)优先:该目录有图就直接用; - 否则按子目录名排序,取第一个有图的目录; - 全空返回 (None, [])。 返回 (dir_name or None, images: List[Path])。 """ root = Path(material_root) if not root.exists(): return None, [] candidates = [] if preferred: d = root / preferred if d.is_dir(): candidates.append(d) candidates += [d for d in sorted(root.iterdir()) if d.is_dir()] seen = set() for d in candidates: if d in seen: continue seen.add(d) imgs = [f for f in sorted(d.iterdir()) if f.is_file() and f.suffix.lower() in IMG_EXTS] if imgs: return d.name, imgs return None, [] def image_ratio_ok(path, target_ratio: float = 3 / 4, tolerance: float = 0.06) -> bool: """图片宽高比是否接近目标比例(默认 3:4,相对容差 6%)。 相对容差:接受 [target*(1-tol), target*(1+tol)],对 3:4 即 0.705~0.795。 无法解析的图片(损坏/非标准)按不通过处理,避免坏图被当模特。 """ try: with Image.open(path) as im: w, h = im.size if w <= 0 or h <= 0: return False ratio = w / h lo = target_ratio * (1 - tolerance) hi = target_ratio * (1 + tolerance) return lo <= ratio <= hi except Exception: # noqa: BLE001 return False def build_mark_model_map(db_path, material_root) -> Dict[str, str]: """启动任务前检测 spu.mark 字段,建立 {mark: 模特文件夹名} 字典。 规则: - 读取 spu 表全部 mark 值(去重); - 每个 mark 映射到 material_library/ 目录(目录名与 mark 一致); - 目录不存在时回退到 category 默认目录(T-shirt); - 目前库中 mark=1 → 映射到 material_library/T-shirt。 """ mark_map: Dict[str, str] = {} try: root = Path(material_root) if not root.exists(): return mark_map folders = [d.name for d in sorted(root.iterdir()) if d.is_dir()] if not folders: return mark_map # 读 spu.mark 实际值(去重) marks: List[str] = [] try: conn = _connect(db_path) rows = conn.execute("SELECT DISTINCT mark FROM SPU WHERE mark IS NOT NULL AND mark != ''").fetchall() conn.close() marks = [str(r["mark"]).strip() for r in rows if str(r["mark"]).strip()] except Exception: # noqa: BLE001 marks = [] if not marks: marks = ["1"] # 库无 mark 数据时按默认 1 处理 for m in marks: if m in folders: mark_map[m] = m else: # mark 无同名目录 → 回退默认 T-shirt(当前 mark=1 → T-shirt) mark_map[m] = "T-shirt" if "T-shirt" in folders else folders[0] print(f"[product] mark→模特目录映射: {mark_map}") except Exception as e: # noqa: BLE001 print(f"[product] mark→模特目录映射构建失败: {e}") return mark_map def find_model_images_for_mark(db_path, material_root, mark, category: str = "T-shirt", ratio: float = 3 / 4, tolerance: float = 0.06) -> List[Path]: """按 spu.mark 定位模特目录,过滤非目标比例图片,返回合格图片列表。 - mark 有对应目录(material_library/)→ 用该目录; - 否则回退 category(如 T-shirt); - 过滤掉非 3:4 比例(默认容差 6%)的图片; - 返回过滤后的图片列表(供调用方随机抽取)。 """ root = Path(material_root) if not root.exists(): return [] d = None if mark is not None: cand = root / str(mark) if cand.is_dir(): d = cand if d is None: cand = root / category if cand.is_dir(): d = cand if d is None: # 兜底:第一个有图的目录(跳过无图目录) for sub in sorted(root.iterdir()): if not sub.is_dir(): continue if any(f.is_file() and f.suffix.lower() in IMG_EXTS for f in sub.iterdir()): d = sub break if d is None: return [] imgs = [f for f in sorted(d.iterdir()) if f.is_file() and f.suffix.lower() in IMG_EXTS] ok = [f for f in imgs if image_ratio_ok(f, ratio, tolerance)] if len(ok) < len(imgs): print(f"[product] 模特目录 {d.name}/ 过滤非 {int(ratio * 100)}:{int(ratio * 100) + 1} 比例:" f"{len(imgs)} → {len(ok)} 张") return ok def first_available_sku(db_path, basemap_root, spu_code: str) -> Optional[str]: """返回该款号下第一个「本地有底图」的 SKU.code;无则 None。""" for c in list_colors(db_path, spu_code): if find_basemap(basemap_root, spu_code, c["sku_code"]) is not None: return c["sku_code"] return None