Initial commit

1 files changed, 230 insertions, 0 deletions

diff --git a/imgutils.py b/imgutils.py
new file mode 100644
index 0000000..955cf8f
--- /dev/null
+++ b/imgutils.py
@@ -0,0 +1,230 @@
+import os
+import shutil
+import subprocess
+from io import BytesIO
+from tempfile import NamedTemporaryFile
+
+from PIL import Image, ImageOps
+
+from config import IMAGE_STRIP_HEIGHT, IMAGE_QUALITY, IMAGE_CACHE, IMAGE_MAX_WIDTH, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT, \
+    THUMBNAIL_QUALITY
+
+Image.MAX_IMAGE_PIXELS = 1000000000
+from glob import glob
+
+def _copy(src, dst):
+    try:
+        os.link(src, dst)
+    except:
+        shutil.copy(src, dst)
+
+
+def rechunk(im):
+
+    height = im.height
+    width = im.width
+    last_chunk = (height // IMAGE_STRIP_HEIGHT) - 1
+
+    for i, cursor in enumerate(range(0, height, IMAGE_STRIP_HEIGHT)):
+
+        if i == last_chunk:
+            chunk_height = IMAGE_STRIP_HEIGHT + (height % IMAGE_STRIP_HEIGHT)
+        else:
+            chunk_height = min(IMAGE_STRIP_HEIGHT, height)
+
+        yield im.crop((0, cursor, width, cursor + chunk_height))
+
+        if i == last_chunk:
+            break
+
+def webp_convert(input_file, output_file, quality=IMAGE_QUALITY):
+
+    subprocess.run([
+        "ffmpeg", "-y", "-loglevel", "error", "-hide_banner", "-nostats", "-i", input_file,
+        "-c:v", "libwebp", "-preset", "drawing", "-quality", str(quality), "-compression_level", "6",
+        output_file
+    ], text=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
+
+def copy_from_cache(cache_folder, output_folder):
+
+    for filename in glob(os.path.join(cache_folder, "*.webp")):
+        _copy(filename, os.path.join(output_folder, os.path.basename(filename)))
+
+def is_cached(cache_folder):
+    if os.path.exists(os.path.join(cache_folder, "ok")):
+        return True
+
+    return False
+
+def convert_cover(cache_key, filename, output_folder):
+    cache_folder = os.path.join(IMAGE_CACHE, cache_key)
+
+    os.makedirs(output_folder, exist_ok=True)
+    output_filename = os.path.join(output_folder, "cover.webp")
+    if os.path.exists(output_filename):
+        os.remove(output_filename)
+
+    if is_cached(cache_folder):
+        copy_from_cache(cache_folder, output_folder)
+        return
+
+    shutil.rmtree(cache_folder, ignore_errors=True)
+    os.makedirs(cache_folder)
+
+    im = open_image(filename)
+
+    if im.width > IMAGE_MAX_WIDTH:
+        aspect_ratio = im.height / im.width
+        new_height = int(IMAGE_MAX_WIDTH * aspect_ratio)
+
+        im = im.resize((IMAGE_MAX_WIDTH, new_height), resample=Image.Resampling.LANCZOS)
+
+    if im.mode != "RGB":
+        im = im.convert("RGB")
+
+    try:
+        with NamedTemporaryFile(delete=False, suffix=".png") as f:
+            im.save(f)
+
+        webp_convert(f.name, os.path.join(cache_folder, f"cover.webp"), quality=90)
+
+    finally:
+        os.remove(f.name)
+
+    with open(os.path.join(cache_folder, "ok"), "w"):
+        pass
+
+    copy_from_cache(cache_folder, output_folder)
+
+def convert_rgb(im):
+    if im.mode == "RGBA":
+        background = Image.new('RGBA', im.size, color=(0, 0, 0))
+        alpha_composite = Image.alpha_composite(background, im)
+        im = alpha_composite.convert("RGB")
+
+    if im.mode != "RGB":
+        im = im.convert("RGB")
+
+    return im
+
+def convert_rgba(im):
+    if im.mode != "RGBA":
+        im = im.convert("RGBA")
+
+    return im
+
+
+def convert_chapter(cache_key, images, output_folder):
+
+    # TODO: for non-long strips, stack horizontally
+    # Limit height to
+
+    cache_folder = os.path.join(IMAGE_CACHE, cache_key)
+
+    shutil.rmtree(output_folder, ignore_errors=True)
+    os.makedirs(output_folder)
+
+    if is_cached(cache_folder):
+        copy_from_cache(cache_folder, output_folder)
+        return
+
+    shutil.rmtree(cache_folder, ignore_errors=True)
+    os.makedirs(cache_folder)
+
+    resized = []
+
+    # Resize images
+    for filename in images:
+        im = open_image(filename)
+
+        im = convert_rgb(im)
+
+        if im.width > IMAGE_MAX_WIDTH:
+
+            aspect_ratio = im.height / im.width
+            new_height = int(IMAGE_MAX_WIDTH * aspect_ratio)
+
+            if new_height == 0:
+                continue
+
+            im = im.resize((IMAGE_MAX_WIDTH, new_height), resample=Image.Resampling.LANCZOS)
+
+        resized.append(im)
+
+    # Prepare canvases
+    canvases = [
+        [resized[0]]
+    ]
+    current_canvas = canvases[0]
+    last_image_width = resized[0].width
+    for i in range(1, len(resized)):
+
+        if last_image_width != resized[i].width:
+            canvases.append([])
+            current_canvas = canvases[-1]
+
+        current_canvas.append(resized[i])
+        last_image_width = resized[i].width
+
+    # Stack
+    stacked_images = []
+    for canvas in canvases:
+        width = canvas[0].width
+        height = sum(im.height for im in canvas)
+
+        new_image = Image.new("RGB", (width, height), color=(0, 0, 0, 0))
+
+        cursor = 0
+        for img in canvas:
+            new_image.paste(img, (0, cursor))
+            cursor += img.height
+
+        stacked_images.append(new_image)
+
+    # Rechunk
+    chunks = []
+    for im in stacked_images:
+        for im_chunk in rechunk(im):
+            chunks.append(im_chunk)
+
+    # Convert
+    for i, chunk in enumerate(chunks):
+
+        try:
+            with NamedTemporaryFile(delete=False, suffix=".png") as f:
+                chunk.save(f)
+
+            webp_convert(f.name, os.path.join(cache_folder, f"{i:03d}.webp"))
+
+        finally:
+            os.remove(f.name)
+
+    with open(os.path.join(cache_folder, "ok"), "w"):
+        pass
+
+    copy_from_cache(cache_folder, output_folder)
+
+def open_image(filename):
+
+    with open(filename, "rb") as f:
+        data = f.read()
+
+    io = BytesIO(data)
+    io.name = filename
+    return Image.open(io)
+
+def create_thumbnail_sheet(images, output_filename):
+    images = [open_image(image["cover"]) for image in images]
+
+    width = len(images) * THUMBNAIL_WIDTH
+    height = THUMBNAIL_HEIGHT
+
+    new_image = Image.new("RGBA", (width, height), color=(0, 0, 0, 0))
+
+    for i, im in enumerate(images):
+        im = convert_rgba(im)
+
+        im = ImageOps.pad(im, (THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT), color=(0, 0, 0, 0), method=Image.Resampling.LANCZOS)
+        new_image.paste(im, (i * THUMBNAIL_WIDTH, 0))
+
+    new_image.save(output_filename, quality=THUMBNAIL_QUALITY)