gradationPicker.py

# -*- coding: utf-8 -*-

import sys
import math

from PySide2.QtWidgets import (QWidget,
                               QMainWindow,
                               QPushButton,
                               QDialog,
                               QApplication,
                               QVBoxLayout,
                               QGraphicsView,
                               QGraphicsScene)
from PySide2.QtCore import (Qt,
                            Signal,
                            QRect,
                            QLineF,
                            QPoint,
                            QPointF,
                            QEvent)
from PySide2.QtGui import (QColor,
                           QFont,
                           QPainter,
                           QCursor,
                           QPen,
                           QBrush,
                           QPixmap,
                           QColor,
                           QGradient,
                           QLinearGradient)


class ColorPick(QMainWindow):

    getGradation = Signal(list)

    def __init__(self, parent=None):
        super().__init__(parent)

        self.setMouseTracking(True)

        self.screen = QApplication.primaryScreen()
        self.originalPixmap = self.screen.grabWindow(QApplication.desktop().winId())

        self.isDrag = False
        self.currentPos = None
        self.threshold = 0.5  # 前のPickとの色の差をどの程度までとるか
        self.totalLength = 0
        self.gradation = []  # [距離,R,G,B]の配列を入れる
        self.startColor = None

    def mousePressEvent(self, event):
        self.isDrag = True
        self.currentPos = QCursor().pos()
        # クリックしたタイミングを0としてグラデーションをピックする
        color = self.getCurrentColor(self.currentPos)
        self.startColor = (color.red(), color.green(), color.blue())

    def mouseReleaseEvent(self, event):

        self.isDrag = False
        # 1つ目は0なので、そのまま取り出す
        retVal = [[0, self.startColor]]

        # 正規化
        for i in self.gradation:
            retVal.append([i[0] / self.totalLength, i[1]])
        # マウスを話したらSignalで結果を返す
        self.getGradation.emit(retVal)
        self.close()

    def mouseMoveEvent(self, event):
        # GradColorを処理する
        if self.isDrag:
            pos = QCursor().pos()
            line = QLineF(QPointF(pos), QPointF(self.currentPos))
            # マウスの移動距離を保存（あとで正規化）
            self.totalLength += line.length()
            color = self.getCurrentColor(pos)
            currentColor = (color.red(), color.green(), color.blue())

            if len(self.gradation) == 0:
                if colorDifference(self.startColor, currentColor) > self.threshold:
                    self.gradation.append([self.totalLength, currentColor])
            else:
                lastColor = self.gradation[-1][1]
                # 色の変化が指定以上だった場合はセットする
                if colorDifference(lastColor, currentColor) > self.threshold:
                    self.gradation.append([self.totalLength, currentColor])

            self.currentPos = pos
            self.repaint()

    def paintEvent(self, event):

        painter = QPainter()
        painter.begin(self)
        rectSize = QApplication.desktop().screenGeometry()
        painter.drawPixmap(rectSize, self.originalPixmap)

        x = QCursor().pos().x()
        y = QCursor().pos().y()
        offset = 30

        if self.isDrag:
            pen = QPen(Qt.red, 7)
            painter.setPen(pen)
            painter.drawPoint(x, y)

            color = self.getCurrentColor(QCursor().pos())
            painter.setBrush(color)
            pen = QPen(Qt.black, 1)
            painter.setPen(pen)
            rect = QRect(x + offset, y + offset, 20, 20)
            painter.drawRect(rect)

            painter.setFont(QFont(u'メイリオ', 8, QFont.Bold, False))
            pen = QPen(Qt.white, 1)
            painter.setPen(pen)
            painter.drawText(x + offset + 30, y + offset + 17, f"RGB({color.red()},{color.green()},{color.blue()})")

        painter.end()

    def getCurrentColor(self, pos: QPoint):
        x = pos.x()
        y = pos.y()
        return self.originalPixmap.toImage().pixelColor(x, y)


def colorDifference(src, dst):
    # 超シンプルな色差を求める関数
    rd = (src[0] / 255.0) - (dst[0] / 255.0)
    gd = (src[1] / 255.0) - (dst[1] / 255.0)
    bd = (src[2] / 255.0) - (dst[2] / 255.0)
    return math.sqrt(rd * rd + gd * gd + bd * bd) / math.sqrt(3)


class Gradation(QGraphicsScene):

    def __init__(self, parent=None):
        super().__init__(parent)

        self.colorList = []

    @property
    def width(self):
        return self.sceneRect().width()

    @property
    def height(self):
        return self.sceneRect().height()

    def drawBackground(self, painter, rect):

        grad1 = QLinearGradient(0.0, 0.0, self.width, 0.0)
        for ratio, color in self.colorList:
            grad1.setColorAt(ratio, QColor(*color))
        painter.setBrush(QBrush(grad1))
        painter.drawRect(0, 0, self.width, self.height)

    def setColor(self, colorList):

        self.colorList = colorList


# グラーデーション確認用GUIを作る
class UISample(QDialog):

    def __init__(self, parent=None):
        super().__init__(parent)

        self.resize(350, 100)

        layout = QVBoxLayout()
        self.setLayout(layout)

        self.scene = Gradation()
        self.scene.setSceneRect(0, 0, 300, 50)

        self.view  = QGraphicsView(self.scene, self)
        layout.addWidget(self.view)

        btn = QPushButton("Pickする")
        layout.addWidget(btn)

        btn.clicked.connect(self.showPick)

    def showPick(self):
        # ボタンを押すと、ピッカーGUIが全画面に表示される
        a = ColorPick(self)
        a.showFullScreen()
        a.getGradation.connect(self.setGradationColor)

    def setGradationColor(self, colorList):

        self.scene.setColor(colorList)
        self.view.repaint()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    # a = ColorPick()
    # a.show()

    win = UISample()
    win.show()

    sys.exit(app.exec_())