Cell magic to save image output as a .png file

It would be great if there was an easy option to save Jupyter cell output to .png files.
Let’s say I have:

import matplotlib.pyplot as plt
x= range(0,10)
fig, ax = plt.subplots()
ax.plot(x, x)
plt.show()

In this example, I could use fig.savefig("img/foo.png"), however this syntax is tied to Matplotlib, and there are many other libraries that produce images in Jupyter. In my most comfortable scenario, one would only have to add the cell magic %%capture_png img/foo.png which would work for any library. Memorize it once, use it everywhere.

There was a similar discussion (Is there a way to save the output of a cell? - #8 by George_Gaines) already two years ago, which was suggesting the %%capture command (@mgeier) and the project GitHub - nteract/scrapbook: A library for recording and reading data in notebooks.. But with both approaches did not solve the problem.
Currently, I am using a cumbersome workaround to extract all images from the Jupyter notebook while running the notebook (can be seen here: matplotlib - Jupyter Notebook: automatically saves all images - Stack Overflow)
Maybe by 2021 there is a new solution? Or would it be possible to make this a new Jupyter feature?

IPython allows to define custom magics. A very rough draft would be:

from base64 import b64decode
from io import BytesIO

from IPython import get_ipython
from IPython.core.magic import register_cell_magic
import PIL


@register_cell_magic
def capture_png(line, cell):
    get_ipython().run_cell_magic(
        'capture',
        ' --no-stderr --no-stdout result',
        cell
    )
    out_paths = line.strip().split(' ')
    for output in result.outputs:
        data = output.data
        if 'image/png' in data:
            path = out_paths.pop(0)
            if not path:
                raise ValueError('Too few paths given!')
            png_bytes = data['image/png']
            if isinstance(png_bytes, str):
                png_bytes = b64decode(png_bytes)
            assert isinstance(png_bytes, bytes)
            bytes_io = BytesIO(png_bytes)
            image = PIL.Image.open(bytes_io)
            image.save(path, 'png')

And you could use it like:

%%capture_png test_copy.png hist.png
from IPython.display import Image
from pandas import Series

display(Image('test.png'))
Series([1, 2, 2, 3, 3, 3]).hist()
4 Likes

Hey @krassowski, thanks for taking time out of your day and making this implementation, I really appreciate your effort!

I think having a cell magic like this has really the potential to improve the workflow in Jupyter for lots of users when it comes to saving images.

Therefore, some thoughts from my side:

  • Would it be possible to make a package out of your implementation?
  • Or would it be possible to add this feature to an already existing library (e.g. IPython or Ipywidgets)?
  • Saving the output with %%capture_png does not display the image in Jupyter. That is also the case when using the %%capture magic.
    I think it would be more intuitive behavior if cell output is still shown by default, but can be hidden with an extra parameter like --hide_cell_output. I did not find an option to hide/show cell output in Jupyter here Built-in magic commands — IPython 7.30.0 documentation, but maybe this would be also possible to implement?
    EDIT: yes, that is possible by simply adding:
    for output in result.outputs:
        display(output)
  • the current implementation has a dependency to PIL, maybe it is also possible to save the BytesIO object to a png file without PIL?

  • EDIT 2: to avoid overwriting of images, one could also add an option that includes the current time of the cell execution to this magic:

 path = out_paths.pop(0)
 path = path.split(".png")[0] + str(time.time_ns()) + ".png"
1 Like

Would it be possible to make a package out of your implementation?

Yes. Personally, I would not install it as-is though, too little benefit (equivalent to 15-line of code) but maybe after expanding to cover more cases it could be useful.

Or would it be possible to add this feature to an already existing library (e.g. IPython or Ipywidgets)?

I could see a case for including something more generic (%capture_images) in IPython-contrib, but I am not a maintainer there so the decision on accepting it would be up to somebody else :wink:

Saving the output with %%capture_png does not display the image in Jupyter. That is also the case when using the %%capture magic.

Well, it’s called capture, this is what capture functions do. If it should display output it ought to be called tee_png (after tee) or save_png. But, as you correctly noticed the modifications needed to display the outputs are minimal (though the order can be lost if you also have non-image outputs - I don’t have a solution for that of top of my head).

Possibly - it’s barely 200 lines of code (+ maybe few hundreds more in helper methods); personally I would accept the dependency (on PIL or pillow) :slight_smile:

1 Like

I could see a case for including something more generic ( %capture_images ) in IPython-contrib, but I am not a maintainer there so the decision on accepting it would be up to somebody else :wink:

I’ve just made an issue at Ipython, as they have already similar cell magic implemented, maybe the maintainers might be interested.

Well, it’s called capture , this is what capture functions do. If it should display output, it ought to be called tee_png (after tee ) or save_png

Good point, I like save_png ! My first association for capture was like in “to capture a picture with a camera”, where the photo object is still there after taking the photo, but I think with save_png this is clear.

1 Like

I’ve just written a more fleshed out version of this magic, that will normally save to png, but also allows compressed jpg files when using the --compression flag :

from base64 import b64decode
from io import BytesIO
import matplotlib.pyplot as plt
import PIL
from IPython import get_ipython
from IPython.core import magic_arguments
from IPython.core.magic import (Magics, cell_magic, magics_class)
from IPython.display import display
from IPython.utils.capture import capture_output


@magics_class
class MyMagic(Magics):

    @cell_magic
    @magic_arguments.magic_arguments()
    @magic_arguments.argument(
        "--path",
        "-p",
        default=None,
        help=("The path where the image will be saved to. When there is more then one image, multiple paths have to be defined"),
    )
    @magic_arguments.argument(
        "--compression",
        "-c",
        default=None,
        help=("Defines the amout of compression,  quality can be from 0.1 - 100 , images must be .jpg"),
    )
    def polaroid_camera(self, line, cell):
        args = magic_arguments.parse_argstring(MyMagic.polaroid_camera, line)
        paths = args.path.strip('"').split(' ')
        with capture_output(stdout=False, stderr=False, display=True) as result:
            self.shell.run_cell(cell) # thanks @krassowski for this idea!

        for output in result.outputs:
            display(output)
            data = output.data
            if 'image/png' in data:
                path = paths.pop(0)
                if not path:
                    raise ValueError('Too few paths given!')
                png_bytes = data['image/png']
                if isinstance(png_bytes, str):
                    png_bytes = b64decode(png_bytes)
                assert isinstance(png_bytes, bytes)
                bytes_io = BytesIO(png_bytes)
                img = PIL.Image.open(bytes_io)
                if args.compression:
                    if img.mode != "RGB":
                        img = img.convert("RGB")
                    img.save(path, "JPEG", optimize=True,
                             quality=int(args.compression))
                else:
                    img.save(path, 'png')


ipy = get_ipython()
ipy.register_magics(MyMagic)

It can be used like this:

%%polaroid_camera --path "foo.png bar.png"
plt.plot([1,2],[10,20])
plt.show()
plt.plot([3,4],[-10,-20])
plt.show()
%%polaroid_camera --path "foo.jpg bar.jpg" --compression 50
plt.plot([1,2],[10,20], color = "r")
plt.show()
plt.plot([3,4],[-10,-20],color = "r")
plt.show()
2 Likes