mirror of
https://github.com/Ultimaker/Cura.git
synced 2025-12-26 17:38:36 -07:00
240 lines
10 KiB
Python
240 lines
10 KiB
Python
# Copyright (c) 2025 UltiMaker
|
|
# Cura is released under the terms of the LGPLv3 or higher.
|
|
|
|
import os
|
|
from PyQt6.QtCore import QRect
|
|
from typing import Optional, List, Tuple, Dict, cast
|
|
|
|
from PyQt6.QtGui import QImage, QColor, QPainter
|
|
|
|
from cura.CuraApplication import CuraApplication
|
|
from cura.BuildVolume import BuildVolume
|
|
from cura.CuraView import CuraView
|
|
from Machines.Models.ExtrudersModel import ExtrudersModel
|
|
from UM.PluginRegistry import PluginRegistry
|
|
from UM.View.GL.ShaderProgram import ShaderProgram
|
|
from UM.View.GL.Texture import Texture
|
|
from UM.Scene.Iterator.DepthFirstIterator import DepthFirstIterator
|
|
from UM.Scene.Selection import Selection
|
|
from UM.View.GL.OpenGL import OpenGL
|
|
from UM.i18n import i18nCatalog
|
|
from UM.Math.Color import Color
|
|
|
|
catalog = i18nCatalog("cura")
|
|
|
|
|
|
class PaintView(CuraView):
|
|
"""View for model-painting."""
|
|
|
|
UNDO_STACK_SIZE = 1024
|
|
|
|
class PaintType:
|
|
def __init__(self, display_color: Color, value: int):
|
|
self.display_color: Color = display_color
|
|
self.value: int = value
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__(use_empty_menu_placeholder = True)
|
|
self._paint_shader: Optional[ShaderProgram] = None
|
|
self._current_paint_texture: Optional[Texture] = None
|
|
self._current_bits_ranges: tuple[int, int] = (0, 0)
|
|
self._current_paint_type = ""
|
|
self._paint_modes: Dict[str, Dict[str, "PaintView.PaintType"]] = {}
|
|
|
|
self._stroke_undo_stack: List[Tuple[QImage, int, int]] = []
|
|
self._stroke_redo_stack: List[Tuple[QImage, int, int]] = []
|
|
|
|
self._force_opaque_mask = QImage(2, 2, QImage.Format.Format_Mono)
|
|
self._force_opaque_mask.fill(1)
|
|
|
|
application = CuraApplication.getInstance()
|
|
application.engineCreatedSignal.connect(self._makePaintModes)
|
|
self._scene = application.getController().getScene()
|
|
|
|
self._extruders_model: Optional[ExtrudersModel] = None
|
|
|
|
def _makePaintModes(self):
|
|
application = CuraApplication.getInstance()
|
|
|
|
self._extruders_model = application.getExtrudersModel()
|
|
self._extruders_model.modelChanged.connect(self._onExtrudersChanged)
|
|
|
|
theme = application.getTheme()
|
|
usual_types = {"none": self.PaintType(Color(*theme.getColor("paint_normal_area").getRgb()), 0),
|
|
"preferred": self.PaintType(Color(*theme.getColor("paint_preferred_area").getRgb()), 1),
|
|
"avoid": self.PaintType(Color(*theme.getColor("paint_avoid_area").getRgb()), 2)}
|
|
self._paint_modes = {
|
|
"seam": usual_types,
|
|
"support": usual_types,
|
|
"extruder": self._makeExtrudersColors(),
|
|
}
|
|
|
|
self._current_paint_type = "seam"
|
|
|
|
def _makeExtrudersColors(self) -> Dict[str, "PaintView.PaintType"]:
|
|
extruders_colors: Dict[str, "PaintView.PaintType"] = {}
|
|
|
|
for extruder_item in self._extruders_model.items:
|
|
if "color" in extruder_item:
|
|
material_color = extruder_item["color"]
|
|
else:
|
|
material_color = self._extruders_model.defaultColors[0]
|
|
|
|
index = extruder_item["index"]
|
|
extruders_colors[str(index)] = self.PaintType(Color(*QColor(material_color).getRgb()), index)
|
|
|
|
return extruders_colors
|
|
|
|
def _onExtrudersChanged(self) -> None:
|
|
if self._paint_modes is None:
|
|
return
|
|
|
|
self._paint_modes["extruder"] = self._makeExtrudersColors()
|
|
|
|
controller = CuraApplication.getInstance().getController()
|
|
if controller.getActiveView() != self:
|
|
return
|
|
|
|
selected_objects = Selection.getAllSelectedObjects()
|
|
if len(selected_objects) != 1:
|
|
return
|
|
|
|
controller.getScene().sceneChanged.emit(selected_objects[0])
|
|
|
|
def _checkSetup(self):
|
|
if not self._paint_shader:
|
|
shader_filename = os.path.join(PluginRegistry.getInstance().getPluginPath("PaintTool"), "paint.shader")
|
|
self._paint_shader = OpenGL.getInstance().createShaderProgram(shader_filename)
|
|
|
|
def _forceOpaqueDeepCopy(self, image: QImage):
|
|
res = QImage(image.width(), image.height(), QImage.Format.Format_RGBA8888)
|
|
res.fill(QColor(255, 255, 255, 255))
|
|
painter = QPainter(res)
|
|
painter.setRenderHint(QPainter.RenderHint.Antialiasing, False)
|
|
painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_SourceOver)
|
|
painter.drawImage(0, 0, image)
|
|
painter.end()
|
|
res.setAlphaChannel(self._force_opaque_mask.scaled(image.width(), image.height()))
|
|
return res
|
|
|
|
def addStroke(self, stroke_mask: QImage, start_x: int, start_y: int, brush_color: str) -> None:
|
|
if self._current_paint_texture is None or self._current_paint_texture.getImage() is None:
|
|
return
|
|
|
|
self._prepareDataMapping()
|
|
|
|
actual_image = self._current_paint_texture.getImage()
|
|
|
|
bit_range_start, bit_range_end = self._current_bits_ranges
|
|
set_value = self._paint_modes[self._current_paint_type][brush_color].value << self._current_bits_ranges[0]
|
|
full_int32 = 0xffffffff
|
|
clear_mask = full_int32 ^ (((full_int32 << (32 - 1 - (bit_range_end - bit_range_start))) & full_int32) >> (32 - 1 - bit_range_end))
|
|
image_rect = QRect(0, 0, stroke_mask.width(), stroke_mask.height())
|
|
|
|
clear_bits_image = stroke_mask.copy()
|
|
clear_bits_image.invertPixels()
|
|
painter = QPainter(clear_bits_image)
|
|
painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Lighten)
|
|
painter.fillRect(image_rect, clear_mask)
|
|
painter.end()
|
|
|
|
set_value_image = stroke_mask.copy()
|
|
painter = QPainter(set_value_image)
|
|
painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Multiply)
|
|
painter.fillRect(image_rect, set_value)
|
|
painter.end()
|
|
|
|
stroked_image = actual_image.copy(start_x, start_y, stroke_mask.width(), stroke_mask.height())
|
|
painter = QPainter(stroked_image)
|
|
painter.setCompositionMode(QPainter.CompositionMode.RasterOp_SourceAndDestination)
|
|
painter.drawImage(0, 0, clear_bits_image)
|
|
painter.setCompositionMode(QPainter.CompositionMode.RasterOp_SourceOrDestination)
|
|
painter.drawImage(0, 0, set_value_image)
|
|
painter.end()
|
|
|
|
self._stroke_redo_stack.clear()
|
|
if len(self._stroke_undo_stack) >= PaintView.UNDO_STACK_SIZE:
|
|
self._stroke_undo_stack.pop(0)
|
|
undo_image = self._forceOpaqueDeepCopy(self._current_paint_texture.setSubImage(stroked_image, start_x, start_y))
|
|
if undo_image is not None:
|
|
self._stroke_undo_stack.append((undo_image, start_x, start_y))
|
|
|
|
def _applyUndoStacksAction(self, from_stack: List[Tuple[QImage, int, int]], to_stack: List[Tuple[QImage, int, int]]) -> bool:
|
|
if len(from_stack) <= 0 or self._current_paint_texture is None:
|
|
return False
|
|
from_image, x, y = from_stack.pop()
|
|
to_image = self._forceOpaqueDeepCopy(self._current_paint_texture.setSubImage(from_image, x, y))
|
|
if to_image is None:
|
|
return False
|
|
if len(to_stack) >= PaintView.UNDO_STACK_SIZE:
|
|
to_stack.pop(0)
|
|
to_stack.append((to_image, x, y))
|
|
return True
|
|
|
|
def undoStroke(self) -> bool:
|
|
return self._applyUndoStacksAction(self._stroke_undo_stack, self._stroke_redo_stack)
|
|
|
|
def redoStroke(self) -> bool:
|
|
return self._applyUndoStacksAction(self._stroke_redo_stack, self._stroke_undo_stack)
|
|
|
|
def getUvTexDimensions(self):
|
|
if self._current_paint_texture is not None:
|
|
return self._current_paint_texture.getWidth(), self._current_paint_texture.getHeight()
|
|
return 0, 0
|
|
|
|
def getPaintType(self) -> str:
|
|
return self._current_paint_type
|
|
|
|
def setPaintType(self, paint_type: str) -> None:
|
|
self._current_paint_type = paint_type
|
|
self._prepareDataMapping()
|
|
|
|
def _prepareDataMapping(self):
|
|
node = Selection.getAllSelectedObjects()[0]
|
|
if node is None:
|
|
return
|
|
|
|
paint_data_mapping = node.callDecoration("getTextureDataMapping")
|
|
|
|
if self._current_paint_type not in paint_data_mapping:
|
|
new_mapping = self._add_mapping(paint_data_mapping, len(self._paint_modes[self._current_paint_type]))
|
|
paint_data_mapping[self._current_paint_type] = new_mapping
|
|
node.callDecoration("setTextureDataMapping", paint_data_mapping)
|
|
|
|
self._current_bits_ranges = paint_data_mapping[self._current_paint_type]
|
|
|
|
@staticmethod
|
|
def _add_mapping(actual_mapping: Dict[str, tuple[int, int]], nb_storable_values: int) -> tuple[int, int]:
|
|
start_index = 0
|
|
if actual_mapping:
|
|
start_index = max(end_index for _, end_index in actual_mapping.values()) + 1
|
|
|
|
end_index = start_index + int.bit_length(nb_storable_values - 1) - 1
|
|
|
|
return start_index, end_index
|
|
|
|
def beginRendering(self) -> None:
|
|
if self._current_paint_type not in self._paint_modes:
|
|
return
|
|
|
|
self._checkSetup()
|
|
renderer = self.getRenderer()
|
|
|
|
for node in DepthFirstIterator(self._scene.getRoot()):
|
|
if isinstance(node, BuildVolume):
|
|
node.render(renderer)
|
|
|
|
paint_batch = renderer.createRenderBatch(shader=self._paint_shader)
|
|
renderer.addRenderBatch(paint_batch)
|
|
|
|
for node in Selection.getAllSelectedObjects():
|
|
paint_batch.addItem(node.getWorldTransformation(copy=False), node.getMeshData(), normal_transformation=node.getCachedNormalMatrix())
|
|
self._current_paint_texture = node.callDecoration("getPaintTexture")
|
|
self._paint_shader.setTexture(0, self._current_paint_texture)
|
|
|
|
self._paint_shader.setUniformValue("u_bitsRangesStart", self._current_bits_ranges[0])
|
|
self._paint_shader.setUniformValue("u_bitsRangesEnd", self._current_bits_ranges[1])
|
|
|
|
colors = [paint_type_obj.display_color for paint_type_obj in self._paint_modes[self._current_paint_type].values()]
|
|
colors_values = [[int(color_part * 255) for color_part in [color.r, color.g, color.b]] for color in colors]
|
|
self._paint_shader.setUniformValueArray("u_renderColors", colors_values)
|