A quick script to compress PNG images in a Jupyter notebook

compress_png_in_notebook.py

#!/usr/bin/env python3
import sys
import json
from subprocess import run
from base64 import decodebytes, encodebytes
from tempfile import NamedTemporaryFile

try:
    from tqdm import tqdm
except ImportError:
    tqdm = lambda x: x

input_path = sys.argv[1]

print(f'Will compress PNG images in the {input_path} notebook')

with open(input_path) as f:
    nb_in = json.load(f)

    new_cells = []

    for cell in tqdm(nb_in['cells']):
        if 'outputs' in cell:
            if not cell['outputs']:
                continue
            outputs = []
            for output in cell['outputs']:
                if 'data' not in output:
                    continue
                data = output['data']
                if 'image/png' in data:
                    base_64_png_data = data['image/png']
                    with NamedTemporaryFile(mode='wb', suffix='.png', delete=False) as p:
                        binary_png = decodebytes(''.join(base_64_png_data.split('\n')).encode())
                        p.write(binary_png)
                        p.flush()
                        run(['trimage', '-f', p.name])
                        with open(p.name, 'rb') as p2:
                            compressed = encodebytes(p2.read()).decode()
                    data['image/png'] = compressed
                outputs.append(output)
            cell['outputs'] = outputs
        new_cells.append(cell)

with open(input_path, 'w') as f:
    nb_in['cells'] = new_cells
    json.dump(nb_in, f)

print(f'Saved as {input_path}')