Skip to content

Commit

Permalink
Make image as part of the widget model
Browse files Browse the repository at this point in the history
  • Loading branch information
martinRenou committed Oct 8, 2021
1 parent c5f8b5a commit 1a83e34
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 53 deletions.
76 changes: 59 additions & 17 deletions ipympl/backend_nbagg.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,29 @@
import json
from base64 import b64encode

import matplotlib
from PIL import Image

import numpy as np

from traitlets import (
Any,
Bool,
CaselessStrEnum,
CInt,
Enum,
Instance,
List,
Unicode,
default,
)

from IPython import get_ipython
from IPython import version_info as ipython_version_info
from IPython.display import HTML, display

from ipywidgets import DOMWidget, widget_serialization

import matplotlib
from matplotlib import is_interactive, rcParams
from matplotlib._pylab_helpers import Gcf
from matplotlib.backend_bases import NavigationToolbar2, _Backend, cursors
Expand All @@ -18,16 +36,6 @@
NavigationToolbar2WebAgg,
TimerTornado,
)
from traitlets import (
Bool,
CaselessStrEnum,
CInt,
Enum,
Instance,
List,
Unicode,
default,
)

from ._version import js_semver

Expand All @@ -40,6 +48,12 @@
}


def image_serialization(image, owner):
if image is None:
return None
return memoryview(image)


def connection_info():
"""
Return a string showing the figure and connection status for
Expand Down Expand Up @@ -142,15 +156,15 @@ class Canvas(DOMWidget, FigureCanvasWebAggCore):
resizable = Bool(True).tag(sync=True)
capture_scroll = Bool(False).tag(sync=True)

_image = Any().tag(sync=True, to_json=image_serialization)

_width = CInt().tag(sync=True)
_height = CInt().tag(sync=True)

_figure_label = Unicode('Figure').tag(sync=True)
_message = Unicode().tag(sync=True)
_cursor = Unicode('pointer').tag(sync=True)

_image_mode = Unicode('full').tag(sync=True)

_rubberband_x = CInt(0).tag(sync=True)
_rubberband_y = CInt(0).tag(sync=True)
_rubberband_width = CInt(0).tag(sync=True)
Expand All @@ -170,6 +184,8 @@ def __init__(self, figure, *args, **kwargs):
DOMWidget.__init__(self, *args, **kwargs)
FigureCanvasWebAggCore.__init__(self, figure, *args, **kwargs)

self.set_image_mode('full')

self.on_msg(self._handle_message)

def _handle_message(self, object, content, buffers):
Expand Down Expand Up @@ -206,15 +222,34 @@ def send_json(self, content):
# Send resize message anyway
self.send({'data': json.dumps(content)})

elif content['type'] == 'image_mode':
self._image_mode = content['mode']

else:
# Default: send the message to the front-end
self.send({'data': json.dumps(content)})

def send_binary(self, data):
self.send({'data': '{"type": "binary"}'}, buffers=[data])
if self._image == data:
# We need this, otherwise the front-end is stuck waiting
# for new image data that never comes
self.send({'data': '{"type": "stop_waiting"}'})
else:
self._image = data

def get_image(self):
if self._png_is_old:
renderer = self.get_renderer()

# The buffer is created as type uint32 so that entire
# pixels can be compared in one numpy call, rather than
# needing to compare each plane separately.
output = (np.frombuffer(renderer.buffer_rgba(), dtype=np.uint32)
.reshape((renderer.height, renderer.width)))

self._png_is_old = False

data = output.view(dtype=np.uint8).reshape((*output.shape, 4))
with io.BytesIO() as png:
Image.fromarray(data).save(png, format="png")
return png.getvalue()

def new_timer(self, *args, **kwargs):
return TimerTornado(*args, **kwargs)
Expand Down Expand Up @@ -337,6 +372,13 @@ def show(self):
def destroy(self):
self.canvas.close()

def refresh_all(self):
if self.web_sockets:
image = self.canvas.get_image()
if image is not None:
for s in self.web_sockets:
s.send_binary(image)


@_Backend.export
class _Backend_ipympl(_Backend):
Expand Down
107 changes: 71 additions & 36 deletions src/mpl_widget.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,38 @@ import * as utils from './utils';

import { MODULE_VERSION } from './version';

function deserialize_image(image: DataView | null): Uint8Array | null {
// No image
if (image === null) {
return null;
}
// Base64 image
if (typeof image === 'string') {
return new Uint8Array(
Array.from(atob(image)).map(c => c.charCodeAt(0))
)
}
// Array of pixels
return new Uint8Array(image.buffer);
}

// Useful for saving widget state in Notebook, we save as base64
function serialize_image(image: Uint8Array | null): string | null {
if (image === null) {
return null;
}
return btoa(String.fromCharCode.apply(null, Array.from(image)));
}

export class MPLCanvasModel extends DOMWidgetModel {
offscreen_canvas: HTMLCanvasElement;
offscreen_context: CanvasRenderingContext2D;
requested_size: Array<number> | null;
resize_requested: boolean;
ratio: number;
waiting: any;
waiting_for_image: boolean;
image: HTMLImageElement;

defaults() {
return {
...super.defaults(),
Expand All @@ -32,12 +56,12 @@ export class MPLCanvasModel extends DOMWidgetModel {
toolbar_position: 'horizontal',
resizable: true,
capture_scroll: false,
_image: null,
_width: 0,
_height: 0,
_figure_label: 'Figure',
_message: '',
_cursor: 'pointer',
_image_mode: 'full',
_rubberband_x: 0,
_rubberband_y: 0,
_rubberband_width: 0,
Expand All @@ -48,6 +72,7 @@ export class MPLCanvasModel extends DOMWidgetModel {
static serializers: ISerializers = {
...DOMWidgetModel.serializers,
toolbar: { deserialize: unpack_models as any },
_image: { deserialize: deserialize_image, serialize: serialize_image },
};

initialize(attributes: any, options: any) {
Expand All @@ -71,13 +96,16 @@ export class MPLCanvasModel extends DOMWidgetModel {
this.resize_requested = false;
this.ratio = (window.devicePixelRatio || 1) / backingStore;
this._init_image();
this.resize_canvas(this.get('_width'), this.get('_height'));
this.on_image_change();

this.on('msg:custom', this.on_comm_message.bind(this));
this.on('change:resizable', () => {
this._for_each_view((view: MPLCanvasView) => {
view.update_canvas();
});
});
this.on('change:_image', this.on_image_change.bind(this));

this.send_initialization_message();
}
Expand All @@ -96,15 +124,14 @@ export class MPLCanvasModel extends DOMWidgetModel {
});
}

this.send_message('send_image_mode');
this.send_message('refresh');

this.send_message('initialized');
}

send_draw_message() {
if (!this.waiting) {
this.waiting = true;
if (!this.waiting_for_image) {
this.waiting_for_image = true;
this.send_message('draw');
}
}
Expand Down Expand Up @@ -190,23 +217,6 @@ export class MPLCanvasModel extends DOMWidgetModel {
this.send_draw_message();
}

handle_binary(msg: any, dataviews: any) {
const url_creator = window.URL || window.webkitURL;

const buffer = new Uint8Array(dataviews[0].buffer);
const blob = new Blob([buffer], { type: 'image/png' });
const image_url = url_creator.createObjectURL(blob);

// Free the memory for the previous frames
if (this.image.src) {
url_creator.revokeObjectURL(this.image.src);
}

this.image.src = image_url;

this.waiting = false;
}

handle_history_buttons(msg: any) {
// No-op
}
Expand All @@ -217,6 +227,10 @@ export class MPLCanvasModel extends DOMWidgetModel {
// button to toggle?
}

handle_stop_waiting(msg: any) {
this.waiting_for_image = false;
}

on_comm_message(evt: any, dataviews: any) {
const msg = JSON.parse(evt.data);
const msg_type = msg['type'];
Expand All @@ -239,20 +253,44 @@ export class MPLCanvasModel extends DOMWidgetModel {
}
}

on_image_change() {
const image = this.get('_image') as Uint8Array | null;

if (image === null) {
this.offscreen_context.clearRect(
0,
0,
this.offscreen_canvas.width,
this.offscreen_canvas.height
);

return;
}

const url_creator = window.URL || window.webkitURL;

const blob = new Blob([image], { type: 'image/png' });
const image_url = url_creator.createObjectURL(blob);

// Free the memory for the previous frames
if (this.image.src) {
url_creator.revokeObjectURL(this.image.src);
}

this.image.src = image_url;

this.waiting_for_image = false;
}

_init_image() {
this.image = document.createElement('img');
this.image.onload = () => {
if (this.get('_image_mode') === 'full') {
// Full images could contain transparency (where diff images
// almost always do), so we need to clear the canvas so that
// there is no ghosting.
this.offscreen_context.clearRect(
0,
0,
this.offscreen_canvas.width,
this.offscreen_canvas.height
);
}
this.offscreen_context.clearRect(
0,
0,
this.offscreen_canvas.width,
this.offscreen_canvas.height
);
this.offscreen_context.drawImage(this.image, 0, 0);

this._for_each_view((view: MPLCanvasView) => {
Expand Down Expand Up @@ -285,7 +323,6 @@ export class MPLCanvasView extends DOMWidgetView {
context: CanvasRenderingContext2D;
top_canvas: HTMLCanvasElement;
top_context: CanvasRenderingContext2D;
waiting: boolean;
footer: HTMLDivElement;
model: MPLCanvasModel;
private _key: string | null;
Expand Down Expand Up @@ -313,8 +350,6 @@ export class MPLCanvasView extends DOMWidgetView {
window.addEventListener('mousemove', this._resize_event);
window.addEventListener('mouseup', this._stop_resize_event);

this.waiting = false;

return this.create_child_view(this.model.get('toolbar')).then(
(toolbar_view) => {
this.toolbar_view = toolbar_view;
Expand Down

0 comments on commit 1a83e34

Please sign in to comment.