Markov Chain for image generation

markovscroll.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-

import json
import itertools
import collections
import random
import pathlib

from PIL import Image

class BlurBucket:

    def __init__(self, blur=1):
        self.bucket = collections.defaultdict(list)
        self.blur = blur

    def __getitem__(self, key):
        return self.bucket[key]

    def __len__(self):
        return len(self.bucket)

    def keys(self):
        return self.bucket.keys()

    def pick(self, key):
        return random.choice(self[key])


class MarkovScroll:

    def __init__(self, origin: pathlib.Path, window_size=5, blur=1):
        self.buckets = BlurBucket(blur)
        self.window_size = window_size
        self.img = Image.open(origin)

        lines = tuple(tuple(self.img.getpixel((x, y)) for y in range(self.img.height)) for x in range(self.img.width))

        for i in range(window_size, len(lines)):
            window = lines[i-window_size:i]
            succ = lines[i]
            self.buckets[window].append(succ)

        print(f"Generated {len(self.buckets)} different windows")

    @property
    def size(self):
        return self.img.height

    def generator(self):
        value = tuple(random.choice(list(self.buckets.keys())))
        while True:
            succ = self.buckets.pick(value)
            result, *value = *value, succ
            value = tuple(value)
            yield result

    def generate(self, fname:pathlib.Path, lines=100):
        img = Image.new("P", (lines, self.size))
        img.putpalette(self.img.getpalette())
        gen = self.generator()
        for x in range(lines):
            line = next(gen)
            for y in range(self.size):
                img.putpixel((x,y), line[y])
        return img.save(fname)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("source", type=pathlib.Path)
    parser.add_argument("output", type=pathlib.Path)
    parser.add_argument("--size", type=int, default=1000)
    parser.add_argument("--window-size", type=int, default=5)
    parser.add_argument("--blur", type=int, default=1)
    args = parser.parse_args()

    s = MarkovScroll(args.source, window_size=args.window_size, blur=args.blur)
    s.generate(args.output, args.size)