Label images dataset in Jupyter notebook
By: Rostyslav Neskorozhenyi
Tags: Jupyter Notebook, Data Science, Tips, Python, Labeling

If you are working with computer vision, then probably you have been dealing with the task of image dataset labeling. Many great utilities were created to help with this task, such as CVAT, but if you need to label several attributes quickly, CVAT functionality may be redundant. In this case, you can deal with image labeling without leaving the familiar Jupyter notebook.

We will consider the process of image labeling in a Jupyter notebook with a specific example. We will need basic knowledge of pandas DataFrames, Jupyter notebook widgets and some JavaScript magic.

As an example, we will label a dataset with images of people wearing face masks. This dataset is just a folder with images of people. In the process of labeling we will create a csv file with the first column containing the names of the images, and the second — the presence (1) or absence (0) of a face mask in each specific image.

We will need to install the Jupyter notebook widget qgrid, which adds the ability of pandas DataFrames interactive filtering and editing directly to the Jupyter notebook. To install qgrid, execute the following commands in the terminal:

pip install qgrid jupyter nbextension enable --py qgrid

Qgrid is a convenient and very fast widget. In my case, it displayed datasets with tens of thousands of records without any problems and delays, but by default it lacks the function we need, it does not know how to display images in cells. To solve this problem, we need a little bit of Javascript.

We will also need to install pandas, if you have not installed it yet:

pip install pandas

After qgrid and pandas are installed, we can create a Jupyter notebook next to the image folder and start the dataset labeling.

To get started, we will import the necessary modules for image labeling in a Jupyter notebook:

import os

import pandas as pd

import numpy as np

import qgrid

Next we will specify the name of the folder with images and create a DataFrame in two columns: image — the name of the image, mask — the presence or absence of a mask on the image.

fold = 'images/'

files = os.listdir(fold)

df = pd.DataFrame({'image': files,

'mask': 0

}

)

df.head()

In the next step, we will create and display a qgrid widget with certain parameters, which seem convenient for image labeling for me. You can read more about these parameters here.

qgrid_widget = qgrid.show_grid(df, show_toolbar = True, grid_options={'forceFitColumns': True, 'rowHeight': 100,

'highlightSelectedCell': True, 'maxVisibleRows': 5}, column_options = {'maxWidth': 150}) qgrid_widget

It is time for JavaScript magic. qgrid widget is based on the SlickGrid JavaScript component, which supports custom cell formatters. Let’s execute a snippet in the next cell, which will add a formatter to display images in the image column.

from IPython.display import Javascript, HTML display(Javascript("""\

window.slick_grid.setColumns(window.slick_grid.getColumns().map(x => {

if (x.id === 'image') {

x.formatter = (r, c, v, cd, dc) =>

(v = "<img src = images/" + v + "></img>");

}

return x

}));

"""))

Done! If we did everything correctly, then the images were displayed directly in the qgrid widget. Now we can quickly label the necessary attributes (in our case, only one mask attribute)

After the attributes are labeled, it is important not to forget to save our DataFrame. All changes made in qgrid widget are not saved in the original DataFrame, so firstly we need to get a copy of the DataFrame with the changes we entered, and then save it to csv file. To do this, execute the following commands:

labels_df = qgrid_widget.get_changed_df()

labels_df.to_csv(‘dataset_masks.csv’, index=False)

Congratulations! Through simple manipulations, we turned Jupyter notebook into an image labeling tool and labeled our dataset in it.