Compress images in notebooks

nb_image_compressor.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "nbformat==5.10.4",
#     "pillow==11.0.0",
# ]
# ///

from pathlib import Path
import nbformat
from PIL import Image
import io
import base64

# Function to compress images in a notebook using WebP
def compress_images_in_notebook(input_notebook_path, output_notebook_path):
    # Load the notebook
    with open(input_notebook_path, 'r') as f:
        notebook = nbformat.read(f, as_version=4)
    
    # Iterate over cells in the notebook
    for cell in notebook.cells:
        # Check if the cell has an image attachment
        if cell.cell_type == "markdown" and "attachments" in cell:
            for attachment_name, attachment_data in cell['attachments'].items():
                if 'image/png' in attachment_data:
                    # Decode the base64 image data
                    png_data = base64.b64decode(attachment_data['image/png'])
                    
                    # Convert the PNG to WebP and compress
                    with Image.open(io.BytesIO(png_data)) as img:
                        webp_buffer = io.BytesIO()
                        img.convert("RGB").save(webp_buffer, format="WEBP", quality=70)  # Adjust quality as needed
                        webp_data = webp_buffer.getvalue()
                    
                    # Encode the WebP back to base64
                    cell['attachments'][attachment_name]['image/webp'] = base64.b64encode(webp_data).decode('utf-8')
                    
                    # Remove the original PNG data
                    del cell['attachments'][attachment_name]['image/png']
    
    # Save the modified notebook
    with open(output_notebook_path, 'w') as f:
        nbformat.write(notebook, f)

# Scan all notebooks in the current folder and compress images
for notebook_path in Path('.').glob('*.ipynb'):
    output_path = notebook_path.with_stem(notebook_path.stem + '_compressed_webp')
    compress_images_in_notebook(notebook_path, output_path)