glTF exporter: Performance improvment on image export
This commit is contained in:
parent
73b85949a0
commit
289fb2b8b8
|
|
@ -15,7 +15,7 @@
|
|||
bl_info = {
|
||||
'name': 'glTF 2.0 format',
|
||||
'author': 'Julien Duroure, Norbert Nopper, Urs Hanselmann, Moritz Becher, Benjamin Schmithüsen, Jim Eckerlein, and many external contributors',
|
||||
"version": (1, 1, 30),
|
||||
"version": (1, 1, 31),
|
||||
'blender': (2, 81, 6),
|
||||
'location': 'File > Import-Export',
|
||||
'description': 'Import-Export as glTF 2.0',
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from io_scene_gltf2.blender.exp import gltf2_blender_search_node_tree
|
|||
from io_scene_gltf2.io.exp import gltf2_io_binary_data
|
||||
from io_scene_gltf2.io.exp import gltf2_io_image_data
|
||||
from io_scene_gltf2.io.com import gltf2_io_debug
|
||||
from io_scene_gltf2.blender.exp import gltf2_blender_image
|
||||
from io_scene_gltf2.blender.exp.gltf2_blender_image import Channel, ExportImage
|
||||
from io_scene_gltf2.blender.exp.gltf2_blender_gather_cache import cached
|
||||
from io_scene_gltf2.io.exp.gltf2_io_user_extensions import export_user_extensions
|
||||
|
||||
|
|
@ -38,8 +38,8 @@ def gather_image(
|
|||
return None
|
||||
|
||||
image_data = __get_image_data(blender_shader_sockets_or_texture_slots, export_settings)
|
||||
if image_data is None:
|
||||
# The blender image has no data
|
||||
if image_data.empty():
|
||||
# The export image has no data
|
||||
return None
|
||||
|
||||
mime_type = __gather_mime_type(blender_shader_sockets_or_texture_slots, export_settings)
|
||||
|
|
@ -144,13 +144,13 @@ def __is_slot(sockets_or_slots):
|
|||
return isinstance(sockets_or_slots[0], bpy.types.MaterialTextureSlot)
|
||||
|
||||
|
||||
def __get_image_data(sockets_or_slots, export_settings) -> gltf2_blender_image.ExportImage:
|
||||
def __get_image_data(sockets_or_slots, 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.
|
||||
if __is_socket(sockets_or_slots):
|
||||
results = [__get_tex_from_socket(socket, export_settings) for socket in sockets_or_slots]
|
||||
composed_image = None
|
||||
composed_image = ExportImage()
|
||||
for result, socket in zip(results, sockets_or_slots):
|
||||
if result.shader_node.image.channels == 0:
|
||||
gltf2_io_debug.print_console("WARNING",
|
||||
|
|
@ -159,44 +159,45 @@ def __get_image_data(sockets_or_slots, export_settings) -> gltf2_blender_image.E
|
|||
continue
|
||||
|
||||
# rudimentarily try follow the node tree to find the correct image data.
|
||||
source_channel = 0
|
||||
src_chan = Channel.R
|
||||
for elem in result.path:
|
||||
if isinstance(elem.from_node, bpy.types.ShaderNodeSeparateRGB):
|
||||
source_channel = {
|
||||
'R': 0,
|
||||
'G': 1,
|
||||
'B': 2
|
||||
src_chan = {
|
||||
'R': Channel.R,
|
||||
'G': Channel.G,
|
||||
'B': Channel.B,
|
||||
}[elem.from_socket.name]
|
||||
|
||||
image = gltf2_blender_image.ExportImage.from_blender_image(result.shader_node.image)
|
||||
|
||||
target_channel = None
|
||||
dst_chan = None
|
||||
|
||||
# some sockets need channel rewriting (gltf pbr defines fixed channels for some attributes)
|
||||
if socket.name == 'Metallic':
|
||||
target_channel = 2
|
||||
dst_chan = Channel.B
|
||||
elif socket.name == 'Roughness':
|
||||
target_channel = 1
|
||||
dst_chan = Channel.G
|
||||
elif socket.name == 'Occlusion' and len(sockets_or_slots) > 1 and sockets_or_slots[1] is not None:
|
||||
target_channel = 0
|
||||
dst_chan = Channel.R
|
||||
elif socket.name == 'Alpha' and len(sockets_or_slots) > 1 and sockets_or_slots[1] is not None:
|
||||
composed_image.set_alpha(True)
|
||||
target_channel = 3
|
||||
dst_chan = Channel.A
|
||||
|
||||
if target_channel is not None:
|
||||
if composed_image is None:
|
||||
composed_image = gltf2_blender_image.ExportImage.white_image(image.width, image.height)
|
||||
if dst_chan is not None:
|
||||
composed_image.fill_image(result.shader_node.image, dst_chan, src_chan)
|
||||
|
||||
composed_image[target_channel] = image[source_channel]
|
||||
# Since metal/roughness are always used together, make sure
|
||||
# the other channel is filled.
|
||||
if socket.name == 'Metallic' and not composed_image.is_filled(Channel.G):
|
||||
composed_image.fill_white(Channel.G)
|
||||
elif socket.name == 'Roughness' and not composed_image.is_filled(Channel.B):
|
||||
composed_image.fill_white(Channel.B)
|
||||
else:
|
||||
# copy full image...eventually following sockets might overwrite things
|
||||
composed_image = image
|
||||
composed_image = ExportImage.from_blender_image(result.shader_node.image)
|
||||
|
||||
return composed_image
|
||||
|
||||
elif __is_slot(sockets_or_slots):
|
||||
texture = __get_tex_from_slot(sockets_or_slots[0])
|
||||
image = gltf2_blender_image.ExportImage.from_blender_image(texture.image)
|
||||
image = ExportImage.from_blender_image(texture.image)
|
||||
return image
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
|
|
|||
|
|
@ -14,122 +14,206 @@
|
|||
|
||||
import bpy
|
||||
import os
|
||||
import typing
|
||||
from typing import Optional
|
||||
import numpy as np
|
||||
import tempfile
|
||||
import enum
|
||||
|
||||
|
||||
class Channel(enum.IntEnum):
|
||||
R = 0
|
||||
G = 1
|
||||
B = 2
|
||||
A = 3
|
||||
|
||||
# These describe how an ExportImage's channels should be filled.
|
||||
|
||||
class FillImage:
|
||||
"""Fills a channel with the channel src_chan from a Blender image."""
|
||||
def __init__(self, image: bpy.types.Image, src_chan: Channel):
|
||||
self.image = image
|
||||
self.src_chan = src_chan
|
||||
|
||||
class FillWhite:
|
||||
"""Fills a channel with all ones (1.0)."""
|
||||
pass
|
||||
|
||||
|
||||
class ExportImage:
|
||||
"""Custom image class that allows manipulation and encoding of images"""
|
||||
# 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
|
||||
"""Custom image class.
|
||||
|
||||
def __init__(self, img: typing.Union[np.ndarray, typing.List[np.ndarray]], max_channels: int = 4,\
|
||||
blender_image: bpy.types.Image = None, has_alpha: bool = False):
|
||||
if isinstance(img, list):
|
||||
np.stack(img, axis=2)
|
||||
An image is represented by giving a description of how to fill its red,
|
||||
green, blue, and alpha channels. For example:
|
||||
|
||||
if len(img.shape) == 2:
|
||||
# images must always have a channels dimension
|
||||
img = np.expand_dims(img, axis=2)
|
||||
self.fills = {
|
||||
Channel.R: FillImage(image=bpy.data.images['Im1'], src_chan=Channel.B),
|
||||
Channel.G: FillWhite(),
|
||||
}
|
||||
|
||||
if not len(img.shape) == 3 or img.shape[2] > 4:
|
||||
raise RuntimeError("Cannot construct an export image from an array of shape {}".format(img.shape))
|
||||
This says that the ExportImage's R channel should be filled with the B
|
||||
channel of the Blender image 'Im1', and the ExportImage's G channel
|
||||
should be filled with all 1.0s. Undefined channels mean we don't care
|
||||
what values that channel has.
|
||||
|
||||
self._img = img
|
||||
self._max_channels = max_channels
|
||||
self._blender_image = blender_image
|
||||
self._has_alpha = has_alpha
|
||||
This is flexible enough to handle the case where eg. the user used the R
|
||||
channel of one image as the metallic value and the G channel of another
|
||||
image as the roughness, and we need to synthesize an ExportImage that
|
||||
packs those into the B and G channels for glTF.
|
||||
|
||||
def set_alpha(self, alpha: bool):
|
||||
self._has_alpha = alpha
|
||||
Storing this description (instead of raw pixels) lets us make more
|
||||
intelligent decisions about how to encode the image.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def from_blender_image(cls, blender_image: bpy.types.Image):
|
||||
img = np.array(blender_image.pixels[:])
|
||||
img = img.reshape((blender_image.size[0], blender_image.size[1], blender_image.channels))
|
||||
has_alpha = blender_image.depth == 32
|
||||
return ExportImage(img=img, blender_image=blender_image, has_alpha=has_alpha)
|
||||
def __init__(self):
|
||||
self.fills = {}
|
||||
|
||||
@classmethod
|
||||
def white_image(cls, width, height, num_channels: int = 4):
|
||||
img = np.ones((width, height, num_channels))
|
||||
return ExportImage(img=img)
|
||||
@staticmethod
|
||||
def from_blender_image(image: bpy.types.Image):
|
||||
export_image = ExportImage()
|
||||
for chan in range(image.channels):
|
||||
export_image.fill_image(image, dst_chan=chan, src_chan=chan)
|
||||
return export_image
|
||||
|
||||
def split_channels(self):
|
||||
"""return a list of numpy arrays where each list element corresponds to one image channel (r,g?,b?,a?)"""
|
||||
return np.split(self._img, self._img.shape[2], axis=2)
|
||||
def fill_image(self, image: bpy.types.Image, dst_chan: Channel, src_chan: Channel):
|
||||
self.fills[dst_chan] = FillImage(image, src_chan)
|
||||
|
||||
@property
|
||||
def img(self) -> np.ndarray:
|
||||
return self._img
|
||||
def fill_white(self, dst_chan: Channel):
|
||||
self.fills[dst_chan] = FillWhite()
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._img.shape
|
||||
def is_filled(self, chan: Channel) -> bool:
|
||||
return chan in self.fills
|
||||
|
||||
@property
|
||||
def width(self):
|
||||
return self.shape[0]
|
||||
def empty(self) -> bool:
|
||||
return not self.fills
|
||||
|
||||
@property
|
||||
def height(self):
|
||||
return self.shape[1]
|
||||
def __on_happy_path(self) -> bool:
|
||||
# Whether there is an existing Blender image we can use for this
|
||||
# ExportImage because all the channels come from the matching
|
||||
# channel of that image, eg.
|
||||
#
|
||||
# self.fills = {
|
||||
# Channel.R: FillImage(image=im, src_chan=Channel.R),
|
||||
# Channel.G: FillImage(image=im, src_chan=Channel.G),
|
||||
# }
|
||||
return (
|
||||
all(isinstance(fill, FillImage) for fill in self.fills.values()) and
|
||||
all(dst_chan == fill.src_chan for dst_chan, fill in self.fills.items()) and
|
||||
len(set(fill.image.name for fill in self.fills.values())) == 1
|
||||
)
|
||||
|
||||
@property
|
||||
def channels(self):
|
||||
return self.shape[2]
|
||||
|
||||
def __getitem__(self, key):
|
||||
"""returns a new ExportImage with only the selected channels"""
|
||||
return ExportImage(self._img[:, :, key])
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
"""set the selected channels to a new value"""
|
||||
if isinstance(key, slice):
|
||||
self._img[:, :, key] = value.img
|
||||
else:
|
||||
self._img[:, :, key] = value.img[:, :, 0]
|
||||
|
||||
def append(self, other):
|
||||
if self.channels + other.channels > self._max_channels:
|
||||
raise RuntimeError("Cannot append image data to this image "
|
||||
"because the maximum number of channels is exceeded.")
|
||||
|
||||
self._img = np.concatenate([self.img, other.img], axis=2)
|
||||
|
||||
def __add__(self, other):
|
||||
self.append(other)
|
||||
|
||||
def encode(self, mime_type: typing.Optional[str]) -> bytes:
|
||||
file_format = {
|
||||
def encode(self, mime_type: Optional[str]) -> bytes:
|
||||
self.file_format = {
|
||||
"image/jpeg": "JPEG",
|
||||
"image/png": "PNG"
|
||||
}.get(mime_type, "PNG")
|
||||
|
||||
if self._blender_image is not None and file_format == self._blender_image.file_format:
|
||||
src_path = bpy.path.abspath(self._blender_image.filepath_raw)
|
||||
if os.path.isfile(src_path):
|
||||
with open(src_path, "rb") as f:
|
||||
encoded_image = f.read()
|
||||
return encoded_image
|
||||
# Happy path = we can just use an existing Blender image
|
||||
if self.__on_happy_path():
|
||||
return self.__encode_happy()
|
||||
|
||||
image = bpy.data.images.new("TmpImage", width=self.width, height=self.height, alpha=self._has_alpha)
|
||||
pixels = self._img.flatten().tolist()
|
||||
image.pixels = pixels
|
||||
# Unhappy path = we need to create the image self.fills describes.
|
||||
return self.__encode_unhappy()
|
||||
|
||||
# we just use blenders built in save mechanism, this can be considered slightly dodgy but currently is the only
|
||||
# way to support
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
tmpfilename = tmpdirname + "/img"
|
||||
image.filepath_raw = tmpfilename
|
||||
image.file_format = file_format
|
||||
image.save()
|
||||
def __encode_happy(self) -> bytes:
|
||||
for fill in self.fills.values():
|
||||
return self.__encode_from_image(fill.image)
|
||||
|
||||
with open(tmpfilename, "rb") as f:
|
||||
encoded_image = f.read()
|
||||
def __encode_unhappy(self) -> bytes:
|
||||
# This will be a numpy array we fill in with pixel data.
|
||||
result = None
|
||||
|
||||
bpy.data.images.remove(image, do_unlink=True)
|
||||
img_fills = {
|
||||
chan: fill
|
||||
for chan, fill in self.fills.items()
|
||||
if isinstance(fill, FillImage)
|
||||
}
|
||||
# Loop over images instead of dst_chans; ensures we only decode each
|
||||
# image once even if it's used in multiple channels.
|
||||
image_names = list(set(fill.image.name for fill in img_fills.values()))
|
||||
for image_name in image_names:
|
||||
image = bpy.data.images[image_name]
|
||||
|
||||
return encoded_image
|
||||
if result is None:
|
||||
result = np.ones((image.size[0], image.size[1], 4), np.float32)
|
||||
# Images should all be the same size (should be guaranteed by
|
||||
# gather_texture_info).
|
||||
assert (image.size[0], image.size[1]) == result.shape[:2]
|
||||
|
||||
# Slow and eats all your memory.
|
||||
pixels = np.array(image.pixels[:])
|
||||
|
||||
pixels = pixels.reshape((image.size[0], image.size[1], image.channels))
|
||||
|
||||
for dst_chan, img_fill in img_fills.items():
|
||||
if img_fill.image == image:
|
||||
result[:, :, dst_chan] = pixels[:, :, img_fill.src_chan]
|
||||
|
||||
pixels = None # GC this please
|
||||
|
||||
if result is None:
|
||||
# No ImageFills; use a 1x1 white pixel
|
||||
result = np.array([1.0, 1.0, 1.0, 1.0])
|
||||
result = result.reshape((1, 1, 4))
|
||||
|
||||
return self.__encode_from_numpy_array(result)
|
||||
|
||||
def __encode_from_numpy_array(self, array: np.ndarray) -> bytes:
|
||||
tmp_image = None
|
||||
try:
|
||||
tmp_image = bpy.data.images.new(
|
||||
"##gltf-export:tmp-image##",
|
||||
width=array.shape[0],
|
||||
height=array.shape[1],
|
||||
alpha=Channel.A in self.fills,
|
||||
)
|
||||
assert tmp_image.channels == 4 # 4 regardless of the alpha argument above.
|
||||
|
||||
# Also slow and eats all your memory.
|
||||
tmp_image.pixels = array.flatten().tolist()
|
||||
|
||||
return _encode_temp_image(tmp_image, self.file_format)
|
||||
|
||||
finally:
|
||||
if tmp_image is not None:
|
||||
bpy.data.images.remove(tmp_image, do_unlink=True)
|
||||
|
||||
def __encode_from_image(self, image: bpy.types.Image) -> bytes:
|
||||
# See if there is an existing file we can use.
|
||||
if image.source == 'FILE' and image.file_format == self.file_format and \
|
||||
not image.is_dirty:
|
||||
if image.packed_file is not None:
|
||||
return image.packed_file.data
|
||||
else:
|
||||
src_path = bpy.path.abspath(image.filepath_raw)
|
||||
if os.path.isfile(src_path):
|
||||
with open(src_path, 'rb') as f:
|
||||
return f.read()
|
||||
|
||||
# Copy to a temp image and save.
|
||||
tmp_image = None
|
||||
try:
|
||||
tmp_image = image.copy()
|
||||
if image.is_dirty:
|
||||
tmp_image.pixels = image.pixels[:]
|
||||
|
||||
return _encode_temp_image(tmp_image, self.file_format)
|
||||
finally:
|
||||
if tmp_image is not None:
|
||||
bpy.data.images.remove(tmp_image, do_unlink=True)
|
||||
|
||||
|
||||
def _encode_temp_image(tmp_image: bpy.types.Image, file_format: str) -> bytes:
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
tmpfilename = tmpdirname + '/img'
|
||||
tmp_image.filepath_raw = tmpfilename
|
||||
|
||||
# NOT A TYPO!!! If you delete this line, the
|
||||
# assignment on the next line will not work.
|
||||
tmp_image.file_format
|
||||
tmp_image.file_format = file_format
|
||||
|
||||
tmp_image.save()
|
||||
|
||||
with open(tmpfilename, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue