taisei/scripts/gen-atlas.py
2020-11-08 10:04:40 +02:00

696 lines
20 KiB
Python
Executable file

#!/usr/bin/env python3
import rectpack
import argparse
import shutil
import subprocess
import re
import functools
import itertools
from pathlib import (
Path,
)
from tempfile import (
TemporaryDirectory,
)
from contextlib import (
ExitStack,
suppress,
)
from concurrent.futures import (
ThreadPoolExecutor,
ProcessPoolExecutor,
)
from PIL import (
Image,
)
from taiseilib.common import (
run_main,
update_text_file,
TaiseiError,
wait_for_futures,
)
#
# To the reasonable person wondering why this god-awful script invokes
# an ImageMagick subprocess even though I've just imported PIL:
#
# It used to just use PIL to handle image composition. That is, until
# I discovered the fact that PIL/Pillow does not support multichannel
# 16-bits per pixel images.
# https://github.com/python-pillow/Pillow/issues/2107#issuecomment-246007526
#
# Now PIL is only used to extract the image size, though I guess I might
# as well just parse the output of `identify` at this point.
#
# Fuck software tbqh.
#
texture_formats = ['png', 'webp']
re_comment = re.compile(r'#.*')
re_keyval = re.compile(r'([a-z0-9_-]+)\s*=\s*(.+)', re.I)
class ConfigSyntaxError(Exception):
pass
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, texture, region, texture_dimensions, overrides=None):
dst.parent.mkdir(exist_ok=True, parents=True)
text = (
'# Autogenerated by the atlas packer, do not modify\n\n'
f'texture = {texture}\n'
'region_x = {region_x:.17f}\n'
'region_y = {region_y:.17f}\n'
'region_w = {region_w:.17f}\n'
'region_h = {region_h:.17f}\n'
).format(
region_x=region[0] / texture_dimensions[0],
region_y=region[1] / texture_dimensions[1],
region_w=region[2] / texture_dimensions[0],
region_h=region[3] / texture_dimensions[1],
)
if overrides is not None:
text += f'\n# -- Pasted from the override file --\n\n{overrides.strip()}\n'
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)
text = (
'# Autogenerated by the atlas packer, do not modify\n\n'
f'source = res/gfx/{texture}.{texture_fmt}\n'
)
if alphamap_tex_fmt is not None:
text += f'alphamap = res/gfx/{texture}.alphamap.{alphamap_tex_fmt}\n'
if global_overrides is not None:
text += f'\n# -- Pasted from the global override file --\n\n{global_overrides.strip()}\n'
if local_overrides is not None:
text += f'\n# -- Pasted from the local override file --\n\n{local_overrides.strip()}\n'
update_text_file(dst, text)
def write_override_template(dst, size):
dst.parent.mkdir(exist_ok=True, parents=True)
text = (
'# This file was generated automatically, because this sprite doesn\'t have a custom override file.\n'
'# To override this sprite\'s parameters, edit this file, remove the .renameme suffix and this comment, then commit it to the repository.\n\n'
'# Modify these to change the virtual size of the sprite\n'
f'w = {size[0]}\n'
f'h = {size[1]}\n'
)
update_text_file(dst.with_suffix(dst.suffix + '.renameme'), text)
def parse_sprite_conf(path):
conf = {}
with suppress(FileNotFoundError):
with open(path, 'r') as f:
for line in f.readlines():
line = re_comment.sub('', line)
line = line.strip()
if not line:
continue
try:
key, val = re_keyval.findall(line)[0]
except IndexError:
raise ConfigSyntaxError(line)
conf[key] = val
return conf
def get_override_file_name(basename):
if re.match(r'.*\.frame\d{4}$', basename):
basename, _ = basename.rsplit('.', 1)
basename += '.framegroup'
return f'{basename}.spr'
def find_alphamap(basepath):
for fmt in texture_formats:
mpath = basepath.with_suffix(f'.alphamap.{fmt}')
if mpath.is_file():
return mpath
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))
results = zip(variants, ex.map(bruteforcer_pack, variants, chunksize=4))
print('results', results)
for (algo_class, (sort_func, sort_name)), result in results:
verbose(f'Trying {algo_class} with {sort_name} sort...')
if best is None or result < best:
verbose(f'\tResult: {result} (best yet)')
best = result
best_algos = (algo_class, sort_name)
else:
verbose(f'\tResult: {result}')
verbose('*' * 64)
verbose(f'WINNER: {best_algos[0]} with {best_algos[1]} sort')
verbose(f'\tBest result: {best}')
verbose('*' * 64)
return best
def gen_atlas(overrides, src, dst, binsize, atlasname, tex_format=texture_formats[0], border=1, force_single=False, crop=True, leanify=True):
overrides = Path(overrides).resolve()
src = Path(src).resolve()
dst = Path(dst).resolve()
sprite_configs = {}
def get_border(sprite, default_border=border):
return max(default_border, int(sprite_configs[sprite].get('border', default_border)))
try:
texture_local_overrides = (src / 'atlas.tex').read_text()
except FileNotFoundError:
texture_local_overrides = None
try:
texture_global_overrides = (overrides / 'atlas.tex').read_text()
except FileNotFoundError:
texture_global_overrides = None
total_images = 0
packed_images = 0
rects = []
def imgfilter(path):
return (
path.is_file() and path.suffix[1:].lower() in texture_formats and
path.with_suffix('').suffix != '.alphamap'
)
for path in sorted(filter(imgfilter, src.glob('**/*.*'))):
img = Image.open(path)
sprite_name = path.relative_to(src).with_suffix('').as_posix()
sprite_config_path = overrides / (get_override_file_name(sprite_name) + '.conf')
sprite_configs[sprite_name] = parse_sprite_conf(sprite_config_path)
border = get_border(sprite_name)
rects.append((img.size[0]+border*2, img.size[1]+border*2, (path, sprite_name)))
img.close()
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}'
# dstfile = temp_dst / f'{textureid}.{tex_format}'
# 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'
print(dstfile)
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)
base_composite_cmd = [
'convert',
'-verbose',
'-size', f'{actual_size[0]}x{actual_size[1]}',
]
composite_cmd = base_composite_cmd.copy() + ['xc:none']
alphamap_composite_cmd = None
for rect in bin:
rect.move(rect.x + bin_offset[0], rect.y + bin_offset[1])
img_path, name = rect.rid
border = get_border(name)
alphamap_path = find_alphamap(img_path)
composite_cmd += [
str(img_path), '-geometry', '{:+}{:+}'.format(rect.x + border, rect.y + border), '-composite'
]
if alphamap_path:
if alphamap_composite_cmd is None:
alphamap_composite_cmd = base_composite_cmd.copy() + [
'xc:white',
'-colorspace', 'Gray',
]
alphamap_composite_cmd += [
str(alphamap_path),
'-geometry', '{:+}{:+}'.format(rect.x + border, rect.y + border),
'-channel', 'R',
'-separate',
'-composite',
]
override_path = overrides / get_override_file_name(name)
if override_path.exists():
override_contents = override_path.read_text()
else:
override_contents = None
write_override_template(override_path, (rect.width - border * 2, rect.height - border * 2))
write_sprite_def(
temp_dst / f'{name}.spr',
textureid,
(rect.x + border, rect.y + border, rect.width - border * 2, rect.height - border * 2),
actual_size,
overrides=override_contents
)
composite_cmd += [str(dstfile)]
write_texture_def(
dstfile_meta,
textureid,
tex_format,
texture_global_overrides,
texture_local_overrides,
alphamap_tex_fmt=None if alphamap_composite_cmd is None else 'png'
)
def process(composite_cmd, dstfile, tex_format=tex_format):
subprocess.check_call(composite_cmd)
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(lambda: process(composite_cmd, dstfile)))
if alphamap_composite_cmd is not None:
alphamap_composite_cmd += [str(dstfile_alphamap)]
futures.append(executor.submit(lambda: process(alphamap_composite_cmd, dstfile_alphamap, tex_format='png')))
# 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 main(args):
parser = argparse.ArgumentParser(description='Generate texture atlases and sprite definitions', prog=args[0])
parser.add_argument('overrides_dir',
help='Directory containing per-sprite override files; templates are created automatically for missing files',
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],
)
args = parser.parse_args()
if args.name is None:
args.name = args.source_dir.name
gen_atlas(
args.overrides_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)