ldivider/ld_utils.py

import random
import string

import numpy as np

import pytoshop
from pytoshop import layers


def randomname(n):
    randlst = [random.choice(string.ascii_letters + string.digits) for i in range(n)]
    return "".join(randlst)


def add_psd(psd: pytoshop.core.PsdFile, img, name, mode):
    layer_1 = layers.ChannelImageData(image=img[:, :, 3], compression=1)
    layer0 = layers.ChannelImageData(image=img[:, :, 0], compression=1)
    layer1 = layers.ChannelImageData(image=img[:, :, 1], compression=1)
    layer2 = layers.ChannelImageData(image=img[:, :, 2], compression=1)

    new_layer = layers.LayerRecord(
        channels={-1: layer_1, 0: layer0, 1: layer1, 2: layer2},
        top=0,
        bottom=img.shape[0],
        left=0,
        right=img.shape[1],
        blend_mode=mode,
        name=name,
        opacity=255,
    )
    psd.layer_and_mask_info.layer_info.layer_records.append(new_layer)
    return psd


def save_psd(input_image: np.ndarray, layers, names, modes, output_dir, layer_mode):
    psd = pytoshop.core.PsdFile(
        num_channels=3, height=input_image.shape[0], width=input_image.shape[1]
    )
    if layer_mode == "normal":
        for idx, output in enumerate(layers[0]):
            psd = add_psd(psd, layers[0][idx], names[0] + str(idx), modes[0])
            psd = add_psd(psd, layers[1][idx], names[1] + str(idx), modes[1])
            psd = add_psd(psd, layers[2][idx], names[2] + str(idx), modes[2])
    else:
        for idx, output in enumerate(layers[0]):
            psd = add_psd(psd, layers[0][idx], names[0] + str(idx), modes[0])
            psd = add_psd(psd, layers[1][idx], names[1] + str(idx), modes[1])
            psd = add_psd(psd, layers[2][idx], names[2] + str(idx), modes[2])
            psd = add_psd(psd, layers[3][idx], names[3] + str(idx), modes[3])
            psd = add_psd(psd, layers[4][idx], names[4] + str(idx), modes[4])

    name = randomname(10)

    with open(f"{output_dir}/output_{name}.psd", "wb") as fd2:
        psd.write(fd2)

    return f"{output_dir}/output_{name}.psd"