feat: EXRInput, EXROutput을 추가하다

EXRInput, EXROutput : EXR File을 읽고 쓰기 위한 Class

Author: 404Vector

openexr_python/__init__.py

@@ -10,4 +10,6 @@
 __version__ = "0.0.0"  # dynamic
 
 from .exr_header import *
-from .exr_window import *
+from .exr_input_file import *
+from .exr_output_file import *
+from .util import *

openexr_python/exr_input_file.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+import os
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+from openexr_python.exr_header import EXRHeader
+import OpenEXR
+from openexr_python.util import (
+    get_size_from_window,
+    convert_numpy_type_to_imath_type,
+    convert_imath_type_to_numpy_type,
+)
+
+
+class EXRInputFile:
+    def __init__(self, file_path: str):
+        assert OpenEXR.isOpenExrFile(file_path)
+        self.file_path = file_path
+        self.raw_input_file: Any = OpenEXR.InputFile(file_path)
+
+    def get_single_channel_buffer(self, channel_key: str) -> bytes:
+        return self.raw_input_file.channel(channel_key)
+
+    def get_multi_channel_buffer(self, channel_keys: List[str]) -> List[bytes]:
+        assert len(channel_keys) > 0
+        assert all([len(c) == 1 for c in channel_keys])
+        return self.raw_input_file.channels(channel_keys)
+
+    def read_numpy(self, channel_key: str) -> np.ndarray:
+        header = self.header
+        buffer = self.get_single_channel_buffer(channel_key=channel_key)
+        channel_info = header.channels[channel_key]
+        dtype = convert_imath_type_to_numpy_type(channel_info.type)
+        np_buffer = np.frombuffer(buffer=buffer, dtype=dtype)
+        size = get_size_from_window(header.data_window)
+        return np.reshape(np_buffer, newshape=(size[1], size[0]))
+
+    def close(self):
+        self.raw_input_file.close()
+
+    @property
+    def header(self) -> EXRHeader:
+        raw_header = self.raw_input_file.header()
+        return EXRHeader(raw_header=raw_header)
+
+    @property
+    def is_complete(self) -> bool:
+        """
+        Returns whether this exr file contains all data and is not missing.
+
+        Returns:
+            bool
+        """
+        return self.raw_input_file.isComplete()

openexr_python/exr_output_file.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+import os
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+from openexr_python.exr_header import EXRHeader
+import OpenEXR
+from openexr_python.util import (
+    get_size_from_window,
+    convert_numpy_type_to_imath_type,
+    convert_imath_type_to_numpy_type,
+)
+
+
+class EXROutputFile:
+    def __init__(self, file_path: str, header: EXRHeader) -> None:
+        assert os.path.isdir(os.path.dirname(file_path))
+        self.file_path = file_path
+        self.header = header
+        self.raw_output_file: Any = OpenEXR.OutputFile(file_path, header.raw_header)
+        self._is_closed = False
+
+    def write_pixels(self, pixels: Dict[str, bytes]):
+        assert self._is_closed is not True, f"The EXR File already has been closed."
+        self.raw_output_file.writePixels(pixels)
+        self._close()
+
+    def _convert_bytes_to_numpy(self, numpy_buffer: np.ndarray):
+        # check the buffer is 2 dim shape(h x w)
+        shape = numpy_buffer.shape
+        assert len(shape) == 2
+        # check the buffer is fit to the window size
+        data_height, data_width = shape
+        window_width, window_height = get_size_from_window(self.header.data_window)
+        assert window_width == data_width
+        assert window_height == data_height
+        return numpy_buffer.tobytes()
+
+    def write_pixels_with_numpy(self, pixels: Dict[str, np.ndarray]):
+        pixels_bytes: Dict[str, bytes] = {}
+        for channel_name, numpy_buffer in pixels.items():
+            # check the channel is exist
+            assert channel_name in self.header.channels
+            pixels_bytes[channel_name] = self._convert_bytes_to_numpy(
+                numpy_buffer=numpy_buffer
+            )
+            # check the buffer type is fit to the window type
+            assert (
+                convert_numpy_type_to_imath_type(numpy_buffer.dtype)
+                == self.header.channels[channel_name].type
+            )
+        self.write_pixels(pixels=pixels_bytes)
+
+    def current_scan_line(self):
+        return self.raw_output_file.currentScanLine()
+
+    def _close(self):
+        assert self._is_closed is not True, f"The EXR File already has been closed."
+        self.raw_output_file.close()
+        self._is_closed = True

openexr_python/util.py

@@ -0,0 +1,42 @@
+from typing import Tuple
+import Imath
+import numpy as np
+
+PIXEL_TYPE_CONVERT_MAP = {
+    np.uint32: Imath.PixelType.UINT,
+    np.float16: Imath.PixelType.HALF,
+    np.float32: Imath.PixelType.FLOAT,
+    Imath.PixelType.UINT: np.uint32,
+    Imath.PixelType.HALF: np.float16,
+    Imath.PixelType.FLOAT: np.float32,
+}
+
+
+def convert_numpy_type_to_imath_type(data_type: np.dtype) -> Imath.PixelType:
+    assert data_type in [np.float16, np.float32, np.uint32]
+    keys = [k for k in PIXEL_TYPE_CONVERT_MAP.keys()]
+    idx = -1
+    for i, key in enumerate(keys):
+        if key != data_type:
+            continue
+        idx = i
+        break
+    assert idx != -1
+    return Imath.PixelType(PIXEL_TYPE_CONVERT_MAP[keys[idx]])
+
+
+def convert_imath_type_to_numpy_type(data_type: Imath.PixelType):
+    assert type(data_type) is Imath.PixelType
+    assert data_type.v in PIXEL_TYPE_CONVERT_MAP
+    return PIXEL_TYPE_CONVERT_MAP[data_type.v]
+
+
+def create_channel(data_type: np.dtype, sampling_x: int = 1, sampling_y: int = 1):
+    assert sampling_x > 0
+    assert sampling_y > 0
+    im_data_type = convert_numpy_type_to_imath_type(data_type)
+    return Imath.Channel(type=im_data_type, xSampling=1, ySampling=1)
+
+
+def get_size_from_window(window: Imath.Box2i) -> Tuple[int, int]:
+    return (1 + window.max.x - window.min.x, 1 + window.max.y - window.min.y)