Source code for locpix.img_processing.watershed

"""Module which allows for the watershed algorithm to
be implemented"""

# imports for widget
from PyQt5.QtWidgets import QGridLayout, QWidget, QGraphicsScene, QGraphicsView
from PyQt5.QtWidgets import QGraphicsPixmapItem, QMessageBox, QPushButton
from PyQt5.QtGui import QPixmap, QPen, QBrush
from PyQt5.QtGui import QImage, QColor
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication

from skimage.segmentation import watershed
import numpy as np
import warnings


[docs] class WatershedWidget(QWidget): """WatershedWidget Widget to visualise an image as Grayscale 8 bit, then annotate with the seeds used for watershed. Child of QWidget. Left click places a seed, right click removes the seed. On closing the seed coordinates are saved Args: None Attributes: scene(QGraphicsScene): This holds all the 2D items (img, all ellipses added) view(QGraphicsView): Visualise the scene pixmap_item(QGraphicsPixmapItem): Image representation that can be used as a paint device help_button (QPushButton): Displays help for user pen(QPen): How the QPainter draws lines and outlines of shapes brush(QBrush): How to fill the shapes drawn by the painter marker_coords(list): List containing coordinates of the seeds for watershed """
[docs] def __init__(self, img, coords=[], file_name="Image"): """Constructor Args: img (int8 numpy array): Numpy array range [0 255] which will be converted to correct format for pixmap coords (list): Seed coordinate list which will pass by reference file_name (img name): Will display the image name """ super().__init__() # create scene to add to and visualise it self.scene = QGraphicsScene() self.view = QGraphicsView(self.scene) self.setWindowTitle(file_name) # create layout manager add the widget to it self.layout = QGridLayout() self.layout.addWidget(self.view, 0, 0) self.setLayout(self.layout) # set layout manager for widget # image in format pyqt needs # create pixmap item and add to scene img = img.copy() image = QImage(img, img.shape[1], img.shape[0], QImage.Format.Format_Grayscale8) self.pixmap_item = QGraphicsPixmapItem(QPixmap(image)) self.scene.addItem(self.pixmap_item) # instructions button self.help_button = QPushButton("Help") self.help_button.setCheckable(True) self.help_button.toggle() self.help_button.clicked.connect(self.help_button_state) self.layout.addWidget(self.help_button, 0, 1) # overload mouse press event for pixmap item to our function self.pixmap_item.mousePressEvent = self.label_cell # set pen and brush for labels self.pen = QPen(QColor("red")) self.brush = QBrush(QColor("red")) # create empty set of marker coordinates self.marker_coords = coords
[docs] def help_button_state(self): """If help button is clicked display alert widget which details the help""" alert = QMessageBox() alert.setWindowTitle("Help") alert.setText( "Left click: add a seed \n" "Right click: remove a seed\n" "When closing all seed coordinates will be saved!" ) alert.exec()
[docs] def label_cell(self, event): """On mouse click this definition will be called Left click: add seed Right click: remove seed Args: event (QEvent): Event triggered by click""" # add label on left click if event.button() == Qt.MouseButton.LeftButton: self.scene.addEllipse( event.pos().x(), event.pos().y(), 10, 10, self.pen, self.brush ) # delete label on right click if event.button() == Qt.MouseButton.RightButton: # get list of items at that location items = self.scene.items( event.pos().x(), event.pos().y(), 1, 1, Qt.ItemSelectionMode.IntersectsItemShape, Qt.SortOrder.AscendingOrder, ) # if item is ellipse (type 4) remove it (don't remove image!) for item in items: if item.type() == 4: self.scene.removeItem(item)
[docs] def get_coords(self): """On closing the widget the coordinates of the remaining seeds are extracted""" # perform when closed items = self.scene.items() warnings.warn( "Not accurate - needs to be adjusted before exact coordinate is returned" ) # if ellipse add item coordinates to list for item in items: # if ellipse if item.type() == 4: # note not sure if lines up correct position - not issue here # i.e. if click at 100,240 is marker placed at 100,240? # and if its the centre of the marker? shouldn't we return # centre of marker # as the location # flag as issue!!! # but would be for accurate segmentation pos = item.pos() pos += item.rect().topLeft() # marker_coords is [H,W] relative to top left of image # i.e. if click bottom left of image marker_coords = [470,10] # this is what is returned to user self.marker_coords.append((int(pos.y()), int(pos.x())))
[docs] def closeEvent(self, event): """When closing the widget this function will be overloaded. It will ask the user if they are sure, on closing the coords of the markers will be extracted Args: event (QEvent): Event triggered by click""" # Confirm user choice reply = QMessageBox() reply.setText("Are you sure you want to close? (Markers are saved on closing)") reply.setStandardButtons( QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No ) reply = reply.exec() if reply == QMessageBox.StandardButton.Yes: # get the coordinates of the markers self.get_coords() event.accept() else: event.ignore()
[docs] def get_markers(img, file_name="Image") -> list: """Get markers for the image for watershed_segment using PyQt6 widget Image will be converted to [0 255] greyscale image for compatibility with PyQt6 widget Args: img (np.ndarray): Image for which need markers for watershed file_name (string): Name of the file Returns: marker_coords (list): List of tuples containing coordinates of markers (h,w) relative to top left of image""" app = QApplication([]) # sys.argv if need command line inputs # scale img to [0 255] and convert to int8 img = (img - np.min(img)) / (np.max(img) - np.min(img)) * 255 img = img.astype("uint8") # marker coords marker_coords = [] # create widget widget = WatershedWidget(img, coords=marker_coords, file_name=file_name) widget.show() app.exec() # note marker coords are (h,w) relative to top left of image i.e. # if click bottom left of image # marker_coords = [470,10] return marker_coords
[docs] def watershed_segment(img, file_name="Image", coords=None) -> np.ndarray: """Perform watershed segmentation on image - with option either to provide coordinates of markers (coords) or obtain annotation using a widget. Args: img (np.ndarray): Image which performing watershed on file_name (string): Name of the image coords (list): List of tuples where each tuple represents coordinate of a marker (x,y) Returns: labels (np.ndarray): Numpy array containing integer labels, each representing different segmented region of the input """ markers = np.zeros(img.shape, dtype="int32") # get markers if not provided if coords is None: coords = get_markers(img, file_name) # coords are (h,w) in image space # markers[row,col] where row is h-3:h+3 and col is w-3:w+3 # i.e. imagine if select marker in bottom left of image # in image space [h,w] with origin at top this would be e.g. (470,10) # return coordinate (470,10) from get_coords # markers [467:473,7:13] is populated i.e. height of 470 ish # and width 10 ish is populated # this ensures img and marker coords are in same space for index, coord in enumerate(coords): markers[coord[0] - 3 : coord[0] + 3, coord[1] - 3 : coord[1] + 3] = int( index + 1 ) # perform watershed labels = watershed(img, markers=markers) return labels