Python script to blend several images from a parametrized striped pattern

blend_images.py

import glob
import argparse
import logging
import math
import tqdm
import itertools
from PIL import Image, ImageDraw


class Vector:

    def __init__(self, x, y):
        self.x = x
        self.y = y
    
    def __str__(self):
        return f"({ self.x }, { self.y })"

    def __sub__(self, other):
        return Vector(self.x - other.x, self.y - other.y)
    
    def __rmul__(self, other):
        return Vector(self.x * other, self.y * other)
    
    def dot(self, other):
        return self.x * other.x + self.y * other.y
    
    def norm2(self):
        return math.sqrt(self.dot(self))


class GradientStop:

    def __init__(self, t, color):
        self.t = t
        self.color = color
    
    def __str__(self):
        return "#%02x%02x%02x%02x %d%%" % (*self.color, self.t * 100)


class Gradient:

    def __init__(self, start, end, stops):
        self.start = Vector(*start)
        self.end = Vector(*end)
        self.support = self.end - self.start
        self.support_norm_square = math.pow(self.support.norm2(), 2)
        self.stops = stops
    
    def __str__(self):
        return f"{ self.start }, { self.end }, " + ", ".join(map(str, self.stops))
    
    def t_at(self, j, i):
        return (Vector(j, i) - self.start).dot(self.support) / self.support_norm_square

    def color_at(self, t):
        if t < 0 or t > 1:
            return (0, 0, 0, 0)
        for stop_prev, stop_next in zip(self.stops[:-1], self.stops[1:]):
            if stop_prev.t <= t and stop_next.t >= t:
                s = (t - stop_prev.t) / (stop_next.t - stop_prev.t)
                r = stop_prev.color[0] * (1 - s) + stop_next.color[0] * s
                g = stop_prev.color[1] * (1 - s) + stop_next.color[1] * s
                b = stop_prev.color[2] * (1 - s) + stop_next.color[2] * s
                a = stop_prev.color[3] * (1 - s) + stop_next.color[3] * s
                return (int(r), int(g), int(b), int(a))
        return (0, 0, 0, 0)


def create_gradient(k, n, width, height, angle=0, feather=0, horizontal_symmetry=False, vertical_symmetry=False):
    alpha = angle * math.pi * 2 / 360
    if math.tan(alpha) <= height / width:
        length = width / math.cos(alpha) + (height - width * math.tan(alpha)) * math.sin(alpha)
    else:
        length = height / math.sin(alpha) + (width - height / math.tan(alpha)) * math.cos(alpha)
    if feather == 0:
        stops = [
            GradientStop(0, (255, 255, 255, 255)),
            GradientStop(1, (255, 255, 255, 255)),
        ]
    else:
        stops = [
            GradientStop(0, (255, 255, 255, 0)),
            GradientStop(feather / (1 + 2 * feather), (255, 255, 255, 255)),
            GradientStop((feather + 1) / (1 + 2 * feather), (255, 255, 255, 255)),
            GradientStop(1, (255, 255, 255, 0))
        ]
    x_step = math.cos(alpha) * length / n
    y_step = math.sin(alpha) * length / n
    start = [x_step * (k - feather), y_step * (k - feather)]
    end = [x_step * (k + 1 + feather), y_step * (k + 1 + feather)]
    if horizontal_symmetry:
        start[0] = width - start[0]
        end[0] = width - end[0]
    if vertical_symmetry:
        start[1] = height - start[1]
        end[1] = height - end[1]
    gradient = Gradient(start, end, stops)
    logging.info("Created gradient %s", gradient)
    return gradient


def create_mask(k, n, width, height, options):
    image = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    context = ImageDraw.Draw(image)
    gradient = create_gradient(k, n, width, height, angle=options.angle,
                               feather=options.feather,
                               horizontal_symmetry=options.horizontal_symmetry,
                               vertical_symmetry=options.vertical_symmetry)
    for i in range(height):
        for j in range(width):
            t = gradient.t_at(j, i)
            if 0 <= t <= 1:
                context.point((j, i), gradient.color_at(t))
    return image


def blend_images(input_paths, options):
    width = None
    height = None
    images = []
    for input_path in input_paths:
        for path in glob.glob(input_path):
            logging.info("Loading %s", path)
            image = Image.open(path)
            if width is None or height is None:
                width = image.width
                height = image.height
            elif image.width != width or image.height != height:
                logging.warning("Expecting shape %dx%d but got %dx%d", width, height, image.width, image.height)
                continue
            images.append(image)
    if options.n == "auto":
        n = len(images)
    else:
        n = int(options.n)
    logging.info("Loaded %d images", n)
    if len(images) == 0:
        logging.warning("No image to blend")
        return None
    blended_image = Image.new("RGB", (width, height))
    images_iterator = itertools.cycle(images)
    for k in tqdm.tqdm(range(n), unit="strip"):
        image = next(images_iterator)
        mask = create_mask(k, n, width, height, options)
        blended_image.paste(image, (0, 0), mask)
    return blended_image


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("input", nargs="+", type=str)
    parser.add_argument("-o", "--output", type=str, default="composite.jpg", help="output path")
    parser.add_argument("-s", "--offset", type=int, default=0, help="start using images from given index")
    parser.add_argument("-hs", "--horizontal-symmetry", action="store_true", help="horizontally mirror stripes")
    parser.add_argument("-vs", "--vertical-symmetry", action="store_true", help="vertically mirror stripes")
    parser.add_argument("-a", "--angle", type=float, default=0, help="stripes angle in degree, from 0° to 90°")
    parser.add_argument("-f", "--feather", type=float, default=0, help="blending amount between stripes, from 0 to 1")
    parser.add_argument("-n", "--n", type=str, default="auto", help="number of stripes; if 'auto', there will be one stripe per photo")
    args = parser.parse_args()
    blended_image = blend_images(args.input, args)
    blended_image.save(args.output)


if __name__ == "__main__":
    main()