| from typing import Tuple, List, Union | |
| import numpy as np | |
| import cv2 | |
| from qtpy.QtCore import QRectF, Qt, QPointF, QSize | |
| from qtpy.QtWidgets import QStyleOptionGraphicsItem, QGraphicsPixmapItem, QWidget, QGraphicsItem | |
| from qtpy.QtGui import QPen, QPainter, QPixmap, QImage, QBrush | |
| from .misc import pixmap2ndarray | |
| SIZE_MAX = 2147483647 | |
| class ImageEditMode: | |
| NONE = 0 | |
| HandTool = 0 | |
| InpaintTool = 1 | |
| PenTool = 2 | |
| RectTool = 3 | |
| class PenShape: | |
| Circle = 0 | |
| Rectangle = 1 | |
| Triangle = 2 | |
| class StrokeImgItem(QGraphicsItem): | |
| def __init__(self, pen: QPen, point: QPointF, size: QSize, format: QImage.Format = QImage.Format.Format_ARGB32, shape=PenShape.Circle): | |
| super().__init__() | |
| self._img = QImage(size, format) | |
| self._img.fill(Qt.GlobalColor.transparent) | |
| pen = QPen(pen) | |
| if shape == PenShape.Rectangle: | |
| pen.setJoinStyle(Qt.PenJoinStyle.MiterJoin) | |
| self.pen = pen | |
| self._d = d = pen.widthF() | |
| self._d_rect = d // 32 | |
| self._r = d / 2 | |
| self.clipped_rect = None | |
| self.shape = shape | |
| self._line_to = [self._line_to_circle, self._line_to_rectangle][shape] | |
| self.painter = QPainter(self._img) | |
| self.painter.setRenderHint(QPainter.RenderHint.Antialiasing) | |
| self.painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Source) | |
| if shape != PenShape.Circle: | |
| pen.setWidthF(0) | |
| self.painter.setPen(pen) | |
| self.painter.setBrush(pen.color()) | |
| self.setBoundingRegionGranularity(0) | |
| self.cur_point = point | |
| self._br = QRectF(0, 0, size.width(), size.height()) | |
| self.is_painting = True | |
| min_x = self.cur_point.x() - self._r | |
| min_y = self.cur_point.y() - self._r | |
| if shape == PenShape.Circle: | |
| self._line_to(self.cur_point, None) | |
| else: | |
| self._line_to(self.cur_point, None) | |
| rect = QRectF(min_x, min_y, self._d, self._d) | |
| self.init_rect = rect | |
| self.update(rect) | |
| def finishPainting(self): | |
| self.painter.end() | |
| self.is_painting = False | |
| def clip(self, mask_only=False, format=QImage.Format.Format_ARGB32_Premultiplied) -> Tuple[List, np.ndarray, QImage]: | |
| img_array = pixmap2ndarray(self._img, True) | |
| ar = cv2.boundingRect(cv2.findNonZero(img_array[..., -1])) | |
| img_array = img_array[ar[1]: ar[1] + ar[3], ar[0]: ar[0] + ar[2]] | |
| if not (ar[2] > 0 and ar[3] > 0): | |
| return None, None, None | |
| if mask_only: | |
| img_array = img_array[..., -1] | |
| img_array[img_array > 0] = 255 | |
| return ar, img_array, self._img.copy(*ar).convertToFormat(format) | |
| def startNewPoint(self, pos: QPointF): | |
| self.is_painting = True | |
| self.painter.begin(self._img) | |
| self.painter.setPen(self.pen) | |
| self.painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Source) | |
| self.cur_point = pos | |
| self.lineTo(pos) | |
| def boundingRect(self) -> QRectF: | |
| return self._br | |
| def _line_to_circle(self, pnt1: QPointF, pnt2: QPointF): | |
| if pnt2 is not None: | |
| self.painter.drawLine(pnt1, pnt2) | |
| else: | |
| pen = QPen(self.pen) | |
| pen.setWidthF(0) | |
| self.painter.setPen(pen) | |
| self.painter.setBrush(self.pen.color()) | |
| rect = QRectF(pnt1.x() - self._r, pnt1.y() - self._r, self._d, self._d) | |
| self.painter.drawEllipse(rect) | |
| self.painter.setPen(self.pen) | |
| def _line_to_rectangle(self, pnt1: QPointF, pnt2: QPointF): | |
| shape_rect = QRectF(pnt1.x() - self._r, pnt1.y() - self._r, self._d, self._d) | |
| self.painter.drawRect(shape_rect) | |
| def lineTo(self, new_pnt: QPointF, update=True) -> QRectF: | |
| delta = self.cur_point - new_pnt | |
| delta_w, delta_h = abs(delta.x()), abs(delta.y()) | |
| rect = None | |
| if delta_w + delta_h > 1: | |
| min_x = min(self.cur_point.x(), new_pnt.x()) - self._r | |
| min_y = min(self.cur_point.y(), new_pnt.y()) - self._r | |
| delta_w += self._d | |
| delta_h += self._d | |
| rect = QRectF(min_x, min_y, delta_w, delta_h) | |
| self._line_to(self.cur_point, new_pnt) | |
| self.cur_point = new_pnt | |
| if update: | |
| self.update(rect) | |
| return rect | |
| def paint(self, painter: QPainter, option: QStyleOptionGraphicsItem, widget: QWidget) -> None: | |
| painter.drawImage(0, 0, self._img) | |
| class PixmapItem(QGraphicsPixmapItem): | |
| def __init__(self, border_pen: QPen, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.border_pen = border_pen | |
| def paint(self, painter: QPainter, option: 'QStyleOptionGraphicsItem', widget: QWidget) -> None: | |
| pen = painter.pen() | |
| painter.setPen(self.border_pen) | |
| painter.drawRect(self.boundingRect()) | |
| painter.setPen(pen) | |
| return super().paint(painter, option, widget) | |
| class DrawingLayer(QGraphicsPixmapItem): | |
| def __init__(self): | |
| super().__init__() | |
| self.qimg_dict = {} | |
| self.drawing_items_info = {} | |
| self.drawed_pixmap = None | |
| def addQImage(self, x: int, y: int, qimg: QImage, compose_mode, key: str): | |
| self.qimg_dict[key] = qimg | |
| self.drawing_items_info[key] = {'pos': [x, y], 'compose': compose_mode} | |
| self.update() | |
| def removeQImage(self, key: str): | |
| if key in self.qimg_dict: | |
| self.qimg_dict.pop(key) | |
| self.drawing_items_info.pop(key) | |
| def paint(self, painter: QPainter, option: QStyleOptionGraphicsItem, widget: QWidget): | |
| pixmap = self.pixmap() | |
| if pixmap.isNull(): | |
| self.drawed_pixmap = None | |
| return | |
| p = QPainter() | |
| p.begin(pixmap) | |
| for key in self.qimg_dict: | |
| item = self.qimg_dict[key] | |
| info = self.drawing_items_info[key] | |
| if isinstance(item, QImage): | |
| p.setCompositionMode(info['compose']) | |
| p.drawImage(info['pos'][0], info['pos'][1], item) | |
| p.end() | |
| painter.drawPixmap(self.offset(), pixmap) | |
| self.drawed_pixmap = pixmap | |
| def get_drawed_pixmap(self, format=QImage.Format.Format_ARGB32) -> QPixmap: | |
| pixmap = self.pixmap() if self.drawed_pixmap is None else self.drawed_pixmap | |
| return pixmap | |
| def drawed(self) -> bool: | |
| return len(self.qimg_dict) > 0 | |
| def clearAllDrawings(self): | |
| self.qimg_dict.clear() | |
| self.drawing_items_info.clear() | |