__init__.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699
  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. from __future__ import annotations
  50. import collections
  51. import contextlib
  52. import copy
  53. import dataclasses
  54. import datetime
  55. import itertools
  56. import locale
  57. import re
  58. import struct
  59. import typing
  60. import numpy
  61. try:
  62. from freesurfer_surface.version import __version__
  63. except ModuleNotFoundError:
  64. # package is not installed
  65. __version__ = None
  66. class UnsupportedLocaleSettingError(locale.Error):
  67. pass
  68. @contextlib.contextmanager
  69. def setlocale(temporary_locale):
  70. primary_locale = locale.setlocale(locale.LC_ALL)
  71. try:
  72. yield locale.setlocale(locale.LC_ALL, temporary_locale)
  73. except locale.Error as exc:
  74. if str(exc) == "unsupported locale setting":
  75. raise UnsupportedLocaleSettingError(temporary_locale) from exc
  76. raise exc # pragma: no cover
  77. finally:
  78. locale.setlocale(locale.LC_ALL, primary_locale)
  79. class Vertex(numpy.ndarray):
  80. def __new__(cls, right: float, anterior: float, superior: float):
  81. return numpy.array((right, anterior, superior), dtype=float).view(cls)
  82. @property
  83. def right(self) -> float:
  84. return self[0]
  85. @property
  86. def anterior(self) -> float:
  87. return self[1]
  88. @property
  89. def superior(self) -> float:
  90. return self[2]
  91. @property
  92. def __dict__(self) -> typing.Dict[str, typing.Any]: # type: ignore
  93. # type hint: https://github.com/python/mypy/issues/6523#issuecomment-470733447
  94. return {
  95. "right": self.right,
  96. "anterior": self.anterior,
  97. "superior": self.superior,
  98. }
  99. def __format_coords(self) -> str:
  100. return ", ".join(
  101. f"{name}={getattr(self, name)}"
  102. for name in ["right", "anterior", "superior"]
  103. )
  104. def __repr__(self) -> str:
  105. return f"{type(self).__name__}({self.__format_coords()})"
  106. def distance_mm(
  107. self, others: typing.Union[Vertex, typing.Iterable[Vertex], numpy.ndarray]
  108. ) -> numpy.ndarray:
  109. if isinstance(others, Vertex):
  110. others = others.reshape((1, 3))
  111. return numpy.linalg.norm(self - others, axis=1)
  112. class PolygonalCircuit:
  113. def __init__(self, vertex_indices: typing.Iterable[int]):
  114. self._vertex_indices = tuple(vertex_indices)
  115. assert all(isinstance(idx, int) for idx in self._vertex_indices)
  116. @property
  117. def vertex_indices(self):
  118. return self._vertex_indices
  119. def _normalize(self) -> PolygonalCircuit:
  120. vertex_indices = collections.deque(self.vertex_indices)
  121. vertex_indices.rotate(int(-numpy.argmin(self.vertex_indices)))
  122. if len(vertex_indices) > 2 and vertex_indices[-1] < vertex_indices[1]:
  123. vertex_indices.reverse()
  124. vertex_indices.rotate(1)
  125. return type(self)(vertex_indices)
  126. def __eq__(self, other: object) -> bool:
  127. # pylint: disable=protected-access
  128. return (
  129. isinstance(other, PolygonalCircuit)
  130. and self._normalize().vertex_indices == other._normalize().vertex_indices
  131. )
  132. def __hash__(self) -> int:
  133. # pylint: disable=protected-access
  134. return hash(self._normalize()._vertex_indices)
  135. def adjacent_vertex_indices(
  136. self, vertices_num: int = 2
  137. ) -> typing.Iterable[typing.Tuple[int, ...]]:
  138. vertex_indices_cycle = list(
  139. itertools.islice(
  140. itertools.cycle(self.vertex_indices),
  141. 0,
  142. len(self.vertex_indices) + vertices_num - 1,
  143. )
  144. )
  145. return zip(
  146. *(
  147. itertools.islice(
  148. vertex_indices_cycle, offset, len(self.vertex_indices) + offset
  149. )
  150. for offset in range(vertices_num)
  151. )
  152. )
  153. class LineSegment(PolygonalCircuit):
  154. def __init__(self, indices: typing.Iterable[int]):
  155. super().__init__(indices)
  156. assert len(self.vertex_indices) == 2
  157. def __repr__(self) -> str:
  158. return f"LineSegment(vertex_indices={self.vertex_indices})"
  159. class Triangle(PolygonalCircuit):
  160. def __init__(self, indices: typing.Iterable[int]):
  161. super().__init__(indices)
  162. assert len(self.vertex_indices) == 3
  163. def __repr__(self) -> str:
  164. return f"Triangle(vertex_indices={self.vertex_indices})"
  165. class PolygonalChainsNotOverlapingError(ValueError):
  166. pass
  167. class PolygonalChain:
  168. def __init__(self, vertex_indices: typing.Iterable[int]):
  169. self.vertex_indices: typing.Deque[int] = collections.deque(vertex_indices)
  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 f"PolygonalChain(vertex_indices={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. @dataclasses.dataclass
  211. class Label:
  212. index: int
  213. name: str
  214. red: int
  215. green: int
  216. blue: int
  217. transparency: int
  218. @property
  219. def color_code(self) -> int:
  220. if self.index == 0: # unknown
  221. return 0
  222. return int.from_bytes(
  223. (self.red, self.green, self.blue, self.transparency),
  224. byteorder="little",
  225. signed=False,
  226. )
  227. @property
  228. def hex_color_code(self) -> str:
  229. return f"#{self.red:02x}{self.green:02x}{self.blue:02x}"
  230. def __str__(self) -> str:
  231. return (
  232. f"Label(name={self.name}, index={self.index}, color={self.hex_color_code})"
  233. )
  234. def __repr__(self) -> str:
  235. return str(self)
  236. class Annotation:
  237. # pylint: disable=too-few-public-methods
  238. _TAG_OLD_COLORTABLE = b"\0\0\0\x01"
  239. def __init__(self):
  240. self.vertex_label_index: typing.Dict[int, int] = {}
  241. self.colortable_path: typing.Optional[bytes] = None
  242. self.labels: typing.Dict[int, Label] = {}
  243. @staticmethod
  244. def _read_label(stream: typing.BinaryIO) -> Label:
  245. index, name_length = struct.unpack(">II", stream.read(4 * 2))
  246. name = stream.read(name_length - 1).decode()
  247. assert stream.read(1) == b"\0"
  248. red, green, blue, transparency = struct.unpack(">IIII", stream.read(4 * 4))
  249. return Label(
  250. index=index,
  251. name=name,
  252. red=red,
  253. green=green,
  254. blue=blue,
  255. transparency=transparency,
  256. )
  257. def _read(self, stream: typing.BinaryIO) -> None:
  258. # https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles
  259. (annotations_num,) = struct.unpack(">I", stream.read(4))
  260. annotations = [
  261. struct.unpack(">II", stream.read(4 * 2)) for _ in range(annotations_num)
  262. ]
  263. assert stream.read(4) == self._TAG_OLD_COLORTABLE
  264. colortable_version, _, filename_length = struct.unpack(
  265. ">III", stream.read(4 * 3)
  266. )
  267. assert colortable_version > 0 # new version
  268. self.colortable_path = stream.read(filename_length - 1)
  269. assert stream.read(1) == b"\0"
  270. (labels_num,) = struct.unpack(">I", stream.read(4))
  271. self.labels = {
  272. label.index: label
  273. for label in (self._read_label(stream) for _ in range(labels_num))
  274. }
  275. label_index_by_color_code = {
  276. label.color_code: label.index for label in self.labels.values()
  277. }
  278. self.vertex_label_index = {
  279. vertex_index: label_index_by_color_code[color_code]
  280. for vertex_index, color_code in annotations
  281. }
  282. assert not stream.read(1)
  283. @classmethod
  284. def read(cls, annotation_file_path: str) -> "Annotation":
  285. annotation = cls()
  286. with open(annotation_file_path, "rb") as annotation_file:
  287. # pylint: disable=protected-access
  288. annotation._read(annotation_file)
  289. return annotation
  290. class Surface:
  291. # pylint: disable=too-many-instance-attributes
  292. _MAGIC_NUMBER = b"\xff\xff\xfe"
  293. _TAG_CMDLINE = b"\x00\x00\x00\x03"
  294. _TAG_OLD_SURF_GEOM = b"\x00\x00\x00\x14"
  295. _TAG_OLD_USEREALRAS = b"\x00\x00\x00\x02"
  296. _DATETIME_FORMAT = "%a %b %d %H:%M:%S %Y"
  297. def __init__(self):
  298. self.creator: bytes = b"pypi.org/project/freesurfer-surface/"
  299. self.creation_datetime: typing.Optional[datetime.datetime] = None
  300. self.vertices: typing.List[Vertex] = []
  301. self.triangles: typing.List[Triangle] = []
  302. self.using_old_real_ras: bool = False
  303. self.volume_geometry_info: typing.Optional[typing.Tuple[bytes, ...]] = None
  304. self.command_lines: typing.List[bytes] = []
  305. self.annotation: typing.Optional[Annotation] = None
  306. @classmethod
  307. def _read_cmdlines(cls, stream: typing.BinaryIO) -> typing.Iterator[bytes]:
  308. while True:
  309. tag = stream.read(4)
  310. if not tag:
  311. return
  312. assert tag == cls._TAG_CMDLINE # might be TAG_GROUP_AVG_SURFACE_AREA
  313. # TAGwrite
  314. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/tags.c#L94
  315. (str_length,) = struct.unpack(">Q", stream.read(8))
  316. yield stream.read(str_length - 1)
  317. assert stream.read(1) == b"\x00"
  318. def _read_triangular(self, stream: typing.BinaryIO):
  319. assert stream.read(3) == self._MAGIC_NUMBER
  320. creation_match = re.match(
  321. rb"^created by (\w+) on (.* \d{4})\n", stream.readline()
  322. )
  323. assert creation_match
  324. self.creator, creation_dt_str = creation_match.groups()
  325. with setlocale("C"):
  326. self.creation_datetime = datetime.datetime.strptime(
  327. creation_dt_str.decode(), self._DATETIME_FORMAT
  328. )
  329. assert stream.read(1) == b"\n"
  330. # fwriteInt
  331. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/fio.c#L290
  332. vertices_num, triangles_num = struct.unpack(">II", stream.read(4 * 2))
  333. self.vertices = [
  334. Vertex(*struct.unpack(">fff", stream.read(4 * 3)))
  335. for _ in range(vertices_num)
  336. ]
  337. self.triangles = [
  338. Triangle(struct.unpack(">III", stream.read(4 * 3)))
  339. for _ in range(triangles_num)
  340. ]
  341. assert all(
  342. vertex_idx < vertices_num
  343. for triangle in self.triangles
  344. for vertex_idx in triangle.vertex_indices
  345. )
  346. assert stream.read(4) == self._TAG_OLD_USEREALRAS
  347. (using_old_real_ras,) = struct.unpack(">I", stream.read(4))
  348. assert using_old_real_ras in [0, 1], using_old_real_ras
  349. self.using_old_real_ras = bool(using_old_real_ras)
  350. assert stream.read(4) == self._TAG_OLD_SURF_GEOM
  351. # writeVolGeom
  352. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/transform.c#L368
  353. self.volume_geometry_info = tuple(stream.readline() for _ in range(8))
  354. self.command_lines = list(self._read_cmdlines(stream))
  355. @classmethod
  356. def read_triangular(cls, surface_file_path: str) -> "Surface":
  357. surface = cls()
  358. with open(surface_file_path, "rb") as surface_file:
  359. # pylint: disable=protected-access
  360. surface._read_triangular(surface_file)
  361. return surface
  362. @classmethod
  363. def _triangular_strftime(cls, creation_datetime: datetime.datetime) -> bytes:
  364. padded_day = f"{creation_datetime.day:>2}"
  365. fmt = cls._DATETIME_FORMAT.replace("%d", padded_day)
  366. with setlocale("C"):
  367. return creation_datetime.strftime(fmt).encode()
  368. def write_triangular(
  369. self,
  370. surface_file_path: str,
  371. creation_datetime: typing.Optional[datetime.datetime] = None,
  372. ):
  373. if creation_datetime is None:
  374. creation_datetime = datetime.datetime.now()
  375. with open(surface_file_path, "wb") as surface_file:
  376. surface_file.write(
  377. self._MAGIC_NUMBER
  378. + b"created by "
  379. + self.creator
  380. + b" on "
  381. + self._triangular_strftime(creation_datetime)
  382. + b"\n\n"
  383. + struct.pack(">II", len(self.vertices), len(self.triangles))
  384. )
  385. for vertex in self.vertices:
  386. surface_file.write(struct.pack(">fff", *vertex))
  387. for triangle in self.triangles:
  388. assert all(
  389. vertex_index < len(self.vertices)
  390. for vertex_index in triangle.vertex_indices
  391. )
  392. surface_file.write(struct.pack(">III", *triangle.vertex_indices))
  393. surface_file.write(
  394. self._TAG_OLD_USEREALRAS
  395. + struct.pack(">I", 1 if self.using_old_real_ras else 0)
  396. )
  397. if not self.volume_geometry_info:
  398. raise ValueError(
  399. "Missing geometry information (set attribute `volume_geometry_info`)"
  400. )
  401. surface_file.write(
  402. self._TAG_OLD_SURF_GEOM + b"".join(self.volume_geometry_info)
  403. )
  404. for command_line in self.command_lines:
  405. surface_file.write(
  406. self._TAG_CMDLINE
  407. + struct.pack(">Q", len(command_line) + 1)
  408. + command_line
  409. + b"\0"
  410. )
  411. def load_annotation_file(self, annotation_file_path: str) -> None:
  412. annotation = Annotation.read(annotation_file_path)
  413. assert len(annotation.vertex_label_index) <= len(self.vertices)
  414. assert max(annotation.vertex_label_index.keys()) < len(self.vertices)
  415. self.annotation = annotation
  416. def add_vertex(self, vertex: Vertex) -> int:
  417. self.vertices.append(vertex)
  418. return len(self.vertices) - 1
  419. def add_rectangle(self, vertex_indices: typing.Iterable[int]) -> None:
  420. vertex_indices = list(vertex_indices)
  421. if len(vertex_indices) == 3:
  422. vertex_indices.append(
  423. self.add_vertex(
  424. self.vertices[vertex_indices[0]]
  425. + self.vertices[vertex_indices[2]]
  426. - self.vertices[vertex_indices[1]]
  427. )
  428. )
  429. assert len(vertex_indices) == 4
  430. self.triangles.append(Triangle(vertex_indices[:3]))
  431. self.triangles.append(Triangle(vertex_indices[2:] + vertex_indices[:1]))
  432. def _triangle_count_by_adjacent_vertex_indices(
  433. self,
  434. ) -> typing.Dict[int, typing.DefaultDict[int, int]]:
  435. counts: typing.Dict[int, typing.DefaultDict[int, int]] = {
  436. vertex_index: collections.defaultdict(lambda: 0)
  437. for vertex_index in range(len(self.vertices))
  438. }
  439. for triangle in self.triangles:
  440. for vertex_index_pair in triangle.adjacent_vertex_indices(2):
  441. counts[vertex_index_pair[0]][vertex_index_pair[1]] += 1
  442. counts[vertex_index_pair[1]][vertex_index_pair[0]] += 1
  443. return counts
  444. def find_borders(self) -> typing.Iterator[PolygonalCircuit]:
  445. border_neighbours = {}
  446. for (
  447. vertex_index,
  448. neighbour_counts,
  449. ) in self._triangle_count_by_adjacent_vertex_indices().items():
  450. if not neighbour_counts:
  451. yield PolygonalCircuit((vertex_index,))
  452. else:
  453. neighbours = [
  454. neighbour_index
  455. for neighbour_index, counts in neighbour_counts.items()
  456. if counts != 2
  457. ]
  458. if neighbours:
  459. assert len(neighbours) % 2 == 0, (vertex_index, neighbour_counts)
  460. border_neighbours[vertex_index] = neighbours
  461. while border_neighbours:
  462. vertex_index, neighbour_indices = border_neighbours.popitem()
  463. cycle_indices = [vertex_index]
  464. border_neighbours[vertex_index] = neighbour_indices[1:]
  465. vertex_index = neighbour_indices[0]
  466. while vertex_index != cycle_indices[0]:
  467. neighbour_indices = border_neighbours.pop(vertex_index)
  468. neighbour_indices.remove(cycle_indices[-1])
  469. cycle_indices.append(vertex_index)
  470. if len(neighbour_indices) > 1:
  471. border_neighbours[vertex_index] = neighbour_indices[1:]
  472. vertex_index = neighbour_indices[0]
  473. assert vertex_index in border_neighbours, (
  474. vertex_index,
  475. cycle_indices,
  476. border_neighbours,
  477. )
  478. final_neighbour_indices = border_neighbours.pop(vertex_index)
  479. assert final_neighbour_indices == [cycle_indices[-1]], (
  480. vertex_index,
  481. final_neighbour_indices,
  482. cycle_indices,
  483. )
  484. yield PolygonalCircuit(cycle_indices)
  485. def _get_vertex_label_index(self, vertex_index: int) -> typing.Optional[int]:
  486. if not self.annotation:
  487. raise RuntimeError(
  488. "Missing annotation (call method `load_annotation_file` first)."
  489. )
  490. return self.annotation.vertex_label_index.get(vertex_index, None)
  491. def _find_label_border_segments(self, label: Label) -> typing.Iterator[LineSegment]:
  492. for triangle in self.triangles:
  493. border_vertex_indices = tuple(
  494. filter(
  495. lambda i: self._get_vertex_label_index(i) == label.index,
  496. triangle.vertex_indices,
  497. )
  498. )
  499. if len(border_vertex_indices) == 2:
  500. yield LineSegment(border_vertex_indices)
  501. _VertexSubindex = typing.Tuple[int, int]
  502. @classmethod
  503. def _duplicate_border(
  504. cls,
  505. neighbour_indices: typing.DefaultDict[
  506. _VertexSubindex, typing.Set[_VertexSubindex]
  507. ],
  508. previous_index: _VertexSubindex,
  509. current_index: _VertexSubindex,
  510. junction_counter: int,
  511. ) -> None:
  512. split_index = (current_index[0], junction_counter)
  513. neighbour_indices[previous_index].add(split_index)
  514. neighbour_indices[split_index].add(previous_index)
  515. next_index, *extra_indices = filter(
  516. lambda i: i != previous_index, neighbour_indices[current_index]
  517. )
  518. if extra_indices:
  519. neighbour_indices[next_index].add(split_index)
  520. neighbour_indices[split_index].add(next_index)
  521. neighbour_indices[next_index].remove(current_index)
  522. neighbour_indices[current_index].remove(next_index)
  523. return
  524. cls._duplicate_border(
  525. neighbour_indices=neighbour_indices,
  526. previous_index=split_index,
  527. current_index=next_index,
  528. junction_counter=junction_counter,
  529. )
  530. def find_label_border_polygonal_chains(
  531. self, label: Label
  532. ) -> typing.Iterator[PolygonalChain]:
  533. neighbour_indices: ( # type: ignore
  534. typing.DefaultDict[self._VertexSubindex, typing.Set[self._VertexSubindex]]
  535. ) = collections.defaultdict(set)
  536. for segment in self._find_label_border_segments(label):
  537. vertex_indices = [(i, 0) for i in segment.vertex_indices]
  538. neighbour_indices[vertex_indices[0]].add(vertex_indices[1])
  539. neighbour_indices[vertex_indices[1]].add(vertex_indices[0])
  540. junction_counter = 0
  541. found_leaf = True
  542. while found_leaf:
  543. found_leaf = False
  544. for leaf_index, leaf_neighbour_indices in neighbour_indices.items():
  545. if len(leaf_neighbour_indices) == 1:
  546. found_leaf = True
  547. junction_counter += 1
  548. self._duplicate_border(
  549. neighbour_indices=neighbour_indices,
  550. previous_index=leaf_index,
  551. # pylint: disable=stop-iteration-return; false positive, has 1 item
  552. current_index=next(iter(leaf_neighbour_indices)),
  553. junction_counter=junction_counter,
  554. )
  555. break
  556. assert all(len(n) == 2 for n in neighbour_indices.values()), neighbour_indices
  557. while neighbour_indices:
  558. # pylint: disable=stop-iteration-return; has >= 1 item
  559. chain = collections.deque([next(iter(neighbour_indices.keys()))])
  560. chain.append(neighbour_indices[chain[0]].pop())
  561. neighbour_indices[chain[1]].remove(chain[0])
  562. while chain[0] != chain[-1]:
  563. previous_index = chain[-1]
  564. next_index = neighbour_indices[previous_index].pop()
  565. neighbour_indices[next_index].remove(previous_index)
  566. chain.append(next_index)
  567. assert not neighbour_indices[previous_index], neighbour_indices[
  568. previous_index
  569. ]
  570. del neighbour_indices[previous_index]
  571. assert not neighbour_indices[chain[0]], neighbour_indices[chain[0]]
  572. del neighbour_indices[chain[0]]
  573. chain.pop()
  574. yield PolygonalChain(v[0] for v in chain)
  575. def _unused_vertices(self) -> typing.Set[int]:
  576. vertex_indices = set(range(len(self.vertices)))
  577. for triangle in self.triangles:
  578. for vertex_index in triangle.vertex_indices:
  579. vertex_indices.discard(vertex_index)
  580. return vertex_indices
  581. def remove_unused_vertices(self) -> None:
  582. vertex_index_conversion = numpy.zeros(len(self.vertices), dtype=int)
  583. for vertex_index in sorted(self._unused_vertices(), reverse=True):
  584. del self.vertices[vertex_index]
  585. vertex_index_conversion[vertex_index] -= 1
  586. vertex_index_conversion = numpy.cumsum(vertex_index_conversion)
  587. for triangle_index, triangle in enumerate(self.triangles):
  588. self.triangles[triangle_index] = Triangle(
  589. map(
  590. lambda i: i + int(vertex_index_conversion[i]),
  591. triangle.vertex_indices,
  592. )
  593. )
  594. def select_vertices(
  595. self, vertex_indices: typing.Iterable[int]
  596. ) -> typing.List[Vertex]:
  597. return [self.vertices[idx] for idx in vertex_indices]
  598. @staticmethod
  599. def unite(surfaces: typing.Iterable["Surface"]) -> "Surface":
  600. surfaces_iter = iter(surfaces)
  601. union = copy.deepcopy(next(surfaces_iter))
  602. for surface in surfaces_iter:
  603. vertex_index_offset = len(union.vertices)
  604. union.vertices.extend(surface.vertices)
  605. union.triangles.extend(
  606. Triangle(
  607. vertex_idx + vertex_index_offset
  608. for vertex_idx in triangle.vertex_indices
  609. )
  610. for triangle in surface.triangles
  611. )
  612. return union