576 lines
20 KiB
Python
Executable File
576 lines
20 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""Combine per-material Substance exports with a Blender Material ID map.
|
||
|
||
The companion Blender extension writes a PNG and JSON legend. This tool uses
|
||
those exact ID colors as masks, discovers material names in Substance Painter
|
||
filenames, and combines matching channels without trusting texture alpha.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import sys
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any, Iterable, Sequence
|
||
|
||
try:
|
||
from PIL import Image, ImageChops
|
||
except ImportError: # pragma: no cover - exercised by the friendly CLI error.
|
||
Image = None
|
||
ImageChops = None
|
||
|
||
|
||
SUPPORTED_SUFFIXES = {
|
||
".bmp",
|
||
".jpeg",
|
||
".jpg",
|
||
".png",
|
||
".tga",
|
||
".tif",
|
||
".tiff",
|
||
".webp",
|
||
}
|
||
MATERIAL_TOKEN = "{material}"
|
||
|
||
|
||
class CombineError(RuntimeError):
|
||
"""A user-correctable texture discovery or composition error."""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class LegendMaterial:
|
||
slot: int
|
||
name: str
|
||
rgb: tuple[int, int, int]
|
||
filename_name: str
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class TextureCandidate:
|
||
path: Path
|
||
material: LegendMaterial
|
||
relative_parent: Path
|
||
stem_pattern: str
|
||
suffix: str
|
||
|
||
|
||
@dataclass
|
||
class TextureGroup:
|
||
relative_parent: Path
|
||
stem_pattern: str
|
||
suffix: str
|
||
textures: dict[str, Path]
|
||
|
||
@property
|
||
def output_name(self) -> str:
|
||
stem = self.stem_pattern.replace(MATERIAL_TOKEN, "")
|
||
stem = re.sub(r"[\s_.-]+", "_", stem).strip("_")
|
||
return f"{stem or 'Combined'}{self.suffix.lower()}"
|
||
|
||
@property
|
||
def output_parent(self) -> Path:
|
||
parts = []
|
||
for part in self.relative_parent.parts:
|
||
cleaned = part.replace(MATERIAL_TOKEN, "")
|
||
cleaned = re.sub(r"[\s_.-]+", "_", cleaned).strip("_")
|
||
if cleaned:
|
||
parts.append(cleaned)
|
||
return Path(*parts) if parts else Path(".")
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
parent = "" if self.relative_parent == Path(".") else f"{self.relative_parent}/"
|
||
return f"{parent}{self.stem_pattern}{self.suffix}"
|
||
|
||
|
||
def _require_pillow() -> None:
|
||
if Image is None or ImageChops is None:
|
||
raise CombineError(
|
||
"Pillow is required. Install it with "
|
||
"'python -m pip install -r tools/requirements.txt'."
|
||
)
|
||
|
||
|
||
def _parse_aliases(values: Sequence[str]) -> dict[str, str]:
|
||
aliases: dict[str, str] = {}
|
||
for value in values:
|
||
if "=" not in value:
|
||
raise CombineError(
|
||
f"Invalid material alias {value!r}; expected 'LEGEND NAME=FILENAME NAME'."
|
||
)
|
||
legend_name, filename_name = (part.strip() for part in value.split("=", 1))
|
||
if not legend_name or not filename_name:
|
||
raise CombineError(f"Invalid material alias {value!r}; both names are required.")
|
||
aliases[legend_name.casefold()] = filename_name
|
||
return aliases
|
||
|
||
|
||
def _rgb8(value: Any, material_name: str) -> tuple[int, int, int]:
|
||
if not isinstance(value, list) or len(value) < 3:
|
||
raise CombineError(f"Legend material {material_name!r} has no valid RGB triplet.")
|
||
try:
|
||
components = tuple(float(component) for component in value[:3])
|
||
except (TypeError, ValueError) as exc:
|
||
raise CombineError(
|
||
f"Legend material {material_name!r} has a non-numeric RGB value."
|
||
) from exc
|
||
if any(component < 0.0 or component > 1.0 for component in components):
|
||
raise CombineError(
|
||
f"Legend material {material_name!r} has RGB values outside 0.0–1.0."
|
||
)
|
||
return tuple(round(component * 255.0) for component in components)
|
||
|
||
|
||
def load_legend(path: Path, aliases: dict[str, str] | None = None) -> list[LegendMaterial]:
|
||
aliases = aliases or {}
|
||
try:
|
||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||
except OSError as exc:
|
||
raise CombineError(f"Could not read legend {path}: {exc}") from exc
|
||
except json.JSONDecodeError as exc:
|
||
raise CombineError(f"Legend {path} is not valid JSON: {exc}") from exc
|
||
|
||
entries = payload.get("materials") if isinstance(payload, dict) else None
|
||
if not isinstance(entries, list) or not entries:
|
||
raise CombineError(f"Legend {path} does not contain a non-empty 'materials' list.")
|
||
|
||
materials: list[LegendMaterial] = []
|
||
seen_names: set[str] = set()
|
||
seen_colors: set[tuple[int, int, int]] = set()
|
||
for fallback_slot, entry in enumerate(entries, start=1):
|
||
if not isinstance(entry, dict):
|
||
raise CombineError(f"Legend entry {fallback_slot} is not an object.")
|
||
name = str(entry.get("material", "")).strip()
|
||
if not name:
|
||
raise CombineError(f"Legend entry {fallback_slot} has no material name.")
|
||
folded_name = name.casefold()
|
||
if folded_name in seen_names:
|
||
raise CombineError(f"Legend contains duplicate material name {name!r}.")
|
||
rgb = _rgb8(entry.get("rgb"), name)
|
||
if rgb in seen_colors:
|
||
raise CombineError(
|
||
f"Legend color {rgb} is shared by multiple materials; masks would be ambiguous."
|
||
)
|
||
try:
|
||
slot = int(entry.get("slot", fallback_slot))
|
||
except (TypeError, ValueError) as exc:
|
||
raise CombineError(f"Legend material {name!r} has an invalid slot.") from exc
|
||
materials.append(
|
||
LegendMaterial(
|
||
slot=slot,
|
||
name=name,
|
||
rgb=rgb,
|
||
filename_name=aliases.get(folded_name, name),
|
||
)
|
||
)
|
||
seen_names.add(folded_name)
|
||
seen_colors.add(rgb)
|
||
return materials
|
||
|
||
|
||
def _name_pattern(name: str) -> re.Pattern[str]:
|
||
tokens = re.findall(r"[A-Za-z0-9]+", name)
|
||
if not tokens:
|
||
raise CombineError(f"Material filename name {name!r} has no letters or digits.")
|
||
expression = r"[\s_.-]*".join(re.escape(token) for token in tokens)
|
||
return re.compile(
|
||
rf"(?<![A-Za-z0-9]){expression}(?![A-Za-z0-9])",
|
||
flags=re.IGNORECASE,
|
||
)
|
||
|
||
|
||
def _match_candidate(
|
||
path: Path,
|
||
input_dir: Path,
|
||
materials: Sequence[LegendMaterial],
|
||
) -> TextureCandidate | None:
|
||
relative = path.relative_to(input_dir)
|
||
searchable = [("stem", relative.stem)] + [
|
||
(f"parent:{index}", part)
|
||
for index, part in enumerate(relative.parent.parts)
|
||
]
|
||
matches: list[
|
||
tuple[int, int, LegendMaterial, re.Match[str], str]
|
||
] = []
|
||
for material in materials:
|
||
pattern = _name_pattern(material.filename_name)
|
||
for location, text in searchable:
|
||
match = pattern.search(text)
|
||
if match:
|
||
matches.append(
|
||
(
|
||
match.end() - match.start(),
|
||
len(material.filename_name),
|
||
material,
|
||
match,
|
||
location,
|
||
)
|
||
)
|
||
if not matches:
|
||
return None
|
||
|
||
_, _, material, match, location = max(
|
||
matches,
|
||
key=lambda item: (item[0], item[1], item[4] == "stem"),
|
||
)
|
||
stem_pattern = relative.stem
|
||
parent_parts = list(relative.parent.parts)
|
||
if location == "stem":
|
||
stem_pattern = (
|
||
f"{stem_pattern[:match.start()]}{MATERIAL_TOKEN}{stem_pattern[match.end():]}"
|
||
)
|
||
else:
|
||
parent_index = int(location.split(":", 1)[1])
|
||
parent_part = parent_parts[parent_index]
|
||
parent_parts[parent_index] = (
|
||
f"{parent_part[:match.start()]}{MATERIAL_TOKEN}{parent_part[match.end():]}"
|
||
)
|
||
relative_parent = Path(*parent_parts) if parent_parts else Path(".")
|
||
return TextureCandidate(
|
||
path=path,
|
||
material=material,
|
||
relative_parent=relative_parent,
|
||
stem_pattern=stem_pattern,
|
||
suffix=path.suffix,
|
||
)
|
||
|
||
|
||
def discover_groups(
|
||
input_dir: Path,
|
||
materials: Sequence[LegendMaterial],
|
||
recursive: bool = False,
|
||
excluded_paths: Iterable[Path] = (),
|
||
) -> tuple[list[TextureGroup], list[Path]]:
|
||
excluded = {path.resolve() for path in excluded_paths}
|
||
iterator = input_dir.rglob("*") if recursive else input_dir.glob("*")
|
||
files = sorted(
|
||
path
|
||
for path in iterator
|
||
if path.is_file()
|
||
and path.suffix.lower() in SUPPORTED_SUFFIXES
|
||
and path.resolve() not in excluded
|
||
)
|
||
|
||
grouped: dict[tuple[Path, str, str], TextureGroup] = {}
|
||
unmatched: list[Path] = []
|
||
for path in files:
|
||
candidate = _match_candidate(path, input_dir, materials)
|
||
if candidate is None:
|
||
unmatched.append(path)
|
||
continue
|
||
key = (
|
||
candidate.relative_parent,
|
||
candidate.stem_pattern.casefold(),
|
||
candidate.suffix.casefold(),
|
||
)
|
||
group = grouped.setdefault(
|
||
key,
|
||
TextureGroup(
|
||
relative_parent=candidate.relative_parent,
|
||
stem_pattern=candidate.stem_pattern,
|
||
suffix=candidate.suffix,
|
||
textures={},
|
||
),
|
||
)
|
||
material_key = candidate.material.name.casefold()
|
||
if material_key in group.textures:
|
||
raise CombineError(
|
||
f"Both {group.textures[material_key]} and {candidate.path} match material "
|
||
f"{candidate.material.name!r} in group {group.description!r}."
|
||
)
|
||
group.textures[material_key] = candidate.path
|
||
|
||
return sorted(grouped.values(), key=lambda group: group.description.casefold()), unmatched
|
||
|
||
|
||
def _mask_for_color(id_rgb: Any, rgb: tuple[int, int, int], tolerance: int) -> Any:
|
||
channel_masks = []
|
||
for channel, expected in zip(id_rgb.split(), rgb, strict=True):
|
||
lookup = [255 if abs(value - expected) <= tolerance else 0 for value in range(256)]
|
||
channel_masks.append(channel.point(lookup, mode="L"))
|
||
mask = ImageChops.multiply(
|
||
ImageChops.multiply(channel_masks[0], channel_masks[1]),
|
||
channel_masks[2],
|
||
)
|
||
# One bit is sufficient for exact material membership and keeps 4K/8K
|
||
# projects from retaining a full byte per material per pixel.
|
||
return mask.convert("1")
|
||
|
||
|
||
def _common_mode(images: Sequence[Any]) -> str:
|
||
modes = {image.mode for image in images}
|
||
if len(modes) == 1:
|
||
mode = next(iter(modes))
|
||
if mode in {"1", "L", "LA", "RGB", "RGBA", "I", "F"}:
|
||
return mode
|
||
if any("A" in mode or mode == "P" for mode in modes):
|
||
return "RGBA"
|
||
if modes <= {"1", "L", "I", "F"}:
|
||
return "L"
|
||
return "RGB"
|
||
|
||
|
||
def combine_group(
|
||
group: TextureGroup,
|
||
materials: Sequence[LegendMaterial],
|
||
masks: dict[str, Any],
|
||
expected_size: tuple[int, int],
|
||
output_path: Path,
|
||
overwrite: bool = False,
|
||
) -> list[str]:
|
||
if output_path.exists() and not overwrite:
|
||
raise CombineError(f"Output already exists: {output_path} (use --overwrite).")
|
||
|
||
opened: dict[str, Any] = {}
|
||
try:
|
||
for material in materials:
|
||
path = group.textures.get(material.name.casefold())
|
||
if path is None:
|
||
continue
|
||
image = Image.open(path)
|
||
image.load()
|
||
if image.size != expected_size:
|
||
raise CombineError(
|
||
f"{path} is {image.size[0]}x{image.size[1]}, but the ID map is "
|
||
f"{expected_size[0]}x{expected_size[1]}."
|
||
)
|
||
opened[material.name.casefold()] = image
|
||
|
||
if not opened:
|
||
raise CombineError(f"Texture group {group.description!r} contains no usable images.")
|
||
|
||
output_mode = _common_mode(list(opened.values()))
|
||
output = Image.new(output_mode, expected_size, 0)
|
||
try:
|
||
missing: list[str] = []
|
||
for material in materials:
|
||
key = material.name.casefold()
|
||
source = opened.get(key)
|
||
if source is None:
|
||
missing.append(material.name)
|
||
continue
|
||
converted = None
|
||
if source.mode != output_mode:
|
||
converted = source.convert(output_mode)
|
||
source = converted
|
||
try:
|
||
output.paste(source, (0, 0), masks[key])
|
||
finally:
|
||
if converted is not None:
|
||
converted.close()
|
||
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
output.save(output_path)
|
||
return missing
|
||
finally:
|
||
output.close()
|
||
finally:
|
||
for image in opened.values():
|
||
image.close()
|
||
|
||
|
||
def combine_all(
|
||
id_map_path: Path,
|
||
legend_path: Path,
|
||
input_dir: Path,
|
||
output_dir: Path,
|
||
aliases: dict[str, str] | None = None,
|
||
recursive: bool = False,
|
||
tolerance: int | None = None,
|
||
overwrite: bool = False,
|
||
dry_run: bool = False,
|
||
) -> tuple[list[Path], list[Path]]:
|
||
_require_pillow()
|
||
if not input_dir.is_dir():
|
||
raise CombineError(f"Input directory does not exist: {input_dir}")
|
||
if tolerance is not None and (tolerance < 0 or tolerance > 32):
|
||
raise CombineError("Color tolerance must be between 0 and 32.")
|
||
|
||
materials = load_legend(legend_path, aliases)
|
||
groups, unmatched = discover_groups(
|
||
input_dir,
|
||
materials,
|
||
recursive=recursive,
|
||
excluded_paths=(id_map_path,),
|
||
)
|
||
if not groups:
|
||
names = ", ".join(material.filename_name for material in materials)
|
||
raise CombineError(
|
||
f"No supported textures in {input_dir} contained a legend material name "
|
||
f"({names}). Use --material-alias when Substance texture-set names differ."
|
||
)
|
||
|
||
outputs: list[Path] = []
|
||
if dry_run:
|
||
for group in groups:
|
||
outputs.append(output_dir / group.output_parent / group.output_name)
|
||
return outputs, unmatched
|
||
|
||
try:
|
||
with Image.open(id_map_path) as loaded_id_map:
|
||
id_map = loaded_id_map.convert("RGB")
|
||
id_map.load()
|
||
except OSError as exc:
|
||
raise CombineError(f"Could not load ID map {id_map_path}: {exc}") from exc
|
||
|
||
try:
|
||
masks = {}
|
||
for material in materials:
|
||
material_tolerance = tolerance if tolerance is not None else 0
|
||
mask = _mask_for_color(id_map, material.rgb, material_tolerance)
|
||
if tolerance is None and mask.getbbox() is None:
|
||
# Blender stores byte images after float-to-byte conversion.
|
||
# Values exactly between bytes can differ from ordinary Python
|
||
# rounding by one, so only widen a mask when exact pixels do
|
||
# not exist at all.
|
||
material_tolerance = 1
|
||
mask = _mask_for_color(id_map, material.rgb, material_tolerance)
|
||
if mask.getbbox() is not None:
|
||
print(
|
||
f"notice: matched {material.name!r} with automatic ±1 byte tolerance",
|
||
file=sys.stderr,
|
||
)
|
||
masks[material.name.casefold()] = mask
|
||
occupied = Image.new("1", id_map.size, 0)
|
||
for material in materials:
|
||
mask = masks[material.name.casefold()]
|
||
if ImageChops.multiply(occupied, mask).getbbox() is not None:
|
||
raise CombineError(
|
||
f"The mask for {material.name!r} overlaps another material mask. "
|
||
"Lower --tolerance (exact matching is --tolerance 0)."
|
||
)
|
||
occupied = ImageChops.lighter(occupied, mask)
|
||
empty_masks = [
|
||
material.name
|
||
for material in materials
|
||
if masks[material.name.casefold()].getbbox() is None
|
||
]
|
||
if empty_masks:
|
||
print(
|
||
"warning: no ID-map pixels matched: " + ", ".join(empty_masks),
|
||
file=sys.stderr,
|
||
)
|
||
|
||
output_paths = [
|
||
output_dir / group.output_parent / group.output_name
|
||
for group in groups
|
||
]
|
||
if len(set(output_paths)) != len(output_paths):
|
||
duplicates = sorted(
|
||
{path for path in output_paths if output_paths.count(path) > 1},
|
||
key=str,
|
||
)
|
||
raise CombineError(
|
||
"Multiple discovered groups resolve to the same output: "
|
||
+ ", ".join(str(path) for path in duplicates)
|
||
)
|
||
|
||
for group, output_path in zip(groups, output_paths, strict=True):
|
||
missing = combine_group(
|
||
group,
|
||
materials,
|
||
masks,
|
||
id_map.size,
|
||
output_path,
|
||
overwrite=overwrite,
|
||
)
|
||
if missing:
|
||
print(
|
||
f"warning: {group.description} has no texture for: {', '.join(missing)}",
|
||
file=sys.stderr,
|
||
)
|
||
outputs.append(output_path)
|
||
finally:
|
||
id_map.close()
|
||
|
||
return outputs, unmatched
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(
|
||
description=(
|
||
"Combine Substance Painter textures exported as separate material/texture "
|
||
"sets, using a Material ID bake and JSON legend as exact masks."
|
||
)
|
||
)
|
||
parser.add_argument("--id-map", required=True, type=Path, help="Material ID PNG")
|
||
parser.add_argument(
|
||
"--legend",
|
||
required=True,
|
||
type=Path,
|
||
help="JSON legend written beside the Material ID PNG",
|
||
)
|
||
parser.add_argument(
|
||
"--input-dir",
|
||
required=True,
|
||
type=Path,
|
||
help="Directory containing Substance Painter exports",
|
||
)
|
||
parser.add_argument(
|
||
"--output-dir",
|
||
required=True,
|
||
type=Path,
|
||
help="Directory for combined channel textures",
|
||
)
|
||
parser.add_argument(
|
||
"--material-alias",
|
||
action="append",
|
||
default=[],
|
||
metavar="LEGEND=FILENAME",
|
||
help="Map a Blender material name to a different Substance filename token",
|
||
)
|
||
parser.add_argument("--recursive", action="store_true", help="Search input subdirectories")
|
||
parser.add_argument(
|
||
"--tolerance",
|
||
type=int,
|
||
default=None,
|
||
help=(
|
||
"Fixed per-channel ID color tolerance in 8-bit values; by default the tool "
|
||
"tries exact matching and a ±1 fallback only when exact pixels are absent"
|
||
),
|
||
)
|
||
parser.add_argument("--overwrite", action="store_true", help="Replace existing outputs")
|
||
parser.add_argument(
|
||
"--dry-run",
|
||
action="store_true",
|
||
help="Print discovered groups and output paths without reading or writing images",
|
||
)
|
||
return parser
|
||
|
||
|
||
def main(argv: Sequence[str] | None = None) -> int:
|
||
parser = build_parser()
|
||
args = parser.parse_args(argv)
|
||
try:
|
||
aliases = _parse_aliases(args.material_alias)
|
||
outputs, unmatched = combine_all(
|
||
id_map_path=args.id_map.resolve(),
|
||
legend_path=args.legend.resolve(),
|
||
input_dir=args.input_dir.resolve(),
|
||
output_dir=args.output_dir.resolve(),
|
||
aliases=aliases,
|
||
recursive=args.recursive,
|
||
tolerance=args.tolerance,
|
||
overwrite=args.overwrite,
|
||
dry_run=args.dry_run,
|
||
)
|
||
except CombineError as exc:
|
||
parser.exit(2, f"error: {exc}\n")
|
||
|
||
action = "Would write" if args.dry_run else "Wrote"
|
||
for output in outputs:
|
||
print(f"{action}: {output}")
|
||
if unmatched:
|
||
print(f"Skipped {len(unmatched)} image(s) with no material-name match:")
|
||
for path in unmatched:
|
||
print(f" {path}")
|
||
print(f"{len(outputs)} texture channel(s) {'discovered' if args.dry_run else 'combined'}.")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|