Image resizer for Android based on PIL

buckethit.py

#!/usr/bin/python3
# @author cab404
# Thing that resizes images for android

import os
import argparse
import random
from multiprocessing import Pool
from PIL import Image

random.seed()

args = argparse.ArgumentParser(description="Authomatic resizer for android resource images")
args.add_argument("source", type=str, help="Folder with source images", default="./")
args.add_argument("target", type=str, help="Target folder for image buckets", default="./res")
args.add_argument("--dp", "-d", type=int, help="Default size for images without dp in their name", default=48)
args.add_argument("--verbose", "-v", help="Prints a little more info", action='count')
args = args.parse_args()
S = "/"

replies = [
    "is a marshmallow", "is an alien spacecraft", "is a cow", "is a beautiful pony", "is an owl in disguise", "is an enemy spy",
    "[REDACTED]", "is not a folder", "is in another castle", "has ascended beyond super folder 2", "is just messing with ya"
]

if not os.path.isdir(args.source):
    print("error: specified source folder does not exist, or it actially {0}.".format(random.choice(replies)))
    exit(1)

if not os.path.isdir(args.target):
    print("error: specified target folder does not exist, or it actially {0}.".format(random.choice(replies)))
    exit(1)

buckets = {
    "mdpi": 1,
    "hdpi": 1.5,
    "xhdpi": 2,
    "xxhdpi": 3,
    "xxxhdpi": 4,
}


logformat = "{0:<40}  {1:>10}   {2:<40} {3:60}"

class ImgData:

    def __init__(self, filename, dpsize, extension, name) -> None:
        super().__init__()
        self.extension = extension
        self.name = name
        self.dpsize = dpsize
        self.filename = filename

    def __str__(self) -> str:
        return self.name

    def get_size(self, w, h, bucket):
        side = int(self.dpsize * buckets[bucket])
        return side, int(side * (h / w))

    def get_target_path(self, bucket: [str, None]):
        bucket_path = args.target + S + "drawable"
        if bucket:
            bucket_path += "-" + bucket

        if not os.path.isdir(bucket_path):
            os.mkdir(bucket_path)
        return bucket_path + S + self.name

    def log_resize(self, bucket, src_size, dst_size):
        warn = []
        if max(src_size) < max(dst_size):
            warn.append("upscaled")
        if src_size[0] / src_size[1] != dst_size[0] / dst_size[1]:
            warn.append("aspect ratio lost")
        warn = ", ".join(warn)

        print(
            logformat.format(
                "{name}".format(
                    name=self.name,
                ),
                "{src_size[0]}x{src_size[1]}".format(
                    src_size=src_size,
                ),
                "{dp}dp {bucket} ({dst_size[0]}x{dst_size[1]})".format(
                    bucket=bucket, dst_size=dst_size, dp=self.dpsize
                ),
                "{warn}".format(
                    warn=warn,
                )
            )
        )


def module_nine_patch_resize(data: ImgData):
    for bucket in buckets:
        target_file = data.get_target_path(bucket)
        bmp = Image.open(data.filename)
        if not bmp:
            raise FileNotFoundError("this bitmap is bad .-.")
        assert isinstance(bmp, Image.Image)

        w = bmp.width
        h = bmp.height
        size = data.get_size(w - 2, h - 2, bucket)

        # width and height of center part
        ow = size[0]
        oh = size[1]

        data.log_resize(bucket, bmp.size, size)

        bmp = bmp.convert("RGBA")

        def extract(source, dest, resample=Image.NONE):
            return bmp.crop(source).resize(dest, resample=resample)

        canvas = Image.new("RGBA", (ow + 2, oh + 2))

        #left
        canvas.paste(extract((0, 1, 1, h), (1, oh)), (0, 1))
        #bottom
        canvas.paste(extract((1, 0, w, 1), (ow, 1)), (1, 0))
        #right
        canvas.paste(extract((w - 1, 1, w, h), (1, oh)), (ow - 1, 1))
        #top
        canvas.paste(extract((1, h - 1, w, h), (ow, 1)), (1, oh - 1))

        #core
        canvas.paste(extract((1, 1, w - 1, h - 1), (ow, oh), resample=Image.LANCZOS), (1, 1))

        canvas.save(target_file)


def module_img_resize(data: ImgData):
    for bucket in buckets:
        target_file = data.get_target_path(bucket)
        bmp = Image.open(data.filename)
        if not bmp:
            raise FileNotFoundError("this bitmap is bad .-.")
        assert isinstance(bmp, Image.Image)
        size = data.get_size(bmp.width, bmp.height, bucket)

        data.log_resize(bucket, bmp.size, size)

        bmp.load()
        bmp.resize(size, resample=Image.LANCZOS).save(target_file)


def module_copy(data: ImgData):
    target_file = data.get_target_path(None)
    print("copying", data, "to", target_file)
    drawable_dir = os.path.dirname(target_file)

    if not os.path.isdir(drawable_dir):
        os.mkdir(drawable_dir)

    img_output = open(target_file, mode="wb")
    img_input = open(data.filename, mode="rb")

    while True:
        buf = img_input.read(1024)
        if len(buf):
            img_output.write(buf)
        else:
            break


modules = {
    "9.png": module_nine_patch_resize,
    "png": module_img_resize,
    "jpg": module_img_resize,
    "jpeg": module_img_resize,
    "xml": module_copy
}


def process_image(name: str):
    imgpath = args.source + str(name)

    if not os.path.isfile(imgpath):
        return False

    if "." not in name:
        print("cannot extract extension, skipping file {name}".format(name=name))

    ext_dot_index = name.index(".")
    ext = name[ext_dot_index + 1:]

    if ext not in modules:
        print("couldn't find module for {ext}, skipping file {name}".format(ext=ext, name=name))

    if args.verbose:
        print("using", modules[ext].__name__, "for", name)

    data = ImgData(imgpath, args.dp, ext, name)

    try:
        size_end_index = name.rfind("dp", 0, ext_dot_index)
        size_start_index = name.rfind("_", 0, size_end_index - 1) + 1
        data.dpsize = int(name[size_start_index:size_end_index])
    except (IndexError, ValueError):
        if args.verbose:
            print("could not extract dp size for image {name}, defaulting to {dp}dp".format(dp=args.dp, name=name))

    modules[ext](data)

pool = Pool()
print("using up to {0} threads".format(os.cpu_count()))
print()
header = logformat.format("IMAGE", "SOURCE", "DEST", "WARNINGS")
print(header)
print("=" * len(header))
pool.map(process_image, os.listdir(args.source))