__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. 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 ModuleNotFoundError:
  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. f"{name}={getattr(self, name)}"
  100. for name in ["right", "anterior", "superior"]
  101. )
  102. def __repr__(self) -> str:
  103. return f"{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 f"LineSegment(vertex_indices={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 f"Triangle(vertex_indices={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: typing.Deque[int] = collections.deque(vertex_indices)
  168. def normalized(self) -> "PolygonalChain":
  169. vertex_indices = list(self.vertex_indices)
  170. min_index = vertex_indices.index(min(vertex_indices))
  171. indices_min_first = vertex_indices[min_index:] + vertex_indices[:min_index]
  172. if indices_min_first[1] < indices_min_first[-1]:
  173. return PolygonalChain(indices_min_first)
  174. return PolygonalChain(indices_min_first[0:1] + indices_min_first[-1:0:-1])
  175. def __eq__(self, other: object) -> bool:
  176. return isinstance(other, PolygonalChain) and (
  177. self.vertex_indices == other.vertex_indices
  178. or self.normalized().vertex_indices == other.normalized().vertex_indices
  179. )
  180. def __repr__(self) -> str:
  181. return f"PolygonalChain(vertex_indices={tuple(self.vertex_indices)})"
  182. def connect(self, other: "PolygonalChain") -> None:
  183. if self.vertex_indices[-1] == other.vertex_indices[0]:
  184. self.vertex_indices.pop()
  185. self.vertex_indices.extend(other.vertex_indices)
  186. elif self.vertex_indices[-1] == other.vertex_indices[-1]:
  187. self.vertex_indices.pop()
  188. self.vertex_indices.extend(reversed(other.vertex_indices))
  189. elif self.vertex_indices[0] == other.vertex_indices[0]:
  190. self.vertex_indices.popleft()
  191. self.vertex_indices.extendleft(other.vertex_indices)
  192. elif self.vertex_indices[0] == other.vertex_indices[-1]:
  193. self.vertex_indices.popleft()
  194. self.vertex_indices.extendleft(reversed(other.vertex_indices))
  195. else:
  196. raise PolygonalChainsNotOverlapingError()
  197. def adjacent_vertex_indices(
  198. self, vertices_num: int = 2
  199. ) -> typing.Iterator[typing.Tuple[int, ...]]:
  200. return zip(
  201. *(
  202. itertools.islice(self.vertex_indices, offset, len(self.vertex_indices))
  203. for offset in range(vertices_num)
  204. )
  205. )
  206. def segments(self) -> typing.Iterable[LineSegment]:
  207. return map(LineSegment, self.adjacent_vertex_indices(2))
  208. class Label:
  209. # pylint: disable=too-many-arguments
  210. def __init__(
  211. self, index: int, name: str, red: int, green: int, blue: int, transparency: int
  212. ):
  213. self.index: int = index
  214. self.name: str = name
  215. self.red: int = red
  216. self.green: int = green
  217. self.blue: int = blue
  218. self.transparency: int = transparency
  219. @property
  220. def color_code(self) -> int:
  221. if self.index == 0: # unknown
  222. return 0
  223. return int.from_bytes(
  224. (self.red, self.green, self.blue, self.transparency),
  225. byteorder="little",
  226. signed=False,
  227. )
  228. @property
  229. def hex_color_code(self) -> str:
  230. return f"#{self.red:02x}{self.green:02x}{self.blue:02x}"
  231. def __str__(self) -> str:
  232. return (
  233. f"Label(name={self.name}, index={self.index}, color={self.hex_color_code})"
  234. )
  235. def __repr__(self) -> str:
  236. return str(self)
  237. class Annotation:
  238. # pylint: disable=too-few-public-methods
  239. _TAG_OLD_COLORTABLE = b"\0\0\0\x01"
  240. def __init__(self):
  241. self.vertex_label_index: typing.Dict[int, int] = {}
  242. self.colortable_path: typing.Optional[bytes] = None
  243. self.labels: typing.Dict[int, Label] = {}
  244. @staticmethod
  245. def _read_label(stream: typing.BinaryIO) -> Label:
  246. index, name_length = struct.unpack(">II", stream.read(4 * 2))
  247. name = stream.read(name_length - 1).decode()
  248. assert stream.read(1) == b"\0"
  249. red, green, blue, transparency = struct.unpack(">IIII", stream.read(4 * 4))
  250. return Label(
  251. index=index,
  252. name=name,
  253. red=red,
  254. green=green,
  255. blue=blue,
  256. transparency=transparency,
  257. )
  258. def _read(self, stream: typing.BinaryIO) -> None:
  259. # https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles
  260. (annotations_num,) = struct.unpack(">I", stream.read(4))
  261. annotations = [
  262. struct.unpack(">II", stream.read(4 * 2)) for _ in range(annotations_num)
  263. ]
  264. assert stream.read(4) == self._TAG_OLD_COLORTABLE
  265. colortable_version, _, filename_length = struct.unpack(
  266. ">III", stream.read(4 * 3)
  267. )
  268. assert colortable_version > 0 # new version
  269. self.colortable_path = stream.read(filename_length - 1)
  270. assert stream.read(1) == b"\0"
  271. (labels_num,) = struct.unpack(">I", stream.read(4))
  272. self.labels = {
  273. label.index: label
  274. for label in (self._read_label(stream) for _ in range(labels_num))
  275. }
  276. label_index_by_color_code = {
  277. label.color_code: label.index for label in self.labels.values()
  278. }
  279. self.vertex_label_index = {
  280. vertex_index: label_index_by_color_code[color_code]
  281. for vertex_index, color_code in annotations
  282. }
  283. assert not stream.read(1)
  284. @classmethod
  285. def read(cls, annotation_file_path: str) -> "Annotation":
  286. annotation = cls()
  287. with open(annotation_file_path, "rb") as annotation_file:
  288. # pylint: disable=protected-access
  289. annotation._read(annotation_file)
  290. return annotation
  291. class Surface:
  292. # pylint: disable=too-many-instance-attributes
  293. _MAGIC_NUMBER = b"\xff\xff\xfe"
  294. _TAG_CMDLINE = b"\x00\x00\x00\x03"
  295. _TAG_OLD_SURF_GEOM = b"\x00\x00\x00\x14"
  296. _TAG_OLD_USEREALRAS = b"\x00\x00\x00\x02"
  297. _DATETIME_FORMAT = "%a %b %d %H:%M:%S %Y"
  298. def __init__(self):
  299. self.creator: bytes = b"pypi.org/project/freesurfer-surface/"
  300. self.creation_datetime: typing.Optional[datetime.datetime] = None
  301. self.vertices: typing.List[Vertex] = []
  302. self.triangles: typing.List[Triangle] = []
  303. self.using_old_real_ras: bool = False
  304. self.volume_geometry_info: typing.Optional[typing.Tuple[bytes, ...]] = None
  305. self.command_lines: typing.List[bytes] = []
  306. self.annotation: typing.Optional[Annotation] = None
  307. @classmethod
  308. def _read_cmdlines(cls, stream: typing.BinaryIO) -> typing.Iterator[bytes]:
  309. while True:
  310. tag = stream.read(4)
  311. if not tag:
  312. return
  313. assert tag == cls._TAG_CMDLINE # might be TAG_GROUP_AVG_SURFACE_AREA
  314. # TAGwrite
  315. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/tags.c#L94
  316. (str_length,) = struct.unpack(">Q", stream.read(8))
  317. yield stream.read(str_length - 1)
  318. assert stream.read(1) == b"\x00"
  319. def _read_triangular(self, stream: typing.BinaryIO):
  320. assert stream.read(3) == self._MAGIC_NUMBER
  321. creation_match = re.match(
  322. rb"^created by (\w+) on (.* \d{4})\n", stream.readline()
  323. )
  324. assert creation_match
  325. self.creator, creation_dt_str = creation_match.groups()
  326. with setlocale("C"):
  327. self.creation_datetime = datetime.datetime.strptime(
  328. creation_dt_str.decode(), self._DATETIME_FORMAT
  329. )
  330. assert stream.read(1) == b"\n"
  331. # fwriteInt
  332. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/fio.c#L290
  333. vertices_num, triangles_num = struct.unpack(">II", stream.read(4 * 2))
  334. self.vertices = [
  335. Vertex(*struct.unpack(">fff", stream.read(4 * 3)))
  336. for _ in range(vertices_num)
  337. ]
  338. self.triangles = [
  339. Triangle(struct.unpack(">III", stream.read(4 * 3)))
  340. for _ in range(triangles_num)
  341. ]
  342. assert all(
  343. vertex_idx < vertices_num
  344. for triangle in self.triangles
  345. for vertex_idx in triangle.vertex_indices
  346. )
  347. assert stream.read(4) == self._TAG_OLD_USEREALRAS
  348. (using_old_real_ras,) = struct.unpack(">I", stream.read(4))
  349. assert using_old_real_ras in [0, 1], using_old_real_ras
  350. self.using_old_real_ras = bool(using_old_real_ras)
  351. assert stream.read(4) == self._TAG_OLD_SURF_GEOM
  352. # writeVolGeom
  353. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/transform.c#L368
  354. self.volume_geometry_info = tuple(stream.readline() for _ in range(8))
  355. self.command_lines = list(self._read_cmdlines(stream))
  356. @classmethod
  357. def read_triangular(cls, surface_file_path: str) -> "Surface":
  358. surface = cls()
  359. with open(surface_file_path, "rb") as surface_file:
  360. # pylint: disable=protected-access
  361. surface._read_triangular(surface_file)
  362. return surface
  363. @classmethod
  364. def _triangular_strftime(cls, creation_datetime: datetime.datetime) -> bytes:
  365. padded_day = f"{creation_datetime.day:>2}"
  366. fmt = cls._DATETIME_FORMAT.replace("%d", padded_day)
  367. with setlocale("C"):
  368. return creation_datetime.strftime(fmt).encode()
  369. def write_triangular(
  370. self,
  371. surface_file_path: str,
  372. creation_datetime: typing.Optional[datetime.datetime] = None,
  373. ):
  374. if creation_datetime is None:
  375. creation_datetime = datetime.datetime.now()
  376. with open(surface_file_path, "wb") as surface_file:
  377. surface_file.write(
  378. self._MAGIC_NUMBER
  379. + b"created by "
  380. + self.creator
  381. + b" on "
  382. + self._triangular_strftime(creation_datetime)
  383. + b"\n\n"
  384. + struct.pack(">II", len(self.vertices), len(self.triangles))
  385. )
  386. for vertex in self.vertices:
  387. surface_file.write(struct.pack(">fff", *vertex))
  388. for triangle in self.triangles:
  389. assert all(
  390. vertex_index < len(self.vertices)
  391. for vertex_index in triangle.vertex_indices
  392. )
  393. surface_file.write(struct.pack(">III", *triangle.vertex_indices))
  394. surface_file.write(
  395. self._TAG_OLD_USEREALRAS
  396. + struct.pack(">I", 1 if self.using_old_real_ras else 0)
  397. )
  398. if not self.volume_geometry_info:
  399. raise ValueError(
  400. "Missing geometry information (set attribute `volume_geometry_info`)"
  401. )
  402. surface_file.write(
  403. self._TAG_OLD_SURF_GEOM + b"".join(self.volume_geometry_info)
  404. )
  405. for command_line in self.command_lines:
  406. surface_file.write(
  407. self._TAG_CMDLINE
  408. + struct.pack(">Q", len(command_line) + 1)
  409. + command_line
  410. + b"\0"
  411. )
  412. def load_annotation_file(self, annotation_file_path: str) -> None:
  413. annotation = Annotation.read(annotation_file_path)
  414. assert len(annotation.vertex_label_index) <= len(self.vertices)
  415. assert max(annotation.vertex_label_index.keys()) < len(self.vertices)
  416. self.annotation = annotation
  417. def add_vertex(self, vertex: Vertex) -> int:
  418. self.vertices.append(vertex)
  419. return len(self.vertices) - 1
  420. def add_rectangle(self, vertex_indices: typing.Iterable[int]) -> None:
  421. vertex_indices = list(vertex_indices)
  422. if len(vertex_indices) == 3:
  423. vertex_indices.append(
  424. self.add_vertex(
  425. self.vertices[vertex_indices[0]]
  426. + self.vertices[vertex_indices[2]]
  427. - self.vertices[vertex_indices[1]]
  428. )
  429. )
  430. assert len(vertex_indices) == 4
  431. self.triangles.append(Triangle(vertex_indices[:3]))
  432. self.triangles.append(Triangle(vertex_indices[2:] + vertex_indices[:1]))
  433. def _triangle_count_by_adjacent_vertex_indices(
  434. self,
  435. ) -> typing.Dict[int, typing.DefaultDict[int, int]]:
  436. counts: typing.Dict[int, typing.DefaultDict[int, int]] = {
  437. vertex_index: collections.defaultdict(lambda: 0)
  438. for vertex_index in range(len(self.vertices))
  439. }
  440. for triangle in self.triangles:
  441. for vertex_index_pair in triangle.adjacent_vertex_indices(2):
  442. counts[vertex_index_pair[0]][vertex_index_pair[1]] += 1
  443. counts[vertex_index_pair[1]][vertex_index_pair[0]] += 1
  444. return counts
  445. def find_borders(self) -> typing.Iterator[PolygonalCircuit]:
  446. border_neighbours = {}
  447. for (
  448. vertex_index,
  449. neighbour_counts,
  450. ) in self._triangle_count_by_adjacent_vertex_indices().items():
  451. if not neighbour_counts:
  452. yield PolygonalCircuit((vertex_index,))
  453. else:
  454. neighbours = [
  455. neighbour_index
  456. for neighbour_index, counts in neighbour_counts.items()
  457. if counts != 2
  458. ]
  459. if neighbours:
  460. assert len(neighbours) % 2 == 0, (vertex_index, neighbour_counts)
  461. border_neighbours[vertex_index] = neighbours
  462. while border_neighbours:
  463. vertex_index, neighbour_indices = border_neighbours.popitem()
  464. cycle_indices = [vertex_index]
  465. border_neighbours[vertex_index] = neighbour_indices[1:]
  466. vertex_index = neighbour_indices[0]
  467. while vertex_index != cycle_indices[0]:
  468. neighbour_indices = border_neighbours.pop(vertex_index)
  469. neighbour_indices.remove(cycle_indices[-1])
  470. cycle_indices.append(vertex_index)
  471. if len(neighbour_indices) > 1:
  472. border_neighbours[vertex_index] = neighbour_indices[1:]
  473. vertex_index = neighbour_indices[0]
  474. assert vertex_index in border_neighbours, (
  475. vertex_index,
  476. cycle_indices,
  477. border_neighbours,
  478. )
  479. final_neighbour_indices = border_neighbours.pop(vertex_index)
  480. assert final_neighbour_indices == [cycle_indices[-1]], (
  481. vertex_index,
  482. final_neighbour_indices,
  483. cycle_indices,
  484. )
  485. yield PolygonalCircuit(cycle_indices)
  486. def _get_vertex_label_index(self, vertex_index: int) -> typing.Optional[int]:
  487. if not self.annotation:
  488. raise RuntimeError(
  489. "Missing annotation (call method `load_annotation_file` first)."
  490. )
  491. return self.annotation.vertex_label_index.get(vertex_index, None)
  492. def _find_label_border_segments(self, label: Label) -> typing.Iterator[LineSegment]:
  493. for triangle in self.triangles:
  494. border_vertex_indices = tuple(
  495. filter(
  496. lambda i: self._get_vertex_label_index(i) == label.index,
  497. triangle.vertex_indices,
  498. )
  499. )
  500. if len(border_vertex_indices) == 2:
  501. yield LineSegment(border_vertex_indices)
  502. _VertexSubindex = typing.Tuple[int, int]
  503. @classmethod
  504. def _duplicate_border(
  505. cls,
  506. neighbour_indices: typing.DefaultDict[
  507. _VertexSubindex, typing.Set[_VertexSubindex]
  508. ],
  509. previous_index: _VertexSubindex,
  510. current_index: _VertexSubindex,
  511. junction_counter: int,
  512. ) -> None:
  513. split_index = (current_index[0], junction_counter)
  514. neighbour_indices[previous_index].add(split_index)
  515. neighbour_indices[split_index].add(previous_index)
  516. next_index, *extra_indices = filter(
  517. lambda i: i != previous_index, neighbour_indices[current_index]
  518. )
  519. if extra_indices:
  520. neighbour_indices[next_index].add(split_index)
  521. neighbour_indices[split_index].add(next_index)
  522. neighbour_indices[next_index].remove(current_index)
  523. neighbour_indices[current_index].remove(next_index)
  524. return
  525. cls._duplicate_border(
  526. neighbour_indices=neighbour_indices,
  527. previous_index=split_index,
  528. current_index=next_index,
  529. junction_counter=junction_counter,
  530. )
  531. def find_label_border_polygonal_chains(
  532. self, label: Label
  533. ) -> typing.Iterator[PolygonalChain]:
  534. neighbour_indices: ( # type: ignore
  535. typing.DefaultDict[self._VertexSubindex, typing.Set[self._VertexSubindex]]
  536. ) = collections.defaultdict(set)
  537. for segment in self._find_label_border_segments(label):
  538. vertex_indices = [(i, 0) for i in segment.vertex_indices]
  539. neighbour_indices[vertex_indices[0]].add(vertex_indices[1])
  540. neighbour_indices[vertex_indices[1]].add(vertex_indices[0])
  541. junction_counter = 0
  542. found_leaf = True
  543. while found_leaf:
  544. found_leaf = False
  545. for leaf_index, leaf_neighbour_indices in neighbour_indices.items():
  546. if len(leaf_neighbour_indices) == 1:
  547. found_leaf = True
  548. junction_counter += 1
  549. self._duplicate_border(
  550. neighbour_indices=neighbour_indices,
  551. previous_index=leaf_index,
  552. # pylint: disable=stop-iteration-return; false positive, has 1 item
  553. current_index=next(iter(leaf_neighbour_indices)),
  554. junction_counter=junction_counter,
  555. )
  556. break
  557. assert all(len(n) == 2 for n in neighbour_indices.values()), neighbour_indices
  558. while neighbour_indices:
  559. # pylint: disable=stop-iteration-return; has >= 1 item
  560. chain = collections.deque([next(iter(neighbour_indices.keys()))])
  561. chain.append(neighbour_indices[chain[0]].pop())
  562. neighbour_indices[chain[1]].remove(chain[0])
  563. while chain[0] != chain[-1]:
  564. previous_index = chain[-1]
  565. next_index = neighbour_indices[previous_index].pop()
  566. neighbour_indices[next_index].remove(previous_index)
  567. chain.append(next_index)
  568. assert not neighbour_indices[previous_index], neighbour_indices[
  569. previous_index
  570. ]
  571. del neighbour_indices[previous_index]
  572. assert not neighbour_indices[chain[0]], neighbour_indices[chain[0]]
  573. del neighbour_indices[chain[0]]
  574. chain.pop()
  575. yield PolygonalChain(v[0] for v in chain)
  576. def _unused_vertices(self) -> typing.Set[int]:
  577. vertex_indices = set(range(len(self.vertices)))
  578. for triangle in self.triangles:
  579. for vertex_index in triangle.vertex_indices:
  580. vertex_indices.discard(vertex_index)
  581. return vertex_indices
  582. def remove_unused_vertices(self) -> None:
  583. vertex_index_conversion = [0] * len(self.vertices)
  584. for vertex_index in sorted(self._unused_vertices(), reverse=True):
  585. del self.vertices[vertex_index]
  586. vertex_index_conversion[vertex_index] -= 1
  587. vertex_index_conversion = numpy.cumsum(vertex_index_conversion)
  588. for triangle_index, triangle in enumerate(self.triangles):
  589. self.triangles[triangle_index] = Triangle(
  590. map(
  591. lambda i: i + int(vertex_index_conversion[i]),
  592. triangle.vertex_indices,
  593. )
  594. )
  595. def select_vertices(
  596. self, vertex_indices: typing.Iterable[int]
  597. ) -> typing.List[Vertex]:
  598. return [self.vertices[idx] for idx in vertex_indices]
  599. @staticmethod
  600. def unite(surfaces: typing.Iterable["Surface"]) -> "Surface":
  601. surfaces_iter = iter(surfaces)
  602. union = copy.deepcopy(next(surfaces_iter))
  603. for surface in surfaces_iter:
  604. vertex_index_offset = len(union.vertices)
  605. union.vertices.extend(surface.vertices)
  606. union.triangles.extend(
  607. Triangle(
  608. vertex_idx + vertex_index_offset
  609. for vertex_idx in triangle.vertex_indices
  610. )
  611. for triangle in surface.triangles
  612. )
  613. return union