__init__.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513
  1. """
  2. Python Library to Read and Write Surface Files in Freesurfer's TriangularSurface Format
  3. compatible with Freesurfer's MRISwriteTriangularSurface()
  4. https://github.com/freesurfer/freesurfer/blob/release_6_0_0/include/mrisurf.h#L1281
  5. https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/mrisurf.c
  6. https://raw.githubusercontent.com/freesurfer/freesurfer/release_6_0_0/utils/mrisurf.c
  7. Freesurfer
  8. https://surfer.nmr.mgh.harvard.edu/
  9. Edit Surface File
  10. >>> from freesurfer_surface import Surface, Vertex, Triangle
  11. >>>
  12. >>> surface = Surface.read_triangular('bert/surf/lh.pial'))
  13. >>>
  14. >>> vertex_a = surface.add_vertex(Vertex(0.0, 0.0, 0.0))
  15. >>> vertex_b = surface.add_vertex(Vertex(1.0, 1.0, 1.0))
  16. >>> vertex_c = surface.add_vertex(Vertex(2.0, 2.0, 2.0))
  17. >>> surface.triangles.append(Triangle((vertex_a, vertex_b, vertex_c)))
  18. >>>
  19. >>> surface.write_triangular('somewhere/else/lh.pial')
  20. List Labels in Annotation File
  21. >>> from freesurfer_surface import Annotation
  22. >>>
  23. >>> annotation = Annotation.read('bert/label/lh.aparc.annot')
  24. >>> for label in annotation.labels.values():
  25. >>> print(label.index, label.hex_color_code, label.name)
  26. Find Border of Labelled Region
  27. >>> surface = Surface.read_triangular('bert/surf/lh.pial'))
  28. >>> surface.load_annotation_file('bert/label/lh.aparc.annot')
  29. >>> region, = filter(lambda l: l.name == 'precentral',
  30. >>> annotation.labels.values())
  31. >>> print(surface.find_label_border_polygonal_chains(region))
  32. """
  33. import collections
  34. import contextlib
  35. import datetime
  36. import itertools
  37. import locale
  38. import re
  39. import struct
  40. import typing
  41. import numpy
  42. from freesurfer_surface.version import __version__
  43. class UnsupportedLocaleSettingError(locale.Error):
  44. pass
  45. @contextlib.contextmanager
  46. def setlocale(temporary_locale):
  47. primary_locale = locale.setlocale(locale.LC_ALL)
  48. try:
  49. yield locale.setlocale(locale.LC_ALL, temporary_locale)
  50. except locale.Error as exc:
  51. if str(exc) == 'unsupported locale setting':
  52. raise UnsupportedLocaleSettingError(temporary_locale)
  53. raise exc # pragma: no cover
  54. finally:
  55. locale.setlocale(locale.LC_ALL, primary_locale)
  56. class Vertex(numpy.ndarray):
  57. def __new__(cls, right: float, anterior: float, superior: float):
  58. return numpy.array((right, anterior, superior),
  59. dtype=float).view(cls)
  60. @property
  61. def right(self) -> float:
  62. return self[0]
  63. @property
  64. def anterior(self) -> float:
  65. return self[1]
  66. @property
  67. def superior(self) -> float:
  68. return self[2]
  69. @property
  70. def __dict__(self) -> typing.Dict[str, float]:
  71. return {'right': self.right,
  72. 'anterior': self.anterior,
  73. 'superior': self.superior}
  74. def __format_coords(self) -> str:
  75. return ', '.join('{}={}'.format(name, getattr(self, name))
  76. for name in ['right', 'anterior', 'superior'])
  77. def __repr__(self) -> str:
  78. return '{}({})'.format(type(self).__name__, self.__format_coords())
  79. def distance_mm(self, others: typing.Union['Vertex',
  80. typing.Iterable['Vertex'],
  81. numpy.ndarray],
  82. ) -> numpy.ndarray:
  83. if isinstance(others, Vertex):
  84. others = others.reshape((1, 3))
  85. return numpy.linalg.norm(self - others, axis=1)
  86. class PolygonalCircuit:
  87. _VERTEX_INDICES_TYPE = typing.Tuple[int]
  88. def __init__(self, vertex_indices: _VERTEX_INDICES_TYPE):
  89. self._vertex_indices = tuple(vertex_indices)
  90. assert all(isinstance(idx, int) for idx in self._vertex_indices)
  91. @property
  92. def vertex_indices(self):
  93. return self._vertex_indices
  94. def _normalize(self) -> 'PolygonalCircuit':
  95. vertex_indices = collections.deque(self.vertex_indices)
  96. vertex_indices.rotate(-numpy.argmin(self.vertex_indices))
  97. if len(vertex_indices) > 2 and vertex_indices[-1] < vertex_indices[1]:
  98. vertex_indices.reverse()
  99. vertex_indices.rotate(1)
  100. return type(self)(vertex_indices)
  101. def __eq__(self, other: 'PolygonalCircuit') -> bool:
  102. # pylint: disable=protected-access
  103. return self._normalize().vertex_indices == other._normalize().vertex_indices
  104. def __hash__(self) -> int:
  105. # pylint: disable=protected-access
  106. return hash(self._normalize()._vertex_indices)
  107. def adjacent_vertex_indices(self, vertices_num: int = 2
  108. ) -> typing.Iterable[typing.Tuple[int]]:
  109. vertex_indices_cycle = list(itertools.islice(
  110. itertools.cycle(self.vertex_indices),
  111. 0,
  112. len(self.vertex_indices) + vertices_num - 1,
  113. ))
  114. return zip(*(itertools.islice(vertex_indices_cycle,
  115. offset,
  116. len(self.vertex_indices) + offset)
  117. for offset in range(vertices_num)))
  118. class LineSegment(PolygonalCircuit):
  119. def __init__(self, indices: PolygonalCircuit._VERTEX_INDICES_TYPE):
  120. super().__init__(indices)
  121. assert len(self.vertex_indices) == 2
  122. def __repr__(self) -> str:
  123. return 'LineSegment(vertex_indices={})'.format(self.vertex_indices)
  124. class Triangle(PolygonalCircuit):
  125. def __init__(self, indices: PolygonalCircuit._VERTEX_INDICES_TYPE):
  126. super().__init__(indices)
  127. assert len(self.vertex_indices) == 3
  128. def __repr__(self) -> str:
  129. return 'Triangle(vertex_indices={})'.format(self.vertex_indices)
  130. class PolygonalChainsNotOverlapingError(ValueError):
  131. pass
  132. class PolygonalChain:
  133. def __init__(self, vertex_indices: typing.Iterable[int]):
  134. self.vertex_indices \
  135. = collections.deque(vertex_indices) # type: Deque[int]
  136. def __eq__(self, other: 'PolygonalChain') -> bool:
  137. return self.vertex_indices == other.vertex_indices
  138. def __repr__(self) -> str:
  139. return 'PolygonalChain(vertex_indices={})'.format(tuple(self.vertex_indices))
  140. def connect(self, other: 'PolygonalChain') -> None:
  141. if self.vertex_indices[-1] == other.vertex_indices[0]:
  142. self.vertex_indices.pop()
  143. self.vertex_indices.extend(other.vertex_indices)
  144. elif self.vertex_indices[-1] == other.vertex_indices[-1]:
  145. self.vertex_indices.pop()
  146. self.vertex_indices.extend(reversed(other.vertex_indices))
  147. elif self.vertex_indices[0] == other.vertex_indices[0]:
  148. self.vertex_indices.popleft()
  149. self.vertex_indices.extendleft(other.vertex_indices)
  150. elif self.vertex_indices[0] == other.vertex_indices[-1]:
  151. self.vertex_indices.popleft()
  152. self.vertex_indices.extendleft(reversed(other.vertex_indices))
  153. else:
  154. raise PolygonalChainsNotOverlapingError()
  155. def adjacent_vertex_indices(self, vertices_num: int = 2
  156. ) -> typing.Iterable[typing.Tuple[int]]:
  157. return zip(*(itertools.islice(self.vertex_indices,
  158. offset,
  159. len(self.vertex_indices))
  160. for offset in range(vertices_num)))
  161. def segments(self) -> typing.Iterable[LineSegment]:
  162. return map(LineSegment, self.adjacent_vertex_indices(2))
  163. class Label:
  164. # pylint: disable=too-many-arguments
  165. def __init__(self, index: int, name: str, red: int,
  166. green: int, blue: int, transparency: int):
  167. self.index = index # type: int
  168. self.name = name # type: str
  169. self.red = red # type: int
  170. self.green = green # type: int
  171. self.blue = blue # type: int
  172. self.transparency = transparency # type: int
  173. @property
  174. def color_code(self) -> int:
  175. if self.index == 0: # unknown
  176. return 0
  177. return int.from_bytes((self.red, self.green, self.blue, self.transparency),
  178. byteorder='little', signed=False)
  179. @property
  180. def hex_color_code(self) -> str:
  181. return '#{:02x}{:02x}{:02x}'.format(self.red, self.green, self.blue)
  182. def __str__(self) -> str:
  183. return 'Label(name={}, index={}, color={})'.format(
  184. self.name, self.index, self.hex_color_code)
  185. def __repr__(self) -> str:
  186. return str(self)
  187. class Annotation:
  188. # pylint: disable=too-few-public-methods
  189. _TAG_OLD_COLORTABLE = b'\0\0\0\x01'
  190. def __init__(self):
  191. self.vertex_label_index = {} # type: Dict[int, int]
  192. self.colortable_path = None # type: Optional[bytes]
  193. self.labels = {} # type: Dict[int, Label]
  194. @staticmethod
  195. def _read_label(stream: typing.BinaryIO) -> Label:
  196. index, name_length = struct.unpack('>II', stream.read(4 * 2))
  197. name = stream.read(name_length - 1).decode()
  198. assert stream.read(1) == b'\0'
  199. red, green, blue, transparency \
  200. = struct.unpack('>IIII', stream.read(4 * 4))
  201. return Label(index=index, name=name, red=red, green=green,
  202. blue=blue, transparency=transparency)
  203. def _read(self, stream: typing.BinaryIO) -> None:
  204. # https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles
  205. annotations_num, = struct.unpack('>I', stream.read(4))
  206. annotations = [struct.unpack('>II', stream.read(4 * 2))
  207. for _ in range(annotations_num)]
  208. assert stream.read(4) == self._TAG_OLD_COLORTABLE
  209. colortable_version, _, filename_length \
  210. = struct.unpack('>III', stream.read(4 * 3))
  211. assert colortable_version > 0 # new version
  212. self.colortable_path = stream.read(filename_length - 1)
  213. assert stream.read(1) == b'\0'
  214. labels_num, = struct.unpack('>I', stream.read(4))
  215. self.labels = {label.index: label for label
  216. in (self._read_label(stream) for _ in range(labels_num))}
  217. label_index_by_color_code = {label.color_code: label.index
  218. for label in self.labels.values()}
  219. self.vertex_label_index = {vertex_index: label_index_by_color_code[color_code]
  220. for vertex_index, color_code in annotations}
  221. assert not stream.read(1)
  222. @classmethod
  223. def read(cls, annotation_file_path: str) -> 'Annotation':
  224. annotation = cls()
  225. with open(annotation_file_path, 'rb') as annotation_file:
  226. # pylint: disable=protected-access
  227. annotation._read(annotation_file)
  228. return annotation
  229. class Surface:
  230. # pylint: disable=too-many-instance-attributes
  231. _MAGIC_NUMBER = b'\xff\xff\xfe'
  232. _TAG_CMDLINE = b'\x00\x00\x00\x03'
  233. _TAG_OLD_SURF_GEOM = b'\x00\x00\x00\x14'
  234. _TAG_OLD_USEREALRAS = b'\x00\x00\x00\x02'
  235. _DATETIME_FORMAT = '%a %b %d %H:%M:%S %Y'
  236. def __init__(self):
  237. self.creator = None # type: Optional[bytes]
  238. self.creation_datetime = None # type: Optional[datetime.datetime]
  239. self.vertices = [] # type: List[Vertex]
  240. self.triangles = [] # type: List[Triangle]
  241. self.using_old_real_ras = False # type: bool
  242. self.volume_geometry_info = None # type: Optional[Tuple[bytes]]
  243. self.command_lines = [] # type: List[bytes]
  244. self.annotation = None # type: Optional[Annotation]
  245. @classmethod
  246. def _read_cmdlines(cls, stream: typing.BinaryIO) -> typing.Iterator[str]:
  247. while True:
  248. tag = stream.read(4)
  249. if not tag:
  250. return
  251. assert tag == cls._TAG_CMDLINE # might be TAG_GROUP_AVG_SURFACE_AREA
  252. # TAGwrite
  253. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/tags.c#L94
  254. str_length, = struct.unpack('>Q', stream.read(8))
  255. yield stream.read(str_length - 1)
  256. assert stream.read(1) == b'\x00'
  257. def _read_triangular(self, stream: typing.BinaryIO):
  258. assert stream.read(3) == self._MAGIC_NUMBER
  259. self.creator, creation_dt_str = re.match(rb'^created by (\w+) on (.* \d{4})\n',
  260. stream.readline()).groups()
  261. with setlocale('C'):
  262. self.creation_datetime = datetime.datetime.strptime(creation_dt_str.decode(),
  263. self._DATETIME_FORMAT)
  264. assert stream.read(1) == b'\n'
  265. # fwriteInt
  266. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/fio.c#L290
  267. vertices_num, triangles_num = struct.unpack('>II', stream.read(4 * 2))
  268. self.vertices = [Vertex(*struct.unpack('>fff', stream.read(4 * 3)))
  269. for _ in range(vertices_num)]
  270. self.triangles = [Triangle(struct.unpack('>III', stream.read(4 * 3)))
  271. for _ in range(triangles_num)]
  272. assert all(vertex_idx < vertices_num
  273. for triangle in self.triangles
  274. for vertex_idx in triangle.vertex_indices)
  275. assert stream.read(4) == self._TAG_OLD_USEREALRAS
  276. using_old_real_ras, = struct.unpack('>I', stream.read(4))
  277. assert using_old_real_ras in [0, 1], using_old_real_ras
  278. self.using_old_real_ras = bool(using_old_real_ras)
  279. assert stream.read(4) == self._TAG_OLD_SURF_GEOM
  280. # writeVolGeom
  281. # https://github.com/freesurfer/freesurfer/blob/release_6_0_0/utils/transform.c#L368
  282. self.volume_geometry_info = tuple(stream.readline() for _ in range(8))
  283. self.command_lines = list(self._read_cmdlines(stream))
  284. @classmethod
  285. def read_triangular(cls, surface_file_path: str) -> 'Surface':
  286. surface = cls()
  287. with open(surface_file_path, 'rb') as surface_file:
  288. # pylint: disable=protected-access
  289. surface._read_triangular(surface_file)
  290. return surface
  291. @classmethod
  292. def _triangular_strftime(cls, creation_datetime: datetime.datetime) -> bytes:
  293. padded_day = '{:>2}'.format(creation_datetime.day)
  294. fmt = cls._DATETIME_FORMAT.replace('%d', padded_day)
  295. with setlocale('C'):
  296. return creation_datetime.strftime(fmt).encode()
  297. def write_triangular(self, surface_file_path: str,
  298. creation_datetime: typing.Optional[datetime.datetime] = None):
  299. if creation_datetime is None:
  300. creation_datetime = datetime.datetime.now()
  301. with open(surface_file_path, 'wb') as surface_file:
  302. surface_file.write(
  303. self._MAGIC_NUMBER
  304. + b'created by ' + self.creator
  305. + b' on ' + self._triangular_strftime(creation_datetime)
  306. + b'\n\n'
  307. + struct.pack('>II', len(self.vertices), len(self.triangles))
  308. )
  309. for vertex in self.vertices:
  310. surface_file.write(struct.pack('>fff', *vertex))
  311. for triangle in self.triangles:
  312. assert all(vertex_index < len(self.vertices)
  313. for vertex_index in triangle.vertex_indices)
  314. surface_file.write(struct.pack('>III',
  315. *triangle.vertex_indices))
  316. surface_file.write(self._TAG_OLD_USEREALRAS
  317. + struct.pack('>I', 1 if self.using_old_real_ras else 0))
  318. surface_file.write(self._TAG_OLD_SURF_GEOM
  319. + b''.join(self.volume_geometry_info))
  320. for command_line in self.command_lines:
  321. surface_file.write(self._TAG_CMDLINE + struct.pack('>Q', len(command_line) + 1)
  322. + command_line + b'\0')
  323. def load_annotation_file(self, annotation_file_path: str) -> None:
  324. annotation = Annotation.read(annotation_file_path)
  325. assert len(annotation.vertex_label_index) <= len(self.vertices)
  326. assert max(annotation.vertex_label_index.keys()) < len(self.vertices)
  327. self.annotation = annotation
  328. def add_vertex(self, vertex: Vertex) -> int:
  329. self.vertices.append(vertex)
  330. return len(self.vertices) - 1
  331. def add_rectangle(self, vertex_indices: typing.Iterable[int]) -> typing.Iterable[int]:
  332. vertex_indices = list(vertex_indices)
  333. if len(vertex_indices) == 3:
  334. vertex_indices.append(self.add_vertex(
  335. self.vertices[vertex_indices[0]]
  336. + self.vertices[vertex_indices[2]]
  337. - self.vertices[vertex_indices[1]]
  338. ))
  339. assert len(vertex_indices) == 4
  340. self.triangles.append(Triangle(vertex_indices[:3]))
  341. self.triangles.append(Triangle(vertex_indices[2:]
  342. + vertex_indices[:1]))
  343. def _triangle_count_by_adjacent_vertex_indices(self) \
  344. -> typing.Dict[int, typing.Dict[int, int]]:
  345. counts = {vertex_index: collections.defaultdict(lambda: 0)
  346. for vertex_index in range(len(self.vertices))}
  347. for triangle in self.triangles:
  348. for vertex_index_pair in triangle.adjacent_vertex_indices(2):
  349. counts[vertex_index_pair[0]][vertex_index_pair[1]] += 1
  350. counts[vertex_index_pair[1]][vertex_index_pair[0]] += 1
  351. return counts
  352. def find_borders(self) -> typing.Iterator[PolygonalCircuit]:
  353. border_neighbours = {}
  354. for vertex_index, neighbour_counts \
  355. in self._triangle_count_by_adjacent_vertex_indices().items():
  356. if not neighbour_counts:
  357. yield PolygonalCircuit((vertex_index,))
  358. else:
  359. neighbours = [neighbour_index for neighbour_index, counts
  360. in neighbour_counts.items()
  361. if counts != 2]
  362. if neighbours:
  363. assert len(neighbours) == 2, (vertex_index, neighbours)
  364. border_neighbours[vertex_index] = neighbours
  365. while border_neighbours:
  366. vertex_index, neighbour_indices = border_neighbours.popitem()
  367. cycle_indices = [vertex_index]
  368. vertex_index = neighbour_indices[0]
  369. while vertex_index != cycle_indices[0]:
  370. neighbour_indices = border_neighbours.pop(vertex_index)
  371. neighbour_indices.remove(cycle_indices[-1])
  372. cycle_indices.append(vertex_index)
  373. vertex_index, = neighbour_indices
  374. yield PolygonalCircuit(cycle_indices)
  375. def _get_vertex_label_index(self, vertex_index: int) -> typing.Optional[int]:
  376. return self.annotation.vertex_label_index.get(vertex_index, None)
  377. def _find_label_border_segments(self, label: Label) -> typing.Iterator[LineSegment]:
  378. for triangle in self.triangles:
  379. border_vertex_indices = tuple(filter(
  380. lambda i: self._get_vertex_label_index(i) == label.index,
  381. triangle.vertex_indices,
  382. ))
  383. if len(border_vertex_indices) == 2:
  384. yield LineSegment(border_vertex_indices)
  385. def find_label_border_polygonal_chains(self, label: Label) -> typing.Iterator[PolygonalChain]:
  386. segments = set(self._find_label_border_segments(label))
  387. available_chains = collections.deque(PolygonalChain(segment.vertex_indices)
  388. for segment in segments)
  389. # irrespective of its poor performance,
  390. # we keep this approach since it's easy to read and fast enough
  391. while available_chains:
  392. chain = available_chains.pop()
  393. last_chains_len = None
  394. while last_chains_len != len(available_chains):
  395. last_chains_len = len(available_chains)
  396. checked_chains = collections.deque()
  397. while available_chains:
  398. potential_neighbour = available_chains.pop()
  399. try:
  400. chain.connect(potential_neighbour)
  401. except PolygonalChainsNotOverlapingError:
  402. checked_chains.append(potential_neighbour)
  403. available_chains = checked_chains
  404. assert all((segment in segments) for segment in chain.segments())
  405. yield chain
  406. def _unused_vertices(self) -> typing.Set[int]:
  407. vertex_indices = set(range(len(self.vertices)))
  408. for triangle in self.triangles:
  409. for vertex_index in triangle.vertex_indices:
  410. vertex_indices.discard(vertex_index)
  411. return vertex_indices
  412. def remove_unused_vertices(self) -> None:
  413. vertex_index_conversion = [0] * len(self.vertices)
  414. for vertex_index in sorted(self._unused_vertices(), reverse=True):
  415. del self.vertices[vertex_index]
  416. vertex_index_conversion[vertex_index] -= 1
  417. vertex_index_conversion = numpy.cumsum(vertex_index_conversion)
  418. for triangle_index in range(len(self.triangles)):
  419. self.triangles[triangle_index] \
  420. = Triangle(map(lambda i: i + int(vertex_index_conversion[i]),
  421. self.triangles[triangle_index].vertex_indices))
  422. def select_vertices(self, vertex_indices: typing.Iterable[int]) \
  423. -> numpy.ndarray:
  424. if not hasattr(vertex_indices, '__getitem__'):
  425. vertex_indices = list(vertex_indices)
  426. return numpy.take(self.vertices, indices=vertex_indices, axis=0)