KDTree.cpp 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321
  1. #include <algorithm>
  2. #include <iostream>
  3. #include "common/utils/KDTree.h"
  4. KDTree::Triangle::Triangle(const Vector3& a, const Vector3& b, const Vector3& c) {
  5. v[0] = a;
  6. v[1] = b;
  7. v[2] = c;
  8. mid = (a + b + c) * (1.0f / 3.0f);
  9. }
  10. const Array<Vector3, 3>& KDTree::Triangle::data() const {
  11. return v;
  12. }
  13. const Vector3& KDTree::Triangle::operator[](int index) const {
  14. return v[index];
  15. }
  16. const Vector3& KDTree::Triangle::getMid() const {
  17. return mid;
  18. }
  19. KDTree::Node::Node() : splitDim(0), splitValue(0.0f), color(0xFFFFFF) {
  20. childs[0] = nullptr;
  21. childs[1] = nullptr;
  22. }
  23. KDTree::~KDTree() {
  24. clean(&root);
  25. }
  26. void KDTree::clean(Node* n) {
  27. if(n->childs[0] != nullptr) {
  28. clean(n->childs[0]);
  29. }
  30. if(n->childs[1] != nullptr) {
  31. clean(n->childs[1]);
  32. }
  33. delete n->childs[0];
  34. delete n->childs[1];
  35. }
  36. void KDTree::build(std::vector<KDTree::Triangle>& data) {
  37. if(data.size() == 0) {
  38. return;
  39. }
  40. min = data[0][0];
  41. max = data[0][0];
  42. for(const Triangle& t : data) {
  43. for(const Vector3& v : t.data()) {
  44. min.set(std::min(min[0], v[0]), std::min(min[1], v[1]), std::min(min[2], v[2]));
  45. max.set(std::max(max[0], v[0]), std::max(max[1], v[1]), std::max(max[2], v[2]));
  46. }
  47. }
  48. build(&root, data);
  49. }
  50. float KDTree::median(std::vector<KDTree::Triangle>& data, int dim) const {
  51. auto compare = [dim](const Triangle& a, const Triangle & b) {
  52. return a.getMid()[dim] < b.getMid()[dim];
  53. };
  54. size_t length = data.size();
  55. if((length & 1) == 0) {
  56. std::nth_element(data.begin(), data.begin() + (length / 2 - 1), data.end(), compare);
  57. float tmp = data[length / 2 - 1].getMid()[dim];
  58. std::nth_element(data.begin(), data.begin() + (length / 2), data.end(), compare);
  59. return (tmp + data[length / 2].getMid()[dim]) / 2;
  60. }
  61. std::nth_element(data.begin(), data.begin() + (length / 2), data.end(), compare);
  62. return data[length / 2].getMid()[dim];
  63. }
  64. void KDTree::build(Node* n, std::vector<KDTree::Triangle>& data) {
  65. if(data.size() == 0) {
  66. return;
  67. } else if(data.size() == 1) {
  68. n->data.push_back(data[0]);
  69. return;
  70. }
  71. // find min and max coordinates
  72. Vector3 min = data[0][0];
  73. Vector3 max = data[0][0];
  74. for(const Triangle& t : data) {
  75. for(const Vector3& v : t.data()) {
  76. min.set(std::min(min[0], v[0]), std::min(min[1], v[1]), std::min(min[2], v[2]));
  77. max.set(std::max(max[0], v[0]), std::max(max[1], v[1]), std::max(max[2], v[2]));
  78. }
  79. }
  80. // find biggest span and its dimension
  81. int splitDim = 0;
  82. float maxSpan = max[0] - min[0];
  83. for(int i = 1; i < 3; i++) {
  84. float span = max[i] - min[i];
  85. if(span > maxSpan) {
  86. splitDim = i;
  87. maxSpan = span;
  88. }
  89. }
  90. // assign data to node
  91. n->splitDim = splitDim;
  92. n->splitValue = median(data, splitDim);
  93. // storage for split data
  94. std::vector<KDTree::Triangle> lessEqualData;
  95. std::vector<KDTree::Triangle> greaterData;
  96. // actually split the data
  97. for(const Triangle& t : data) {
  98. // count points on each split side
  99. int lessEqualCounter = 0;
  100. int greaterCount = 0;
  101. for(const Vector3& v : t.data()) {
  102. if(v[n->splitDim] <= n->splitValue) {
  103. lessEqualCounter++;
  104. } else {
  105. greaterCount++;
  106. }
  107. }
  108. // put the data in the correct container
  109. if(lessEqualCounter == 3) {
  110. lessEqualData.push_back(t);
  111. } else if(greaterCount == 3) {
  112. greaterData.push_back(t);
  113. } else {
  114. n->data.push_back(t);
  115. }
  116. }
  117. // recursive calls
  118. if(lessEqualData.size() > 0) {
  119. n->childs[0] = new Node();
  120. build(n->childs[0], lessEqualData);
  121. }
  122. if(greaterData.size() > 0) {
  123. n->childs[1] = new Node();
  124. build(n->childs[1], greaterData);
  125. }
  126. }
  127. void KDTree::fillLines(Lines& lines) const {
  128. lines.clear();
  129. lines.add(Vector3(min[0], min[1], min[2]), Vector3(max[0], min[1], min[2]), root.color);
  130. lines.add(Vector3(min[0], min[1], min[2]), Vector3(min[0], min[1], max[2]), root.color);
  131. lines.add(Vector3(max[0], min[1], min[2]), Vector3(max[0], min[1], max[2]), root.color);
  132. lines.add(Vector3(min[0], min[1], max[2]), Vector3(max[0], min[1], max[2]), root.color);
  133. lines.add(Vector3(min[0], min[1], min[2]), Vector3(min[0], max[1], min[2]), root.color);
  134. lines.add(Vector3(max[0], min[1], min[2]), Vector3(max[0], max[1], min[2]), root.color);
  135. lines.add(Vector3(min[0], min[1], max[2]), Vector3(min[0], max[1], max[2]), root.color);
  136. lines.add(Vector3(max[0], min[1], max[2]), Vector3(max[0], max[1], max[2]), root.color);
  137. lines.add(Vector3(min[0], max[1], min[2]), Vector3(max[0], max[1], min[2]), root.color);
  138. lines.add(Vector3(min[0], max[1], min[2]), Vector3(min[0], max[1], max[2]), root.color);
  139. lines.add(Vector3(max[0], max[1], min[2]), Vector3(max[0], max[1], max[2]), root.color);
  140. lines.add(Vector3(min[0], max[1], max[2]), Vector3(max[0], max[1], max[2]), root.color);
  141. fillLines(lines, &root, min, max);
  142. lines.build();
  143. }
  144. void KDTree::fillLines(Lines& lines, const Node* n, const Vector3& min, const Vector3& max) const {
  145. if(n->childs[0] == nullptr && n->childs[1] == nullptr) {
  146. return;
  147. }
  148. switch(n->splitDim) {
  149. case 0:
  150. lines.add(Vector3(n->splitValue, min[1], min[2]), Vector3(n->splitValue, max[1], min[2]), n->color);
  151. lines.add(Vector3(n->splitValue, max[1], min[2]), Vector3(n->splitValue, max[1], max[2]), n->color);
  152. lines.add(Vector3(n->splitValue, max[1], max[2]), Vector3(n->splitValue, min[1], max[2]), n->color);
  153. lines.add(Vector3(n->splitValue, min[1], max[2]), Vector3(n->splitValue, min[1], min[2]), n->color);
  154. if(n->childs[0] != nullptr) {
  155. fillLines(lines, n->childs[0], min, Vector3(n->splitValue, max[1], max[2]));
  156. }
  157. if(n->childs[1] != nullptr) {
  158. fillLines(lines, n->childs[1], Vector3(n->splitValue, min[1], min[2]), max);
  159. }
  160. break;
  161. case 1:
  162. lines.add(Vector3(min[0], n->splitValue, min[2]), Vector3(max[0], n->splitValue, min[2]), n->color);
  163. lines.add(Vector3(min[0], n->splitValue, min[2]), Vector3(min[0], n->splitValue, max[2]), n->color);
  164. lines.add(Vector3(max[0], n->splitValue, min[2]), Vector3(max[0], n->splitValue, max[2]), n->color);
  165. lines.add(Vector3(min[0], n->splitValue, max[2]), Vector3(max[0], n->splitValue, max[2]), n->color);
  166. if(n->childs[0] != nullptr) {
  167. fillLines(lines, n->childs[0], min, Vector3(max[0], n->splitValue, max[2]));
  168. }
  169. if(n->childs[1] != nullptr) {
  170. fillLines(lines, n->childs[1], Vector3(min[0], n->splitValue, min[2]), max);
  171. }
  172. break;
  173. case 2:
  174. lines.add(Vector3(min[0], min[1], n->splitValue), Vector3(min[0], max[1], n->splitValue), n->color);
  175. lines.add(Vector3(min[0], max[1], n->splitValue), Vector3(max[0], max[1], n->splitValue), n->color);
  176. lines.add(Vector3(max[0], max[1], n->splitValue), Vector3(max[0], min[1], n->splitValue), n->color);
  177. lines.add(Vector3(max[0], min[1], n->splitValue), Vector3(min[0], min[1], n->splitValue), n->color);
  178. if(n->childs[0] != nullptr) {
  179. fillLines(lines, n->childs[0], min, Vector3(max[0], max[1], n->splitValue));
  180. }
  181. if(n->childs[1] != nullptr) {
  182. fillLines(lines, n->childs[1], Vector3(min[0], min[1], n->splitValue), max);
  183. }
  184. break;
  185. }
  186. }
  187. void KDTree::cleanVisited(Node* n) {
  188. if(n->childs[0] != nullptr) {
  189. cleanVisited(n->childs[0]);
  190. }
  191. if(n->childs[1] != nullptr) {
  192. cleanVisited(n->childs[1]);
  193. }
  194. n->color = 0xFFFFFF;
  195. }
  196. bool KDTree::findIntersection(const Vector3& pos, const Vector3& direction) {
  197. minDistance = 9999999.0f;
  198. original = pos;
  199. cleanVisited(&root);
  200. float t = 0.0f;
  201. for(int i = 0; i < 3; i++) {
  202. if(direction[i] == 0.0f) {
  203. continue;
  204. }
  205. t = std::max(t, std::abs((min[i] - pos[i]) / direction[i]));
  206. t = std::max(t, std::abs((max[i] - pos[i]) / direction[i]));
  207. }
  208. return findIntersection(&root, pos, direction, t);
  209. }
  210. Vector3 KDTree::getIntersection() const {
  211. return intersection;
  212. }
  213. KDTree::Triangle KDTree::getIntersectedTriangle() const {
  214. return intersectionTriangle;
  215. }
  216. bool KDTree::findIntersection(Node* n, const Vector3& pos, const Vector3& direction, float tMax) {
  217. if(n == nullptr) {
  218. return false;
  219. }
  220. n->color = 0x00FFFF;
  221. int dim = n->splitDim;
  222. // lessEqual = 0
  223. // greater = 1
  224. bool r = false;
  225. bool first = pos[dim] > n->splitValue;
  226. if(direction[dim] == 0.0f) {
  227. r = findIntersection(n->childs[first], pos, direction, tMax);
  228. } else {
  229. float t = (n->splitValue - pos[dim]) / direction[dim];
  230. if(t >= 0.0f && t < tMax) {
  231. r = findIntersection(n->childs[first], pos, direction, t);
  232. r = findIntersection(n->childs[!first], pos + t * direction, direction, tMax - t) || r;
  233. } else {
  234. r = findIntersection(n->childs[first], pos, direction, tMax);
  235. }
  236. }
  237. for(KDTree::Triangle& tri : n->data) {
  238. float t = testIntersection(pos, direction, tri);
  239. if(t < 0.0f) {
  240. continue;
  241. }
  242. Vector3 i = pos + t * direction;
  243. float distance = static_cast<Vector3> (original - i).squareLength();
  244. if(distance > minDistance) {
  245. continue;
  246. }
  247. intersectionTriangle = tri;
  248. intersection = i;
  249. minDistance = distance;
  250. r = true;
  251. }
  252. return r;
  253. }
  254. float KDTree::orient(const Vector3& a, const Vector3& b, const Vector3& c, const Vector3& d) const {
  255. return static_cast<Vector3> (a - d).dot(static_cast<Vector3> (b - d).cross(c - d));
  256. }
  257. float KDTree::testIntersection(const Vector3& pos, const Vector3& direction, const KDTree::Triangle& t) const {
  258. const float eps = 0.0001f;
  259. Vector3 h1 = t[1] - t[0];
  260. Vector3 h2 = t[2] - t[0];
  261. Vector3 abc = h1.cross(h2);
  262. abc.normalize();
  263. float d = -abc.dot(t[0]);
  264. float check = abc.dot(direction);
  265. if(std::abs(check) < eps) {
  266. return -1.0f;
  267. }
  268. float factor = -(abc.dot(pos) + d) / check;
  269. Vector3 inter = pos + factor * direction;
  270. int sideA = 0;
  271. int sideB = 0;
  272. for(int i = 0; i < 3; i++) {
  273. const Vector3& a = t[i];
  274. const Vector3& b = t[(i + 1) % 3];
  275. const Vector3 c = t[i] + abc;
  276. float o = orient(a, b, c, inter);
  277. if(o < -eps) {
  278. sideA++;
  279. } else if(o > eps) {
  280. sideB++;
  281. } else {
  282. sideA++;
  283. sideB++;
  284. }
  285. }
  286. if(sideA >= 3 || sideB >= 3) {
  287. return factor;
  288. }
  289. return -1.0f;
  290. }