ImageReader.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  1. # Copyright (c) 2015 Ultimaker B.V.
  2. # Cura is released under the terms of the LGPLv3 or higher.
  3. import numpy
  4. from PyQt5.QtGui import QImage, qRed, qGreen, qBlue
  5. from PyQt5.QtCore import Qt
  6. from UM.Mesh.MeshReader import MeshReader
  7. from UM.Mesh.MeshBuilder import MeshBuilder
  8. from UM.Math.Vector import Vector
  9. from UM.Job import Job
  10. from UM.Logger import Logger
  11. from .ImageReaderUI import ImageReaderUI
  12. from cura.Scene.CuraSceneNode import CuraSceneNode as SceneNode
  13. class ImageReader(MeshReader):
  14. def __init__(self):
  15. super(ImageReader, self).__init__()
  16. self._supported_extensions = [".jpg", ".jpeg", ".bmp", ".gif", ".png"]
  17. self._ui = ImageReaderUI(self)
  18. def preRead(self, file_name, *args, **kwargs):
  19. img = QImage(file_name)
  20. if img.isNull():
  21. Logger.log("e", "Image is corrupt.")
  22. return MeshReader.PreReadResult.failed
  23. width = img.width()
  24. depth = img.height()
  25. largest = max(width, depth)
  26. width = width / largest * self._ui.default_width
  27. depth = depth / largest * self._ui.default_depth
  28. self._ui.setWidthAndDepth(width, depth)
  29. self._ui.showConfigUI()
  30. self._ui.waitForUIToClose()
  31. if self._ui.getCancelled():
  32. return MeshReader.PreReadResult.cancelled
  33. return MeshReader.PreReadResult.accepted
  34. def read(self, file_name):
  35. size = max(self._ui.getWidth(), self._ui.getDepth())
  36. return self._generateSceneNode(file_name, size, self._ui.peak_height, self._ui.base_height, self._ui.smoothing, 512, self._ui.image_color_invert)
  37. def _generateSceneNode(self, file_name, xz_size, peak_height, base_height, blur_iterations, max_size, image_color_invert):
  38. scene_node = SceneNode()
  39. mesh = MeshBuilder()
  40. img = QImage(file_name)
  41. if img.isNull():
  42. Logger.log("e", "Image is corrupt.")
  43. return None
  44. width = max(img.width(), 2)
  45. height = max(img.height(), 2)
  46. aspect = height / width
  47. if img.width() < 2 or img.height() < 2:
  48. img = img.scaled(width, height, Qt.IgnoreAspectRatio)
  49. base_height = max(base_height, 0)
  50. peak_height = max(peak_height, -base_height)
  51. xz_size = max(xz_size, 1)
  52. scale_vector = Vector(xz_size, peak_height, xz_size)
  53. if width > height:
  54. scale_vector = scale_vector.set(z=scale_vector.z * aspect)
  55. elif height > width:
  56. scale_vector = scale_vector.set(x=scale_vector.x / aspect)
  57. if width > max_size or height > max_size:
  58. scale_factor = max_size / width
  59. if height > width:
  60. scale_factor = max_size / height
  61. width = int(max(round(width * scale_factor), 2))
  62. height = int(max(round(height * scale_factor), 2))
  63. img = img.scaled(width, height, Qt.IgnoreAspectRatio)
  64. width_minus_one = width - 1
  65. height_minus_one = height - 1
  66. Job.yieldThread()
  67. texel_width = 1.0 / (width_minus_one) * scale_vector.x
  68. texel_height = 1.0 / (height_minus_one) * scale_vector.z
  69. height_data = numpy.zeros((height, width), dtype=numpy.float32)
  70. for x in range(0, width):
  71. for y in range(0, height):
  72. qrgb = img.pixel(x, y)
  73. avg = float(qRed(qrgb) + qGreen(qrgb) + qBlue(qrgb)) / (3 * 255)
  74. height_data[y, x] = avg
  75. Job.yieldThread()
  76. if image_color_invert:
  77. height_data = 1 - height_data
  78. for _ in range(0, blur_iterations):
  79. copy = numpy.pad(height_data, ((1, 1), (1, 1)), mode= "edge")
  80. height_data += copy[1:-1, 2:]
  81. height_data += copy[1:-1, :-2]
  82. height_data += copy[2:, 1:-1]
  83. height_data += copy[:-2, 1:-1]
  84. height_data += copy[2:, 2:]
  85. height_data += copy[:-2, 2:]
  86. height_data += copy[2:, :-2]
  87. height_data += copy[:-2, :-2]
  88. height_data /= 9
  89. Job.yieldThread()
  90. height_data *= scale_vector.y
  91. height_data += base_height
  92. heightmap_face_count = 2 * height_minus_one * width_minus_one
  93. total_face_count = heightmap_face_count + (width_minus_one * 2) * (height_minus_one * 2) + 2
  94. mesh.reserveFaceCount(total_face_count)
  95. # initialize to texel space vertex offsets.
  96. # 6 is for 6 vertices for each texel quad.
  97. heightmap_vertices = numpy.zeros((width_minus_one * height_minus_one, 6, 3), dtype = numpy.float32)
  98. heightmap_vertices = heightmap_vertices + numpy.array([[
  99. [0, base_height, 0],
  100. [0, base_height, texel_height],
  101. [texel_width, base_height, texel_height],
  102. [texel_width, base_height, texel_height],
  103. [texel_width, base_height, 0],
  104. [0, base_height, 0]
  105. ]], dtype = numpy.float32)
  106. offsetsz, offsetsx = numpy.mgrid[0: height_minus_one, 0: width - 1]
  107. offsetsx = numpy.array(offsetsx, numpy.float32).reshape(-1, 1) * texel_width
  108. offsetsz = numpy.array(offsetsz, numpy.float32).reshape(-1, 1) * texel_height
  109. # offsets for each texel quad
  110. heightmap_vertex_offsets = numpy.concatenate([offsetsx, numpy.zeros((offsetsx.shape[0], offsetsx.shape[1]), dtype=numpy.float32), offsetsz], 1)
  111. heightmap_vertices += heightmap_vertex_offsets.repeat(6, 0).reshape(-1, 6, 3)
  112. # apply height data to y values
  113. heightmap_vertices[:, 0, 1] = heightmap_vertices[:, 5, 1] = height_data[:-1, :-1].reshape(-1)
  114. heightmap_vertices[:, 1, 1] = height_data[1:, :-1].reshape(-1)
  115. heightmap_vertices[:, 2, 1] = heightmap_vertices[:, 3, 1] = height_data[1:, 1:].reshape(-1)
  116. heightmap_vertices[:, 4, 1] = height_data[:-1, 1:].reshape(-1)
  117. heightmap_indices = numpy.array(numpy.mgrid[0:heightmap_face_count * 3], dtype=numpy.int32).reshape(-1, 3)
  118. mesh._vertices[0:(heightmap_vertices.size // 3), :] = heightmap_vertices.reshape(-1, 3)
  119. mesh._indices[0:(heightmap_indices.size // 3), :] = heightmap_indices
  120. mesh._vertex_count = heightmap_vertices.size // 3
  121. mesh._face_count = heightmap_indices.size // 3
  122. geo_width = width_minus_one * texel_width
  123. geo_height = height_minus_one * texel_height
  124. # bottom
  125. mesh.addFaceByPoints(0, 0, 0, 0, 0, geo_height, geo_width, 0, geo_height)
  126. mesh.addFaceByPoints(geo_width, 0, geo_height, geo_width, 0, 0, 0, 0, 0)
  127. # north and south walls
  128. for n in range(0, width_minus_one):
  129. x = n * texel_width
  130. nx = (n + 1) * texel_width
  131. hn0 = height_data[0, n]
  132. hn1 = height_data[0, n + 1]
  133. hs0 = height_data[height_minus_one, n]
  134. hs1 = height_data[height_minus_one, n + 1]
  135. mesh.addFaceByPoints(x, 0, 0, nx, 0, 0, nx, hn1, 0)
  136. mesh.addFaceByPoints(nx, hn1, 0, x, hn0, 0, x, 0, 0)
  137. mesh.addFaceByPoints(x, 0, geo_height, nx, 0, geo_height, nx, hs1, geo_height)
  138. mesh.addFaceByPoints(nx, hs1, geo_height, x, hs0, geo_height, x, 0, geo_height)
  139. # west and east walls
  140. for n in range(0, height_minus_one):
  141. y = n * texel_height
  142. ny = (n + 1) * texel_height
  143. hw0 = height_data[n, 0]
  144. hw1 = height_data[n + 1, 0]
  145. he0 = height_data[n, width_minus_one]
  146. he1 = height_data[n + 1, width_minus_one]
  147. mesh.addFaceByPoints(0, 0, y, 0, 0, ny, 0, hw1, ny)
  148. mesh.addFaceByPoints(0, hw1, ny, 0, hw0, y, 0, 0, y)
  149. mesh.addFaceByPoints(geo_width, 0, y, geo_width, 0, ny, geo_width, he1, ny)
  150. mesh.addFaceByPoints(geo_width, he1, ny, geo_width, he0, y, geo_width, 0, y)
  151. mesh.calculateNormals(fast=True)
  152. scene_node.setMeshData(mesh.build())
  153. return scene_node