How to Build a Custom Labeler in Python with IPyWidgets and Plotly | by Jacky Kaub | Jan, 2023

By Jessie Hobb On Jan 14, 2023

Create a segmentation tool in a Jupyter environment

Glomerulli segmented from kidney. Original images provided by the HuBMAP’s tissue mapping center at Vanderbilt University

Did you know that you can turn a simple Jupyter Notebook into a powerful labeler in a few line of codes ?

Turning a simple notebook to a labeling tool can be really time saving at the beginning of a project when you are still trying to assess the potential of an idea and you don’t want to spend much in building a robust web app.

In this article I will show you how to build a segmentation tool with plotly and ipywidget to generate binary masks out of images in python. In image segmentation, binary masks can be used to train deep learning models to efficiently recognize and isolate area of interest automatically.

I assume that you are already familiar with the basis of combining ipywidgets and plotly. If this is not the case, I strongly recommend having a look first at this article where we go through the basics.

A word about the usecase

I am using the HuBMAP dataset from a Kaggle competition hosted in 2021.

I chose this dataset as it illustrates perfectly real life applications: a client reach out to you with a bunch of images and asks if you could find a way to automatically isolate within those images specific areas of interests. Those areas(in this example) are called glomeruli, tiny parts of the kidney that helps to filter out waste and extra water from the blood, see the illustration below.

The illustration below shows what your customer is expected here: a model to quickly identify those areas of interest.

Project objective: providing segmented area of the kidney images. Original images provided by the HuBMAP’s tissue mapping center at Vanderbilt University

In exploration phases, time is of the essence. Often, we only have access to raw data and need to quickly generate labels to train and evaluate a model and assess the feasibility of the project.

Instead of spending hours searching for an open-source tool that may still require customization, I will guide you through building your own custom tool. You’ll see that with a bit of effort, you can easily create any type of widget to aid in your daily tasks.

Building the widget technical specifications

Before jumping directly to the coding part, we need to take a step back and think about what we want exactly from our widget. Among the questions we need to answer:

What is the format of the inputs ? Are they jpeg images ? Numpy array ? rasters ? etc..
Where are they stored ? Locally ? In multiple area ? Or even in the cloud ?
What are you expecting as an output ? Another jpeg images or numpy array ? Maybe annotations stored in a json ?
What are the interactions you are expecting ? Clicking in a button do to something, having some UI elements to navigate efficiently from data to data, interacting directly with a figure, etc…
For each interaction, how, practically will you manage it ? How this will affect the inner state of the widget ?
Try to draw somewhere, even on paper, the way you want your widget to looks like. This will help you structure your widget’s layout.

In our case:

I assume that the main images are already divided into 256×256 pixels sub-images and stored as .npy files in an imgs/ folder locally.

main/
|-- imgs/
|     |-- 1.npy
|     |-- 2.npy
|     |-- etc..
|-- masks/
|-- widget_notebook.ipynb

The generated masks will be stored in the masks/ folder, also as .npy, with the same name as the original image.

Among the functionalities desired by our tool:

The widget will display two figures: the raw image with a transparent overlay of the mask (work ongoing), and the currently saved mask, that will allow a user to visualize the current work done and complete it if required.
A user can click multiple times on the image to generate a polygon defining an area. This will be used inside the widget to generate a mask with 0 for pixels outside of the area, and 1 for pixels inside.
A button should be used to save the mask when the user is happy with its label, so that the result is stored and can be retrieved by him or another service (like a deep learning modele)
A button should be used to delete the current polygon, which will be convenient to fix drawing mistakes.
A button should be used to reset the locally saved mask, allowing to revert to the original state if needed without having to delete the mask manually.
The user can navigate from image to image via a dropdown menu

A rough view of the widget layout that we want to achieve. Original images provided by the HuBMAP’s tissue mapping center at Vanderbilt University

Some good practices before starting

I recommend you to build your widget within a custom Class. That will facilitate your life in term of state management as everything will be held internally, and be way easier to package.

Also, as a general advice,you should build the main layout first, and build the interactions only in a second phase.

Finally, keep things as much separated as possible ! It does not cost more to create a dedicated function for creating each button, or a function to handle separately each callback, but it might save you a precious time for testing and debugging.

To summarize, for building efficiently all kind of widgets:

Create a class to hold all the internal states used in the widget
Build functions to initialize the figures and the components, without interactions in a first phase
Add the interactions and test them step by step
Keep your code clean and organized

Widget internal state logic

Our widget needs to keeps in memory and modify several information:

The _current_id is used to store the current name of the image and is also use to access and save the binary mask.
The _polygon_coordinates list will contain the coordinates of the points clicked by the user and will be used to generate the binary mask (using a skimage function)
the _initialize_widget function will generate all the other states used to build the layout such as the buttons, figures and intermediate masks.

class SegmentWidget:def __init__(self, path_imgs, path_masks):
self._path_imgs = path_imgs
self._path_masks = path_masks
self._ids = sorted([os.path.splitext(e)[0] for e in  os.listdir(path_imgs)], key=int)
self._current_id = self.ids[0]
#This list will be used later to save in memory the coordinates 
#of the clicks by the user
self._polygon_coordinates = []
#TODO: Initiation of all the layout components
self._initialize_widget()

Note: Python does not manage public and private variables, but as a good pep8 practice, non-public variable and function can be prefixed with “_”.

Read the images

We need to develop a function for reading both the image and its associated mask, when it exists.

This function will be called during the initialization of the widget, but also every time the user navigate to a new image.

    def _load_images(self):
'''This method will be used to load image and mask when we select another image'''
img_path = os.path.join(self._path_imgs,f"{self._current_id}.npy")
self._current_img = np.load(img_path)
h,w, _ = self._current_img.shape#There is not always a mask saved. When no mask is saved, we create an empty one.
mask_path = os.path.join(self._path_masks,f"{self._current_id}.npy")
if os.path.exists(mask_path):
self._current_mask = np.load(mask_path)
else:
self._current_mask = np.zeros((h,w))
#initiate an intermediate mask which will be used to store ongoing work
self._intermediate_mask = self._current_mask.copy()

Note: I prefer to package a dedicated function here. I could have put everything directly in my __init__, but it would have been much less digest. If you identify a task that make sense as a whole, use a function.

Initialize the graphs

The next step is to initialize our FigureWidgets.

We are going to display the RGB image using the plotly.express.imshow method while we are using plotly.graph_objects.Heatmap for the binary mask.

We are currently giving a focus to the general layout of our widget so we are leaving the definitions of the callback for later.

def _initialize_figures(self):
'''This function is called to initialize the figure and its callback'''
self._image_fig = go.FigureWidget()
self._mask_fig = go.FigureWidget()self._load_images() #Update the state loading the images
#We use plotly express to generate the RGB image from the 3D array loaded
img_trace = px.imshow(self._current_img).data[0]
#We use plotly HeatMap for the 2D mask array
mask_trace = go.Heatmap(z=self._current_mask, showscale=False, zmin=0, zmax=1)
#Add the traces
self._image_fig.add_trace(img_trace)
self._image_fig.add_trace(mask_trace)
self._mask_fig.add_trace(mask_trace)
#A bit of chart formating
self._image_fig.data[1].opacity = 0.3 #make the mask transparent on image 1
self._image_fig.data[1].zmax = 2 #the overlayed mask above the image can have values in range 0..2
self._image_fig.update_xaxes(visible=False)
self._image_fig.update_yaxes(visible=False)
self._image_fig.update_layout(margin={"l": 10, "r": 10, "b": 10, "t": 50}, 
title = "Define your Polygon Here",
title_x = 0.5, title_y = 0.95)
self._mask_fig.update_layout(yaxis=dict(autorange='reversed'), margin={"l": 0, "r": 10, "b": 10, "t": 50},)
self._mask_fig.update_xaxes(visible=False)
self._mask_fig.update_yaxes(visible=False)
#Todo: add the callbacks to the two charts

Don’t hesitate to spend a bit of time on the formatting, it might looks annoying to do, but having a tool of the good size with just the good level of information will save you some time when you manipulate it

Buttons, Dropdown, and final layout

We will now add functions to integrate the interface components to the widget. As precised in the specs, we want 3 buttons (displayed horizontally), a dropdown menu, and the two figures next to each other.

def _build_save_button(self):
self._save_button = Button(description="Save Configuration")
#Todo: add the callbackdef _build_delete_current_config_button(self):
self._delete_current_config_button = Button(description="Delete Current Mask")
#Todo: add the callback
def _build_delete_all_button(self):
self._delete_all_button = Button(description="Delete All Mask")
#Todo: add the callback
def _build_dropdown(self):
#The ids are passed as option for the dropdown
self._dropdown = Dropdown(options = self._ids)
#Todo: add the callback
def _initialize_widget(self):
'''Function called during the init phase to initialize all the components
and build the widget layout
'''
#Initialize the components
self._initialize_figures()
self._build_save_button()
self._build_delete_current_config_button()
self._build_delete_all_button()
self._build_dropdown()
#Build the layout
buttons_widgets = HBox([self._save_button,
self._delete_current_config_button,
self._delete_all_button])
figure_widgets = HBox([self._image_fig, self._mask_fig])
self.widget = VBox([self._dropdown, buttons_widgets, figure_widgets])
def display(self):
display(self.widget)

I add a “public” display() method (the only “public” of the widget). It will result in showing the initial state of your widget.

Having such a method allow users to directly display the widget without having to look for your widget name inside its internal state…

Original images provided by the HuBMAP’s tissue mapping center at Vanderbilt University

Now that we have our widget fully initialized, we can start to work on the different interactions.

The dropdown interaction

The dropdown allows the user to navigate through the different images of the dataset. In order to work correctly, the dropdown callback should:

Read the new image and the new mask
Update both figures
Clear the list of points in _polygon_coordinates to start a new polygon

To update the px.imshow() trace, we need to modify its “source” parameter.

To update the go.Heatmap() trace, we modify the “z” parameter of the trace.

Note that each trace is usually displaying as an object containing all the information of the chart. Don’t hesitate to display it to see what you can modify more easily.

def _callback_dropdown(self, change):#Set the new id to the new dropdown value
self._current_id = change['new']
#Load the new image and the new mask, we already have a method to do this
self._load_images()
img_trace = px.imshow(self._current_img).data[0]
#Update both figure
with self._image_fig.batch_update():
#Update the trace 0 and the trace 1 containing respectively
#the image and the mask
self._image_fig.data[0].source = img_trace.source
self._image_fig.data[1].z = self._current_mask
with self._mask_fig.batch_update():
self._mask_fig.data[0].z = self._current_mask
#Reset the list of coordinates used to store current work in progress
self._polygon_coordinates = []
def _build_dropdown(self):
#The ids are passed as option for the dropdown
self._dropdown = Dropdown(options = self._ids)
self._dropdown.observe(self._callback_dropdown, names="value")

We can now navigate through the different images.

The FigureWidget on_click interaction

Clicking on a pixel in the image will store its coordinate in the widget state.

Having 3 or more pixels stored should trigger a function that will:

Generate a new mask array with a value of 0 or 2 if the pixel is located within or outside the polygon defined by the list of coordinates
Create an intermediate mask containing both information of the previous mask and the mask newly generated
Display on the left figure that newly generated mask

We are using skimage.draw to generate efficiently the polygon mask from the list of coordinates.

Scikit-image is a powerful tool for handling all type of image processing. Check their documentation to get a grab on the possibilities!

def _gen_mask_from_polygon(self):
'''This function set to 2 the values inside the polygon defined by the list of points provided'''
h,w = self._current_mask.shape
new_mask = np.zeros((h,w), dtype=int)
#Get coordinates inside the polygon using skimage.draw.polygon function
rr, cc = polygon([e[0] for e in self._polygon_coordinates], 
[e[1] for e in self._polygon_coordinates], shape=new_mask.shape)#Recreate the intermediate_mask and set values inside ongoing polygon
#to 2
self._intermediate_mask = self._current_mask.copy()
self._intermediate_mask[rr,cc]=2
def _on_click_figure(self, trace, points, state):
#Retrieve coordinates of the clicked point
i,j = points.point_inds[0]
#Add the point to the list of points
self._polygon_coordinates.append((i,j))
#If more than 2 click have been done, create the new intermediate polygon
#and update the mask on the image
if len(self._polygon_coordinates)>2:
self._gen_mask_from_polygon()
with self._image_fig.batch_update():
self._image_fig.data[1].z = self._intermediate_mask

We can then update our _initialize_figures method to attach the callback.

Note: As the figure is made of several image layers, we attach the callback to the top layer.

def _initialize_figures(self):#[...Rest of the function...]
self._image_fig.data[-1].on_click(self._on_click_figure)

And this is it ! We can now visualize the effect of our chart interaction:

The buttons interactions

To finish our app, let’s code the button interactions.

The “Save Configuration” button. This button save the mask by:

Copying 0 and 2 pixels from the _intermediate_mask to the _current_mask
Save the _current_mask in the /masks folder
Refresh the figures
Reset the _polygon_coordinates list.

def _callback_save_button(self, button):
self._current_mask[self._intermediate_mask==2]=1
self._current_mask[self._intermediate_mask==0]=0
mask_path = os.path.join(self._path_masks,f"{self._current_id}.npy")
np.save(mask_path,self._current_mask)
self._intermediate_mask = self._current_mask.copy()
with self._image_fig.batch_update():
self._image_fig.data[1].z = self._current_mask     
with self._mask_fig.batch_update():
self._mask_fig.data[0].z = self._current_mask
self._polygon_coordinates = []def _build_save_button(self):
self._save_button = Button(description="Save Configuration")
self._save_button.on_click(self._callback_save_button)

The “Delete Current Mask” button. This button simply reset the _intermediate_mask to the _current_mask value and refresh the figures and the _polygon_coordinates list.

def _callback_delete_current_config_button(self, button):
self._intermediate_mask = self._current_mask.copy()
with self._image_fig.batch_update():
self._image_fig.data[1].z = self._intermediate_mask
self._polygon_coordinates = []def _build_delete_current_config_button(self):
self._delete_current_config_button = Button(description="Delete Current Mask")
self._delete_current_config_button.on_click(self._callback_delete_current_config_button)

The “Delete All Mask” button. This button simply reset the _intermediate_mask to 0 and the _polygon_coordinates to an empty list.

def _callback_delete_all_button(self, button):
self._intermediate_mask[:] = 0
with self._image_fig.batch_update():
self._image_fig.data[1].z = self._intermediate_mask
self._polygon_coordinates = []def _build_delete_all_button(self):
self._delete_all_button = Button(description="Delete All Mask")
self._delete_all_button.on_click(self._callback_delete_all_button)

You want to try it out yourself ? The full code is available here!