add python file

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())