submodel_render.py

import xml.etree.ElementTree as ET
import sys
import os
from PIL import Image, ImageDraw

SCALE = 20   # increase for higher resolution (10–40 works well)


# ---------------------------------------------------------
# Parse CustomModel and Submodels
# ---------------------------------------------------------
def parse_xmodel(path):
    print(f"Loading {path} ...")
    tree = ET.parse(path)
    root = tree.getroot()

    custom = root.get("CustomModel")
    if not custom:
        raise ValueError("❌ No CustomModel attribute found in .xmodel file!")

    # -------------------------
    # Parse custom model grid
    # -------------------------
    rows = custom.split(";")
    grid = [row.split(",") for row in rows]

    pixel_coords = {}

    for y, row in enumerate(grid):
        for x, cell in enumerate(row):
            cell = cell.strip()
            if cell == "" or not cell.isdigit():
                continue
            pixel_index = int(cell)  # 1-based index used by xLights
            pixel_coords[pixel_index] = (x, y)

    if not pixel_coords:
        raise ValueError("❌ CustomModel grid contained no pixel numbers!")

    max_index = max(pixel_coords.keys())

    # Convert dictionary to sorted 0-based list
    coords_list = [None] * max_index
    for idx, pos in pixel_coords.items():
        coords_list[idx - 1] = pos

    # -------------------------
    # Parse submodels
    # -------------------------
    submodels = {}

    for sm in root.findall(".//subModel"):
        name = sm.get("name")
        groups = sm.get("group", None)

        lines = [sm.get(f"line{i}") for i in range(10)]
        lines = [ln for ln in lines if ln]

        pix_indexes = []

        for ln in lines:
            parts = ln.split(",")
            for part in parts:
                part = part.strip()
                if not part:
                    continue

                if "-" in part:
                    a, b = part.split("-")
                    pix_indexes.extend(range(int(a), int(b) + 1))
                else:
                    pix_indexes.append(int(part))

        # convert to 0-based
        submodels[name] = {
            "pixels": [i - 1 for i in pix_indexes],
            "groups": groups
        }

    return coords_list, submodels, len(grid[0]), len(grid)


# ---------------------------------------------------------
# Draw a submodel image
# ---------------------------------------------------------
def render_submodel(name, coords, active_pixels, width, height, outdir):
    print(f"Rendering submodel: {name}")

    img = Image.new("RGB", (width * SCALE, height * SCALE), "black")
    draw = ImageDraw.Draw(img)

    # Draw all pixels (dark gray)
    for idx, pos in enumerate(coords):
        if pos is None:
            continue
        x, y = pos
        x0 = x * SCALE
        y0 = y * SCALE
        x1 = x0 + SCALE
        y1 = y0 + SCALE

        color = (75, 75, 75)
        draw.rectangle([x0, y0, x1, y1], fill=color)

    # Draw active submodel pixels (red)
    for idx in active_pixels:
        if idx < 0 or idx >= len(coords) or coords[idx] is None:
            continue
        x, y = coords[idx]
        x0 = x * SCALE
        y0 = y * SCALE
        x1 = x0 + SCALE
        y1 = y0 + SCALE

        color = (255, 0, 0)
        draw.rectangle([x0, y0, x1, y1], fill=color)

    # Save
    outpath = os.path.join(outdir, f"{name}.png")
    img.save(outpath)
    print(f"→ Saved {outpath}")


# ---------------------------------------------------------
# Parse model groups
# ---------------------------------------------------------
def parse_model_groups(path):
    tree = ET.parse(path)
    root = tree.getroot()

    model_groups = {}

    for mg in root.findall(".//modelGroup"):
        name = mg.get("name")
        models_raw = mg.get("models", "")

        # Split models attribute: "EXPORTEDMODEL/Arm 1,EXPORTEDMODEL/Arm 2"
        models = [
            m.replace("EXPORTEDMODEL/", "").strip()
            for m in models_raw.split(",")
            if m.strip() != ""
        ]

        model_groups[name] = models

    return model_groups


# ---------------------------------------------------------
# Create README.md with embedded images
# ---------------------------------------------------------
def create_readme(outdir, model_name, submodels, model_groups):
    readme_path = os.path.join(outdir, "README.md")

    def enc(name):
        """URL-encode spaces for markdown images."""
        return name.replace(" ", "%20")

    with open(readme_path, "w", encoding="utf-8") as f:
        f.write(f"# Submodel Renders for **{model_name}**\n\n")
        f.write("Automatically generated by `submodel.py`.\n\n")

        f.write("## Model Groups Overview\n\n")

        for group_name, model_list in model_groups.items():
            f.write(f"### {group_name}\n\n")
            f.write("| " + " | ".join(model_list) + " |\n")
            f.write("| " + " | ".join(["---"] * len(model_list)) + " |\n")

            row = "| "
            for sm in model_list:
                filename = enc(f"images/{sm}.png")  # points to images folder
                row += f'<img src="{filename}" width="150"><br>{sm} | '
            f.write(row + "\n\n")

        f.write("\n---\n\n## Individual Submodels\n\n")
        for name in submodels:
            filename = enc(f"images/{name}.png")  # points to images folder
            f.write(f"### {name}\n")
            f.write(f"![{name}]({filename})\n\n")
            f.write("---\n\n")

    print(f"→ README created at {readme_path}")


# ---------------------------------------------------------
# MAIN
# ---------------------------------------------------------
def main():
    if len(sys.argv) < 2:
        print("Usage: python submodel.py <file.xmodel>")
        return

    path = sys.argv[1]

    coords, submodels, width, height = parse_xmodel(path)

    # Parse model groups
    model_groups = parse_model_groups(path)

    # Output directories
    model_name = os.path.splitext(os.path.basename(path))[0]
    outdir = os.path.join("renders", model_name)
    images_dir = os.path.join(outdir, "images")
    os.makedirs(images_dir, exist_ok=True)

    # Render each submodel into images_dir
    for name, info in submodels.items():
        render_submodel(
            name,
            coords,
            info["pixels"],
            width,
            height,
            images_dir
        )

    # Create README.md pointing to images/
    create_readme(outdir, model_name, submodels, model_groups)

    print("\nAll submodels rendered!")


if __name__ == "__main__":
    main()