Use up-to-date loader.

Author: kwokcb

plugins/OIIO/OIIOLoader.py

@@ -27,7 +27,6 @@
     logger.error("Required modules not found. Please install OpenImageIO and numpy.")
     raise
 
-
 have_matplot = False
 try:
     import matplotlib.pyplot as plt
@@ -62,19 +61,48 @@ def __init__(self):
                 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):
         """
-        Return a set of supported image file extensions.
+        Derived method to return a set of supported image file extensions.
         """
-        logger.info(f"Supported OIIO supported extensions: {self._extensions}")
+        logger.info(f"OIIO supported extensions: {self._extensions}")
         return self._extensions
 
-    def previewImage(self, data, width, height, nchannels):
+    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:
@@ -99,19 +127,23 @@ def previewImage(self, data, width, height, nchannels):
             else:
                 rgb_disp = rgb
 
-            plt.imshow(rgb_disp)
-            plt.axis('off')
+            # 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).
-        
-        Args:
-            filePath (MaterialX.FilePath): Path to the image file
-            
-        Returns:
-            MaterialX.ImagePtr: MaterialX Image object or None if loading fails
+
+        @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}")
@@ -129,8 +161,14 @@ def loadImage(self, filePath):
 
             # Get image specifications
             spec = img_input.spec()
-            self.last_spec = spec
-            self.last_loaded_path = file_path_str
+            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)
@@ -140,25 +178,28 @@ def loadImage(self, filePath):
                 return None
 
             # Create MaterialX image
-            mx_image = mx_render.Image.create(spec.width, spec.height, spec.nchannels, base_type)
+            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}")
+            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, spec.nchannels, spec.format)
+            data = mx_image.getResourceBuffer()
+            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 * spec.nchannels * bytes_per_channel
-            logger.info(f"Total bytes read in: {total_bytes} (width: {spec.width}, height: {spec.height}, channels: {spec.nchannels}, format: {spec.format})")
+            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:
@@ -186,19 +227,26 @@ def saveImage(self, filePath, image, verticalFlip=False):
         filename = filePath.asString()
         width = image.getWidth()
         height = image.getHeight()
-        channels = image.getChannelCount()
+
+        # 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"Error: Unsupported MaterialX base type for OIIO: {mx_basetype}")
+            logger.error(f"Unsupported MaterialX base type for OIIO: {mx_basetype}")
             return False
-        
-        buffer = image.getResourceBuffer()
-        
+
+        buffer_addr = image.getResourceBuffer()
         np_type = self._materialx_type_to_np_type(mx_basetype)
-        pixels = np.zeros((height, width, channels), dtype=np_type)
-        # Copy from buffer to pixels
+        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.
@@ -207,37 +255,64 @@ def saveImage(self, filePath, image, verticalFlip=False):
             # - Optionally previews the image for debugging.
             # - Creates an OIIO ImageOutput and writes the image to disk.
             #
-            base_stride = image.getBaseStride()
-            total_bytes = width * height * channels * base_stride
+            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)
-            np_buffer = np.frombuffer(buf, dtype=np_type).reshape((height, width, channels))
-            np.copyto(pixels, np_buffer)
+            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()
 
-            # Handle vertical flip if requested
             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...")
-            self.previewImage(pixels, width, height, channels)
+            # 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 out:
-            if np_type is None:
-                logger.error(f"Error: Unsupported NumPy type for OIIO: {mx_basetype}")
-                return False
-            spec = oiio.ImageSpec(width, height, channels, np_type)
+        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} with width: {width}, height: {height}, channels: {channels}, base type: {mx_basetype}")
+            logger.info(f"Image saved to {filename} (w={width}, h={height}, c={channels}, type={mx_basetype})")
             out.close()
             return True
-        return False
+        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."""
@@ -268,14 +343,14 @@ def _materialx_to_oiio_type(self, mx_basetype):
         return return_val
 
     def _materialx_type_to_np_type(self, mx_basetype):
-        """Map MaterialX base type to NumPy dtype."""
+        """Map MaterialX base type to NumPy dtype with explicit widths."""
         type_mapping = {
-            mx_render.BaseType.UINT8: 'uint8',
-            mx_render.BaseType.UINT16: 'uint16',
-            mx_render.BaseType.INT8: 'int8',
-            mx_render.BaseType.INT16: 'int16',
-            mx_render.BaseType.HALF: 'half',
-            mx_render.BaseType.FLOAT: 'float',
+            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)
 
@@ -287,6 +362,7 @@ def test_load_save():
     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
@@ -296,15 +372,15 @@ def test_load_save():
 
     # Create MaterialX handler with custom OIIO image loader
     loader = OiioImageLoader()
-    #handler = mx_render.ImageHandler.create(loader)
-    manager = mx_render.getPluginManager()
-    handler = manager.getImageHandler()
-    logger.info(f"Got handler from plugin manager {handler}")
+    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:
@@ -319,12 +395,18 @@ def test_load_save():
         logger.info(f"  Base type: {mx_image.getBaseType()}")
 
         # Save image using handler API (to a new file)
-        logger.info('*'*45)
-        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()}")
+        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 save image.")
+            logger.error("Failed to acquire image for saving.")
     else:
         logger.error("Failed to load image.")