__init__.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693
  1. # freesurfer-surface - Read and Write Surface Files in Freesurfer’s TriangularSurface Format
  2. #
  3. # Copyright (C) 2020 Fabian Peter Hammerle <fabian@hammerle.me>
  4. #
  5. # This program is free software: you can redistribute it and/or modify
  6. # it under the terms of the GNU General Public License as published by
  7. # the Free Software Foundation, either version 3 of the License, or
  8. # any later version.
  9. #
  10. # This program is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. # GNU General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU General Public License
  16. # along with this program. If not, see <https://www.gnu.org/licenses/>.
  17. """
  18. Python Library to Read and Write Surface Files in Freesurfer's TriangularSurface Format
  19. compatible with Freesurfer's MRISwriteTriangularSurface()
  20. https://github.com/freesurfer/freesurfer/blob/release_6_0_0/include/mrisurf.h#L1281
  21. https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/mrisurf.c
  22. https://raw.githubusercontent.com/freesurfer/freesurfer/release_6_0_0/utils/mrisurf.c
  23. Freesurfer
  24. https://surfer.nmr.mgh.harvard.edu/
  25. Edit Surface File
  26. >>> from freesurfer_surface import Surface, Vertex, Triangle
  27. >>>
  28. >>> surface = Surface.read_triangular('bert/surf/lh.pial'))
  29. >>>
  30. >>> vertex_a = surface.add_vertex(Vertex(0.0, 0.0, 0.0))
  31. >>> vertex_b = surface.add_vertex(Vertex(1.0, 1.0, 1.0))
  32. >>> vertex_c = surface.add_vertex(Vertex(2.0, 2.0, 2.0))
  33. >>> surface.triangles.append(Triangle((vertex_a, vertex_b, vertex_c)))
  34. >>>
  35. >>> surface.write_triangular('somewhere/else/lh.pial')
  36. List Labels in Annotation File
  37. >>> from freesurfer_surface import Annotation
  38. >>>
  39. >>> annotation = Annotation.read('bert/label/lh.aparc.annot')
  40. >>> for label in annotation.labels.values():
  41. >>> print(label.index, label.hex_color_code, label.name)
  42. Find Border of Labelled Region
  43. >>> surface = Surface.read_triangular('bert/surf/lh.pial'))
  44. >>> surface.load_annotation_file('bert/label/lh.aparc.annot')
  45. >>> region, = filter(lambda l: l.name == 'precentral',
  46. >>> annotation.labels.values())
  47. >>> print(surface.find_label_border_polygonal_chains(region))
  48. """
  49. import collections
  50. import contextlib
  51. import copy
  52. import datetime
  53. import itertools
  54. import locale
  55. import re
  56. import struct
  57. import typing
  58. import numpy
  59. try:
  60. from freesurfer_surface.version import __version__
  61. except ImportError: # ModuleNotFoundError not available in python<3.6
  62. # package is not installed
  63. __version__ = None
  64. class UnsupportedLocaleSettingError(locale.Error):
  65. pass
  66. @contextlib.contextmanager
  67. def setlocale(temporary_locale):
  68. primary_locale = locale.setlocale(locale.LC_ALL)
  69. try:
  70. yield locale.setlocale(locale.LC_ALL, temporary_locale)
  71. except locale.Error as exc:
  72. if str(exc) == "unsupported locale setting":
  73. raise UnsupportedLocaleSettingError(temporary_locale) from exc
  74. raise exc # pragma: no cover
  75. finally:
  76. locale.setlocale(locale.LC_ALL, primary_locale)
  77. class Vertex(numpy.ndarray):
  78. def __new__(cls, right: float, anterior: float, superior: float):
  79. return numpy.array((right, anterior, superior), dtype=float).view(cls)
  80. @property
  81. def right(self) -> float:
  82. return self[0]
  83. @property
  84. def anterior(self) -> float:
  85. return self[1]
  86. @property
  87. def superior(self) -> float:
  88. return self[2]
  89. @property
  90. def __dict__(self) -> typing.Dict[str, typing.Any]: # type: ignore
  91. # type hint: https://github.com/python/mypy/issues/6523#issuecomment-470733447
  92. return {
  93. "right": self.right,
  94. "anterior": self.anterior,
  95. "superior": self.superior,
  96. }
  97. def __format_coords(self) -> str:
  98. return ", ".join(
  99. "{}={}".format(name, getattr(self, name))
  100. for name in ["right", "anterior", "superior"]
  101. )
  102. def __repr__(self) -> str:
  103. return "{}({})".format(type(self).__name__, self.__format_coords())
  104. def distance_mm(
  105. self, others: typing.Union["Vertex", typing.Iterable["Vertex"], numpy.ndarray]
  106. ) -> numpy.ndarray:
  107. if isinstance(others, Vertex):
  108. others = others.reshape((1, 3))
  109. return numpy.linalg.norm(self - others, axis=1)
  110. class PolygonalCircuit:
  111. def __init__(self, vertex_indices: typing.Iterable[int]):
  112. self._vertex_indices = tuple(vertex_indices)
  113. assert all(isinstance(idx, int) for idx in self._vertex_indices)
  114. @property
  115. def vertex_indices(self):
  116. return self._vertex_indices
  117. def _normalize(self) -> "PolygonalCircuit":
  118. vertex_indices = collections.deque(self.vertex_indices)
  119. vertex_indices.rotate(-numpy.argmin(self.vertex_indices))
  120. if len(vertex_indices) > 2 and vertex_indices[-1] < vertex_indices[1]:
  121. vertex_indices.reverse()
  122. vertex_indices.rotate(1)
  123. return type(self)(vertex_indices)
  124. def __eq__(self, other: object) -> bool:
  125. # pylint: disable=protected-access
  126. return (
  127. isinstance(other, PolygonalCircuit)
  128. and self._normalize().vertex_indices == other._normalize().vertex_indices
  129. )
  130. def __hash__(self) -> int:
  131. # pylint: disable=protected-access
  132. return hash(self._normalize()._vertex_indices)
  133. def adjacent_vertex_indices(
  134. self, vertices_num: int = 2
  135. ) -> typing.Iterable[typing.Tuple[int]]:
  136. vertex_indices_cycle = list(
  137. itertools.islice(
  138. itertools.cycle(self.vertex_indices),
  139. 0,
  140. len(self.vertex_indices) + vertices_num - 1,
  141. )
  142. )
  143. return zip(
  144. *(
  145. itertools.islice(
  146. vertex_indices_cycle, offset, len(self.vertex_indices) + offset
  147. )
  148. for offset in range(vertices_num)
  149. )
  150. )
  151. class LineSegment(PolygonalCircuit):
  152. def __init__(self, indices: typing.Iterable[int]):
  153. super().__init__(indices)
  154. assert len(self.vertex_indices) == 2
  155. def __repr__(self) -> str:
  156. return "LineSegment(vertex_indices={})".format(self.vertex_indices)
  157. class Triangle(PolygonalCircuit):
  158. def __init__(self, indices: typing.Iterable[int]):
  159. super().__init__(indices)
  160. assert len(self.vertex_indices) == 3
  161. def __repr__(self) -> str:
  162. return "Triangle(vertex_indices={})".format(self.vertex_indices)
  163. class PolygonalChainsNotOverlapingError(ValueError):
  164. pass
  165. class PolygonalChain:
  166. def __init__(self, vertex_indices: typing.Iterable[int]):
  167. self.vertex_indices = collections.deque(
  168. vertex_indices
  169. ) # type: typing.Deque[int]
  170. def normalized(self) -> "PolygonalChain":
  171. vertex_indices = list(self.vertex_indices)
  172. min_index = vertex_indices.index(min(vertex_indices))
  173. indices_min_first = vertex_indices[min_index:] + vertex_indices[:min_index]
  174. if indices_min_first[1] < indices_min_first[-1]:
  175. return PolygonalChain(indices_min_first)
  176. return PolygonalChain(indices_min_first[0:1] + indices_min_first[-1:0:-1])
  177. def __eq__(self, other: object) -> bool:
  178. return isinstance(other, PolygonalChain) and (
  179. self.vertex_indices == other.vertex_indices
  180. or self.normalized().vertex_indices == other.normalized().vertex_indices
  181. )
  182. def __repr__(self) -> str:
  183. return "PolygonalChain(vertex_indices={})".format(tuple(self.vertex_indices))
  184. def connect(self, other: "PolygonalChain") -> None:
  185. if self.vertex_indices[-1] == other.vertex_indices[0]:
  186. self.vertex_indices.pop()
  187. self.vertex_indices.extend(other.vertex_indices)
  188. elif self.vertex_indices[-1] == other.vertex_indices[-1]:
  189. self.vertex_indices.pop()
  190. self.vertex_indices.extend(reversed(other.vertex_indices))
  191. elif self.vertex_indices[0] == other.vertex_indices[0]:
  192. self.vertex_indices.popleft()
  193. self.vertex_indices.extendleft(other.vertex_indices)
  194. elif self.vertex_indices[0] == other.vertex_indices[-1]:
  195. self.vertex_indices.popleft()
  196. self.vertex_indices.extendleft(reversed(other.vertex_indices))
  197. else:
  198. raise PolygonalChainsNotOverlapingError()
  199. def adjacent_vertex_indices(
  200. self, vertices_num: int = 2
  201. ) -> typing.Iterator[typing.Tuple[int, ...]]:
  202. return zip(
  203. *(
  204. itertools.islice(self.vertex_indices, offset, len(self.vertex_indices))
  205. for offset in range(vertices_num)
  206. )
  207. )
  208. def segments(self) -> typing.Iterable[LineSegment]:
  209. return map(LineSegment, self.adjacent_vertex_indices(2))
  210. class Label:
  211. # pylint: disable=too-many-arguments
  212. def __init__(
  213. self, index: int, name: str, red: int, green: int, blue: int, transparency: int
  214. ):
  215. self.index = index # type: int
  216. self.name = name # type: str
  217. self.red = red # type: int
  218. self.green = green # type: int
  219. self.blue = blue # type: int
  220. self.transparency = transparency # type: int
  221. @property
  222. def color_code(self) -> int:
  223. if self.index == 0: # unknown
  224. return 0
  225. return int.from_bytes(
  226. (self.red, self.green, self.blue, self.transparency),
  227. byteorder="little",
  228. signed=False,
  229. )
  230. @property
  231. def hex_color_code(self) -> str:
  232. return "#{:02x}{:02x}{:02x}".format(self.red, self.green, self.blue)
  233. def __str__(self) -> str:
  234. return "Label(name={}, index={}, color={})".format(
  235. self.name, self.index, self.hex_color_code
  236. )
  237. def __repr__(self) -> str:
  238. return str(self)
  239. class Annotation:
  240. # pylint: disable=too-few-public-methods
  241. _TAG_OLD_COLORTABLE = b"\0\0\0\x01"
  242. def __init__(self):
  243. self.vertex_label_index = {} # type: Dict[int, int]
  244. self.colortable_path = None # type: Optional[bytes]
  245. self.labels = {} # type: Dict[int, Label]
  246. @staticmethod
  247. def _read_label(stream: typing.BinaryIO) -> Label:
  248. index, name_length = struct.unpack(">II", stream.read(4 * 2))
  249. name = stream.read(name_length - 1).decode()
  250. assert stream.read(1) == b"\0"
  251. red, green, blue, transparency = struct.unpack(">IIII", stream.read(4 * 4))
  252. return Label(
  253. index=index,
  254. name=name,
  255. red=red,
  256. green=green,
  257. blue=blue,
  258. transparency=transparency,
  259. )
  260. def _read(self, stream: typing.BinaryIO) -> None:
  261. # https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles
  262. (annotations_num,) = struct.unpack(">I", stream.read(4))
  263. annotations = [
  264. struct.unpack(">II", stream.read(4 * 2)) for _ in range(annotations_num)
  265. ]
  266. assert stream.read(4) == self._TAG_OLD_COLORTABLE
  267. colortable_version, _, filename_length = struct.unpack(
  268. ">III", stream.read(4 * 3)
  269. )
  270. assert colortable_version > 0 # new version
  271. self.colortable_path = stream.read(filename_length - 1)
  272. assert stream.read(1) == b"\0"
  273. (labels_num,) = struct.unpack(">I", stream.read(4))
  274. self.labels = {
  275. label.index: label
  276. for label in (self._read_label(stream) for _ in range(labels_num))
  277. }
  278. label_index_by_color_code = {
  279. label.color_code: label.index for label in self.labels.values()
  280. }
  281. self.vertex_label_index = {
  282. vertex_index: label_index_by_color_code[color_code]
  283. for vertex_index, color_code in annotations
  284. }
  285. assert not stream.read(1)
  286. @classmethod
  287. def read(cls, annotation_file_path: str) -> "Annotation":
  288. annotation = cls()
  289. with open(annotation_file_path, "rb") as annotation_file:
  290. # pylint: disable=protected-access
  291. annotation._read(annotation_file)
  292. return annotation
  293. class Surface:
  294. # pylint: disable=too-many-instance-attributes
  295. _MAGIC_NUMBER = b"\xff\xff\xfe"
  296. _TAG_CMDLINE = b"\x00\x00\x00\x03"
  297. _TAG_OLD_SURF_GEOM = b"\x00\x00\x00\x14"
  298. _TAG_OLD_USEREALRAS = b"\x00\x00\x00\x02"
  299. _DATETIME_FORMAT = "%a %b %d %H:%M:%S %Y"
  300. def __init__(self):
  301. self.creator = None # type: Optional[bytes]
  302. self.creation_datetime = None # type: Optional[datetime.datetime]
  303. self.vertices = [] # type: List[Vertex]
  304. self.triangles = [] # type: List[Triangle]
  305. self.using_old_real_ras = False # type: bool
  306. self.volume_geometry_info = None # type: Optional[Tuple[bytes]]
  307. self.command_lines = [] # type: List[bytes]
  308. self.annotation = None # type: Optional[Annotation]
  309. @classmethod
  310. def _read_cmdlines(cls, stream: typing.BinaryIO) -> typing.Iterator[bytes]:
  311. while True:
  312. tag = stream.read(4)
  313. if not tag:
  314. return
  315. assert tag == cls._TAG_CMDLINE # might be TAG_GROUP_AVG_SURFACE_AREA
  316. # TAGwrite
  317. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/tags.c#L94
  318. (str_length,) = struct.unpack(">Q", stream.read(8))
  319. yield stream.read(str_length - 1)
  320. assert stream.read(1) == b"\x00"
  321. def _read_triangular(self, stream: typing.BinaryIO):
  322. assert stream.read(3) == self._MAGIC_NUMBER
  323. creation_match = re.match(
  324. rb"^created by (\w+) on (.* \d{4})\n", stream.readline()
  325. )
  326. assert creation_match
  327. self.creator, creation_dt_str = creation_match.groups()
  328. with setlocale("C"):
  329. self.creation_datetime = datetime.datetime.strptime(
  330. creation_dt_str.decode(), self._DATETIME_FORMAT
  331. )
  332. assert stream.read(1) == b"\n"
  333. # fwriteInt
  334. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/fio.c#L290
  335. vertices_num, triangles_num = struct.unpack(">II", stream.read(4 * 2))
  336. self.vertices = [
  337. Vertex(*struct.unpack(">fff", stream.read(4 * 3)))
  338. for _ in range(vertices_num)
  339. ]
  340. self.triangles = [
  341. Triangle(struct.unpack(">III", stream.read(4 * 3)))
  342. for _ in range(triangles_num)
  343. ]
  344. assert all(
  345. vertex_idx < vertices_num
  346. for triangle in self.triangles
  347. for vertex_idx in triangle.vertex_indices
  348. )
  349. assert stream.read(4) == self._TAG_OLD_USEREALRAS
  350. (using_old_real_ras,) = struct.unpack(">I", stream.read(4))
  351. assert using_old_real_ras in [0, 1], using_old_real_ras
  352. self.using_old_real_ras = bool(using_old_real_ras)
  353. assert stream.read(4) == self._TAG_OLD_SURF_GEOM
  354. # writeVolGeom
  355. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/transform.c#L368
  356. self.volume_geometry_info = tuple(stream.readline() for _ in range(8))
  357. self.command_lines = list(self._read_cmdlines(stream))
  358. @classmethod
  359. def read_triangular(cls, surface_file_path: str) -> "Surface":
  360. surface = cls()
  361. with open(surface_file_path, "rb") as surface_file:
  362. # pylint: disable=protected-access
  363. surface._read_triangular(surface_file)
  364. return surface
  365. @classmethod
  366. def _triangular_strftime(cls, creation_datetime: datetime.datetime) -> bytes:
  367. padded_day = "{:>2}".format(creation_datetime.day)
  368. fmt = cls._DATETIME_FORMAT.replace("%d", padded_day)
  369. with setlocale("C"):
  370. return creation_datetime.strftime(fmt).encode()
  371. def write_triangular(
  372. self,
  373. surface_file_path: str,
  374. creation_datetime: typing.Optional[datetime.datetime] = None,
  375. ):
  376. if creation_datetime is None:
  377. creation_datetime = datetime.datetime.now()
  378. with open(surface_file_path, "wb") as surface_file:
  379. surface_file.write(
  380. self._MAGIC_NUMBER
  381. + b"created by "
  382. + self.creator
  383. + b" on "
  384. + self._triangular_strftime(creation_datetime)
  385. + b"\n\n"
  386. + struct.pack(">II", len(self.vertices), len(self.triangles))
  387. )
  388. for vertex in self.vertices:
  389. surface_file.write(struct.pack(">fff", *vertex))
  390. for triangle in self.triangles:
  391. assert all(
  392. vertex_index < len(self.vertices)
  393. for vertex_index in triangle.vertex_indices
  394. )
  395. surface_file.write(struct.pack(">III", *triangle.vertex_indices))
  396. surface_file.write(
  397. self._TAG_OLD_USEREALRAS
  398. + struct.pack(">I", 1 if self.using_old_real_ras else 0)
  399. )
  400. surface_file.write(
  401. self._TAG_OLD_SURF_GEOM + b"".join(self.volume_geometry_info)
  402. )
  403. for command_line in self.command_lines:
  404. surface_file.write(
  405. self._TAG_CMDLINE
  406. + struct.pack(">Q", len(command_line) + 1)
  407. + command_line
  408. + b"\0"
  409. )
  410. def load_annotation_file(self, annotation_file_path: str) -> None:
  411. annotation = Annotation.read(annotation_file_path)
  412. assert len(annotation.vertex_label_index) <= len(self.vertices)
  413. assert max(annotation.vertex_label_index.keys()) < len(self.vertices)
  414. self.annotation = annotation
  415. def add_vertex(self, vertex: Vertex) -> int:
  416. self.vertices.append(vertex)
  417. return len(self.vertices) - 1
  418. def add_rectangle(self, vertex_indices: typing.Iterable[int]) -> None:
  419. vertex_indices = list(vertex_indices)
  420. if len(vertex_indices) == 3:
  421. vertex_indices.append(
  422. self.add_vertex(
  423. self.vertices[vertex_indices[0]]
  424. + self.vertices[vertex_indices[2]]
  425. - self.vertices[vertex_indices[1]]
  426. )
  427. )
  428. assert len(vertex_indices) == 4
  429. self.triangles.append(Triangle(vertex_indices[:3]))
  430. self.triangles.append(Triangle(vertex_indices[2:] + vertex_indices[:1]))
  431. def _triangle_count_by_adjacent_vertex_indices(
  432. self,
  433. ) -> typing.Dict[int, typing.DefaultDict[int, int]]:
  434. counts = {
  435. vertex_index: collections.defaultdict(lambda: 0)
  436. for vertex_index in range(len(self.vertices))
  437. } # type: typing.Dict[int, typing.DefaultDict[int, int]]
  438. for triangle in self.triangles:
  439. for vertex_index_pair in triangle.adjacent_vertex_indices(2):
  440. counts[vertex_index_pair[0]][vertex_index_pair[1]] += 1
  441. counts[vertex_index_pair[1]][vertex_index_pair[0]] += 1
  442. return counts
  443. def find_borders(self) -> typing.Iterator[PolygonalCircuit]:
  444. border_neighbours = {}
  445. for (
  446. vertex_index,
  447. neighbour_counts,
  448. ) in self._triangle_count_by_adjacent_vertex_indices().items():
  449. if not neighbour_counts:
  450. yield PolygonalCircuit((vertex_index,))
  451. else:
  452. neighbours = [
  453. neighbour_index
  454. for neighbour_index, counts in neighbour_counts.items()
  455. if counts != 2
  456. ]
  457. if neighbours:
  458. assert len(neighbours) % 2 == 0, (vertex_index, neighbour_counts)
  459. border_neighbours[vertex_index] = neighbours
  460. while border_neighbours:
  461. vertex_index, neighbour_indices = border_neighbours.popitem()
  462. cycle_indices = [vertex_index]
  463. border_neighbours[vertex_index] = neighbour_indices[1:]
  464. vertex_index = neighbour_indices[0]
  465. while vertex_index != cycle_indices[0]:
  466. neighbour_indices = border_neighbours.pop(vertex_index)
  467. neighbour_indices.remove(cycle_indices[-1])
  468. cycle_indices.append(vertex_index)
  469. if len(neighbour_indices) > 1:
  470. border_neighbours[vertex_index] = neighbour_indices[1:]
  471. vertex_index = neighbour_indices[0]
  472. assert vertex_index in border_neighbours, (
  473. vertex_index,
  474. cycle_indices,
  475. border_neighbours,
  476. )
  477. final_neighbour_indices = border_neighbours.pop(vertex_index)
  478. assert final_neighbour_indices == [cycle_indices[-1]], (
  479. vertex_index,
  480. final_neighbour_indices,
  481. cycle_indices,
  482. )
  483. yield PolygonalCircuit(cycle_indices)
  484. def _get_vertex_label_index(self, vertex_index: int) -> typing.Optional[int]:
  485. return self.annotation.vertex_label_index.get(vertex_index, None)
  486. def _find_label_border_segments(self, label: Label) -> typing.Iterator[LineSegment]:
  487. for triangle in self.triangles:
  488. border_vertex_indices = tuple(
  489. filter(
  490. lambda i: self._get_vertex_label_index(i) == label.index,
  491. triangle.vertex_indices,
  492. )
  493. )
  494. if len(border_vertex_indices) == 2:
  495. yield LineSegment(border_vertex_indices)
  496. _VertexSubindex = typing.Tuple[int, int]
  497. @classmethod
  498. def _duplicate_border(
  499. cls,
  500. neighbour_indices: typing.DefaultDict[
  501. _VertexSubindex, typing.Set[_VertexSubindex]
  502. ],
  503. previous_index: _VertexSubindex,
  504. current_index: _VertexSubindex,
  505. junction_counter: int,
  506. ) -> None:
  507. split_index = (current_index[0], junction_counter)
  508. neighbour_indices[previous_index].add(split_index)
  509. neighbour_indices[split_index].add(previous_index)
  510. next_index, *extra_indices = filter(
  511. lambda i: i != previous_index, neighbour_indices[current_index]
  512. )
  513. if extra_indices:
  514. neighbour_indices[next_index].add(split_index)
  515. neighbour_indices[split_index].add(next_index)
  516. neighbour_indices[next_index].remove(current_index)
  517. neighbour_indices[current_index].remove(next_index)
  518. return
  519. cls._duplicate_border(
  520. neighbour_indices=neighbour_indices,
  521. previous_index=split_index,
  522. current_index=next_index,
  523. junction_counter=junction_counter,
  524. )
  525. def find_label_border_polygonal_chains(
  526. self, label: Label
  527. ) -> typing.Iterator[PolygonalChain]:
  528. neighbour_indices = collections.defaultdict(
  529. set
  530. ) # type: typing.DefaultDict[_VertexSubindex, typing.Set[_VertexSubindex]] # type: ignore
  531. for segment in self._find_label_border_segments(label):
  532. vertex_indices = [(i, 0) for i in segment.vertex_indices]
  533. neighbour_indices[vertex_indices[0]].add(vertex_indices[1])
  534. neighbour_indices[vertex_indices[1]].add(vertex_indices[0])
  535. junction_counter = 0
  536. found_leaf = True
  537. while found_leaf:
  538. found_leaf = False
  539. for leaf_index, leaf_neighbour_indices in neighbour_indices.items():
  540. if len(leaf_neighbour_indices) == 1:
  541. found_leaf = True
  542. junction_counter += 1
  543. self._duplicate_border(
  544. neighbour_indices=neighbour_indices,
  545. previous_index=leaf_index,
  546. # pylint: disable=stop-iteration-return; false positive, has 1 item
  547. current_index=next(iter(leaf_neighbour_indices)),
  548. junction_counter=junction_counter,
  549. )
  550. break
  551. assert all(len(n) == 2 for n in neighbour_indices.values()), neighbour_indices
  552. while neighbour_indices:
  553. # pylint: disable=stop-iteration-return; has >= 1 item
  554. chain = collections.deque([next(iter(neighbour_indices.keys()))])
  555. chain.append(neighbour_indices[chain[0]].pop())
  556. neighbour_indices[chain[1]].remove(chain[0])
  557. while chain[0] != chain[-1]:
  558. previous_index = chain[-1]
  559. next_index = neighbour_indices[previous_index].pop()
  560. neighbour_indices[next_index].remove(previous_index)
  561. chain.append(next_index)
  562. assert not neighbour_indices[previous_index], neighbour_indices[
  563. previous_index
  564. ]
  565. del neighbour_indices[previous_index]
  566. assert not neighbour_indices[chain[0]], neighbour_indices[chain[0]]
  567. del neighbour_indices[chain[0]]
  568. chain.pop()
  569. yield PolygonalChain(v[0] for v in chain)
  570. def _unused_vertices(self) -> typing.Set[int]:
  571. vertex_indices = set(range(len(self.vertices)))
  572. for triangle in self.triangles:
  573. for vertex_index in triangle.vertex_indices:
  574. vertex_indices.discard(vertex_index)
  575. return vertex_indices
  576. def remove_unused_vertices(self) -> None:
  577. vertex_index_conversion = [0] * len(self.vertices)
  578. for vertex_index in sorted(self._unused_vertices(), reverse=True):
  579. del self.vertices[vertex_index]
  580. vertex_index_conversion[vertex_index] -= 1
  581. vertex_index_conversion = numpy.cumsum(vertex_index_conversion)
  582. for triangle_index in range(len(self.triangles)):
  583. self.triangles[triangle_index] = Triangle(
  584. map(
  585. lambda i: i + int(vertex_index_conversion[i]),
  586. self.triangles[triangle_index].vertex_indices,
  587. )
  588. )
  589. def select_vertices(
  590. self, vertex_indices: typing.Iterable[int]
  591. ) -> typing.List[Vertex]:
  592. return [self.vertices[idx] for idx in vertex_indices]
  593. @staticmethod
  594. def unite(surfaces: typing.Iterable["Surface"]) -> "Surface":
  595. surfaces_iter = iter(surfaces)
  596. union = copy.deepcopy(next(surfaces_iter))
  597. for surface in surfaces_iter:
  598. vertex_index_offset = len(union.vertices)
  599. union.vertices.extend(surface.vertices)
  600. union.triangles.extend(
  601. Triangle(
  602. vertex_idx + vertex_index_offset
  603. for vertex_idx in triangle.vertex_indices
  604. )
  605. for triangle in surface.triangles
  606. )
  607. return union