(WIP) More flexible image export handling for extensions.

addons/io_scene_gltf2/blender/exp/gltf2_blender_gather_image.py

@@ -27,14 +27,16 @@
 from io_scene_gltf2.io.exp.gltf2_io_user_extensions import export_user_extensions
 
 
-@cached
+# @cached
 def gather_image(
         blender_shader_sockets: typing.Tuple[bpy.types.NodeSocket],
-        export_settings):
+        export_settings,
+        export_image_class = ExportImage
+):
     if not __filter_image(blender_shader_sockets, export_settings):
         return None
 
-    image_data = __get_image_data(blender_shader_sockets, export_settings)
+    image_data = __get_image_data(blender_shader_sockets, export_image_class, export_settings)
     if image_data.empty():
         # The export image has no data
         return None
@@ -93,20 +95,14 @@ def __gather_extras(sockets, export_settings):
 
 
 def __gather_mime_type(sockets, export_image, export_settings):
-    # force png if Alpha contained so we can export alpha
-    for socket in sockets:
-        if socket.name == "Alpha":
-            return "image/png"
-
     if export_settings["gltf_image_format"] == "AUTO":
-        image = export_image.blender_image()
-        if image is not None and __is_blender_image_a_jpeg(image):
-            return "image/jpeg"
-        return "image/png"
+        preferred = export_image.preferred_mime_type()
+        if preferred: return preferred
 
     elif export_settings["gltf_image_format"] == "JPEG":
         return "image/jpeg"
 
+    return "image/png"
 
 def __gather_name(export_image, export_settings):
     # Find all Blender images used in the ExportImage
@@ -122,9 +118,7 @@ def __gather_name(export_image, export_settings):
     if len(filepaths) == 1:
         filename = os.path.basename(list(filepaths)[0])
         name, extension = os.path.splitext(filename)
-        if extension.lower() in ['.png', '.jpg', '.jpeg']:
-            if name:
-                return name
+        return name
 
     # Combine the image names: img1-img2-img3
     names = []
@@ -148,12 +142,12 @@ def __gather_uri(image_data, mime_type, name, export_settings):
     return None
 
 
-def __get_image_data(sockets, export_settings) -> ExportImage:
+def __get_image_data(sockets, export_image_class, export_settings) -> ExportImage:
     # For shared resources, such as images, we just store the portion of data that is needed in the glTF property
     # in a helper class. During generation of the glTF in the exporter these will then be combined to actual binary
     # resources.
     results = [__get_tex_from_socket(socket, export_settings) for socket in sockets]
-    composed_image = ExportImage()
+    composed_image = export_image_class()
     for result, socket in zip(results, sockets):
         if result.shader_node.image.channels == 0:
             gltf2_io_debug.print_console("WARNING",
@@ -200,7 +194,7 @@ def __get_image_data(sockets, export_settings) -> ExportImage:
                 composed_image.fill_white(Channel.B)
         else:
             # copy full image...eventually following sockets might overwrite things
-            composed_image = ExportImage.from_blender_image(result.shader_node.image)
+            composed_image = export_image_class.from_blender_image(result.shader_node.image)
 
     return composed_image
 
@@ -213,10 +207,3 @@ def __get_tex_from_socket(blender_shader_socket: bpy.types.NodeSocket, export_se
     if not result:
         return None
     return result[0]
-
-
-def __is_blender_image_a_jpeg(image: bpy.types.Image) -> bool:
-    if image.source != 'FILE':
-        return False
-    path = image.filepath_raw.lower()
-    return path.endswith('.jpg') or path.endswith('.jpeg') or path.endswith('.jpe')

addons/io_scene_gltf2/blender/exp/gltf2_blender_image.py

@@ -67,9 +67,9 @@ class ExportImage:
     def __init__(self):
         self.fills = {}
 
-    @staticmethod
-    def from_blender_image(image: bpy.types.Image):
-        export_image = ExportImage()
+    @classmethod
+    def from_blender_image(cls, image: bpy.types.Image):
+        export_image = cls()
         for chan in range(image.channels):
             export_image.fill_image(image, dst_chan=chan, src_chan=chan)
         return export_image
@@ -104,12 +104,24 @@ def __on_happy_path(self) -> bool:
             len(set(fill.image.name for fill in self.fills.values())) == 1
         )
 
-    def encode(self, mime_type: Optional[str]) -> bytes:
-        self.file_format = {
+    def preferred_mime_type(self) -> str:
+        image = self.blender_image()
+        if image:
+            return {
+                "JPEG": "image/jpeg",
+                "PNG": "image/png"
+            }.get(image.file_format, None)
+        return None
+
+    def __blender_format_for_mime(self, mime_type: Optional[str]) -> str:
+        return {
             "image/jpeg": "JPEG",
             "image/png": "PNG"
         }.get(mime_type, "PNG")
 
+    def encode(self, mime_type: Optional[str]) -> bytes:
+        self.file_format = self.__blender_format_for_mime(mime_type)
+
         # Happy path = we can just use an existing Blender image
         if self.__on_happy_path():
             return self.__encode_happy()
@@ -176,6 +188,13 @@ def __encode_from_numpy_array(self, pixels: np.ndarray, dim: Tuple[int, int]) ->
 
             return _encode_temp_image(tmp_image, self.file_format)
 
+    def __check_magic(self, data: bytes) -> bool:
+        if self.file_format == 'PNG':
+            return data.startswith(b'\x89PNG')
+        elif self.file_format == 'JPEG':
+            return data.startswith(b'\xff\xd8\xff')
+        return False
+
     def __encode_from_image(self, image: bpy.types.Image) -> bytes:
         # See if there is an existing file we can use.
         data = None
@@ -189,13 +208,8 @@ def __encode_from_image(self, image: bpy.types.Image) -> bytes:
                     with open(src_path, 'rb') as f:
                         data = f.read()
         # Check magic number is right
-        if data:
-            if self.file_format == 'PNG':
-                if data.startswith(b'\x89PNG'):
-                    return data
-            elif self.file_format == 'JPEG':
-                if data.startswith(b'\xff\xd8\xff'):
-                    return data
+        if data and self.__check_magic(data):
+            return data
 
         # Copy to a temp image and save.
         with TmpImageGuard() as guard:

addons/io_scene_gltf2/io/exp/gltf2_io_image_data.py

@@ -19,6 +19,13 @@ class ImageData:
     # FUTURE_WORK: as a method to allow the node graph to be better supported, we could model some of
     # the node graph elements with numpy functions
 
+    extension_for_mime = {
+        "image/jpeg": ".jpg",
+        "image/png": ".png",
+        "image/x-exr": ".exr",
+        "image/vnd.radiance": ".hdr"
+    }
+
     def __init__(self, data: bytes, mime_type: str, name: str):
         self._data = data
         self._mime_type = mime_type
@@ -46,9 +53,7 @@ def name(self):
 
     @property
     def file_extension(self):
-        if self._mime_type == "image/jpeg":
-            return ".jpg"
-        return ".png"
+        return ImageData.extension_for_mime.get(self._mime_type)
 
     @property
     def byte_length(self):