Skip to content

adusachev/jupyter-paint-segment

Repository files navigation

jupyter-paint-segment

Overview

A jupyter widget for interactive drawing segmentation masks.

Brush, rectangle and eraser tools are available. Size of brush and eraser can be changed using slider.


Installation

pip install jupyter_paint_segment

requires: python >= 3.9


Usage

Demo notebook: ./examples/main.ipynb

  1. Load image into numpy array:
# with PIL: 
from PIL import Image
import numpy as np
pilImage = Image.open("./examples/images/sheeps.png")
image = np.array(pilImage)

# or with opencv:
import cv2 as cv
image_bgr = cv.imread("./examples/images/sheeps.png")
image = cv.cvtColor(image_bgr, cv.COLOR_BGR2RGB)
  1. Define labels and colors (optionally) and create widget:
from jupyter_paint_segment import SegmentWidget

widget = SegmentWidget(
    image=image,
    labels=["sheep", "dog"],
    colors=["red", "blue"],
    image_scale=1,
)
widget
  1. Get segmentation results:
labels_array, labels_map = widget.segmentation_result()

labels_map
# {'sheep': 1, 'dog': 2, 'unlabeled_background': 0}
labels_array
# array([[0, 0, 0, ..., 0, 0, 0],
#        [0, 0, 0, ..., 0, 0, 0],
#        [0, 0, 0, ..., 0, 0, 0],
#        ...,
#        [0, 0, 0, ..., 0, 0, 0],
#        [0, 0, 0, ..., 0, 0, 0],
#        [0, 0, 0, ..., 0, 0, 0]])
import matplotlib.pyplot as plt

plt.imshow(labels_array)


Related projects

This project was highly inspired by jupyter-bbox-widget.


License

This project is licensed under the MIT license.