Spaces:
Running
Running
| 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) | |