Krita/krita/pykrita/custompreview/custompreview.py

import math
import random
import re, time, datetime
import numpy
import io
import cv2
from PIL import Image
from pathlib import Path
import gmic
import json
from skimage.color import rgb2lab, rgba2rgb
import qimage2ndarray

from PyQt5.QtCore import Qt, pyqtSignal, QEvent, QBuffer
from PyQt5.QtGui import QImage, QPainter, QPixmap, QIcon, QInputEvent, QTabletEvent, QMouseEvent, \
  QHelpEvent, QPaintEvent, QTransform, QPixelFormat
from PyQt5.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QFileDialog, QLabel, QSizePolicy, \
  QScrollArea, QAction, QToolButton, QMdiArea, QAbstractScrollArea, QCheckBox, QComboBox, QLineEdit
from PyQt5.QtCore import QTimer, QByteArray
from krita import Krita, DockWidget, DockWidgetFactory, DockWidgetFactoryBase, Extension
from .util import dump_tablet_event, dump_mouse_event, get_qview, rgb2hsv

from .gmic_filters import GMIC_CMDS, GMIC_CMDS_UNSUPPORTED

KI = Krita.instance()
PREVIEW_MIN = 800

def clickable(widget, mouse_num):
  class Filter(QWidget):
    clicked = pyqtSignal(name="click")

    def eventFilter(self, obj, event):
      if obj == widget:
        if event.type() == QEvent.MouseButtonRelease:
          # print("b", event.button(), mouse_num)
          if event.button() == mouse_num and obj.rect().contains(event.pos()):
            # print(",")
            self.clicked.emit()
            # The developer can opt for .emit(obj) to get the object within the slot.
            return True
      return False

  filter_ = Filter(widget)
  widget.installEventFilter(filter_)
  return filter_.clicked


def hoverable(widget):
  class Filter(QWidget):
    clicked = pyqtSignal(name="click")

    def eventFilter(self, obj, event):
      if obj == widget:
        if isinstance(event, QPaintEvent):
          self.clicked.emit()
          return True
      return False

  filter_ = Filter(widget)
  widget.installEventFilter(filter_)
  return filter_.clicked

aww = None

class CustomPreview(DockWidget):
  def __init__(self):
    super().__init__()

    self.setWindowTitle(Krita.krita_i18n("Preview"))

    self.greyscale = False
    self.rot = 0
    self.locked = False
    self.locks = {}
    self.lastRender = 0


    layout = QVBoxLayout()
    layout.setAlignment(Qt.AlignCenter)

    self._target_qcanvas = None
    self._window = window = None

    # PREVIEW
    self.previewContainer = QWidget()
    clickable(self.previewContainer, 4).connect(self.toggleFlip)
    clickable(self.previewContainer, 2).connect(self.toggleLock)
    # clickable(self.previewContainer, 1).connect(self.blah)
    # hoverable(self.previewContainer).connect(self.blah)
    layout.addWidget(self.previewContainer)
    self.previewContainer.setContentsMargins(0, 0, 0, 0)
    previewContainerLayout = QHBoxLayout()
    previewContainerLayout.setContentsMargins(0, 0, 0, 0)
    previewContainerLayout.setSpacing(0)
    self.previewContainer.setLayout(previewContainerLayout)
    self.scrollArea = QScrollArea()
    previewContainerLayout.addWidget(self.scrollArea)
    self.scrollArea.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
    self.scrollArea.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
    self.scrollArea.setWidgetResizable(True)
    self.scrollArea.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)

    self.previewContainer.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
    self.previewContainer.setMinimumSize(PREVIEW_MIN, PREVIEW_MIN)

    self.previewLabel = QLabel()
    self.scrollArea.setWidget(self.previewLabel)

    clickable(self.scrollArea, 1).connect(self.blah)
    hoverable(self.scrollArea).connect(self.blah)

    # BUTTONS
    self.buttonLayout = QHBoxLayout()
    layout.addLayout(self.buttonLayout)
    self.buttonLayout.setAlignment(Qt.AlignLeft)

    self.dropdown = QComboBox(self)
    self.dropdown.addItem("Normal")
    self.dropdown.addItem("Lab(L)")
    self.dropdown.addItem("HSV(S)")
    self.dropdown.addItem("HSV(V)")
    self.dropdown.addItem("HSV(H)(1)")
    self.dropdown.addItem("HSV(H)(2)")
    self.dropdown.addItem("HSV(H)(3)")
    self.dropdown.addItem("NTSC")
    for item in GMIC_CMDS:
      self.dropdown.addItem(item.split(" ")[0], item)

    self.buttonLayout.addWidget(self.dropdown)
    self.dropdown.currentIndexChanged.connect(self.blah)

    self.b1 = QCheckBox("G")
    self.b1.setChecked(False)
    self.b1.stateChanged.connect(self.toggleGrey)
    self.buttonLayout.addWidget(self.b1)
    self.b1.setVisible(False)

    self.b2 = QLabel("0°")
    self.buttonLayout.addWidget(self.b2)

    self.b4 = QLabel(str(datetime.timedelta(seconds=0)))
    self.buttonLayout.addWidget(self.b4)

    self.b3 = QCheckBox("L")
    self.b3.setChecked(False)
    # self.b3.setAttribute(Qt.WA_TransparentForMouseEvents)
    # self.b3.setFocusPolicy(Qt.NoFocus)
    self.b3.stateChanged.connect(self.toggleLock)
    self.buttonLayout.addWidget(self.b3)

    self.wInput = QLineEdit()
    self.wInput.setText("800")
    self.wInput.textChanged.connect(self.setPreviewWidth)

    self.buttonLayout.addWidget(self.wInput)

    mainWidget = QWidget(self)
    mainWidget.setLayout(layout)
    self.setWidget(mainWidget)
    self.startTimer(300)  # refresh every 250 ms

    global aww
    aww = self

  def setPreviewWidth(self):
    i = self.wInput.text()
    if i.isnumeric():
      self.previewContainer.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
      self.previewContainer.setMinimumSize(int(i), int(i))
    else:
      self.previewContainer.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)
      self.previewContainer.setMinimumSize(50, 50)

  def msg(self, message: str):
    KI.activeWindow().activeView().showFloatingMessage(message, QIcon(), 1000, 2)

  def set_filter_state(self, new_state):
    if new_state:
      # remove old
      if self._target_qcanvas is not None:
        self._target_qcanvas.removeEventFilter(self)
        self._target_qcanvas = None
      # install new
      view = self._window.activeView()
      if view.document() is not None:
        # view has actual Krita document.
        self._target_qcanvas = self.get_qcanvas(view.canvas())
        if self._target_qcanvas is None:
          return
        self._target_qcanvas.installEventFilter(self)
    else:
      if self._target_qcanvas is not None:
        self._target_qcanvas.removeEventFilter(self)
        self._target_qcanvas = None

  def get_qcanvas(self, canvas):
    qview = get_qview(canvas.view())
    if qview is None:
      return
    children = qview.findChild(QAbstractScrollArea).viewport().children()
    for c in children:
      cls_name = c.metaObject().className()
      # KisOpenGLCanvas2 or KisQPainterCanvas
      if 'Canvas' in cls_name:
        return c

  def on_active_view_changed(self):
    """
    first uninstall event filter from old qcanvas (if old qcanvas is not None)
    second install event filter to currently active view's canvas
    """
    if self._target_qcanvas is not None:
      self._target_qcanvas.removeEventFilter(self)
      self._target_qcanvas = None
    view = self._window.activeView()
    if view.document() is not None:
      # bound only to views that have documents.
      self._target_qcanvas = self.get_qcanvas(view.canvas())
      self._target_qcanvas.installEventFilter(self)

  def eventFilter(self, obj, e):
    # print("e:", e)
    if isinstance(e, QHelpEvent):
      self.blah()
    if isinstance(e, QInputEvent):
      # some sort of input event.
      #   QContextMenuEvent, QHoverEvent, QKeyEvent, QMouseEvent,
      #   QNativeGestureEvent, QTabletEvent, QTouchEvent, QWheelEvent
      if isinstance(e, QTabletEvent):
        pass
        # print("e:", e)
        # self.dump_tablet_event(e)
      elif isinstance(e, QMouseEvent):
        pass
        # print("e:", e)
        # self.dump_mouse_event(e)
    return False  # event was NOT consumed.

  def toggleGrey(self):
    self.greyscale = not self.greyscale
    self.blah()

  def nextGmicFilter(self):
    i = self.dropdown.currentIndex()
    # print('next', i, 0 if i + 1 >= len(self.dropdown) else i + 1 % len(self.dropdown))
    self.dropdown.blockSignals(True)
    self.dropdown.setCurrentIndex(0 if i + 1 >= len(self.dropdown) else i + 1 % len(self.dropdown))
    self.dropdown.blockSignals(False)

  def prevGmicFilter(self):
    i = self.dropdown.currentIndex()
    # print('prev', i, i - 1 if i > 1 else len(self.dropdown) - 1)
    self.dropdown.blockSignals(True)
    self.dropdown.setCurrentIndex(i - 1 if i >= 1 else len(self.dropdown) - 1)
    self.dropdown.blockSignals(False)

  def toggleLock(self):
    doc = KI.activeDocument()
    view = Krita.instance().activeWindow().activeView()
    filename = Path(view.document().fileName()).name
    parent_dir = Path(view.document().fileName()).parent.name
    j = parent_dir+"/"+filename
    if self.locked or j in self.locks:
      if j in self.locks:
        del self.locks[j]
        self.msg("Unlocked preview")
      self.locked = False
      self.blah()
      return
    else:
      self.msg("Locked preview")
      self.locked = True
    canvas = view.canvas()
    qview = get_qview(view)
    kis_canvas_controller = qview.findChild(QAbstractScrollArea)
    try:
      zoom = (canvas.zoomLevel() * 72.0) / doc.resolution()
    except RuntimeError as e:
      print(e)
      return
    factor = (1 / zoom)
    c_x = int(kis_canvas_controller.horizontalScrollBar().value() * factor)
    if c_x < 0: c_x = 0
    c_y = int(kis_canvas_controller.verticalScrollBar().value() * factor)
    if c_y < 0: c_y = 0
    o_w = int(kis_canvas_controller.width() * factor)
    if o_w > doc.width(): o_w = doc.width()
    o_h = int(kis_canvas_controller.height() * factor)
    if o_h > doc.height(): o_h = doc.height()
    self.locks[j] = {
      "mirror": canvas.mirror(),
      "factor": factor,
      "c_x": c_x,
      "c_y": c_y,
      "o_w": o_w,
      "o_h": o_h,
    }
    self.blah()

  def toggleFlip(self):
    self.rot = (int(self.rot) + 1) % 4
    self.b2.setText(f"{self.rot * 90}°")
    # self.msg(f"{self.rot * 90}°")
    self.blah(force=True)

  def canvasChanged(self, canvas):
    self.blah()

  def timerEvent(self, event):
    self.refresh()

  def resizeEvent(self, event):
    self.blah()

  def setup(self):
    print(".................")

    def _later():
      # this happens later, so now activeWindow should already be there.
      # just to be safe keep handle to window alive in python.
      self._window = Application.activeWindow()
      self._window.activeViewChanged.connect(self.on_active_view_changed)

    QTimer.singleShot(0, _later)

  def refresh(self):
    try:
      view = Krita.instance().activeWindow().activeView()
      canvas = view.canvas()
      if round(canvas.zoomLevel() * 0.72, 2) > 1:
        view.showFloatingMessage(
          f"ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ZOOMED IN TOO FAR. ",
          Krita.instance().icon("showColoringOff"),
          200,
          1
        )
      self._window = KI.activeWindow()
      self._window.activeViewChanged.connect(lambda: self.set_filter_state(True))
    except Exception as e:
      print(e)

  def blah(self, force=False):
    if self.lastRender + 1.2 > time.time() and not force:
      return
    self.lastRender = time.time()

    doc = KI.activeDocument()
    # return if no document is open
    if doc is None:
      self.setDisabled(True)
      self.previewLabel.setPixmap(QPixmap())
      return

    self.setDisabled(False)

    # update doc info display
    q = doc.annotation("krita2spine")  #type: QByteArray
    r = re.search(r"<editing-time>(\d+)", doc.documentInfo(), re.I | re.M)
    self.b4.setText(str(datetime.timedelta(seconds=0 if r is None else int(r.group(1))))+f" (🍿️{doc.currentTime()}/{doc.animationLength()}) {'' if q.isEmpty() else '[📺]'}")

    # get current drawing
    # previewImage = doc.projection(0, 0, doc.width(), doc.height())
    view = Krita.instance().activeWindow().activeView()
    canvas = view.canvas()
    qview = get_qview(view)
    if qview is None:
      return
    kis_canvas_controller = qview.findChild(QAbstractScrollArea)
    # transform = QTransform()
    filename = Path(view.document().fileName()).name
    parent_dir = Path(view.document().fileName()).parent.name
    j = parent_dir+"/"+filename
    if self.locked and j in self.locks:
      mirror = self.locks[j]['mirror']
      factor = self.locks[j]['factor']
      c_x = self.locks[j]['c_x']
      c_y = self.locks[j]['c_y']
      o_w = self.locks[j]['o_w']
      o_h = self.locks[j]['o_h']
      self.b3.blockSignals(True)
      self.b3.setChecked(True)
      self.b3.blockSignals(False)
    else:
      try:
        factor = 1 / ((canvas.zoomLevel() * 72.0) / doc.resolution())
      except RuntimeError as e:
        print(e)
        return
      c_x = int(kis_canvas_controller.horizontalScrollBar().value() * factor)
      if c_x < 0:
        c_x = 0
      c_y = int(kis_canvas_controller.verticalScrollBar().value() * factor)
      if c_y < 0:
        c_y = 0
      o_w = int(kis_canvas_controller.width() * factor)
      if o_w > doc.width(): o_w = doc.width()
      o_h = int(kis_canvas_controller.height() * factor)
      if o_h > doc.height(): o_h = doc.height()
      mirror = canvas.mirror()
      self.b3.blockSignals(True)
      self.b3.setChecked(False)
      self.b3.blockSignals(False)

    # msg(Krita.instance(), f"factor{factor} cx{c_x} cy{c_y} ow{o_w} oh{o_h} dw{doc.width()} dh{doc.height()}")
    # msg(Krita.instance(), f"factor{factor} cx{c_x} kccw{int(kis_canvas_controller.width())} ow{o_w} oh{o_h} dw{doc.width()} dh{doc.height()}")
    # msg(Krita.instance(), f"{canvas.rotation()}")
    if mirror:
      previewImage = doc.projection(int(((doc.width()) - o_w - c_x)), c_y, int(kis_canvas_controller.width() * factor), o_h)
    else:
      previewImage = doc.projection(c_x, c_y, o_w, o_h)  # type: QBuffer
      # previewImage = doc.projection(c_x, c_y if canvas.rotation() != 180 else int(((doc.height()) - o_h - c_y)), o_w, o_h)
    # pprint(("flipped:", self.flipped, "locked:", self.locked, "canvas.mirror:", mirror))
    # pprint(("o_w:", o_w, "o_h:", o_h, "c_x", c_x, "c_y", c_y))
    # pprint(self.locks)

    dim = self.previewContainer.contentsRect()
    width = dim.width() - self.scrollArea.contentsMargins().top() * 2
    height = dim.height() - self.scrollArea.contentsMargins().top() * 2
    previewImage = previewImage.scaled(width, height, Qt.KeepAspectRatio, Qt.SmoothTransformation)
    previewImage = previewImage.mirrored(canvas.mirror(), 0)  # type: QImage
    if self.rot:
      rot = canvas.rotation()
      transform = QTransform()
      transform.rotate(math.copysign(1, rot)*float(self.rot * 90))
      previewImage = previewImage.transformed(transform)

    choice = self.dropdown.currentText()
    data = self.dropdown.currentData()
    if choice != "Normal":
      buffer = QBuffer()
      buffer.open(QBuffer.ReadWrite)
      previewImage.save(buffer, "PNG")
      im = Image.open(io.BytesIO(buffer.data()))
      num_im = numpy.array(im)

      if choice == "NTSC":
        # ref = (num_im[:, :, 0] * 0.2126 + num_im[:, :, 1] * 0.7152 + num_im[:, :, 2] * 0.081 + 0.5).astype('uint8')
        ref = (num_im[:, :, 0] * 0.2126 + num_im[:, :, 1] * 0.7152 + num_im[:, :, 2] * 0.0722).astype('uint8')
      elif choice == "Lab(L)":
        num_im = rgb2lab(rgba2rgb(num_im))
        num_im[..., 1] = num_im[..., 2] = 0
        ref = (num_im[:, :, 0] / 100 * 255).astype('uint8')
      elif choice == "HSV(V)":
        # remove alpha
        num_im = numpy.delete(num_im, [3, 3], 2)
        num_im = rgb2hsv(num_im)
        ref = (num_im[:, :, 2]).astype('uint8')
      elif choice == "HSV(S)":
        # remove alpha
        num_im = numpy.delete(num_im, [3, 3], 2)
        num_im = rgb2hsv(num_im)
        ref = (num_im[:, :, 1] * -255).astype('uint8')
      elif choice == "HSV(H)(1)":
        # remove alpha
        num_im = numpy.delete(num_im, [3, 3], 2)
        num_im = rgb2hsv(num_im)
        ref = (num_im[:, :, 0] * (1.4225 + random.random())).astype('uint8')
      elif choice == "HSV(H)(2)":
        # remove alpha
        num_im = numpy.delete(num_im, [3, 3], 2)
        num_im = rgb2hsv(num_im)
        ref = (num_im[:, :, 0] * 2.845).astype('uint8')
      elif choice == "HSV(H)(3)":
        # remove alpha
        num_im = numpy.delete(num_im, [3, 3], 2)
        num_im = rgb2hsv(num_im)
        ref = (num_im[:, :, 0] * 0.71125).astype('uint8')
      elif choice == "G'Mic(Kuwahara)":
        g = gmic.GmicImage.from_numpy_helper(num_im, deinterleave=True)
        gmic.run("kuwahara 4", g)  # horizontal blur+special black&white
        ref = g.to_numpy_helper(astype=numpy.uint8, interleave=True, squeeze_shape=True)
        ref = numpy.swapaxes(ref, 0, 1)
      elif choice == "G'Mic(PWB)":
        g = gmic.GmicImage.from_numpy_helper(num_im, deinterleave=True)
        gmic.run(
          'fx_paint_with_brush 0,"0",1,16,30,100,100,10,80,0.5,3,45,0,6,2,10,0,0,0,60,20,1,0,0,0,30,15,15,15,15,15,15,1,45,0,0,50,50',
          g)  # horizontal blur+special black&white
        ref = g.to_numpy_helper(astype=numpy.uint8, interleave=True, squeeze_shape=True)
        ref = numpy.swapaxes(ref, 0, 1)
      elif choice == "G'Mic(Posterize)":
        g = gmic.GmicImage.from_numpy_helper(num_im, deinterleave=True)
        gmic.run('fx_posterize 67.2,15.9,1,2,0,0,1,0,50,50', g)  # horizontal blur+special black&white
        ref = g.to_numpy_helper(astype=numpy.uint8, interleave=True, squeeze_shape=True)
        ref = numpy.swapaxes(ref, 0, 1)
      elif choice == "G'Mic(Sharp Abstract)":
        g = gmic.GmicImage.from_numpy_helper(num_im, deinterleave=True)
        gmic.run('fx_sharp_abstract 4,10,0.5,0,0,50,50', g)  # horizontal blur+special black&white
        ref = g.to_numpy_helper(astype=numpy.uint8, interleave=True, squeeze_shape=True)
        ref = numpy.swapaxes(ref, 0, 1)
      for item in GMIC_CMDS:
        if data == item:
          g = gmic.GmicImage.from_numpy_helper(num_im, deinterleave=True)
          gmic.run(item, g)  # horizontal blur+special black&white
          ref = g.to_numpy_helper(astype=numpy.uint8, interleave=True, squeeze_shape=True)
          ref = numpy.swapaxes(ref, 0, 1)

      qimg = qimage2ndarray.array2qimage(ref)
      qimg.scaled(width, height)

      previewImage = qimg
      resultImage = qimg.convertToFormat(QImage.Format_ARGB32_Premultiplied)
      resultImage.fill(0)
      self.b1.blockSignals(True)
      self.b1.setChecked(True)
      self.b1.blockSignals(False)
    else:
      resultImage = QImage(previewImage.width(), previewImage.height(), QImage.Format_ARGB32_Premultiplied)
      resultImage.fill(0)
      self.b1.blockSignals(True)
      self.b1.setChecked(False)
      self.b1.blockSignals(False)
    painter = QPainter(resultImage)
    painter.setRenderHint(QPainter.Antialiasing, True)
    painter.drawImage(0, 0, previewImage)
    painter.end()

    self.previewLabel.setPixmap(QPixmap.fromImage(resultImage))
    self.scrollArea.setMaximumSize(previewImage.width() + 4, previewImage.height() + 4)


class CustomPreviewFtw(Extension):
  def __init__(self, parent):
    super(CustomPreviewFtw, self).__init__(parent)

  def setup(self):
    pass

  def initialize(self):
    pass

  def blah(self):
    global aww
    if aww: aww.blah(force=True)

  def flip(self):
    global aww
    if aww: aww.toggleFlip()

  def nextf(self):
    global aww
    if aww: aww.nextGmicFilter()

  def prevf(self):
    global aww
    if aww: aww.prevGmicFilter()

  def lock(self):
    global aww
    if aww: aww.toggleLock()

  def createActions(self, window):
    action = window.createAction("preview_refresh", str(f"Refresh preview"), "")
    action.triggered.connect(self.blah)
    action = window.createAction("preview_flip", str(f"Refresh flip preview"), "")
    action.triggered.connect(self.flip)
    action = window.createAction("preview_next_filter", str(f"Next filter preview"), "")
    action.triggered.connect(self.prevf)
    action = window.createAction("preview_prev_filter", str(f"Prev filter preview"), "")
    action.triggered.connect(self.nextf)
    action = window.createAction("preview_lock", str(f"Lock preview"), "")
    action.triggered.connect(self.lock)

KI.addDockWidgetFactory(DockWidgetFactory("customPreview", DockWidgetFactoryBase.DockRight, CustomPreview))