PickingPass.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. # Copyright (c) 2018 Ultimaker B.V.
  2. # Cura is released under the terms of the LGPLv3 or higher.
  3. import random
  4. from typing import Optional, TYPE_CHECKING
  5. from PyQt5.QtCore import QBuffer
  6. from UM.Mesh.MeshBuilder import MeshBuilder
  7. from UM.Qt.QtApplication import QtApplication
  8. from UM.Math.Vector import Vector
  9. from UM.Math.Color import Color
  10. from UM.Resources import Resources
  11. from UM.View.RenderPass import RenderPass
  12. from UM.View.GL.OpenGL import OpenGL
  13. from UM.View.RenderBatch import RenderBatch
  14. from UM.Scene.Iterator.DepthFirstIterator import DepthFirstIterator
  15. if TYPE_CHECKING:
  16. from UM.View.GL.ShaderProgram import ShaderProgram
  17. ## A RenderPass subclass that renders a the distance of selectable objects from the active camera to a texture.
  18. # The texture is used to map a 2d location (eg the mouse location) to a world space position
  19. #
  20. # Note that in order to increase precision, the 24 bit depth value is encoded into all three of the R,G & B channels
  21. class PickingPass(RenderPass):
  22. def __init__(self, width: int, height: int) -> None:
  23. super().__init__("picking", width, height)
  24. self._selection_map = {}
  25. self._renderer = QtApplication.getInstance().getRenderer()
  26. self._shader = None #type: Optional[ShaderProgram]
  27. self._scene = QtApplication.getInstance().getController().getScene()
  28. def render(self) -> None:
  29. if not self._shader:
  30. self._shader = OpenGL.getInstance().createShaderProgram(Resources.getPath(Resources.Shaders, "camera_distance.shader"))
  31. width, height = self.getSize()
  32. self._gl.glViewport(0, 0, width, height)
  33. self._gl.glClearColor(1.0, 1.0, 1.0, 0.0)
  34. self._gl.glClear(self._gl.GL_COLOR_BUFFER_BIT | self._gl.GL_DEPTH_BUFFER_BIT)
  35. # Create a new batch to be rendered
  36. batch = RenderBatch(self._shader)
  37. # Fill up the batch with objects that can be sliced. `
  38. for node in DepthFirstIterator(self._scene.getRoot()): #type: ignore #Ignore type error because iter() should get called automatically by Python syntax.
  39. if node.callDecoration("isSliceable") and node.getMeshData() and node.isVisible():
  40. faces = node.getMeshData().getIndices()
  41. vertices = node.getMeshData().getVertices()
  42. normals = node.getMeshData().getNormals()
  43. print("Faces:", faces)
  44. print("Vertices:", vertices)
  45. print("Normals:", normals)
  46. for index, face in enumerate(faces):
  47. normal_vertex = normals[face][0]
  48. triangle_mesh = vertices[face]
  49. print(face, normal_vertex, triangle_mesh)
  50. batch.addItem(transformation = node.getWorldTransformation(), mesh = node.getMeshData(), uniforms = { "selection_color": self._getFaceColor(face, normal_vertex)})
  51. self.bind()
  52. batch.render(self._scene.getActiveCamera())
  53. self.release()
  54. def _getFaceColor(self, face: Vector, normal_vertex: Vector) -> Color:
  55. while True:
  56. r = random.randint(0, 255)
  57. g = random.randint(0, 255)
  58. b = random.randint(0, 255)
  59. a = 255
  60. color = Color(r, g, b, a)
  61. if color not in self._selection_map:
  62. break
  63. print("Adding color: {color} - {normal}".format(color = color, normal = normal_vertex))
  64. self._selection_map[color] = {"face": id(face), "normal_vertex": normal_vertex}
  65. return color
  66. ## Get the normal vector at a certain pixel coordinate.
  67. def getPickedNormalVertex(self, x: int, y: int) -> Optional[Vector]:
  68. output = self.getOutput()
  69. print("Creating image")
  70. output.save("thumbnail.png")
  71. window_size = self._renderer.getWindowSize()
  72. px = (0.5 + x / 2.0) * window_size[0]
  73. py = (0.5 + y / 2.0) * window_size[1]
  74. if px < 0 or px > (output.width() - 1) or py < 0 or py > (output.height() - 1):
  75. return None
  76. pixel = output.pixel(px, py)
  77. print("######### ", x, y, pixel, Color.fromARGB(pixel))
  78. face = self._selection_map.get(Color.fromARGB(pixel), None)
  79. if not face:
  80. return None
  81. return face.get("normal_vertex", None)
  82. ## Get the distance in mm from the camera to at a certain pixel coordinate.
  83. def getPickedDepth(self, x: int, y: int) -> float:
  84. output = self.getOutput()
  85. window_size = self._renderer.getWindowSize()
  86. px = (0.5 + x / 2.0) * window_size[0]
  87. py = (0.5 + y / 2.0) * window_size[1]
  88. if px < 0 or px > (output.width() - 1) or py < 0 or py > (output.height() - 1):
  89. return -1
  90. distance = output.pixel(px, py) # distance in micron, from in r, g & b channels
  91. distance = (distance & 0x00ffffff) / 1000. # drop the alpha channel and covert to mm
  92. return distance
  93. ## Get the world coordinates of a picked point
  94. def getPickedPosition(self, x: int, y: int) -> Vector:
  95. distance = self.getPickedDepth(x, y)
  96. camera = self._scene.getActiveCamera()
  97. if camera:
  98. return camera.getRay(x, y).getPointAlongRay(distance)
  99. return Vector()