Spaces:
Running
Running
File size: 3,304 Bytes
13aa528 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import argparse
import logging
import multiprocessing
import os
from typing import Iterable
from psd_tools import PSDImage
from tqdm import tqdm
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
def parse_args():
parser = argparse.ArgumentParser(description="Convert PSD files to PNG.")
parser.add_argument(
"-d",
"--directory",
type=str,
default="./",
help="Directory to search for PSD files.",
)
parser.add_argument("-o", "--output", type=str, default="./", help="Directory to save PNG files.")
parser.add_argument(
"--visible_layers",
default=[],
nargs="+",
type=str,
help="List of layer names to make visible.",
)
parser.add_argument(
"--invisible_layers",
default=[],
nargs="+",
type=str,
help="List of layer names to make invisible.",
)
parser.add_argument("--num_processes", "-n", default=None, type=int, help=" Number of processes to use.")
return parser.parse_args()
def find_psd_files(directory):
psd_files = []
for root, dirs, files in os.walk(directory):
for file in files:
if file.endswith(".psd"):
psd_files.append(os.path.join(root, file))
return psd_files
def set_layer_visibility(layer, visible_layers, invisible_layers):
if layer.name in visible_layers:
layer.visible = True
if layer.name in invisible_layers:
layer.visible = False
if isinstance(layer, Iterable):
for child in layer:
set_layer_visibility(child, visible_layers, invisible_layers)
def process_psd_file(task):
"""
Worker function that processes a single PSD file.
Opens the PSD, sets layer visibility, composites the image and saves it as PNG.
"""
psd_file, output, visible_layers, invisible_layers, force = task
try:
psd = PSDImage.open(psd_file)
if force:
for layer in psd:
set_layer_visibility(layer, visible_layers, invisible_layers)
image = psd.composite(force=force)
fname = os.path.basename(psd_file).replace(".psd", ".png")
output_file = os.path.join(output, fname)
image.save(output_file)
except Exception as e:
logger.error("Error processing file %s: %s", psd_file, e)
def main(args):
# Create output directory if it doesn't exist
if not os.path.exists(args.output):
os.makedirs(args.output)
psd_files = find_psd_files(args.directory)
# force=True when any layer visibility is provided
force = True if len(args.visible_layers) + len(args.invisible_layers) else False
tasks = [(psd_file, args.output, args.visible_layers, args.invisible_layers, force) for psd_file in psd_files]
num_processes = args.num_processes if args.num_processes else multiprocessing.cpu_count() // 2
# Use multiprocessing to process PSD files in parallel
with multiprocessing.Pool(processes=num_processes) as pool:
list(
tqdm(
pool.imap_unordered(process_psd_file, tasks),
total=len(tasks),
desc="Convert PSD to PNG files",
)
)
if __name__ == "__main__":
args = parse_args()
main(args)
|