youtube-fs.py

'''
YouTube FileStorage v0.0
requires: ffmpeg, youtube-dl
'''

import os
import shutil
import sys
import numpy as np
from scipy.spatial import cKDTree
import PIL.Image
import moviepy.editor

# same settings, more colors -> lossy :(
colors = ((0, 0, 0), (0, 0, 255), (0, 255, 0), (255, 0, 0))
tree = cKDTree(colors)
vsize = (1920, 1080, 3)
psize = (8, 8)
isize = [int(vsize[i] / psize[i]) for i in range(2)]
bytes_per_frame = int(vsize[0] * vsize[1] / (psize[0] * psize[1]) / 4)
fps = 30
out_dirname, tmp_dirname = 'out', 'tmp'

def encode(filename):
    def render_frame(t):
        f = np.clip(np.int(np.round(t * fps)), 0, frames - 1)
        image = np.zeros(vsize, dtype=np.uint8)
        fp.seek(f * bytes_per_frame)
        payload = fp.read(bytes_per_frame)
        px = 0
        for byte in payload:
            for shift in (6, 4, 2, 0):
                x = psize[0] * (px % isize[0])
                y = psize[1] * int(px / isize[0])
                image[x:x+psize[0],y:y+psize[1]] = colors[byte >> shift & 3]
                px += 1
        return np.transpose(image, (1, 0, 2))
    filesize = os.path.getsize(filename)
    frames = int(np.ceil(filesize / bytes_per_frame))
    duration = frames / fps
    fp = open(filename, 'rb')
    mp4_filename = f'{os.path.basename(filename)} {filesize}.mp4'
    video_clip = moviepy.editor.VideoClip(render_frame, duration=duration)
    video_clip.write_videofile(mp4_filename, fps=fps, codec='libx264', bitrate='8M')
    fp.close()

def decode(id):
    os.system(f'youtube-dl https://youtu.be/{id}')
    filename = [f for f in os.listdir('.') if id in f][0]
    name, filesize = filename.replace(id, '')[:-5].split(' ')
    filesize = int(filesize)
    last_frame_size = filesize % bytes_per_frame
    if os.path.exists(tmp_dirname):
        shutil.rmtree(tmp_dirname)
    os.mkdir(tmp_dirname)
    os.system(f'ffmpeg -i "{filename}" -q:v 0 {tmp_dirname}/%08d.png')
    os.makedirs(out_dirname, exist_ok=True)
    out_filename = f'{out_dirname}/{name}.part'
    fp = open(out_filename, 'wb')
    filenames = sorted([f for f in os.listdir(tmp_dirname) if f.endswith('.png')])
    for i, filename in enumerate(filenames):
        print(f'\r{i+1}/{len(filenames)}', end='', flush=True)
        image = np.array(PIL.Image.open(f'{tmp_dirname}/{filename}'))
        image = np.transpose(image, (1, 0, 2))
        length = last_frame_size if i == len(filenames) - 1 else bytes_per_frame
        payload = np.zeros(length, dtype=np.uint8)
        idx, px = 0, 0
        while idx < length:
            for shift in (6, 4, 2, 0):
                x = psize[0] * (px % isize[0])
                y = psize[1] * int(px / isize[0])
                pixels = image[x+1:x+psize[0]-2,y+1:y+psize[1]-2]
                color = np.mean(pixels, (0, 1))
                payload[idx] += tree.query(color)[1] << shift
                px += 1
            idx += 1
        fp.write(bytes(payload))
    print('\nDone.')
    fp.close()
    os.rename(out_filename, out_filename[:-5])
    shutil.rmtree(tmp_dirname)

if __name__ == '__main__':
    if len(sys.argv) != 3 or sys.argv[1] not in ('encode', 'decode'):
        print('youtube-fs encode filename\nupload\nyoutube-fs decode id')
    action, filename = sys.argv[1:]
    encode(filename) if action == 'encode' else decode(filename)