ITK/Modules/Bridge/NumPy/wrapping/PyBuffer.i.init at ee6aa99712cbb4a1bb449ecb6926d66f3adaec0e · InsightSoftwareConsortium/ITK · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
%pythoncode %{

import numpy as np
class NDArrayITKBase(np.ndarray):
  """A numpy array that provides a view on the data associated with an optional itk "base" object."""

  def __new__(cls, input_array, itk_base=None):
      obj = np.asarray(input_array).view(cls)
      obj.itk_base = itk_base
      return obj

  def __array_finalize__(self, obj):
      if obj is None: return
      self.itk_base = getattr(obj, 'itk_base', None)

try:
    from distributed.protocol import dask_serialize, dask_deserialize
    from typing import Dict, List, Tuple
except (ImportError, RuntimeError):
    pass
else:
    @dask_serialize.register(NDArrayITKBase)
    def serialize(ndarray_itk_base: NDArrayITKBase) -> Tuple[Dict, List[bytes]]:
        dumps = dask_serialize.dispatch(np.ndarray)
        return dumps(ndarray_itk_base)

    @dask_deserialize.register(NDArrayITKBase)
    def deserialize(header: Dict, frames: List[bytes]) -> NDArrayITKBase:
        loads = dask_deserialize.dispatch(np.ndarray)
        return NDArrayITKBase(loads(header, frames))

def _get_formatstring(itk_Image_type) -> str:
    """Returns the struct format string for a given ITK image type.

    Format characters from Python's struct module:
    - 'b': signed char (int8)
    - 'B': unsigned char (uint8)
    - 'h': short (int16)
    - 'H': unsigned short (uint16)
    - 'i': int (int32)
    - 'I': unsigned int (uint32)
    - 'l': long (platform dependent)
    - 'L': unsigned long (platform dependent)
    - 'q': long long (int64)
    - 'Q': unsigned long long (uint64)
    - 'f': float (float32)
    - 'd': double (float64)
    """

    # Mapping from ITK pixel type codes to struct format strings
    _format_map = {
        "UC": "B",   # unsigned char
        "US": "H",   # unsigned short
        "UI": "I",   # unsigned int
        "UL": "L",   # unsigned long
        "ULL": "Q",  # unsigned long long
        "SC": "b",   # signed char
        "SS": "h",   # signed short
        "SI": "i",   # signed int
        "SL": "l",   # signed long
        "SLL": "q",  # signed long long
        "F": "f",    # float
        "D": "d",    # double
        "PF2": "f",  # Point<float, 2> - use float for components
        "PF3": "f",  # Point<float, 3> - use float for components
    }

    import os
    # Platform-specific adjustments for Windows
    if os.name == 'nt':
        _format_map['UL'] = 'I'  # unsigned int on Windows
        _format_map['SL'] = 'i'  # signed int on Windows

    try:
        return _format_map[itk_Image_type]
    except KeyError as e:
        raise ValueError(f"Unknown ITK image type: {itk_Image_type}") from e

def _get_numpy_pixelid(itk_Image_type) -> np.dtype:
    """Returns a ITK PixelID given a numpy array."""

    # This is a Mapping from numpy array types to itk pixel types.
    _np_itk = {"UC":np.dtype(np.uint8),
               "US":np.dtype(np.uint16),
               "UI":np.dtype(np.uint32),
               "UL":np.dtype(np.uint64),
               "ULL":np.dtype(np.uint64),
               "SC":np.dtype(np.int8),
               "SS":np.dtype(np.int16),
               "SI":np.dtype(np.int32),
               "SL":np.dtype(np.int64),
               "SLL":np.dtype(np.int64),
               "F":np.dtype(np.float32),
               "D":np.dtype(np.float64),
               "PF2":np.dtype(np.float32),
               "PF3":np.dtype(np.float32),
                }
    import os
    if os.name == 'nt':
        _np_itk['UL'] = np.dtype(np.uint32)
        _np_itk['SL'] = np.dtype(np.int32)
    try:
        return _np_itk[itk_Image_type]
    except KeyError as e:
        raise e
%}