mirror of
https://github.com/johndoe6345789/SDL3CPlusPlus.git
synced 2026-04-25 14:15:02 +00:00
430 lines
18 KiB
Python
430 lines
18 KiB
Python
"""
|
|
Sample MaterialX ImageLoader implementation using OpenImageIO package.
|
|
|
|
This module provides a MaterialX-compatible ImageLoader implementation using OpenImageIO (OIIO).
|
|
The test will test loading an image, save it out, and optionally previewing it.
|
|
|
|
Steps:
|
|
1. Create an OIIOLoader which is derived from the ImageLoader interface class.
|
|
2. Create a new ImageHandler and register the loader with it.
|
|
3. Request to acquire an image using the ImageHandler. An EXR image is requested.
|
|
4. OIIOLoader will return supported extensions and match the requested image format.
|
|
5. As such the OIIOLoader will be requested to load in the EXR image, convert the
|
|
data and return a MaterialX Image.
|
|
6. Try to acquire the image again. This should returnt the cached MaterialX Image.
|
|
7. Save the image back to disk in the original format.
|
|
|
|
The image can optionally be previewed after load before save.
|
|
|
|
- Python Dependencies:
|
|
- OpenImageIO (version 3.0.6.1)
|
|
- API Docs can be found here: https://openimageio.readthedocs.io/en/v3.0.6.1/)
|
|
- numpy : For numerical operations on image data
|
|
- matplotlib : If image preview is desired.
|
|
"""
|
|
import ctypes
|
|
import os
|
|
import argparse
|
|
|
|
import logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("OIIOLoad")
|
|
|
|
try:
|
|
import MaterialX as mx
|
|
import MaterialX.PyMaterialXRender as mx_render
|
|
except ImportError:
|
|
logger.error("Required modules not found. Please install MaterialX.")
|
|
raise
|
|
try:
|
|
import OpenImageIO as oiio
|
|
import numpy as np
|
|
except ImportError:
|
|
logger.error("Required modules not found. Please install OpenImageIO and numpy.")
|
|
raise
|
|
|
|
have_matplot = False
|
|
try:
|
|
import matplotlib.pyplot as plt
|
|
have_matplot = True
|
|
except ImportError:
|
|
logger.warning("matplotlib module not found. Image preview display is disabled.")
|
|
|
|
class OiioImageLoader(mx_render.ImageLoader):
|
|
"""
|
|
A MaterialX ImageLoader implementation that uses OpenImageIO to read image files.
|
|
|
|
Inherits from MaterialX.ImageLoader and implements the required interface methods.
|
|
Supports common image formats like PNG, JPEG, TIFF, EXR, HDR, etc.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""
|
|
Initialize the OiioImageLoader and set supported extensions."""
|
|
super().__init__()
|
|
|
|
# Set all extensions supported by OpenImageIO. e.g.
|
|
# openexr:exr,sxr,mxr;tiff:tif,tiff,tx,env,sm,vsm;jpeg:jpg,jpe,jpeg,jif,jfif,jfi;bmp:bmp,dib;cineon:cin;dds:dds;dpx:dpx;fits:fits;hdr:hdr,rgbe;ico:ico;iff:iff,z;null:null,nul;png:png;pnm:ppm,pgm,pbm,pnm,pfm;psd:psd,pdd,psb;rla:rla;sgi:sgi,rgb,rgba,bw,int,inta;softimage:pic;targa:tga,tpic;term:term;webp:webp;zfile:zfile
|
|
self._extensions = set()
|
|
oiio_extensions = oiio.get_string_attribute("extension_list")
|
|
# Split string by ";"
|
|
for group in oiio_extensions.split(";"):
|
|
# Each group is like "openexr:exr,sxr,mxr"
|
|
if ":" in group:
|
|
_, exts = group.split(":", 1)
|
|
self._extensions.update(ext.strip() for ext in exts.split(","))
|
|
else:
|
|
self._extensions.update(ext.strip() for ext in group.split(","))
|
|
logger.debug(f"Cache supported extensions: {self._extensions}")
|
|
|
|
self.preview = False
|
|
self.identifier = "OpenImageIO Custom Image Loader"
|
|
self.color_space = {}
|
|
|
|
def supportedExtensions(self):
|
|
"""
|
|
Derived method to return a set of supported image file extensions.
|
|
"""
|
|
logger.info(f"OIIO supported extensions: {self._extensions}")
|
|
return self._extensions
|
|
|
|
def set_preview(self, value):
|
|
"""
|
|
Set whether to preview images when loading and saving
|
|
|
|
@param value: Boolean indicating whether to enable preview
|
|
"""
|
|
self.preview = value
|
|
|
|
def get_identifier(self):
|
|
return "OIIO Custom Loader"
|
|
|
|
def previewImage(self, title, data, width, height, nchannels, color_space):
|
|
"""
|
|
Utility method to preview an image using matplotlib.
|
|
Handles normalization and dtype for correct display.
|
|
|
|
@param title: Title for the preview window
|
|
@param data: Image data array
|
|
@param width: Image width
|
|
@param height: Image height
|
|
@param nchannels: Number of image channels
|
|
@param color_space: Color space of the image
|
|
"""
|
|
if not self.preview:
|
|
return
|
|
|
|
if have_matplot:
|
|
# If the image is float16 (half), convert to float32
|
|
if data.dtype == np.float16:
|
|
data = data.astype(np.float32)
|
|
|
|
flat = data.reshape(height, width, nchannels)
|
|
# Always display as RGB (first 3 channels or repeat if less)
|
|
if nchannels >= 3:
|
|
rgb = flat[..., :3]
|
|
else:
|
|
rgb = np.repeat(flat[..., :1], 3, axis=-1)
|
|
|
|
# Determine if normalization is needed
|
|
if np.issubdtype(flat.dtype, np.floating):
|
|
# If float, normalize to [0, 1] for display
|
|
rgb_disp = np.clip(rgb, 0.0, 1.0)
|
|
elif np.issubdtype(flat.dtype, np.integer):
|
|
# If integer, assume 8 or 16 bit, scale if needed
|
|
if flat.dtype == np.uint8:
|
|
rgb_disp = rgb # matplotlib expects [0,255] for uint8
|
|
elif flat.dtype == np.uint16:
|
|
# Scale 16-bit to 8-bit for display
|
|
rgb_disp = (rgb / 257).astype(np.uint8)
|
|
else:
|
|
# For other integer types, try to scale to [0,255]
|
|
rgb_disp = np.clip(rgb, 0, 255).astype(np.uint8)
|
|
else:
|
|
rgb_disp = rgb
|
|
|
|
# Set title bar text for the preview window
|
|
fig, ax = plt.subplots()
|
|
ax.imshow(rgb_disp)
|
|
ax.axis("off")
|
|
#fig.patch.set_facecolor("black")
|
|
fig.canvas.manager.set_window_title(title)
|
|
info = f"Dimensions:({width}x{height}), {nchannels} channels, type={data.dtype}, colorspace={color_space}"
|
|
fig.suptitle(title, fontsize=12)
|
|
plt.title(info, fontsize=9)
|
|
plt.show()
|
|
|
|
def loadImage(self, filePath):
|
|
"""
|
|
Load an image from the file system (MaterialX interface method).
|
|
|
|
@param filePath (MaterialX.FilePath): Path to the image file
|
|
@returns MaterialX.ImagePtr: MaterialX Image object or None if loading fails
|
|
"""
|
|
file_path_str = filePath.asString()
|
|
logger.info(f"Load using OIIO loader: {file_path_str}")
|
|
|
|
if not os.path.exists(file_path_str):
|
|
print(f"Error: File '{file_path_str}' does not exist")
|
|
return None
|
|
|
|
try:
|
|
# Open the image file
|
|
img_input = oiio.ImageInput.open(file_path_str)
|
|
if not img_input:
|
|
print(f"Error: Could not open '{file_path_str}' - {oiio.geterror()}")
|
|
return None
|
|
|
|
# Get image specifications
|
|
spec = img_input.spec()
|
|
color_space = spec.getattribute("oiio:ColorSpace")
|
|
logger.info(f"ColorSpace: {color_space}")
|
|
self.color_space[file_path_str] = color_space
|
|
|
|
# Check channel count
|
|
channels = spec.nchannels
|
|
if channels > 4:
|
|
channels = 4
|
|
|
|
# Determine MaterialX base type from OIIO format
|
|
base_type = self._oiio_to_materialx_type(spec.format.basetype)
|
|
if base_type is None:
|
|
img_input.close()
|
|
print(f"Error: Unsupported image format for '{file_path_str}'")
|
|
return None
|
|
|
|
# Create MaterialX image
|
|
mx_image = mx_render.Image.create(spec.width, spec.height, channels, base_type)
|
|
mx_image.createResourceBuffer()
|
|
logger.debug(f"Create buffer with width: {spec.width}, height: {spec.height}, channels: {spec.nchannels} -> {channels}")
|
|
|
|
# Read the image data using the correct OIIO Python API (returns a bytes object)
|
|
logger.debug(f"Reading image data from '{file_path_str}' with spec: {spec}")
|
|
data = img_input.read_image(0, 0, 0, channels, spec.format)
|
|
if len(data) > 0:
|
|
logger.debug(f"Done Reading image data from '{file_path_str}' with spec: {spec}")
|
|
else:
|
|
logger.error(f"Could not read image data.")
|
|
return None
|
|
|
|
self.previewImage("Loaded MaterialX Image", data, spec.width, spec.height, channels, color_space)
|
|
|
|
# Steps:
|
|
# - Copy the OIIO data into the MaterialX image resource buffer
|
|
resource_buffer_ptr = mx_image.getResourceBuffer()
|
|
bytes_per_channel = spec.format.size()
|
|
total_bytes = spec.width * spec.height * channels * bytes_per_channel
|
|
logger.info(f"Total bytes read in: {total_bytes} (width: {spec.width}, height: {spec.height}, channels: {channels}, format: {spec.format})")
|
|
try:
|
|
ctypes.memmove(resource_buffer_ptr, (ctypes.c_char * total_bytes).from_buffer_copy(data), total_bytes)
|
|
except Exception as e:
|
|
logger.error(f"Failed to update image resource buffer: {e}")
|
|
|
|
img_input.close()
|
|
|
|
return mx_image
|
|
|
|
except Exception as e:
|
|
print(f"Error loading image from '{file_path_str}': {str(e)}")
|
|
return None
|
|
|
|
return None
|
|
|
|
def saveImage(self, filePath, image, verticalFlip=False):
|
|
"""
|
|
@brief Saves an image to disk using OpenImageIO (OIIO).
|
|
|
|
@param filePath The file path where the image will be saved. Expected to have an asString() method.
|
|
@param image The MaterialX image object to save.
|
|
@param verticalFlip Whether to vertically flip the image before saving. (Currently unused.)
|
|
@return True if the image was saved successfully, False otherwise.
|
|
"""
|
|
filename = filePath.asString()
|
|
width = image.getWidth()
|
|
height = image.getHeight()
|
|
|
|
# Clamp to RGBA
|
|
src_channels = image.getChannelCount()
|
|
channels = min(src_channels, 4)
|
|
if src_channels > 4:
|
|
logger.warning(f"Image has {src_channels} channels. Saving only first {channels} (RGBA).")
|
|
|
|
mx_basetype = image.getBaseType()
|
|
oiio_format = self._materialx_to_oiio_type(mx_basetype)
|
|
logger.info(f"mx_basetype: {mx_basetype}, oiio_format: {oiio_format}, base_stride: {image.getBaseStride()}")
|
|
if oiio_format is None:
|
|
logger.error(f"Unsupported MaterialX base type for OIIO: {mx_basetype}")
|
|
return False
|
|
|
|
buffer_addr = image.getResourceBuffer()
|
|
np_type = self._materialx_type_to_np_type(mx_basetype)
|
|
if np_type is None:
|
|
logger.error(f"No NumPy dtype mapping for base type: {mx_basetype}")
|
|
return False
|
|
|
|
try:
|
|
# Steps:
|
|
# - Maps the MaterialX base type to OIIO and NumPy types.
|
|
# - Allocates a NumPy array for the pixel data.
|
|
# - Copies the raw buffer from the image into the NumPy array.
|
|
# - Optionally previews the image for debugging.
|
|
# - Creates an OIIO ImageOutput and writes the image to disk.
|
|
#
|
|
base_stride = image.getBaseStride() # bytes per channel element
|
|
total_bytes = width * height * src_channels * base_stride
|
|
|
|
buf_type = (ctypes.c_char * total_bytes)
|
|
buf = buf_type.from_address(buffer_addr)
|
|
|
|
np_buffer = np.frombuffer(buf, dtype=np_type)
|
|
|
|
# Validate total elements before reshape to catch mismatches early
|
|
expected_elems = width * height * src_channels
|
|
if np_buffer.size != expected_elems:
|
|
logger.error(f"Buffer element count mismatch: got {np_buffer.size}, expected {expected_elems}.")
|
|
return False
|
|
|
|
np_buffer = np_buffer.reshape((height, width, src_channels))
|
|
|
|
# Keep only up to RGBA
|
|
pixels = np_buffer[..., :channels].copy()
|
|
|
|
if verticalFlip:
|
|
logger.info("Applying vertical flip before saving image.")
|
|
pixels = np.flipud(pixels)
|
|
|
|
logger.info("Previewing image after load into Image and reload for save...")
|
|
# Remove "saved_" prefix if present
|
|
search_name = filename.replace("saved_", "")
|
|
color_space = "Unknown"
|
|
for key in self.color_space:
|
|
value = self.color_space[key]
|
|
path = os.path.basename(key)
|
|
if path in search_name:
|
|
color_space = value
|
|
logger.info(f"colorspace lookup for: {search_name}. list: {color_space}")
|
|
self.previewImage("OpenImageIO Output Image", pixels, width, height, channels, color_space)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error copying buffer to pixels: {e}")
|
|
return False
|
|
|
|
out = oiio.ImageOutput.create(filename)
|
|
if not out:
|
|
logger.error("Failed to create OIIO ImageOutput.")
|
|
return False
|
|
|
|
try:
|
|
spec = oiio.ImageSpec(width, height, channels, oiio_format)
|
|
out.open(filename, spec)
|
|
out.write_image(pixels)
|
|
logger.info(f"Image saved to {filename} (w={width}, h={height}, c={channels}, type={mx_basetype})")
|
|
out.close()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to write image: {e}")
|
|
try:
|
|
out.close()
|
|
finally:
|
|
pass
|
|
return False
|
|
|
|
def _oiio_to_materialx_type(self, oiio_basetype):
|
|
"""Convert OIIO base type to MaterialX Image base type."""
|
|
type_mapping = {
|
|
oiio.UINT8: mx_render.BaseType.UINT8,
|
|
oiio.INT8: mx_render.BaseType.INT8,
|
|
oiio.UINT16: mx_render.BaseType.UINT16,
|
|
oiio.INT16: mx_render.BaseType.INT16,
|
|
oiio.HALF: mx_render.BaseType.HALF,
|
|
oiio.FLOAT: mx_render.BaseType.FLOAT
|
|
}
|
|
return_val = type_mapping.get(oiio_basetype, None)
|
|
logger.debug(f"OIIO to MaterialX type mapping: {return_val} from {oiio_basetype}")
|
|
return return_val
|
|
|
|
def _materialx_to_oiio_type(self, mx_basetype):
|
|
"""Convert MaterialX Image base type to OIIO type."""
|
|
type_mapping = {
|
|
mx_render.BaseType.UINT8: oiio.UINT8,
|
|
mx_render.BaseType.UINT16: oiio.UINT16,
|
|
mx_render.BaseType.INT8: oiio.INT8,
|
|
mx_render.BaseType.INT16: oiio.INT16,
|
|
mx_render.BaseType.HALF: oiio.HALF,
|
|
mx_render.BaseType.FLOAT: oiio.FLOAT,
|
|
}
|
|
return_val = type_mapping.get(mx_basetype, None)
|
|
logger.debug(f"MaterialX type mapping: {mx_basetype} to {return_val}")
|
|
return return_val
|
|
|
|
def _materialx_type_to_np_type(self, mx_basetype):
|
|
"""Map MaterialX base type to NumPy dtype with explicit widths."""
|
|
type_mapping = {
|
|
mx_render.BaseType.UINT8: np.uint8,
|
|
mx_render.BaseType.UINT16: np.uint16,
|
|
mx_render.BaseType.INT8: np.int8,
|
|
mx_render.BaseType.INT16: np.int16,
|
|
mx_render.BaseType.HALF: np.float16,
|
|
mx_render.BaseType.FLOAT: np.float32,
|
|
}
|
|
return type_mapping.get(mx_basetype, None)
|
|
|
|
|
|
def test_load_save():
|
|
"""
|
|
Example usage of the OiioImageLoader class with MaterialX ImageHandler.
|
|
"""
|
|
parser = argparse.ArgumentParser(description="MaterialX OIIO Image Handler")
|
|
parser.add_argument("path", help="Path to the image file")
|
|
parser.add_argument("--flip", action="store_true", help="Flip the image vertically")
|
|
parser.add_argument("--preview", action="store_true", help="Preview the image before saving")
|
|
args = parser.parse_args()
|
|
|
|
test_image_path = args.path
|
|
if not os.path.exists(test_image_path):
|
|
logger.error(f"Image file not found: {test_image_path}")
|
|
return
|
|
|
|
# Create MaterialX handler with custom OIIO image loader
|
|
loader = OiioImageLoader()
|
|
loader.set_preview(args.preview)
|
|
handler = mx_render.ImageHandler.create(loader)
|
|
logger.info(f"Created image handler with loader ({loader.get_identifier()}): {handler is not None}")
|
|
handler.addLoader(loader)
|
|
|
|
mx_filepath = mx.FilePath(test_image_path)
|
|
|
|
# Load image using handler API
|
|
logger.info('-'*45)
|
|
logger.info(f"Loading image from path: {mx_filepath.asString()}")
|
|
mx_image = handler.acquireImage(mx_filepath)
|
|
if mx_image:
|
|
# Q: How to check for failed image load as you
|
|
# get back a 1x1 pixel image.
|
|
if mx_image.getWidth() == 1 and mx_image.getHeight() == 1:
|
|
logger.warning("Failed to load image. Got 1x1 pixel image returned")
|
|
return
|
|
logger.info(f"MaterialX Image loaded via Image Handler:")
|
|
logger.info(f" Dimensions: {mx_image.getWidth()}x{mx_image.getHeight()}")
|
|
logger.info(f" Channels: {mx_image.getChannelCount()}")
|
|
logger.info(f" Base type: {mx_image.getBaseType()}")
|
|
|
|
# Save image using handler API (to a new file)
|
|
logger.info('-'*45)
|
|
|
|
# Retrieve cached image
|
|
mx_image = handler.acquireImage(mx_filepath)
|
|
if mx_image:
|
|
out_path = mx.FilePath("saved_" + os.path.basename(test_image_path))
|
|
if handler.saveImage(out_path, mx_image, verticalFlip=args.flip):
|
|
logger.info(f"MaterialX Image saved to {out_path.asString()}")
|
|
else:
|
|
logger.error("Failed to save image.")
|
|
else:
|
|
logger.error("Failed to acquire image for saving.")
|
|
else:
|
|
logger.error("Failed to load image.")
|
|
|
|
if __name__ == "__main__":
|
|
test_load_save()
|