glTF exporter: fix / enhancement of texture images export
This commit is contained in:
parent
4f34e011d4
commit
8e72572153
|
@ -107,6 +107,22 @@ class ExportGLTF2_Base:
|
|||
default=''
|
||||
)
|
||||
|
||||
export_image_format: EnumProperty(
|
||||
name='Image Format',
|
||||
items=(('NAME', 'from image name',
|
||||
'Determine the output format from the blender image name'),
|
||||
('JPEG', 'jpeg image format (.jpg)',
|
||||
'encode and save textures as .jpeg files. Be aware of a possible loss in quality.'),
|
||||
('PNG', 'png image format (.png)',
|
||||
'encode and save textures as .png files.')
|
||||
),
|
||||
description=(
|
||||
'Output format for images. PNG is lossless and generally preferred, but JPEG might be preferable for web '
|
||||
'applications due to the smaller file size'
|
||||
),
|
||||
default='NAME'
|
||||
)
|
||||
|
||||
export_texcoords: BoolProperty(
|
||||
name='UVs',
|
||||
description='Export UVs (texture coordinates) with meshes',
|
||||
|
@ -348,6 +364,7 @@ class ExportGLTF2_Base:
|
|||
export_settings['gltf_filedirectory'] = os.path.dirname(export_settings['gltf_filepath']) + '/'
|
||||
|
||||
export_settings['gltf_format'] = self.export_format
|
||||
export_settings['gltf_image_format'] = self.export_image_format
|
||||
export_settings['gltf_copyright'] = self.export_copyright
|
||||
export_settings['gltf_texcoords'] = self.export_texcoords
|
||||
export_settings['gltf_normals'] = self.export_normals
|
||||
|
@ -428,6 +445,7 @@ class ExportGLTF2_Base:
|
|||
col.prop(self, 'export_extras')
|
||||
col.prop(self, 'will_save_settings')
|
||||
col.prop(self, 'export_copyright')
|
||||
col.prop(self, 'export_image_format')
|
||||
|
||||
def draw_mesh_settings(self):
|
||||
col = self.layout.box().column()
|
||||
|
|
|
@ -24,8 +24,11 @@ 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_gather_cache import cached
|
||||
|
||||
|
||||
@cached
|
||||
def gather_image(
|
||||
blender_shader_sockets_or_texture_slots: typing.Union[typing.Tuple[bpy.types.NodeSocket],
|
||||
typing.Tuple[bpy.types.Texture]],
|
||||
|
@ -39,13 +42,11 @@ def gather_image(
|
|||
# The blender image has no data
|
||||
return None
|
||||
|
||||
mime_type = __gather_mime_type(uri.filepath if uri is not None else "")
|
||||
|
||||
image = gltf2_io.Image(
|
||||
buffer_view=buffer_view,
|
||||
extensions=__gather_extensions(blender_shader_sockets_or_texture_slots, export_settings),
|
||||
extras=__gather_extras(blender_shader_sockets_or_texture_slots, export_settings),
|
||||
mime_type=mime_type,
|
||||
mime_type=__gather_mime_type(blender_shader_sockets_or_texture_slots, export_settings),
|
||||
name=__gather_name(blender_shader_sockets_or_texture_slots, export_settings),
|
||||
uri=uri
|
||||
)
|
||||
|
@ -64,7 +65,7 @@ def __gather_buffer_view(sockets_or_slots, export_settings):
|
|||
if image is None:
|
||||
return None
|
||||
return gltf2_io_binary_data.BinaryData(
|
||||
data=image.to_image_data(__gather_mime_type()))
|
||||
data=image.encode(__gather_mime_type(sockets_or_slots, export_settings)))
|
||||
return None
|
||||
|
||||
|
||||
|
@ -76,29 +77,43 @@ def __gather_extras(sockets_or_slots, export_settings):
|
|||
return None
|
||||
|
||||
|
||||
def __gather_mime_type(filepath=""):
|
||||
extension_types = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg'}
|
||||
default_extension = extension_types['.png']
|
||||
def __gather_mime_type(sockets_or_slots, export_settings):
|
||||
if export_settings["gltf_image_format"] == "NAME":
|
||||
image_name = __get_texname_from_slot(sockets_or_slots, export_settings)
|
||||
_, extension = os.path.splitext(image_name)
|
||||
if extension in [".jpeg", ".jpg", ".png"]:
|
||||
return {
|
||||
".jpeg": "image/jpeg",
|
||||
".jpg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
}[extension]
|
||||
return "image/png"
|
||||
|
||||
matches = re.findall(r'\.\w+$', filepath)
|
||||
extension = matches[0] if len(matches) > 0 else default_extension
|
||||
return extension_types[extension] if extension.lower() in extension_types.keys() else default_extension
|
||||
elif export_settings["gltf_image_format"] == "JPEG":
|
||||
return "image/jpeg"
|
||||
else:
|
||||
return "image/png"
|
||||
|
||||
|
||||
def __gather_name(sockets_or_slots, export_settings):
|
||||
if __is_socket(sockets_or_slots):
|
||||
node = __get_tex_from_socket(sockets_or_slots[0])
|
||||
if node is not None:
|
||||
return node.shader_node.image.name
|
||||
elif isinstance(sockets_or_slots[0], bpy.types.MaterialTextureSlot):
|
||||
return sockets_or_slots[0].name
|
||||
return None
|
||||
image_name = __get_texname_from_slot(sockets_or_slots, export_settings)
|
||||
|
||||
name, extension = os.path.splitext(image_name)
|
||||
if extension in [".jpeg", ".jpg", ".png"]:
|
||||
return name
|
||||
return image_name
|
||||
|
||||
|
||||
def __gather_uri(sockets_or_slots, export_settings):
|
||||
if export_settings[gltf2_blender_export_keys.FORMAT] == 'GLTF_SEPARATE':
|
||||
# as usual we just store the data in place instead of already resolving the references
|
||||
return __get_image_data(sockets_or_slots, export_settings)
|
||||
mime_type = __gather_mime_type(sockets_or_slots, export_settings)
|
||||
return gltf2_io_image_data.ImageData(
|
||||
data=__get_image_data(sockets_or_slots, export_settings).encode(mime_type=mime_type),
|
||||
mime_type=mime_type,
|
||||
name=__gather_name(sockets_or_slots, export_settings)
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
@ -110,7 +125,7 @@ def __is_slot(sockets_or_slots):
|
|||
return isinstance(sockets_or_slots[0], bpy.types.MaterialTextureSlot)
|
||||
|
||||
|
||||
def __get_image_data(sockets_or_slots, export_settings):
|
||||
def __get_image_data(sockets_or_slots, export_settings) -> gltf2_blender_image.ExportImage:
|
||||
# For shared ressources, 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
|
||||
# ressources.
|
||||
|
@ -128,8 +143,8 @@ def __get_image_data(sockets_or_slots, export_settings):
|
|||
return channels
|
||||
|
||||
if __is_socket(sockets_or_slots):
|
||||
results = [__get_tex_from_socket(socket) for socket in sockets_or_slots]
|
||||
image = None
|
||||
results = [__get_tex_from_socket(socket, export_settings) for socket in sockets_or_slots]
|
||||
composed_image = None
|
||||
for result, socket in zip(results, sockets_or_slots):
|
||||
if result.shader_node.image.channels == 0:
|
||||
gltf2_io_debug.print_console("WARNING",
|
||||
|
@ -138,9 +153,7 @@ def __get_image_data(sockets_or_slots, export_settings):
|
|||
continue
|
||||
|
||||
# rudimentarily try follow the node tree to find the correct image data.
|
||||
source_channel = None
|
||||
target_channel = None
|
||||
source_channels_length = None
|
||||
source_channel = 0
|
||||
for elem in result.path:
|
||||
if isinstance(elem.from_node, bpy.types.ShaderNodeSeparateRGB):
|
||||
source_channel = {
|
||||
|
@ -149,68 +162,30 @@ def __get_image_data(sockets_or_slots, export_settings):
|
|||
'B': 2
|
||||
}[elem.from_socket.name]
|
||||
|
||||
if source_channel is not None:
|
||||
pixels = [split_pixels_by_channels(result.shader_node.image, export_settings)[source_channel]]
|
||||
target_channel = source_channel
|
||||
source_channel = 0
|
||||
source_channels_length = 1
|
||||
else:
|
||||
pixels = split_pixels_by_channels(result.shader_node.image, export_settings)
|
||||
target_channel = 0
|
||||
source_channel = 0
|
||||
source_channels_length = len(pixels)
|
||||
image = gltf2_blender_image.ExportImage.from_blender_image(result.shader_node.image)
|
||||
|
||||
if composed_image is None:
|
||||
composed_image = gltf2_blender_image.ExportImage.white_image(image.width, image.height)
|
||||
|
||||
# Change target channel for metallic and roughness.
|
||||
if elem.to_socket.name == 'Metallic':
|
||||
target_channel = 2
|
||||
source_channels_length = 1
|
||||
composed_image[2] = image[source_channel]
|
||||
elif elem.to_socket.name == 'Roughness':
|
||||
target_channel = 1
|
||||
source_channels_length = 1
|
||||
|
||||
file_name = os.path.splitext(result.shader_node.image.name)[0]
|
||||
if result.shader_node.image.packed_file is None:
|
||||
file_path = result.shader_node.image.filepath
|
||||
composed_image[1] = image[source_channel]
|
||||
else:
|
||||
# empty path for packed textures, because they are converted to png anyway
|
||||
file_path = ""
|
||||
composed_image.update(image)
|
||||
|
||||
image_data = gltf2_io_image_data.ImageData(
|
||||
file_name,
|
||||
file_path,
|
||||
result.shader_node.image.size[0],
|
||||
result.shader_node.image.size[1],
|
||||
source_channel,
|
||||
target_channel,
|
||||
source_channels_length,
|
||||
pixels)
|
||||
return composed_image
|
||||
|
||||
if image is None:
|
||||
image = image_data
|
||||
else:
|
||||
image.add_to_image(target_channel, image_data)
|
||||
|
||||
return image
|
||||
elif __is_slot(sockets_or_slots):
|
||||
texture = __get_tex_from_slot(sockets_or_slots[0])
|
||||
pixels = split_pixels_by_channels(texture.image, export_settings)
|
||||
|
||||
image_data = gltf2_io_image_data.ImageData(
|
||||
texture.name,
|
||||
texture.image.filepath,
|
||||
texture.image.size[0],
|
||||
texture.image.size[1],
|
||||
0,
|
||||
0,
|
||||
len(pixels),
|
||||
pixels)
|
||||
return image_data
|
||||
image = gltf2_blender_image.ExportImage.from_blender_image(texture.image)
|
||||
return image
|
||||
else:
|
||||
# Texture slots
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def __get_tex_from_socket(blender_shader_socket: bpy.types.NodeSocket):
|
||||
@cached
|
||||
def __get_tex_from_socket(blender_shader_socket: bpy.types.NodeSocket, export_settings):
|
||||
result = gltf2_blender_search_node_tree.from_socket(
|
||||
blender_shader_socket,
|
||||
gltf2_blender_search_node_tree.FilterByType(bpy.types.ShaderNodeTexImage))
|
||||
|
@ -222,3 +197,15 @@ def __get_tex_from_socket(blender_shader_socket: bpy.types.NodeSocket):
|
|||
def __get_tex_from_slot(blender_texture_slot):
|
||||
return blender_texture_slot.texture
|
||||
|
||||
|
||||
@cached
|
||||
def __get_texname_from_slot(sockets_or_slots, export_settings):
|
||||
if __is_socket(sockets_or_slots):
|
||||
node = __get_tex_from_socket(sockets_or_slots[0], export_settings)
|
||||
if node is None:
|
||||
return None
|
||||
return node.shader_node.image.name
|
||||
|
||||
elif isinstance(sockets_or_slots[0], bpy.types.MaterialTextureSlot):
|
||||
return sockets_or_slots[0].name
|
||||
|
||||
|
|
|
@ -11,18 +11,15 @@
|
|||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Optional, List, Dict
|
||||
import re
|
||||
from typing import List
|
||||
|
||||
from io_scene_gltf2.io.com import gltf2_io
|
||||
from io_scene_gltf2.io.com import gltf2_io_extensions
|
||||
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.exp import gltf2_io_buffer
|
||||
from io_scene_gltf2.io.exp import gltf2_io_image_data
|
||||
|
||||
import bpy
|
||||
import os
|
||||
from shutil import copyfile
|
||||
|
||||
class GlTF2Exporter:
|
||||
"""
|
||||
|
@ -65,7 +62,7 @@ class GlTF2Exporter:
|
|||
)
|
||||
|
||||
self.__buffer = gltf2_io_buffer.Buffer()
|
||||
self.__images = []
|
||||
self.__images = {}
|
||||
|
||||
# mapping of all glTFChildOfRootProperty types to their corresponding root level arrays
|
||||
self.__childOfRootPropertyTypeLookup = {
|
||||
|
@ -152,18 +149,10 @@ class GlTF2Exporter:
|
|||
:param output_path:
|
||||
:return:
|
||||
"""
|
||||
for image in self.__images:
|
||||
dst_path = output_path + image.name + image.get_extension()
|
||||
src_path = bpy.path.abspath(image.filepath)
|
||||
if os.path.isfile(src_path):
|
||||
# Source file exists.
|
||||
if os.path.abspath(dst_path) != os.path.abspath(src_path):
|
||||
# Only copy, if source and destination are not the same.
|
||||
copyfile(src_path, dst_path)
|
||||
else:
|
||||
# Source file does not exist e.g. it is packed or has been generated.
|
||||
with open(dst_path, 'wb') as f:
|
||||
f.write(image.to_png_data())
|
||||
for name, image in self.__images.items():
|
||||
dst_path = output_path + "/" + name + image.file_extension
|
||||
with open(dst_path, 'wb') as f:
|
||||
f.write(image.data)
|
||||
|
||||
def add_scene(self, scene: gltf2_io.Scene, active: bool = True):
|
||||
"""
|
||||
|
@ -220,11 +209,23 @@ class GlTF2Exporter:
|
|||
return index
|
||||
|
||||
def __add_image(self, image: gltf2_io_image_data.ImageData):
|
||||
self.__images.append(image)
|
||||
name = image.adjusted_name()
|
||||
count = 1
|
||||
regex = re.compile(r"\d+$")
|
||||
regex_found = re.findall(regex, name)
|
||||
while name in self.__images.keys():
|
||||
if regex_found:
|
||||
name = re.sub(regex, str(count), name)
|
||||
else:
|
||||
name += " " + str(count)
|
||||
|
||||
count += 1
|
||||
# TODO: we need to know the image url at this point already --> maybe add all options to the constructor of the
|
||||
# exporter
|
||||
# TODO: allow embedding of images (base64)
|
||||
return image.name + image.get_extension()
|
||||
|
||||
self.__images[name] = image
|
||||
return name + image.file_extension
|
||||
|
||||
@classmethod
|
||||
def __get_key_path(cls, d: dict, keypath: List[str], default=[]):
|
||||
|
@ -288,7 +289,8 @@ class GlTF2Exporter:
|
|||
|
||||
# image data needs to be saved to file
|
||||
if isinstance(node, gltf2_io_image_data.ImageData):
|
||||
return self.__add_image(node)
|
||||
image = self.__add_image(node)
|
||||
return image
|
||||
|
||||
# extensions
|
||||
if isinstance(node, gltf2_io_extensions.Extension):
|
||||
|
|
|
@ -0,0 +1,123 @@
|
|||
# Copyright 2018 The glTF-Blender-IO authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import bpy
|
||||
import typing
|
||||
import numpy as np
|
||||
import tempfile
|
||||
|
||||
|
||||
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
|
||||
|
||||
def __init__(self, img: typing.Union[np.ndarray, typing.List[np.ndarray]], max_channels: int = 4):
|
||||
if isinstance(img, list):
|
||||
np.stack(img, axis=2)
|
||||
|
||||
if len(img.shape) == 2:
|
||||
# images must always have a channels dimension
|
||||
img = np.expand_dims(img, axis=2)
|
||||
|
||||
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))
|
||||
|
||||
self._img = img
|
||||
self._max_channels = max_channels
|
||||
|
||||
@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))
|
||||
return ExportImage(img=img)
|
||||
|
||||
@classmethod
|
||||
def white_image(cls, width, height, num_channels: int = 4):
|
||||
img = np.ones((width, height, num_channels))
|
||||
return ExportImage(img=img)
|
||||
|
||||
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)
|
||||
|
||||
@property
|
||||
def img(self) -> np.ndarray:
|
||||
return self._img
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._img.shape
|
||||
|
||||
@property
|
||||
def width(self):
|
||||
return self.shape[0]
|
||||
|
||||
@property
|
||||
def height(self):
|
||||
return self.shape[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 update(self, other):
|
||||
self[:other.channels] = other[:other.channels]
|
||||
|
||||
def __add__(self, other):
|
||||
self.append(other)
|
||||
|
||||
def encode(self, mime_type: typing.Optional[str]) -> bytes:
|
||||
image = bpy.data.images.new("TmpImage", width=self.width, height=self.height)
|
||||
pixels = self._img.flatten().tolist()
|
||||
image.pixels = pixels
|
||||
|
||||
file_format = {
|
||||
"image/jpeg": "JPEG",
|
||||
"image/png": "PNG"
|
||||
}.get(mime_type, "PNG")
|
||||
|
||||
# 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()
|
||||
|
||||
with open(tmpfilename, "rb") as f:
|
||||
encoded_image = f.read()
|
||||
|
||||
bpy.data.images.remove(image, do_unlink=True)
|
||||
|
||||
return encoded_image
|
||||
|
|
@ -11,127 +11,46 @@
|
|||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import typing
|
||||
import struct
|
||||
import re
|
||||
import zlib
|
||||
import numpy as np
|
||||
|
||||
|
||||
class ImageData:
|
||||
"""Contains channels of an image with raw pixel data."""
|
||||
# TODO: refactor to only operate on numpy arrays
|
||||
"""Contains encoded 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
|
||||
|
||||
def __init__(self, name: str, filepath: str, width: int, height: int, source: int, target: int, source_length: int, channels: typing.Optional[typing.List[np.ndarray]] = []):
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError("Image data can not have zero width or height")
|
||||
if source + source_length > 4:
|
||||
raise ValueError("Source image data can not have more than 4 channels")
|
||||
if target + source_length > 4:
|
||||
raise ValueError("Target image data can not have more than 4 channels")
|
||||
self.channels = [None, None, None, None]
|
||||
for index in range(source, source + source_length):
|
||||
self.channels[target + index - source] = channels[index]
|
||||
self.name = name
|
||||
self.filepath = filepath
|
||||
self.width = width
|
||||
self.height = height
|
||||
def __init__(self, data: bytes, mime_type: str, name: str):
|
||||
self._data = data
|
||||
self._mime_type = mime_type
|
||||
self._name = name
|
||||
|
||||
def add_to_image(self, target: int, image_data):
|
||||
if self.width != image_data.width or self.height != image_data.height:
|
||||
raise ValueError("Image dimensions do not match")
|
||||
if target < 0 or target > 3:
|
||||
raise ValueError("Can't insert image: channels out of bounds")
|
||||
if len(image_data.channels) != 4:
|
||||
raise ValueError("Can't insert image: incomplete image")
|
||||
def __eq__(self, other):
|
||||
return self._data == other.data
|
||||
|
||||
if self.name != image_data.name:
|
||||
self.name += image_data.name
|
||||
self.filepath = ""
|
||||
def __hash__(self):
|
||||
return hash(self._data)
|
||||
|
||||
# Replace channel.
|
||||
self.channels[target] = image_data.channels[target]
|
||||
def adjusted_name(self):
|
||||
regex_dot = re.compile("\.")
|
||||
adjusted_name = re.sub(regex_dot, "_", self.name)
|
||||
new_name = "".join([char for char in adjusted_name if char not in "!#$&'()*+,/:;<>?@[\]^`{|}~"])
|
||||
return new_name
|
||||
|
||||
@property
|
||||
def r(self):
|
||||
if len(self.channels) <= 0:
|
||||
return None
|
||||
return self.channels[0]
|
||||
def data(self):
|
||||
return self._data
|
||||
|
||||
@property
|
||||
def g(self):
|
||||
if len(self.channels) <= 1:
|
||||
return None
|
||||
return self.channels[1]
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def b(self):
|
||||
if len(self.channels) <= 2:
|
||||
return None
|
||||
return self.channels[2]
|
||||
def file_extension(self):
|
||||
if self._mime_type == "image/jpeg":
|
||||
return ".jpg"
|
||||
return ".png"
|
||||
|
||||
@property
|
||||
def a(self):
|
||||
if len(self.channels) <= 3:
|
||||
return None
|
||||
return self.channels[3]
|
||||
|
||||
def get_extension(self):
|
||||
allowed_extensions = ['.png', '.jpg', '.jpeg']
|
||||
fallback_extension = allowed_extensions[0]
|
||||
|
||||
matches = re.findall(r'\.\w+$', self.filepath)
|
||||
extension = matches[0] if len(matches) > 0 else fallback_extension
|
||||
return extension if extension.lower() in allowed_extensions else fallback_extension
|
||||
|
||||
def to_image_data(self, mime_type: str) -> bytes:
|
||||
if mime_type == 'image/png':
|
||||
return self.to_png_data()
|
||||
raise ValueError("Unsupported image file type {}".format(mime_type))
|
||||
|
||||
def to_png_data(self) -> bytes:
|
||||
channels = self.channels
|
||||
|
||||
# if there is no data, create a single pixel image
|
||||
if not channels:
|
||||
channels = np.ones((1, 1))
|
||||
# fill all channels of the png
|
||||
for _ in range(4 - len(channels)):
|
||||
channels.append(np.ones_like(channels[0]))
|
||||
else:
|
||||
template_index = None
|
||||
for index in range(0, 4):
|
||||
if channels[index] is not None:
|
||||
template_index = index
|
||||
break
|
||||
for index in range(0, 4):
|
||||
if channels[index] is None:
|
||||
channels[index] = np.ones_like(channels[template_index])
|
||||
|
||||
image = np.concatenate(channels, axis=1)
|
||||
image = image.flatten()
|
||||
image = (image * 255.0).astype(np.uint8)
|
||||
buf = image.tobytes()
|
||||
|
||||
#
|
||||
# Taken from 'blender-thumbnailer.py' in Blender.
|
||||
#
|
||||
|
||||
# reverse the vertical line order and add null bytes at the start
|
||||
width_byte_4 = self.width * 4
|
||||
raw_data = b"".join(
|
||||
b'\x00' + buf[span:span + width_byte_4] for span in range(
|
||||
(self.height - 1) * self.width * 4, -1, - width_byte_4))
|
||||
|
||||
def png_pack(png_tag, data):
|
||||
chunk_head = png_tag + data
|
||||
return struct.pack("!I", len(data)) + chunk_head + struct.pack("!I", 0xFFFFFFFF & zlib.crc32(chunk_head))
|
||||
|
||||
return b"".join([
|
||||
b'\x89PNG\r\n\x1a\n',
|
||||
png_pack(b'IHDR', struct.pack("!2I5B", self.width, self.height, 8, 6, 0, 0, 0)),
|
||||
png_pack(b'IDAT', zlib.compress(raw_data)),
|
||||
png_pack(b'IEND', b'')])
|
||||
def byte_length(self):
|
||||
return len(self._data)
|
||||
|
||||
|
|
Loading…
Reference in New Issue