Skip to content

Instantly share code, notes, and snippets.

@krassowski
Last active April 1, 2020 15:14
Show Gist options
  • Save krassowski/bc537795717a3d61eaa078cb9e393095 to your computer and use it in GitHub Desktop.
Save krassowski/bc537795717a3d61eaa078cb9e393095 to your computer and use it in GitHub Desktop.
A quick script to compress PNG images in a Jupyter notebook
#!/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}')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment