gen-atlas.py

#!/usr/bin/env python3

import argparse
import functools
import itertools
import re
import rectpack
import shutil
import subprocess

from pathlib import (
    Path,
)

from tempfile import (
    TemporaryDirectory,
)

from contextlib import (
    ExitStack,
)

from concurrent.futures import (
    ProcessPoolExecutor,
    ThreadPoolExecutor,
    as_completed,
)

from PIL import (
    Image,
)

from taiseilib.common import (
    TaiseiError,
    run_main,
    update_text_file,
    wait_for_futures,
)

import taiseilib.atlas as atlas
import taiseilib.keyval as keyval


texture_formats = [s[1:] for s in atlas.SPRITE_SUFFIXES]


def ceildiv(x, y):
    return x // y + int(x % y != 0)


def align_size(size, alignment=4):
    return tuple(ceildiv(x, alignment) * alignment for x in size)


def write_sprite_def(dst, data):
    dst.parent.mkdir(exist_ok=True, parents=True)
    text = '# Autogenerated by the atlas packer, do not modify\n\n'
    text += keyval.dump(data)
    update_text_file(dst, text)


def write_texture_def(dst, texture, texture_fmt, global_overrides=None, local_overrides=None, alphamap_tex_fmt=None):
    dst.parent.mkdir(exist_ok=True, parents=True)

    data = {
        'source': f'res/gfx/{texture}.{texture_fmt}',
    }

    if alphamap_tex_fmt is not None:
        data['alphamap'] = f'res/gfx/{texture}.alphamap.{alphamap_tex_fmt}'

    if global_overrides is not None:
        data.update(global_overrides)

    if local_overrides is not None:
        data.update(local_overrides)

    text = '# Autogenerated by the atlas packer, do not modify\n\n'
    text += keyval.dump(data)
    update_text_file(dst, text)


def get_bin_occupied_bounds(bin):
    xmin, xmax, ymin, ymax = 0xffffffffffffffff, 0, 0xffffffffffffffff, 0

    for rect in bin:
        if rect.bottom < ymin:
            ymin = rect.bottom

        if rect.top > ymax:
            ymax = rect.top

        if rect.left < xmin:
            xmin = rect.left

        if rect.right > xmax:
            xmax = rect.right

    w = xmax - xmin
    h = ymax - ymin
    w, h = align_size((w, h))

    return rectpack.geometry.Rectangle(xmin, ymin, w, h)


@functools.total_ordering
class PackResult:
    def __init__(self, packer, rects):
        self.bins = list(packer)
        self.num_images_packed = sum(map(len, self.bins))

        self.success = self.num_images_packed == len(rects)
        assert(self.num_images_packed <= len(rects))

        self.total_area = self._calculate_total_area()

    @property
    def num_bins(self):
        return len(self.bins)

    def __repr__(self):
        return f'PackResult(num_bins={self.num_bins}, num_images_packed={self.num_images_packed}, total_area={self.total_area}, success={self.success})'

    def _calculate_total_area(self):
        return sum(get_bin_occupied_bounds(bin).area() for bin in self.bins)

    def __eq__(self, other):
        return (
            self.num_bins == other.num_bins and
            self.num_images_packed == other.num_images_packed and
            self.total_area == other.total_area
        )

    # less == better-than
    def __lt__(self, other):
        if self.num_images_packed > other.num_images_packed:
            return True

        if self.num_images_packed < other.num_images_packed:
            return False

        if self.num_bins < other.num_bins:
            return True

        if self.num_bins > other.num_bins:
            return False

        if self.total_area < other.total_area:
            return True

        return False


def pack_rects(rects, bin_size, packer_factory, single_bin):
    bin_size = list(bin_size)

    if single_bin:
        while True:
            packer = packer_factory()
            packer.add_bin(*bin_size)

            for rect in rects:
                packer.add_rect(*rect)

            packer.pack()

            if sum(len(bin) for bin in packer) == len(rects):
                break

            if bin_size[1] < bin_size[0]:
                bin_size[1] *= 2
            else:
                bin_size[0] *= 2
    else:
        packer = packer_factory()

        for rect in rects:
            packer.add_rect(*rect)
            packer.add_bin(*bin_size)

        packer.pack()

    return PackResult(packer, rects)


def bruteforcer_pack(params):
    import rectpack
    global subprocess_args
    rects = subprocess_args['rects']
    bin_size = subprocess_args['bin_size']
    single_bin = subprocess_args['single_bin']

    algo_class, (sort_func, sort_name) = params
    algo_class = getattr(rectpack, algo_class)
    sort_func = getattr(rectpack, sort_func)

    def packer_factory(algo_class=algo_class, sort_func=sort_func):
        return rectpack.newPacker(
            rotation=False,   # No rotation support in Taisei yet
            pack_algo=algo_class,
            sort_algo=sort_func,
            bin_algo=rectpack.PackingBin.BFF,
        )

    result = pack_rects(
        rects=rects,
        packer_factory=packer_factory,
        bin_size=bin_size,
        single_bin=single_bin
    )

    return result


def bruteforcer_init(rects, bin_size, single_bin):
    global subprocess_args
    subprocess_args = {
        'rects': rects,
        'bin_size': bin_size,
        'single_bin': single_bin,
    }


def pack_rects_brute_force(rects, bin_size, single_bin):
    algos_guillotine = [
        'GuillotineBafLas',
        'GuillotineBafLlas',
        'GuillotineBafMaxas',
        'GuillotineBafMinas',
        'GuillotineBafSas',
        'GuillotineBafSlas',
        'GuillotineBlsfLas',
        'GuillotineBlsfLlas',
        'GuillotineBlsfMaxas',
        'GuillotineBlsfMinas',
        'GuillotineBlsfSas',
        'GuillotineBlsfSlas',
        'GuillotineBssfLas',
        'GuillotineBssfLlas',
        'GuillotineBssfMaxas',
        'GuillotineBssfMinas',
        'GuillotineBssfSas',
        'GuillotineBssfSlas',
    ]

    algos_maxrects = [
        'MaxRectsBaf',
        'MaxRectsBl',
        'MaxRectsBlsf',
        'MaxRectsBssf',
    ]

    algos_skyline = [
        'SkylineBl',
        'SkylineBlWm',
        'SkylineMwf',
        'SkylineMwfWm',
        'SkylineMwfl',
        'SkylineMwflWm',
    ]

    algos = []
    # MaxRects is the slowest algorithm, but the results blow the other two out of the water (most of the time).
    algos += algos_guillotine
    algos += algos_maxrects
    algos += algos_skyline

    sorts = [
        ('SORT_AREA',  'area'),
        ('SORT_DIFF',  'difference'),
        ('SORT_LSIDE', 'longest-side'),
        ('SORT_NONE',  'no'),
        ('SORT_PERI',  'perimeter'),
        ('SORT_RATIO', 'ratio'),
        ('SORT_SSIDE', 'shortest-side'),
    ]

    best = None
    best_algos = None

    def verbose(*msg):
        print('[bruteforce]', *msg)

    with ProcessPoolExecutor(initializer=bruteforcer_init, initargs=(rects, bin_size, single_bin)) as ex:
        variants = tuple(itertools.product(algos, sorts))
        futures = {}

        for variant in variants:
            f = ex.submit(bruteforcer_pack, variant)
            futures[f] = variant

        for future in as_completed(futures):
            algo_class, (sort_func, sort_name) = futures[future]
            result = future.result()

            if best is None or result < best:
                verbose(' * {:>20} with {:20} | {!s}'.format(algo_class, sort_name + ' sort', result))
                best = result
                best_algos = (algo_class, sort_name)
            else:
                verbose('   {:>20} with {:20} | {!s}'.format(algo_class, sort_name + ' sort', result))

    verbose('*' * 64)
    verbose(f'WINNER: {best_algos[0]} with {best_algos[1]} sort')
    verbose(f'\tBest result: {best}')
    verbose('*' * 64)
    return best


def sprite_sources_iter(srcroot):
    def imgfilter(path):
        return (
            path.is_file() and path.suffix.lower() in atlas.SPRITE_SUFFIXES and
            path.with_suffix('').suffix != '.alphamap'
        )

    for path in sorted(filter(imgfilter, srcroot.glob('**/*.*'))):
        yield path


def gen_atlas(config_dir, src, dst, binsize, atlasname, tex_format=texture_formats[0], border=1, force_single=False, crop=True, leanify=True):
    def get_sprite_border(sprite, default_border=border):
        return max(default_border, int(sprite.config.get('border', default_border)))

    texture_local_overrides = keyval.parse(src / 'atlas.tex', missing_ok=True)
    texture_global_overrides = keyval.parse(config_dir / 'atlas.tex', missing_ok=True)

    total_images = 0
    packed_images = 0
    rects = []

    sprites = {}

    for path in sprite_sources_iter(src):
        sprite_name = path.relative_to(src).with_suffix('').as_posix()
        sprite_config_path = config_dir / atlas.get_sprite_config_file_name(sprite_name)
        sprite = atlas.Sprite.load(sprite_name, path, sprite_config_path)
        sprites[sprite.name] = sprite
        border = get_sprite_border(sprite)
        rects.append((
            sprite.image.size[0] + border*2,
            sprite.image.size[1] + border*2,
            sprite.name,
        ))

    pack_result = pack_rects_brute_force(rects=rects, bin_size=binsize, single_bin=force_single)

    if not pack_result.success:
        missing = len(rects) - pack_result.num_images_packed
        raise TaiseiError(
            f'{missing} sprite{"s were" if missing > 1 else " was"} not packed (bin size is too small?)'
        )

    futures = []

    with ExitStack() as stack:
        # Do everything in a temporary directory first
        temp_dst = Path(stack.enter_context(TemporaryDirectory(prefix=f'taisei-atlas-{atlasname}')))

        # Run multiple leanify processes in parallel, in case we end up with multiple pages
        # Yeah I'm too lazy to use Popen properly
        executor = stack.enter_context(ThreadPoolExecutor())

        for i, bin in enumerate(pack_result.bins):
            textureid = f'atlas_{atlasname}_{i}'
            # NOTE: we always save PNG first and convert with an external tool later if needed.
            dstfile = temp_dst / f'{textureid}.png'
            dstfile_alphamap = temp_dst / f'{textureid}.alphamap.png'
            dstfile_meta = temp_dst / f'{textureid}.tex'

            if crop:
                occupied_bounds = get_bin_occupied_bounds(bin)
                actual_size = (occupied_bounds.width, occupied_bounds.height)
                bin_offset = (-occupied_bounds.x, -occupied_bounds.y)
            else:
                actual_size = (bin.width, bin.height)
                bin_offset = (0, 0)

            bin_image = Image.new('RGBA', actual_size, (0, 0, 0, 0))
            bin_alphamap_image = None

            for rect in bin:
                rect.move(rect.x + bin_offset[0], rect.y + bin_offset[1])
                sprite = sprites[rect.rid]
                border = get_sprite_border(sprite)
                paste_coord = (rect.x + border, rect.y + border)
                bin_image.paste(sprite.image, paste_coord)

                if sprite.alphamap_image is not None:
                    if bin_alphamap_image is None:
                        bin_alphamap_image = Image.new('L', actual_size, 255)
                    bin_alphamap_image.paste(sprite.alphamap_image, paste_coord)

                sprite.texture_id = textureid
                sprite.texture_region = atlas.Geometry(
                    rect.width - border * 2,
                    rect.height - border * 2,
                    rect.x + border,
                    rect.y + border,
                ) / actual_size

                write_sprite_def(temp_dst / f'{sprite.name}.spr', sprite.dump_spritedef_dict())
                sprite.unload()

            write_texture_def(
                dstfile_meta,
                textureid,
                tex_format,
                texture_global_overrides,
                texture_local_overrides,
                alphamap_tex_fmt=None if bin_alphamap_image is None else 'png'
            )

            def process(
                bin_image=bin_image,
                dstfile=dstfile,
                tex_format=tex_format
            ):
                bin_image.save(dstfile)
                bin_image.close()

                oldfmt = dstfile.suffix[1:].lower()

                if oldfmt != tex_format:
                    new_dstfile = dstfile.with_suffix(f'.{tex_format}')

                    if tex_format == 'webp':
                        subprocess.check_call([
                            'cwebp',
                            '-progress',
                            '-preset', 'icon',
                            '-z', '9',
                            '-lossless',
                            '-q', '100',
                            '-m', '6',
                            str(dstfile),
                            '-o', str(new_dstfile),
                        ])
                    else:
                        raise TaiseiError(f'Unhandled conversion {oldfmt} -> {tex_format}')

                    dstfile.unlink()
                    dstfile = new_dstfile

                if leanify:
                    subprocess.check_call(['leanify', '-v', str(dstfile)])

            futures.append(executor.submit(process))

            def process_alphamap(
                bin_alphamap_image=bin_alphamap_image,
                dstfile_alphamap=dstfile_alphamap,
            ):
                bin_alphamap_image.save(dstfile_alphamap)
                bin_alphamap_image.close()

                if leanify:
                    subprocess.check_call(['leanify', '-v', str(dstfile_alphamap)])

            if bin_alphamap_image is not None:
                futures.append(executor.submit(process_alphamap))

        # Wait for subprocesses to complete.
        wait_for_futures(futures)
        executor.shutdown(wait=True)

        # Only now, if everything is ok so far, copy everything to the destination, possibly overwriting previous results
        pattern = re.compile(rf'^atlas_{re.escape(atlasname)}_\d+(?:\.alphamap)?.({"|".join(texture_formats + ["tex"])})$')
        for path in dst.iterdir():
            if pattern.match(path.name):
                path.unlink()

        targets = list(temp_dst.glob('**/*'))

        for dir in (p.relative_to(temp_dst) for p in targets if p.is_dir()):
            (dst / dir).mkdir(parents=True, exist_ok=True)

        for file in (p.relative_to(temp_dst) for p in targets if not p.is_dir()):
            shutil.copyfile(str(temp_dst / file), str(dst / file))


def upgrade_configs(overrides_dir, src_dir, config_dir):
    for path in sprite_sources_iter(src_dir):
        sprite_name = path.relative_to(src_dir).with_suffix('').as_posix()
        config_fname = atlas.get_sprite_config_file_name(sprite_name)
        sprite_config_path = config_dir / config_fname
        sprite_override_path = overrides_dir / config_fname

        if sprite_config_path.is_file():
            continue

        sprite = atlas.Sprite.load(sprite_name, path, sprite_override_path, force_autotrim=False)
        sprite.upgrade_config()

        if sprite.config:
            sprite_config_path.parent.mkdir(parents=True, exist_ok=True)
            update_text_file(sprite_config_path, keyval.dump(sprite.config))
            print('{}:  {!r}'.format(sprite.name, sprite.config))

        sprite.unload()


def main(args):
    parser = argparse.ArgumentParser(description='Generate texture atlases and sprite definitions', prog=args[0])

    parser.add_argument('config_dir',
        help='Directory containing configuration files for individual sprites',
        type=Path,
    )

    parser.add_argument('source_dir',
        help='Directory containing input textures (searched recursively)',
        type=Path,
    )

    parser.add_argument('dest_dir',
        help='Directory to dump the results into',
        type=Path,
    )

    parser.add_argument('--width', '-W',
        help='Base width of a single atlas bin (default: 2048)',
        default=2048,
        type=int
    )

    parser.add_argument('--height', '-H',
        help='Base height of a single atlas bin (default: 2048)',
        default=2048,
        type=int
    )

    parser.add_argument('--name', '-n',
        help='Unique identifier for this atlas (used to form the texture name), default is inferred from directory name',
        default=None,
        type=str,
    )

    parser.add_argument('--border', '-b',
        help='Add a protective border WIDTH pixels wide around each sprite (default: 1)',
        metavar='WIDTH',
        dest='border',
        type=int,
        default=1
    )

    parser.add_argument('--crop', '-c',
        help='Remove unused space from bins (default)',
        dest='crop',
        action='store_true',
        default=True,
    )

    parser.add_argument('--no-crop', '-C',
        help='Do not remove unused space from atlases',
        dest='crop',
        action='store_false',
        default=True,
    )

    parser.add_argument('--single', '-s',
        help='Package everything into a single texture, possibly extending its size; this can be slow (default)',
        dest='single',
        action='store_true',
        default=True
    )

    parser.add_argument('--multiple', '-m',
        help='Split the atlas across multiple textures if the sprites won\'t fit otherwise',
        dest='single',
        action='store_false',
        default=True
    )

    parser.add_argument('--leanify', '-l',
        help='Leanify atlases to save space; very slow (default)',
        dest='leanify',
        action='store_true',
        default=True
    )

    parser.add_argument('--no-leanify', '-L',
        help='Do not leanify atlases',
        dest='leanify',
        action='store_false',
        default=True
    )

    parser.add_argument('--format', '-f',
        help=f'Format of the atlas textures (default: {texture_formats[0]})',
        choices=texture_formats,
        default=texture_formats[0],
    )

    parser.add_argument('--upgrade-configs',
        help='Generate sprite configs from old-style "overrides". In this mode, config_dir specifies the old "overrides" root, and new configs will be dumped into dest_dir. Existing configs in dest_dir will not be replaced',
        action='store_true',
        default=False,
    )

    args = parser.parse_args()

    if args.name is None:
        args.name = args.source_dir.name

    args.config_dir = args.config_dir.resolve()
    args.source_dir = args.source_dir.resolve()
    args.dest_dir = args.dest_dir.resolve()

    if args.upgrade_configs:
        upgrade_configs(
            args.config_dir,
            args.source_dir,
            args.dest_dir,
        )
    else:
        gen_atlas(
            args.config_dir,
            args.source_dir,
            args.dest_dir,
            align_size((args.width, args.height)),
            tex_format=args.format,
            atlasname=args.name,
            border=args.border,
            force_single=args.single,
            crop=args.crop,
            leanify=args.leanify,
        )


if __name__ == '__main__':
    run_main(main)