diff --git a/pc.py b/pc.py new file mode 100644 index 0000000..3ae1265 --- /dev/null +++ b/pc.py @@ -0,0 +1,145 @@ +''' +the_man_with_the_personal_computer.py +Robert Luxemburg, 2020, Public Domain + +The idea is to split each frame of the source video into 16x16 cells, each of +which has a size of 8x6 pixels. Obviously, higher values give more accurate +results, but take longer to render. In the first pass, each frame is resized +accordingly, and the resulting list of 48-integer vectors is saved to disk. +In the second pass, these vectors are used to populate a k-d tree (see +https://en.wikipedia.org/wiki/K-d_tree) that allows for quick lookup of nearest +neighbors. Then each frame is resized so that its cells have a size of 8x6 +pixels, and for each cell, the tree is queried for the best full-frame match. +The above is optimized for speed by using cKDTree instead of KDTree, caching +the cell frames in memory, and allowing multiple instances of the script to run +in parallel. On a dual-core 2015 MacBook Pro, rendering takes around 60 hours. +''' + +import atexit +import json +import os +import sys +import time +import numpy as np +from scipy.spatial import cKDTree +from PIL import Image + +video_fn = 'The Man with the Movie Camera.mkv' +video_dn = '.'.join(video_fn.split('.')[:-1]) +src_dn = f'{video_dn}/src' +os.makedirs(src_dn, exist_ok=True) + +m = 16 # Matrix size +m2 = m ** 2 +cx, cy = 8, 6 # Cell size +cs = cx * cy +w, h = 1920, 1440 # Target frame size +cell_w, cell_h = w//m, h//m +c = 1 # Color channels +fps = 24 +kf = 1 # Keyframes (every kf-th source frame is a cell frame candidate) +start, stop = 3275, 97701 +ts, ti = 1, 0 # Number of threads, current thread index + +cache = {} +cache_keys = [] +cache_size = m2 * fps +cache_hits = 0 + +def get_image(i): + global cache, cache_keys, cache_hits + if i not in cache: + cache[i] = Image.open(f'{video_dn}/src_{cell_w}x{cell_h}/{i:08d}.jpg') + cache_keys.append(i) + if len(cache_keys) > cache_size: + del cache[cache_keys[0]] + cache_keys = cache_keys[1:] + else: + cache_keys.remove(i) + cache_keys.append(i) + cache_hits += 1 + return cache[i] + +def load_image(fn): + image = Image.open(fn) + if c == 1: + image = image.convert('L') + return image + +def read_images(): + def save_data(): + np.save(data_fn, data) + src_fn = f'{src_dn}/{stop:08d}.jpg' + if not os.path.exists(src_fn): + os.makedirs(src_dn, exist_ok=True) + os.system(f'ffmpeg -i "{video_fn}" -q:v 1 "{src_dn}/%08d.jpg"') + cell_dn = f'{src_dn}_{cell_w}x{cell_h}' + os.makedirs(cell_dn, exist_ok=True) + src_fns = [f'src_dn/{fn}' for fn in os.listdir(src_dn) if fn[0] != '.'] + src_fns = sorted(src_fns) + n = len(src_fns) + data_fn = f'{video_dn}/{cx}x{cy}.npy' + if not os.path.exists(data_fn): + data = np.zeros((n, cs) if c == 1 else (n, cs, c), dtype=np.uint8) + else: + data = np.load(data_fn) + atexit.register(save_data) + for i, src_fn in enumerate(src_fns): + t = time.time() + image = None + cell_fn = f'{cell_dn}/{i:08d}.jpg' + if not os.path.exists(cell_fn): + image = load_image(src_fn) + image.resize((cell_w, cell_h), Image.LANCZOS).save(cell_fn) + if sum(data[i]) == 0: + if not image: + image = load_image(src_fn) + data[i] = np.array(image.resize((cx, cy), Image.LANCZOS).getdata()) + t = time.time() - t + if t > 0: + fn = os.path.basename(src_fn) + print(f'\rreading {fn} {1/t:.1f} fps', ' ' * 16, end='') + save_data() + return data.astype(np.int32) + +def write_images(data): + global cache_hits + dst_dn = f'{video_dn}/dst_{m}x{m}x{cx}x{cy}_{kf}' + os.makedirs(dst_dn, exist_ok=True) + src_fns = [f'src_dn/{fn}' for fn in os.listdir(src_dn) if fn[0] != '.'] + src_fns = sorted(src_fns) + n = len(src_fns) + tree = cKDTree(data if kf == 1 else [data[i] for i in range(0, n, kf)]) + for i, src_fn in enumerate(src_fns): + if i % ts != ti: + continue + dst_fn = f'{dst_dn}/{i:08d}.jpg' + if os.path.exists(dst_fn): + continue + t = time.time() + cache_hits = 0 + src_image = load_image(src_fn).resize((cx * m, cy * m), Image.LANCZOS) + dst_image = Image.new('L' if c == 1 else 'RGB', (w, h)) + for y in range(m): + for x in range(m): + crop = (x * cx, y * cy, (x + 1) * cx, (y + 1) * cy) + cell_data = src_image.crop(crop).getdata() + cell_image = get_image(tree.query(cell_data)[1] * kf) + dst_image.paste(cell_image, (x * cell_w, y * cell_h)) + dst_image.save(dst_fn) + t = time.time() - t + fn = os.path.basename(dst_fn) + print(f'\rwriting {fn} {cache_hits} {t:.1f} spf', ' ' * 16, end='') + done = len([f for f in os.listdir(dst_dn) if f[0] != '.']) == n + dst_fn = f'{dst_dn}/{m}x{m}x{cx}x{cy}_{kf}.mp4' + if done and not os.path.exists(dst_fn): + os.system( + f'ffmpeg -r {fps} -start_number 0 -i "{dst_dn}/%08d.jpg" ' + f'-c:v libx264 -profile:v high -level 4.0 -pix_fmt yuv420p ' + f'-r {fps} -crf 20 "{dst_fn}"' + ) + +if __name__ == '__main__': + if len(sys.argv) == 3: + ts, ti = int(sys.argv[1]), int(sys.argv[2]) + write_images(read_images()) \ No newline at end of file