Source code for arthropod_describer.label_editor.label_view_widget

import time
import typing
from typing import Optional

import cv2
import numba
import numpy as np
import qimage2ndarray
from PySide2.QtCore import QRectF, Signal, QTimer
from PySide2.QtGui import QImage, QPainter, QColor, Qt, QBitmap, QRegion, QBrush, QTransform
from PySide2.QtWidgets import QGraphicsItem, QWidget, QStyleOptionGraphicsItem, QGraphicsSceneMouseEvent, \
    QGraphicsObject, QGraphicsSceneHoverEvent
from cv2 import cv2
from scipy import ndimage
from skimage import morphology as M

from arthropod_describer.common.label_change import CommandEntry
from arthropod_describer.common.photo import LabelImg
from arthropod_describer.common.state import State
from arthropod_describer.common.tool import Tool, EditContext, ToolCursor


[docs]class LabelView(QGraphicsObject): label_img_modified = Signal(CommandEntry) label_picked = Signal(int) label_hovered = Signal(int) def __init__(self, state: State, label_name: str): QGraphicsObject.__init__(self) self.viz_clip_region: QRegion = QRegion() self.viz_mask: QBitmap = None self.viz_mask_buffer: QImage = QImage() self.viz_mask_nd: np.ndarray = None self._label_img: LabelImg = None self._label_qimage: typing.Optional[QImage] = None self._nd_img = None # self.setAcceptedMouseButtons(Qt.NoButton) self.setAcceptedMouseButtons(Qt.LeftButton | Qt.RightButton) self.setAcceptHoverEvents(True) self.state: State = state self.showing_full_mask: bool = True self.outline_1px_nd: np.ndarray = None self.outline_nd: np.ndarray = None self.outline_mask: QBitmap = QBitmap() self.outline_mask_buffer: QImage = QImage() self.outline_width: int = 3 self.current_tool: typing.Optional[Tool] = None self.tool_cursor: Optional[ToolCursor] = None self.clip_region: QRegion = QRegion() self._clip_nd: np.ndarray = None self._clip_qimg = QImage() self.clip_mask = QBitmap() self.restricted_clip_region: QRegion = QRegion() #self.constraint_level: int = -1 #self.label_name: str = label_name self.tool_viz_commands = [] self.brush_offset_idx: int = 0 self.offset_time: int = 0 self.marching_ants_timer = QTimer(self) self.marching_ants_timer.setInterval(500) self.marching_ants_timer.timeout.connect(self.update) self.march_ant_texture: QImage = QImage(32, 32, QImage.Format_Mono) #painter = QPainter(self.march_ant_texture) #painter.fillRect(self.march_ant_texture.rect(), Qt.GlobalColor.black) #painter.fillRect(self.march_ant_texture.rect(), QBrush()) #painter.setBrush() self.state.new_label_constraint.connect(self.adapt_to_constraint)
[docs] def adapt_to_constraint(self, _): self.compute_clip_mask() self.compute_viz_mask()
[docs] def set_tool(self, tool: Tool, tool_cursor: Optional[ToolCursor]): self.current_tool = tool self.tool_cursor = tool_cursor
[docs] def set_label_name(self, label_name: str): lbl = self.state.current_photo[label_name] level = self.state.current_label_level self.set_label_image(lbl, level)
[docs] def set_label_image(self, mask: LabelImg, level: int): self.prepareGeometryChange() self._label_img = mask #self.setTransformOriginPoint(mask.size[1] * 0.5, mask.size[0] * 0.5) self._recolor_image() now = time.time() self.compute_clip_mask() print(f'computing clip mask took {time.time() - now} secs') now = time.time() self.outline_nd = 0 * np.ones_like(self.state.current_photo[self.state.current_label_name].label_image, dtype=np.uint8) print(f'outline_nd took {time.time() - now} secs') if not self.showing_full_mask: self.draw_outline() now = time.time() self.compute_viz_mask() print(f'computing viz mask took {time.time() - now} secs')
#def set_color_map(self, colormap: ColormapItemModel): # self._colormap = colormap # self._recolor_image() # self.update() def _recolor_image(self): if self._label_img is None or not self._label_img.is_set: self._label_qimage = None return if self._nd_img is None or self._label_img.label_image.shape != self._nd_img.shape: self._nd_img = np.zeros(self.state.current_photo[self.state.current_label_name].label_image.shape + (1,), np.uint32) self.outline_mask_buffer = QImage(self._nd_img.shape[1], self._nd_img.shape[0], QImage.Format_Mono) else: self._nd_img[:, :] = 0 self.outline_mask_buffer.fill(1) level_img = self.state.current_photo[self.state.current_label_name][self.state.current_label_level] used_labels = np.unique(level_img) cmap = {label: QColor.fromRgb(*self.state.colormap[label]).rgba() for label in used_labels} if 0 in used_labels: cmap[0] = QColor(0, 0, 0, 0).rgba() for label in used_labels: coords = np.nonzero(level_img == label) self._nd_img[coords] = cmap[label] self._label_qimage = QImage(self._nd_img.data, self._nd_img.shape[1], self._nd_img.shape[0], 4 * self._nd_img.shape[1], QImage.Format_ARGB32) # self.outline_mask = QImage(self._nd_img.shape[1], self._nd_img.shape[0], QImage.Format_Grayscale8) self.update()
[docs] def boundingRect(self): if self._label_qimage is None: return QRectF() return self._label_qimage.rect()
[docs] def paint(self, painter: QPainter, option: QStyleOptionGraphicsItem, widget: typing.Optional[QWidget] = ...): if self._label_qimage is not None: painter.save() if self.viz_mask is not None: painter.setClipRegion(self.viz_clip_region, Qt.ReplaceClip) #if not self.showing_full_mask: # reg = QRegion(self.outline_mask) # painter.setClipRegion(reg, Qt.ReplaceClip) painter.drawImage(option.rect, self._label_qimage) #if not self.showing_full_mask: # painter.setBrush(QBrush(QColor.fromRgb(0, 0, 0, 125), Qt.FDiagPattern)) # if time.time() - self.offset_time > 0.1: # self.brush_offset_idx = (self.brush_offset_idx + 1) % 8 # self.offset_time = time.time() # painter.setBrushOrigin(self.brush_offset_idx * 10, 0) # painter.fillRect(self.boundingRect(), painter.brush()) painter.setClipRegion(self.restricted_clip_region, Qt.ReplaceClip) painter.setRenderHint(QPainter.HighQualityAntialiasing, True) painter.fillRect(self.boundingRect(), QBrush(QColor.fromRgb(0, 0, 0, 200))) #, Qt.Dense5Pattern)) painter.setRenderHint(QPainter.HighQualityAntialiasing, False) painter.restore() if len(self.tool_viz_commands) > 0: for viz_cmd in self.tool_viz_commands: viz_cmd(painter)
[docs] def switch_label_level(self, level: int): self._recolor_image() if not self.showing_full_mask: self.draw_outline()
#self.compute_clip_mask() #mask_level = self.state.current_label_level - 1 if self.state.current_label_level > 0 else 0 #if self.state.current_label_level == 0 and False: # clip_mask = 255 * np.ones_like(self.state.current_photo.regions_image.label_img, dtype=np.uint8) # self._clip_nd = np.require(clip_mask, np.uint8, 'C') # self._clip_qimg = QImage(self._clip_nd.data, self._clip_nd.shape[1], self._clip_nd.shape[0], # self._clip_nd.strides[0], QImage.Format_Grayscale8) #else: # mask = self.state.current_photo.regions_image[mask_level] # clip_mask = 255 * (mask > 0).astype(np.uint8) # self._clip_nd = np.require(clip_mask, np.uint8, 'C') #self._clip_qimg = QImage(self._clip_nd.data, self._clip_nd.shape[1], self._clip_nd.shape[0], # self._clip_nd.strides[0], QImage.Format_Grayscale8) #self._clip_qimg.invertPixels() #self.clip_mask = QBitmap.fromImage(self._clip_qimg, Qt.AutoColor) #self.clip_region = QRegion(self.clip_mask)
[docs] def draw_outline(self, bbox: typing.Optional[typing.Tuple[int, int, int, int]] = None, compute_from_viz: bool = True): print('redraw outline - START') now = time.time() level_img = self.state.current_photo[self.state.current_label_name][self.state.current_label_level] #print(f'lv took {time.time() - now}') if bbox is not None: bbox = list(bbox) bbox[0] -= 10 bbox[1] += 10 bbox[2] -= 10 bbox[3] += 10 if compute_from_viz: level_img = qimage2ndarray.raw_view(self._label_qimage.copy(bbox[2], bbox[0], bbox[3] - bbox[2] + 1, bbox[1] - bbox[0] + 1)) else: level_img = level_img[bbox[0]:bbox[1]+1, bbox[2]:bbox[3]+1] outline = compute_outline(level_img).astype(np.uint8) self.outline_1px_nd[bbox[0]+9:bbox[1] - 10 + 1, bbox[2]+9:bbox[3] - 10 + 1] = outline[9:-10, 9:-10] selem = M.diamond(1) if self.outline_width > 1: self.outline_nd = cv2.morphologyEx(self.outline_1px_nd, cv2.MORPH_ERODE, selem, # cv2.getStructuringElement(cv2.MORPH_CROSS, # 2 * (2 * self.outline_width + 1,)), borderType=cv2.BORDER_CONSTANT, borderValue=0, iterations=self.outline_width-1) else: self.outline_nd = self.outline_1px_nd self.outline_mask_buffer = QImage(self.outline_nd.data, self.outline_nd.shape[1], self.outline_nd.shape[0], self.outline_nd.strides[0], QImage.Format_Grayscale8) self.outline_mask = QBitmap.fromImage(self.outline_mask_buffer, Qt.MonoOnly) else: if compute_from_viz: level_img = qimage2ndarray.raw_view(self._label_qimage) self.outline_1px_nd = compute_outline(level_img).astype(np.uint8) selem = M.diamond(1) if self.outline_width > 1: self.outline_nd = cv2.morphologyEx(self.outline_1px_nd, cv2.MORPH_ERODE, selem, #cv2.getStructuringElement(cv2.MORPH_CROSS, # 2 * (2 * self.outline_width + 1,)), borderType=cv2.BORDER_CONSTANT, borderValue=0, iterations=self.outline_width-1) else: self.outline_nd = self.outline_1px_nd self.outline_mask_buffer = QImage(self.outline_nd.data, self.outline_nd.shape[1], self.outline_nd.shape[0], self.outline_nd.strides[0], QImage.Format_Grayscale8) self.outline_mask = QBitmap.fromImage(self.outline_mask_buffer, Qt.MonoOnly) self.compute_viz_mask(bbox) self.update()
[docs] def handle_outline_width_changed(self, value: int): if value == self.outline_width: return selem = M.diamond(1) if value > 1: self.outline_nd = cv2.morphologyEx(self.outline_1px_nd, cv2.MORPH_ERODE, selem, iterations=value-1) else: self.outline_nd = self.outline_1px_nd self.outline_mask_buffer = QImage(self.outline_nd.data, self.outline_nd.shape[1], self.outline_nd.shape[0], self.outline_nd.shape[1], QImage.Format_Grayscale8) self.outline_mask = QBitmap.fromImage(self.outline_mask_buffer, Qt.MonoOnly) self.outline_width = value self.compute_viz_mask() self.update()
[docs] def show_outline(self): self.draw_outline() self.showing_full_mask = False
#self.compute_viz_mask() #if not self.marching_ants_timer.isActive(): # self.marching_ants_timer.start()
[docs] def show_mask(self): self._recolor_image() self.outline_nd = np.zeros_like(self._label_img.label_image, dtype=np.uint8) self.compute_viz_mask() self.showing_full_mask = True
#if self.marching_ants_timer.isActive(): # self.marching_ants_timer.stop()
[docs] def mousePressEvent(self, event: QGraphicsSceneMouseEvent): if event.buttons() & Qt.MiddleButton: QGraphicsItem.mousePressEvent(self, event) return elif event.button() & Qt.LeftButton: if self.current_tool is not None: hovered_label = self.state.current_photo[self.state.current_label_name][self.state.current_label_level][event.pos().toPoint().toTuple()[::-1]] self.current_tool.left_press(None, event.pos().toPoint(), self._create_context()) self.update()
# QGraphicsItem.mousePressEvent(self, event)
[docs] def mouseMoveEvent(self, event: QGraphicsSceneMouseEvent): if event.button() & Qt.RightButton: return elif event.buttons() & Qt.LeftButton: if self.current_tool is not None: ctx = self._create_context() _, rect = self.current_tool.mouse_move(None, event.pos().toPoint(), event.lastPos().toPoint(), ctx) bbox = [rect.top(), rect.top() + rect.height(), rect.left(), rect.left() + rect.width()] self.tool_viz_commands = ctx.tool_viz_commands if not self.showing_full_mask: self.draw_outline(bbox) if self.tool_cursor is not None: self.tool_cursor.setPos(event.pos()) self.tool_cursor.setVisible(False) self.state.viz_layer.update() self.update()
[docs] def mouseReleaseEvent(self, event: QGraphicsSceneMouseEvent): if event.button() & Qt.RightButton: if self.current_tool is not None: self.current_tool.right_release(None, event.pos().toPoint(), self._create_context()) self.label_picked.emit(self._label_img[self.state.current_label_level][event.pos().toPoint().toTuple()[::-1]]) return if self.tool_cursor is not None: cmd, bbox = self.current_tool.left_release(None, event.pos().toPoint(), self._create_context()) self.tool_viz_commands = [] if cmd is not None: self.label_img_modified.emit(cmd) self.tool_cursor.setVisible(True) self.update() self.state.viz_layer.update()
# QGraphicsItem.mouseReleaseEvent(self, event)
[docs] def hoverEnterEvent(self, event: QGraphicsSceneHoverEvent): if self.tool_cursor is not None: self.tool_cursor.setVisible(True)
[docs] def hoverMoveEvent(self, event: QGraphicsSceneHoverEvent): ev = event.pos().toPoint() if not (0 <= ev.x() < self._label_qimage.width()) or not (0 <= ev.y() < self._label_qimage.height()): return if self.current_tool is not None and self.tool_cursor is not None: self.tool_cursor.setPos(event.pos()) self.tool_cursor.setVisible(True) self.update() label = self.state.current_photo[self.state.current_label_name].label_image[ev.toTuple()[::-1]] leveled_label = self.state.label_hierarchy.level_mask(self.state.current_label_level) & label self.label_hovered.emit(leveled_label) lab_img = self.state.current_photo[self.state.current_label_name] if (reg_props := lab_img.get_region_props(label)) is not None: tooltip = '' for reg_prop in reg_props.values(): tooltip = tooltip + str(reg_prop) + '\n' self.setToolTip(tooltip) else: self.setToolTip('')
[docs] def hoverLeaveEvent(self, event: QGraphicsSceneHoverEvent): if self.tool_cursor is not None: self.tool_cursor.setVisible(False) self.label_hovered.emit(-1)
#def set_constraint_level(self, level: int): # #self.constraint_level = self.state.current_constraint.label_level # if self.state.current_photo is None: # return # if self.state.redraw_canvas: # self.compute_clip_mask() # self.compute_viz_mask()
[docs] def handle_primary_label_changed(self): if self.state.redraw_canvas: self.compute_clip_mask() self.compute_viz_mask()
def _create_context(self): return EditContext(self._label_img, self.state.primary_label, self.state.current_photo.image, self.state.colormap, self._label_qimage, self.state.current_photo, self.state.current_label_level, self._clip_nd, self.clip_region, self._clip_nd)
[docs] def compute_clip_mask(self): print('computing now clip mask') if self.state.current_photo is None: return if self.state.constraint_label == 0: clip_mask = 255 * np.ones(self.state.current_photo[self.state.current_label_name].label_image.shape, dtype=np.uint8) else: #lab_hier = self.state.storage.get_label_hierarchy(lbl_type) lab_hier = self.state.label_hierarchy level_img = self.state.current_photo[self.state.current_constraint.label_name][lab_hier.get_level(self.state.constraint_label)]#[self.state.current_constraint.label_level] hier_mask = lab_hier.label_mask(self.state.constraint_label) prefix = self.state.constraint_label & hier_mask clip_mask = 255 * (level_img == prefix).astype(np.uint8) self._clip_nd = np.require(clip_mask, np.uint8, 'C') self._clip_qimg = QImage(self._clip_nd.data, self._clip_nd.shape[1], self._clip_nd.shape[0], self._clip_nd.strides[0], QImage.Format_Grayscale8) #self._clip_qimg.save('C:\\Users\\radoslav\\Desktop\\clip_mask.png') self._clip_qimg.invertPixels() self.clip_mask = QBitmap.fromImage(self._clip_qimg, Qt.AutoColor) self.clip_region = QRegion(self.clip_mask)
[docs] def compute_viz_mask(self, bbox: Optional[typing.Tuple[int, int, int, int]] = None): #print('computing viz mask') #io.imsave('/home/radoslav/arth_clip_nd.png', self._clip_nd) #io.imsave('/home/radoslav/arth_outline_mask.png', self.outline_nd) #io.imsave('/home/radoslav/arth_not_clip.png', ~self._clip_nd) #io.imsave('/home/radoslav/arth_not_outline.png', ~self.outline_nd) #io.imsave('/home/radoslav/arth_nclip_and_noutline.png', ~np.bitwise_and(~self._clip_nd, ~self.outline_nd)) now = time.time() if bbox is not None: self.viz_mask_nd[bbox[0]:bbox[1]+1, bbox[2]:bbox[3]+1] = ~np.bitwise_and( ~self._clip_nd[bbox[0]:bbox[1]+1, bbox[2]:bbox[3]+1], ~self.outline_nd[bbox[0]:bbox[1]+1, bbox[2]:bbox[3]+1] ) else: self.viz_mask_nd = ~np.bitwise_and(~self._clip_nd, ~self.outline_nd) self.viz_mask_buffer = QImage(self.viz_mask_nd.data, self.viz_mask_nd.shape[1], self.viz_mask_nd.shape[0], self.viz_mask_nd.strides[0], QImage.Format_Grayscale8) self.viz_mask = QBitmap.fromImage(self.viz_mask_buffer) #self.viz_mask.save('/home/radoslav/arth_viz_mask.png') self.viz_clip_region = QRegion(self.viz_mask) b = QBitmap(self.viz_mask.size()) b.fill(QColor.fromRgb(0, 0, 0, 0)) #b.save('/home/radoslav/arth_viz_clip_reg.png') reg = QRegion(b) self.restricted_clip_region = self.viz_clip_region ^ reg if self.showing_full_mask else self.clip_region ^ reg #io.imsave('/home/radoslav/arth_clip_and_outline.png', self.viz_mask_nd) #print(f'computing {"with" if bbox is not None else "without"} bbox took {time.time() - now} secs') self.update()
[docs] def rotate(self, ccw: bool = True): transform = QTransform() transform.rotate(90 * (-1 if ccw else 1)) self.prepareGeometryChange() self._label_qimage = self._label_qimage.transformed(transform) self.viz_mask = self.viz_mask.transformed(transform) self.viz_clip_region = QRegion(self.viz_mask) self.viz_mask_buffer = self.viz_mask_buffer.transformed(transform) self.viz_mask_nd = ndimage.rotate(self.viz_mask_nd, 90 if ccw else -90) #T.rotate(self.viz_mask_nd, 90 * (-1 if ccw else 1), order=0) self._clip_nd = ndimage.rotate(self._clip_nd, 90 if ccw else -90) #T.rotate(self._clip_nd, 90 * (-1 if ccw else 1)) #self._clip_qimg = self._clip_qimg.transformed(transform) self.clip_mask = self.clip_mask.transformed(transform) self.clip_region = QRegion(self.clip_mask) b = QBitmap(self.viz_mask.size()) b.fill(QColor.fromRgb(0, 0, 0, 0)) reg = QRegion(b) self.restricted_clip_region = self.viz_clip_region ^ reg if self.showing_full_mask else self.clip_region ^ reg self.update()
[docs]@numba.njit def compute_outline(label_img: np.ndarray) -> np.ndarray: #result = np.zeros(label_img.shape + (1,), dtype=label_img.dtype) result = 255 * np.ones(label_img.shape, dtype=np.uint8) d = [-1, 0, 1] for y in range(label_img.shape[0]): for x in range(label_img.shape[1]): if label_img[y, x] == 0: continue lab = label_img[y, x] for k in d: for j in d: if k == 0 and j == 0 or outside(y + k, x + j, label_img.shape) or abs(k) + abs(j) > 1: continue if label_img[y + k, x + j] != lab: result[y, x] = 0 # label_img[y, x] return result
[docs]@numba.njit def outside(y, x, shape): return y < 0 or y >= shape[0] or x < 0 or x >= shape[1]