packer.py

from PIL import Image
import os
import glob
import ntpath
import sys
import numpy as np

# Bake sbsar into better channels
# Currently using this for UE4 terrain packing

SBSRENDER = '"C:\Program Files\Allegorithmic\Substance Automation Toolkit\sbsrender.exe"'
out_folder = "temp_channel_packer_output"

# true image size = 2**image_size
# for example image size = 11 => 2**11 => dimensions = (2048, 2048)
image_size = 10

try:
    print("Creating folder:", out_folder)
    os.mkdir(out_folder)
except FileExistsError:
    print("Folder already exists:", out_folder)

print("Number of arguments:", len(sys.argv), "arguments.")
arglist = sys.argv[1:]
print("Argument List:", str(arglist))

if len(arglist) == 0:
    print("No arguments given, exit.")
    exit(1)

# Can batch process multiple files
for arg in arglist:
    SBSAR_NAME = ntpath.basename(arg)
    print("Name:", SBSAR_NAME)

    name_base = SBSAR_NAME.split(".")[0]
    print("Base:", name_base)

    command = f"""
    {SBSRENDER} render
        --inputs {SBSAR_NAME}
        --input-graph-output basecolor
        --input-graph-output normal
        --input-graph-output roughness
        --input-graph-output height
        --input-graph-output ambientocclusion
        --set-value $outputsize@{image_size},{image_size}
        --output-name {{inputName}}_{{outputNodeName}}
        --output-path {out_folder}
    """
    command = " ".join(command.split())
    print("Working directory:", os.getcwd())
    print("Running:", command)
    os.system(command)

    images = {}

    for infile in glob.glob(f"{out_folder}/*.png"):
        file, ext = os.path.splitext(infile)
        if name_base in file and "sbsar" not in file and ext == ".png":
            print("Opening:", file)
            channel = file.split("_")[-1]
            images[channel] = Image.open(infile)

    print("Channels:", images.keys())

    # Color(RGB), Roughness(A)
    # Normal(RGB), Height(A)

    R, G, B = images["basecolor"].split()[:3]
    Rg = images["roughness"].split()[-1]
    img0 = Image.merge("RGBA", [R, G, B, Rg])
    print("Image_color data:", img0.format, img0.mode)
    img0.save(out_folder + "\\" + SBSAR_NAME + "_color.png", "PNG")

    # if "ambientocclusion" in images:
    #     Ao = images["ambientocclusion"].split()[-1].convert("L")
    # else:
    #     Ao = Image.new("L", R.size, color=255)
    R, G, B = images["normal"].split()[:3]

    # PIL doesn't understand howto properly convert from 16-bit grayscale to 8-bit
    # Se let's do it manually then...
    Hg_n = np.asarray(images["height"])
    print("Height min/max:", np.min(Hg_n), np.max(Hg_n))
    Hg = Image.fromarray(Hg_n // 256).convert("L")
    # Hg.save(out_folder + "\\" + "hg_test.png", "PNG")

    img1 = Image.merge("RGBA", [R, G, B, Hg])
    print("Image_normal data:", img1.format, img1.mode)
    img1.save(out_folder + "\\" + SBSAR_NAME + "_normal.png", "PNG")