io_scene_gltf2/blender/exp/material/gltf2_blender_gather_image.py

   1 # SPDX-FileCopyrightText: 2018-2021 The glTF-Blender-IO authors
   2 #
   3 # SPDX-License-Identifier: Apache-2.0
   4
   5 import bpy
   6 import typing
   7 import os
   8
   9 from ....io.com import gltf2_io
  10 from ....io.com.gltf2_io_path import path_to_uri
  11 from ....io.exp import gltf2_io_binary_data, gltf2_io_image_data
  12 from ....io.com import gltf2_io_debug
  13 from ....io.exp.gltf2_io_user_extensions import export_user_extensions
  14 from ..gltf2_blender_gather_cache import cached
  15 from .extensions.gltf2_blender_image import Channel, ExportImage, FillImage
  16 from ..gltf2_blender_get import get_tex_from_socket
  17
  18 @cached
  19 def gather_image(
  20         blender_shader_sockets: typing.Tuple[bpy.types.NodeSocket],
  21         default_sockets: typing.Tuple[bpy.types.NodeSocket],
  22         export_settings):
  23     if not __filter_image(blender_shader_sockets, export_settings):
  24         return None, None, None
  25
  26     image_data = __get_image_data(blender_shader_sockets, default_sockets, export_settings)
  27     if image_data.empty():
  28         # The export image has no data
  29         return None, None, None
  30
  31     mime_type = __gather_mime_type(blender_shader_sockets, image_data, export_settings)
  32     name = __gather_name(image_data, export_settings)
  33
  34     factor = None
  35
  36     if image_data.original is None:
  37         uri, factor_uri = __gather_uri(image_data, mime_type, name, export_settings)
  38     else:
  39         # Retrieve URI relative to exported glTF files
  40         uri = __gather_original_uri(image_data.original.filepath, export_settings)
  41         # In case we can't retrieve image (for example packed images, with original moved)
  42         # We don't create invalid image without uri
  43         factor_uri = None
  44         if uri is None: return None, None, None
  45
  46     buffer_view, factor_buffer_view = __gather_buffer_view(image_data, mime_type, name, export_settings)
  47
  48     factor = factor_uri if uri is not None else factor_buffer_view
  49
  50     image = __make_image(
  51         buffer_view,
  52         __gather_extensions(blender_shader_sockets, export_settings),
  53         __gather_extras(blender_shader_sockets, export_settings),
  54         mime_type,
  55         name,
  56         uri,
  57         export_settings
  58     )
  59
  60     export_user_extensions('gather_image_hook', export_settings, image, blender_shader_sockets)
  61
  62     # We also return image_data, as it can be used to generate same file with another extension for webp management
  63     return image, image_data, factor
  64
  65 def __gather_original_uri(original_uri, export_settings):
  66
  67     path_to_image = bpy.path.abspath(original_uri)
  68     if not os.path.exists(path_to_image): return None
  69     try:
  70         rel_path = os.path.relpath(
  71             path_to_image,
  72             start=export_settings['gltf_filedirectory'],
  73         )
  74     except ValueError:
  75         # eg. because no relative path between C:\ and D:\ on Windows
  76         return None
  77     return path_to_uri(rel_path)
  78
  79
  80 @cached
  81 def __make_image(buffer_view, extensions, extras, mime_type, name, uri, export_settings):
  82     return gltf2_io.Image(
  83         buffer_view=buffer_view,
  84         extensions=extensions,
  85         extras=extras,
  86         mime_type=mime_type,
  87         name=name,
  88         uri=uri
  89     )
  90
  91
  92 def __filter_image(sockets, export_settings):
  93     if not sockets:
  94         return False
  95     return True
  96
  97
  98 @cached
  99 def __gather_buffer_view(image_data, mime_type, name, export_settings):
 100     if export_settings['gltf_format'] != 'GLTF_SEPARATE':
 101         data, factor = image_data.encode(mime_type, export_settings)
 102         return gltf2_io_binary_data.BinaryData(data=data), factor
 103     return None, None
 104
 105
 106 def __gather_extensions(sockets, export_settings):
 107     return None
 108
 109
 110 def __gather_extras(sockets, export_settings):
 111     return None
 112
 113
 114 def __gather_mime_type(sockets, export_image, export_settings):
 115     # force png or webp if Alpha contained so we can export alpha
 116     for socket in sockets:
 117         if socket.name == "Alpha":
 118             if export_settings["gltf_image_format"] == "WEBP":
 119                 return "image/webp"
 120             else:
 121                 # If we keep image as is (no channel composition), we need to keep original format (for webp)
 122                 image = export_image.blender_image()
 123                 if image is not None and __is_blender_image_a_webp(image):
 124                     return "image/webp"
 125                 return "image/png"
 126
 127     if export_settings["gltf_image_format"] == "AUTO":
 128         if export_image.original is None: # We are going to create a new image
 129             image = export_image.blender_image()
 130         else:
 131             # Using original image
 132             image = export_image.original
 133
 134         if image is not None and __is_blender_image_a_jpeg(image):
 135             return "image/jpeg"
 136         elif image is not None and __is_blender_image_a_webp(image):
 137             return "image/webp"
 138         return "image/png"
 139
 140     elif export_settings["gltf_image_format"] == "WEBP":
 141         return "image/webp"
 142     elif export_settings["gltf_image_format"] == "JPEG":
 143         return "image/jpeg"
 144
 145
 146 def __gather_name(export_image, export_settings):
 147     if export_image.original is None:
 148         # Find all Blender images used in the ExportImage
 149         imgs = []
 150         for fill in export_image.fills.values():
 151             if isinstance(fill, FillImage):
 152                 img = fill.image
 153                 if img not in imgs:
 154                     imgs.append(img)
 155
 156         # If all the images have the same path, use the common filename
 157         filepaths = set(img.filepath for img in imgs)
 158         if len(filepaths) == 1:
 159             filename = os.path.basename(list(filepaths)[0])
 160             name, extension = os.path.splitext(filename)
 161             if extension.lower() in ['.png', '.jpg', '.jpeg']:
 162                 if name:
 163                     return name
 164
 165         # Combine the image names: img1-img2-img3
 166         names = []
 167         for img in imgs:
 168             name, extension = os.path.splitext(img.name)
 169             names.append(name)
 170         name = '-'.join(names)
 171         return name or 'Image'
 172     else:
 173         return export_image.original.name
 174
 175
 176 @cached
 177 def __gather_uri(image_data, mime_type, name, export_settings):
 178     if export_settings['gltf_format'] == 'GLTF_SEPARATE':
 179         # as usual we just store the data in place instead of already resolving the references
 180         data, factor = image_data.encode(mime_type, export_settings)
 181         return gltf2_io_image_data.ImageData(
 182             data=data,
 183             mime_type=mime_type,
 184             name=name
 185         ), factor
 186
 187     return None, None
 188
 189
 190 def __get_image_data(sockets, default_sockets, export_settings) -> ExportImage:
 191     # For shared resources, such as images, we just store the portion of data that is needed in the glTF property
 192     # in a helper class. During generation of the glTF in the exporter these will then be combined to actual binary
 193     # resources.
 194     results = [get_tex_from_socket(socket) for socket in sockets]
 195
 196     # Check if we need a simple mapping or more complex calculation
 197     if any([socket.name == "Specular" and socket.node.type == "BSDF_PRINCIPLED" for socket in sockets]):
 198         return __get_image_data_specular(sockets, results, export_settings)
 199     else:
 200         return __get_image_data_mapping(sockets, default_sockets, results, export_settings)
 201
 202 def __get_image_data_mapping(sockets, default_sockets, results, export_settings) -> ExportImage:
 203     """
 204     Simple mapping
 205     Will fit for most of exported textures : RoughnessMetallic, Basecolor, normal, ...
 206     """
 207     composed_image = ExportImage()
 208
 209     default_metallic = None
 210     default_roughness = None
 211     if "Metallic" in [s.name for s in default_sockets]:
 212         default_metallic = [s for s in default_sockets if s.name == "Metallic"][0].default_value
 213     if "Roughness" in [s.name for s in default_sockets]:
 214         default_roughness = [s for s in default_sockets if s.name == "Roughness"][0].default_value
 215
 216     for result, socket in zip(results, sockets):
 217         # Assume that user know what he does, and that channels/images are already combined correctly for pbr
 218         # If not, we are going to keep only the first texture found
 219         # Example : If user set up 2 or 3 different textures for Metallic / Roughness / Occlusion
 220         # Only 1 will be used at export
 221         # This Warning is displayed in UI of this option
 222         if export_settings['gltf_keep_original_textures']:
 223             composed_image = ExportImage.from_original(result.shader_node.image)
 224
 225         else:
 226             # rudimentarily try follow the node tree to find the correct image data.
 227             src_chan = Channel.R
 228             for elem in result.path:
 229                 if isinstance(elem.from_node, bpy.types.ShaderNodeSeparateColor):
 230                     src_chan = {
 231                         'Red': Channel.R,
 232                         'Green': Channel.G,
 233                         'Blue': Channel.B,
 234                     }[elem.from_socket.name]
 235                 if elem.from_socket.name == 'Alpha':
 236                     src_chan = Channel.A
 237
 238             dst_chan = None
 239
 240             # some sockets need channel rewriting (gltf pbr defines fixed channels for some attributes)
 241             if socket.name == 'Metallic':
 242                 dst_chan = Channel.B
 243             elif socket.name == 'Roughness' and socket.node.type == "BSDF_PRINCIPLED":
 244                 dst_chan = Channel.G
 245             elif socket.name == 'Occlusion':
 246                 dst_chan = Channel.R
 247             elif socket.name == 'Alpha':
 248                 dst_chan = Channel.A
 249             elif socket.name == 'Coat':
 250                 dst_chan = Channel.R
 251             elif socket.name == 'Coat Roughness':
 252                 dst_chan = Channel.G
 253             elif socket.name == 'Thickness': # For KHR_materials_volume
 254                 dst_chan = Channel.G
 255             elif socket.name == "Specular": # For original KHR_material_specular
 256                 dst_chan = Channel.A
 257             elif socket.name == "Roughness" and socket.node.type == "BSDF_SHEEN": # For KHR_materials_sheen
 258                 dst_chan = Channel.A
 259
 260             if dst_chan is not None:
 261                 composed_image.fill_image(result.shader_node.image, dst_chan, src_chan)
 262
 263                 # Since metal/roughness are always used together, make sure
 264                 # the other channel is filled.
 265                 if socket.name == 'Metallic' and not composed_image.is_filled(Channel.G):
 266                     if default_roughness is not None:
 267                         composed_image.fill_with(Channel.G, default_roughness)
 268                     else:
 269                         composed_image.fill_white(Channel.G)
 270                 elif socket.name == 'Roughness' and not composed_image.is_filled(Channel.B):
 271                     if default_metallic is not None:
 272                         composed_image.fill_with(Channel.B, default_metallic)
 273                     else:
 274                         composed_image.fill_white(Channel.B)
 275             else:
 276                 # copy full image...eventually following sockets might overwrite things
 277                 composed_image = ExportImage.from_blender_image(result.shader_node.image)
 278
 279     # Check that we don't have some empty channels (based on weird images without any size for example)
 280     keys = list(composed_image.fills.keys()) # do not loop on dict, we may have to delete an element
 281     for k in [k for k in keys if isinstance(composed_image.fills[k], FillImage)]:
 282         if composed_image.fills[k].image.size[0] == 0 or composed_image.fills[k].image.size[1] == 0:
 283             gltf2_io_debug.print_console("WARNING",
 284                                          "Image '{}' has no size and cannot be exported.".format(
 285                                              composed_image.fills[k].image))
 286             del composed_image.fills[k]
 287
 288     return composed_image
 289
 290
 291 def __get_image_data_specular(sockets, results, export_settings) -> ExportImage:
 292     """
 293     calculating Specular Texture, settings needed data
 294     """
 295     from .extensions.gltf2_blender_texture_specular import specular_calculation
 296     composed_image = ExportImage()
 297     composed_image.set_calc(specular_calculation)
 298
 299     composed_image.store_data("ior", sockets[4].default_value, type="Data")
 300
 301     results = [get_tex_from_socket(socket) for socket in sockets[:-1]] #Do not retrieve IOR --> No texture allowed
 302
 303     mapping = {
 304         0: "specular",
 305         1: "specular_tint",
 306         2: "base_color",
 307         3: "transmission"
 308     }
 309
 310     for idx, result in enumerate(results):
 311         if get_tex_from_socket(sockets[idx]):
 312
 313             composed_image.store_data(mapping[idx], result.shader_node.image, type="Image")
 314
 315             # rudimentarily try follow the node tree to find the correct image data.
 316             src_chan = None if idx == 2 else Channel.R
 317             for elem in result.path:
 318                 if isinstance(elem.from_node, bpy.types.ShaderNodeSeparateColor):
 319                     src_chan = {
 320                         'Red': Channel.R,
 321                         'Green': Channel.G,
 322                         'Blue': Channel.B,
 323                     }[elem.from_socket.name]
 324                 if elem.from_socket.name == 'Alpha':
 325                     src_chan = Channel.A
 326             # For base_color, keep all channels, as this is a Vec, not scalar
 327             if idx != 2:
 328                 composed_image.store_data(mapping[idx] + "_channel", src_chan, type="Data")
 329             else:
 330                 if src_chan is not None:
 331                     composed_image.store_data(mapping[idx] + "_channel", src_chan, type="Data")
 332
 333         else:
 334             composed_image.store_data(mapping[idx], sockets[idx].default_value, type="Data")
 335
 336     return composed_image
 337
 338
 339 def __is_blender_image_a_jpeg(image: bpy.types.Image) -> bool:
 340     if image.source != 'FILE':
 341         return False
 342     if image.filepath_raw == '' and image.packed_file:
 343         return image.packed_file.data[:3] == b'\xff\xd8\xff'
 344     else:
 345         path = image.filepath_raw.lower()
 346         return path.endswith('.jpg') or path.endswith('.jpeg') or path.endswith('.jpe')
 347
 348 def __is_blender_image_a_webp(image: bpy.types.Image) -> bool:
 349     if image.source != 'FILE':
 350         return False
 351     if image.filepath_raw == '' and image.packed_file:
 352         return image.packed_file.data[8:12] == b'WEBP'
 353     else:
 354         path = image.filepath_raw.lower()
 355         return path.endswith('.webp')