Upload BAM_parser.py

BAM_parser.py

@@ -0,0 +1,399 @@
+from __future__ import annotations
+
+import re
+from typing import Dict, List, Tuple
+
+
+# ============================================================================
+# Output definitions (order matters!)
+# ============================================================================
+_OUTPUT_DEFS: List[Tuple[str, str]] = [
+    # Core identity
+    ("gender_str", "STRING"),
+    ("gender_int", "INT"),
+    ("age_str", "STRING"),
+    ("age_int", "INT"),
+    ("identity_str", "STRING"),
+    ("eyecolor_str", "STRING"),
+    ("hairstyle_str", "STRING"),
+
+    # Equipment (strings)
+    ("topwear_str", "STRING"),
+    ("bellywear_str", "STRING"),
+    ("breastwear_str", "STRING"),
+    ("handwear_left_str", "STRING"),
+    ("handwear_right_str", "STRING"),
+    ("wristwear_left_str", "STRING"),
+    ("wristwear_right_str", "STRING"),
+    ("forearm_left_str", "STRING"),
+    ("forearm_right_str", "STRING"),
+    ("elbow_left_str", "STRING"),
+    ("elbow_right_str", "STRING"),
+    ("upperarm_left_str", "STRING"),
+    ("upperarm_right_str", "STRING"),
+    ("shoulder_left_str", "STRING"),
+    ("shoulder_right_str", "STRING"),
+    ("shank_left_str", "STRING"),
+    ("shank_right_str", "STRING"),
+    ("knee_left_str", "STRING"),
+    ("knee_right_str", "STRING"),
+    ("foot_left_str", "STRING"),
+    ("foot_right_str", "STRING"),
+    ("necklace_str", "STRING"),
+    ("earring_left_str", "STRING"),
+    ("earring_right_str", "STRING"),
+    ("kneewear_str", "STRING"),
+    ("headwear_str", "STRING"),
+    ("facemask_str", "STRING"),
+    ("sunglasses_str", "STRING"),
+    ("glasses_str", "STRING"),
+    ("crotch_str", "STRING"),
+    ("one_piece_str", "STRING"),
+
+    # Tags (strings)
+    ("aesthetic_tag1", "STRING"),
+    ("aesthetic_tag2", "STRING"),
+    ("aesthetic_tag3", "STRING"),
+    ("aesthetic_tag4", "STRING"),
+    ("aesthetic_tag5", "STRING"),
+
+    ("skin_tag1", "STRING"),
+    ("skin_tag2", "STRING"),
+    ("skin_tag3", "STRING"),
+    ("skin_tag4", "STRING"),
+    ("skin_tag5", "STRING"),
+
+    ("expression_tag1", "STRING"),
+    ("expression_tag2", "STRING"),
+    ("expression_tag3", "STRING"),
+    ("expression_tag4", "STRING"),
+    ("expression_tag5", "STRING"),
+
+    # Unique headwear slot AFTER expression tags (your step 23)
+    ("headwear_str_2", "STRING"),
+
+    # Equipment section simplified (values only, comma-separated)
+    ("old_bam", "STRING"),
+]
+
+RETURN_NAMES_TUPLE = tuple(name for name, _t in _OUTPUT_DEFS)
+RETURN_TYPES_TUPLE = tuple(_t for _name, _t in _OUTPUT_DEFS)
+
+
+# ============================================================================
+# Equipment parsing helpers
+# ============================================================================
+def _extract_parenthesized_items(s: str) -> List[str]:
+    """
+    Extract top-level (...) items from a string, handling nested parentheses.
+    Returns inside text of each (...) (outer parens removed).
+    """
+    items: List[str] = []
+    depth = 0
+    buf: List[str] = []
+
+    for ch in s or "":
+        if ch == "(":
+            if depth == 0:
+                buf = []
+            else:
+                buf.append(ch)
+            depth += 1
+        elif ch == ")":
+            if depth > 0:
+                depth -= 1
+            if depth == 0:
+                item = "".join(buf).strip()
+                if item:
+                    items.append(item)
+            else:
+                buf.append(ch)
+        else:
+            if depth > 0:
+                buf.append(ch)
+
+    return items
+
+
+def _normalize_key(k: str) -> str:
+    k = (k or "").strip().lower()
+    k = k.replace(" ", "_").replace("-", "_")
+    k = re.sub(r"_+", "_", k)
+    return k
+
+
+# Map key spellings to a canonical internal key.
+_KEY_CANONICAL: Dict[str, str] = {
+    "belly": "bellywear",
+    "bellywear": "bellywear",
+
+    "topwear": "topwear",
+
+    "breast": "breastwear",
+    "breastwear": "breastwear",
+
+    "hand": "handwear",
+    "handwear": "handwear",
+
+    "wrist": "wristwear",
+    "wristwear": "wristwear",
+
+    "forearm": "forearm",
+    "elbow": "elbow",
+
+    "upperarm": "upperarm",
+    "upper_arm": "upperarm",
+
+    "shoulder": "shoulder",
+    "shank": "shank",
+
+    "knee": "knee",
+
+    "foot": "foot",
+    "footwear": "foot",   # your example uses Footwear:
+    "shoe": "foot",
+    "shoes": "foot",
+
+    "necklace": "necklace",
+
+    "earring": "earring",
+    "earrings": "earring",
+
+    "kneewear": "kneewear",
+    "headwear": "headwear",
+
+    "facemask": "facemask",
+    "face_mask": "facemask",
+    "mask": "facemask",
+
+    "sunglasses": "sunglasses",
+    "glasses": "glasses",
+
+    "crotch": "crotch",
+
+    "onepiece": "one_piece",
+    "one_piece": "one_piece",
+    "one_piecewear": "one_piece",
+}
+
+# Canonical keys that are side-aware (Left/Right)
+_SIDE_FIELDS: Dict[str, Tuple[str, str]] = {
+    "handwear": ("handwear_left_str", "handwear_right_str"),
+    "wristwear": ("wristwear_left_str", "wristwear_right_str"),
+    "forearm": ("forearm_left_str", "forearm_right_str"),
+    "elbow": ("elbow_left_str", "elbow_right_str"),
+    "upperarm": ("upperarm_left_str", "upperarm_right_str"),
+    "shoulder": ("shoulder_left_str", "shoulder_right_str"),
+    "shank": ("shank_left_str", "shank_right_str"),
+    "knee": ("knee_left_str", "knee_right_str"),
+    "foot": ("foot_left_str", "foot_right_str"),
+    "earring": ("earring_left_str", "earring_right_str"),
+}
+
+# Canonical keys that map to a single output
+_SINGLE_FIELDS: Dict[str, str] = {
+    "topwear": "topwear_str",
+    "bellywear": "bellywear_str",
+    "breastwear": "breastwear_str",
+    "necklace": "necklace_str",
+    "kneewear": "kneewear_str",
+    "headwear": "headwear_str",
+    "facemask": "facemask_str",
+    "sunglasses": "sunglasses_str",
+    "glasses": "glasses_str",
+    "crotch": "crotch_str",
+    "one_piece": "one_piece_str",
+}
+
+_ALL_EQUIP_OUTPUTS = set(_SINGLE_FIELDS.values())
+for lf, rf in _SIDE_FIELDS.values():
+    _ALL_EQUIP_OUTPUTS.add(lf)
+    _ALL_EQUIP_OUTPUTS.add(rf)
+
+
+def _parse_equipment(equip: str) -> Tuple[Dict[str, str], str]:
+    """
+    Parse equipment section like:
+      (Knee_Left: Blue Skater Kneepad), (Knee:Green Skater Kneepads), (Belly:), ...
+
+    Rules implemented:
+    - Missing entries => output stays ""
+    - Empty "(Belly:)" => output ""
+    - Side-less "(Knee:...)" => assigns to BOTH left and right outputs
+    - old_bam => values-only, comma-separated, in original order
+    """
+    out: Dict[str, str] = {name: "" for name in _ALL_EQUIP_OUTPUTS}
+    old_values: List[str] = []
+
+    for item in _extract_parenthesized_items(equip or ""):
+        if ":" not in item:
+            continue
+
+        raw_key, raw_val = item.split(":", 1)
+        key = _normalize_key(raw_key)
+        val = (raw_val or "").strip()
+
+        # Trim a trailing comma, if any
+        if val.endswith(","):
+            val = val[:-1].rstrip()
+
+        # old_bam: collect only non-empty values
+        if val:
+            old_values.append(val)
+
+        # Detect side suffix
+        side = None
+        base_key = key
+        if base_key.endswith("_left"):
+            side = "left"
+            base_key = base_key[:-5]
+        elif base_key.endswith("_right"):
+            side = "right"
+            base_key = base_key[:-6]
+        base_key = base_key.rstrip("_")
+
+        canonical = _KEY_CANONICAL.get(base_key, base_key)
+
+        if canonical in _SIDE_FIELDS:
+            left_name, right_name = _SIDE_FIELDS[canonical]
+            if side == "left":
+                out[left_name] = val
+            elif side == "right":
+                out[right_name] = val
+            else:
+                # "(Knee:...)" style => both sides
+                out[left_name] = val
+                out[right_name] = val
+        elif canonical in _SINGLE_FIELDS:
+            out[_SINGLE_FIELDS[canonical]] = val
+        else:
+            # Unknown equipment key => ignored for structured outputs,
+            # but still included in old_bam via old_values.
+            pass
+
+    old_bam = ", ".join(v for v in old_values if v)
+    return out, old_bam
+
+
+# ============================================================================
+# BAM parsing
+# ============================================================================
+def _get_part(parts: List[str], idx: int) -> str:
+    return parts[idx] if idx < len(parts) else ""
+
+
+def _safe_int(s: str, default: int = 0) -> int:
+    try:
+        return int((s or "").strip())
+    except Exception:
+        return default
+
+
+def _parse_bam(bam: str) -> Dict[str, object]:
+    """
+    Expected structure after the first START### marker:
+
+      0  gender_num
+      1  age
+      2  identity
+      3  eyecolor
+      4  hairstyle
+      5  equipment
+      6..10   aesthetic_tag1..5
+      11..15  skin_tag1..5
+      16..20  expression_tag1..5
+      21 headwear_str_2
+      22+ irrelevant
+    """
+    bam = bam or ""
+    marker = "START###"
+    idx = bam.find(marker)
+    payload = bam[idx + len(marker):] if idx != -1 else bam
+
+    parts = payload.split("###")
+
+    gender_token = _get_part(parts, 0).strip()
+    gender_int = 1 if gender_token == "1" else 2
+    gender_str = "boy" if gender_int == 1 else "girl"
+
+    age_str = _get_part(parts, 1).strip()
+    age_int = _safe_int(age_str, default=0)
+
+    identity_str = _get_part(parts, 2).strip()
+    eyecolor_str = _get_part(parts, 3).strip()
+    hairstyle_str = _get_part(parts, 4).strip()
+
+    equipment_raw = _get_part(parts, 5).strip()
+    equip_out, old_bam = _parse_equipment(equipment_raw)
+
+    out: Dict[str, object] = {
+        "gender_str": gender_str,
+        "gender_int": gender_int,
+        "age_str": age_str,
+        "age_int": age_int,
+        "identity_str": identity_str,
+        "eyecolor_str": eyecolor_str,
+        "hairstyle_str": hairstyle_str,
+        "old_bam": old_bam,
+    }
+
+    # Equipment
+    out.update(equip_out)
+
+    # Aesthetic tags
+    for i in range(5):
+        out[f"aesthetic_tag{i+1}"] = _get_part(parts, 6 + i).strip()
+
+    # Skin tags
+    for i in range(5):
+        out[f"skin_tag{i+1}"] = _get_part(parts, 11 + i).strip()
+
+    # Expression tags
+    for i in range(5):
+        out[f"expression_tag{i+1}"] = _get_part(parts, 16 + i).strip()
+
+    # Outer-layer headwear slot
+    out["headwear_str_2"] = _get_part(parts, 21).strip()
+
+    return out
+
+
+# ============================================================================
+# ComfyUI node
+# ============================================================================
+class BAMFormatParser:
+    """
+    ComfyUI custom node: parses your BAM string and exposes each field as outputs.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "bam_string": ("STRING", {"multiline": True, "default": ""}),
+            }
+        }
+
+    RETURN_TYPES = RETURN_TYPES_TUPLE
+    RETURN_NAMES = RETURN_NAMES_TUPLE
+    FUNCTION = "parse"
+    CATEGORY = "BAM"
+
+    def parse(self, bam_string: str):
+        parsed = _parse_bam(bam_string)
+
+        # Guarantee every declared output exists (malformed BAM => defaults)
+        for name, t in _OUTPUT_DEFS:
+            if name not in parsed:
+                parsed[name] = 0 if t == "INT" else ""
+
+        return tuple(parsed[name] for name in RETURN_NAMES_TUPLE)
+
+
+NODE_CLASS_MAPPINGS = {
+    "BAMFormatParser": BAMFormatParser,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BAMFormatParser": "BAM Format Parser",
+}