Pre_Seg_Server/backend/services/default_templates.py

"""Bundled system ontology templates and restore helpers."""

from __future__ import annotations

from copy import deepcopy

from sqlalchemy.orm import Session

from models import Template

RESERVED_UNCLASSIFIED_CLASS = {
    "id": "reserved-unclassified",
    "name": "待分类",
    "color": "#000000",
    "zIndex": 0,
    "maskId": 0,
    "category": "系统保留",
}


def _with_reserved_unclassified_class(classes: list[dict]) -> list[dict]:
    filtered = [
        item for item in classes
        if item.get("id") != RESERVED_UNCLASSIFIED_CLASS["id"]
        and item.get("name") != RESERVED_UNCLASSIFIED_CLASS["name"]
        and item.get("maskId") != 0
    ]
    return [*filtered, dict(RESERVED_UNCLASSIFIED_CLASS)]


def _template_classes(
    template_name: str,
    names: list[str],
    colors: list[tuple[int, int, int]],
    *,
    id_prefix: str,
) -> list[dict]:
    classes = []
    for idx, (rgb, name) in enumerate(zip(colors, names)):
        color_hex = f"#{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}"
        classes.append({
            "id": f"{id_prefix}-{idx}",
            "name": name,
            "color": color_hex,
            "zIndex": (len(names) - idx) * 10,
            "maskId": idx + 1,
            "category": template_name,
        })
    return classes


def bundled_default_template_definitions() -> list[dict]:
    """Return fresh definitions for all bundled system templates."""
    return [
        {
            "name": "腹腔镜胆囊切除术",
            "description": "腹腔镜胆囊切除术（LC）手术器械与解剖结构语义分割模板，共35个分类",
            "color": "#06b6d4",
            "z_index": 0,
            "classes": _with_reserved_unclassified_class(_template_classes(
                "腹腔镜胆囊切除术",
                [
                    "针", "线", "肿瘤", "血管阻断夹", "棉球", "双极电凝",
                    "肝脏", "胆囊", "分离钳", "脂肪", "止血海绵", "肝总管",
                    "吸引器", "剪刀", "超声刀", "止血纱布", "胆总管", "生物夹",
                    "无损伤钳", "钳夹", "喷洒", "胆囊管", "动脉", "电凝",
                    "静脉", "标本袋", "引流管", "纱布", "金属钛夹", "韧带",
                    "肝蒂", "推结器", "乳胶管-血管阻断", "吻合器", "术中超声",
                ],
                [
                    (134, 124, 118), (0, 157, 142), (245, 161, 0), (255, 172, 159), (146, 175, 236), (155, 62, 0),
                    (255, 91, 0), (255, 234, 0), (85, 111, 181), (155, 132, 0), (181, 227, 14), (72, 0, 255),
                    (255, 0, 255), (29, 32, 136), (240, 16, 116), (160, 15, 95), (0, 155, 33), (0, 160, 233),
                    (52, 184, 178), (66, 115, 82), (90, 120, 41), (255, 0, 0), (117, 0, 0), (167, 24, 233),
                    (42, 8, 66), (112, 113, 150), (0, 255, 0), (255, 255, 255), (0, 255, 255), (181, 85, 105),
                    (113, 102, 140), (202, 202, 200), (197, 83, 181), (136, 162, 196), (138, 251, 213),
                ],
                id_prefix="cls-lap",
            )),
        },
        {
            "name": "头颈部CT分割",
            "description": "头颈部CT分割",
            "color": "#ef4444",
            "z_index": 10,
            "classes": _with_reserved_unclassified_class(_template_classes(
                "头颈部CT分割",
                [
                    "肿瘤/结节",
                    "下颌骨",
                    "甲状腺",
                    "气管",
                    "颈椎",
                    "颈动脉",
                    "颈静脉",
                    "腮腺",
                    "下颌下腺",
                    "舌骨",
                ],
                [
                    (255, 0, 0),
                    (0, 255, 0),
                    (0, 0, 255),
                    (255, 255, 0),
                    (255, 0, 255),
                    (0, 255, 255),
                    (255, 128, 0),
                    (128, 0, 128),
                    (0, 128, 128),
                    (128, 128, 0),
                ],
                id_prefix="cls-head-neck-ct",
            )),
        },
    ]


def _has_legacy_head_neck_english_labels(template: Template) -> bool:
    if template.name != "头颈部CT分割":
        return False
    classes = (template.mapping_rules or {}).get("classes") or []
    return any(
        isinstance(item, dict)
        and isinstance(item.get("name"), str)
        and "(" in item["name"]
        and ")" in item["name"]
        for item in classes
    )


def ensure_default_templates(db: Session, *, restore_existing: bool = False) -> list[Template]:
    """Create bundled system templates, optionally restoring existing ones exactly."""
    templates: list[Template] = []
    for definition in bundled_default_template_definitions():
        existing = db.query(Template).filter(
            Template.name == definition["name"],
            Template.owner_user_id.is_(None),
        ).first()
        if existing is None:
            existing = Template(owner_user_id=None)
            db.add(existing)
        elif not restore_existing and not _has_legacy_head_neck_english_labels(existing):
            templates.append(existing)
            continue

        existing.name = definition["name"]
        existing.description = definition["description"]
        existing.color = definition["color"]
        existing.z_index = definition["z_index"]
        existing.owner_user_id = None
        existing.mapping_rules = {
            "classes": deepcopy(definition["classes"]),
            "rules": [],
        }
        templates.append(existing)
    db.commit()
    for template in templates:
        db.refresh(template)
    return templates


def restore_default_templates(db: Session) -> list[Template]:
    """Restore bundled system templates after demo factory reset."""
    return ensure_default_templates(db, restore_existing=True)