picks color under the cursor by shortcut

krita-pickcolor-shortcut.py

from krita import *  # type: ignore
from PyQt5.QtCore import QPointF
from PyQt5.QtGui import QImage, QTransform, QCursor
from PyQt5.QtWidgets import QMdiArea, QAbstractScrollArea
import math

'''
pickcolor shortcut for Krita: 
picks color under the cursor by shortcut
This script is based on https://github.com/seguso/KritaColorPlus

Install:
    1. Save or download this script.
    2. Place it in any preferred folder.
    3. Name it e.g. "hsv-shifter.py".
    4. On Krita use ten script plugin to assing this script to a shortcut.

How to use:
    1. fire the shortcut and it picks the color under the cursor.
'''

g_last_coord_mouse_down = None
g_last_coord_mouse_up = None
g_virtual_fg_color_rgb = None

class XY:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __str__(self):
        return f"{self.x}-{self.y}"

class RGB:
    def __init__(self, r, g, b, a=255):
        self.a = a
        self.r = r
        self.g = g
        self.b = b

    def __str__(self):
        return f"r:{round(self.r)}, g:{round(self.g)}, b:{round(self.b)}, a:{self.a}"

    def print(self, msg):
        print(f"{msg}: {self}")

    def average(self, other):
        return RGB(
            (self.r + other.r) / 2,
            (self.g + other.g) / 2,
            (self.b + other.b) / 2,
            255
        )

    def distance(self, other):
        return math.sqrt(
            (self.r - other.r) ** 2 +
            (self.g - other.g) ** 2 +
            (self.b - other.b) ** 2
        )

    def equals(self, other):
        return self.r == other.r and self.g == other.g and self.b == other.b

    def clone(self):
        return RGB(self.r, self.g, self.b, self.a)

def main():
    global g_last_coord_mouse_down, g_virtual_fg_color_rgb
    g_last_coord_mouse_down = get_cursor_in_document_coords()
    color = get_color_under_cursor_or_at_pos(forced_pos=xy_of_qpoint(g_last_coord_mouse_down))
    set_fg_color(color)
    g_virtual_fg_color_rgb = color

def get_cursor_in_document_coords():
    app = Krita.instance() # type: ignore
    view = app.activeWindow().activeView()
    if view.document():
        q_view = get_q_view(view)
        q_canvas = get_q_canvas(q_view)
        transform = get_transform(view)
        if transform is not None:
            transform_inv, _ = transform.inverted()
            global_pos = QCursor.pos()
            local_pos = q_canvas.mapFromGlobal(global_pos)
            center = q_canvas.rect().center()
            return transform_inv.map(local_pos - QPointF(center))
    return None

def get_q_view(view):
    window = view.window()
    q_window = window.qwindow()
    q_stacked_widget = q_window.centralWidget()
    q_mdi_area = q_stacked_widget.findChild(QMdiArea)
    for v, q_mdi_view in zip(window.views(), q_mdi_area.subWindowList()):
        if v == view:
            return q_mdi_view.widget()

def get_q_canvas(q_view):
    if q_view is None:
        return None

    scroll_area = q_view.findChild(QAbstractScrollArea)
    viewport = scroll_area.viewport()
    for child in viewport.children():
        cls_name = child.metaObject().className()
        if cls_name.startswith('Kis') and ('Canvas' in cls_name):
            return child

def get_transform(view):
    def _offset(scroller):
        mid = (scroller.minimum() + scroller.maximum()) / 2.0
        return -(scroller.value() - mid)

    canvas = view.canvas()
    document = view.document()
    q_view = get_q_view(view)
    if q_view is not None:
        area = q_view.findChild(QAbstractScrollArea)
        zoom = (canvas.zoomLevel() * 72.0) / document.resolution()
        transform = QTransform()
        transform.translate(
            _offset(area.horizontalScrollBar()),
            _offset(area.verticalScrollBar())
        )
        transform.rotate(canvas.rotation())
        transform.scale(zoom, zoom)
        return transform
    return None

def update_label_from_virtual_color(lblActiveColor):
    global g_virtual_fg_color_rgb
    r, g, b = g_virtual_fg_color_rgb.r, g_virtual_fg_color_rgb.g, g_virtual_fg_color_rgb.b
    lblActiveColor.setStyleSheet(f"background-color: rgb({b}, {g}, {r})")

def xy_of_qpoint(q):
    return XY(int(round(q.x())), int(round(q.y())))

def get_color_under_cursor_or_at_pos(forced_pos=None):
    application = Krita.instance() # type: ignore
    document = application.activeDocument()
    if not document:
        return None

    center = QPointF(0.5 * document.width(), 0.5 * document.height())
    if forced_pos is None:
        p = get_cursor_in_document_coords()
        doc_pos = p + center
        doc_posxy = xy_of_qpoint(doc_pos)
    else:
        doc_posxy = XY(forced_pos.x + int(round(center.x())), forced_pos.y + int(round(center.y())))

    pixBytes = document.pixelData(int(doc_posxy.x), int(doc_posxy.y), 1, 1)
    if len(pixBytes) == 4:
        imageData = QImage(pixBytes, 1, 1, QImage.Format_RGBA8888)
    elif len(pixBytes) == 8:
        imageData = QImage(pixBytes, 1, 1, QImage.Format_RGBA64)
    else:
        raise ValueError(f"Unsupported length {len(pixBytes)}")

    pixel_color = imageData.pixelColor(0, 0)
    return RGB(pixel_color.red(), pixel_color.green(), pixel_color.blue(), 255)

def set_fg_color(color):
    app = Krita.instance() # type: ignore
    win = app.activeWindow()
    if not win:
        return

    view = win.activeView()
    if not view:
        return

    fg = view.foregroundColor()
    components = fg.components()
    # clamps to 0 - 1
    components[0] = color.r / 255.0
    components[1] = color.g / 255.0
    components[2] = color.b / 255.0

    fg.setComponents(components)
    view.setForeGroundColor(fg)

if __name__ == "__main__":
    main()