io_scene_gltf2/blender/exp/gltf2_blender_image.py

   1 # SPDX-License-Identifier: Apache-2.0
   2 # Copyright 2018-2021 The glTF-Blender-IO authors.
   3
   4 import bpy
   5 import os
   6 from typing import Optional, Tuple
   7 import numpy as np
   8 import tempfile
   9 import enum
  10
  11
  12 class Channel(enum.IntEnum):
  13     R = 0
  14     G = 1
  15     B = 2
  16     A = 3
  17
  18 # These describe how an ExportImage's channels should be filled.
  19
  20 class FillImage:
  21     """Fills a channel with the channel src_chan from a Blender image."""
  22     def __init__(self, image: bpy.types.Image, src_chan: Channel):
  23         self.image = image
  24         self.src_chan = src_chan
  25
  26 class FillWhite:
  27     """Fills a channel with all ones (1.0)."""
  28     pass
  29
  30
  31 class ExportImage:
  32     """Custom image class.
  33
  34     An image is represented by giving a description of how to fill its red,
  35     green, blue, and alpha channels. For example:
  36
  37         self.fills = {
  38             Channel.R: FillImage(image=bpy.data.images['Im1'], src_chan=Channel.B),
  39             Channel.G: FillWhite(),
  40         }
  41
  42     This says that the ExportImage's R channel should be filled with the B
  43     channel of the Blender image 'Im1', and the ExportImage's G channel
  44     should be filled with all 1.0s. Undefined channels mean we don't care
  45     what values that channel has.
  46
  47     This is flexible enough to handle the case where eg. the user used the R
  48     channel of one image as the metallic value and the G channel of another
  49     image as the roughness, and we need to synthesize an ExportImage that
  50     packs those into the B and G channels for glTF.
  51
  52     Storing this description (instead of raw pixels) lets us make more
  53     intelligent decisions about how to encode the image.
  54     """
  55
  56     def __init__(self, original=None):
  57         self.fills = {}
  58
  59         # In case of keeping original texture images
  60         self.original = original
  61
  62     @staticmethod
  63     def from_blender_image(image: bpy.types.Image):
  64         export_image = ExportImage()
  65         for chan in range(image.channels):
  66             export_image.fill_image(image, dst_chan=chan, src_chan=chan)
  67         return export_image
  68
  69     @staticmethod
  70     def from_original(image: bpy.types.Image):
  71         return ExportImage(image)
  72
  73     def fill_image(self, image: bpy.types.Image, dst_chan: Channel, src_chan: Channel):
  74         self.fills[dst_chan] = FillImage(image, src_chan)
  75
  76     def fill_white(self, dst_chan: Channel):
  77         self.fills[dst_chan] = FillWhite()
  78
  79     def is_filled(self, chan: Channel) -> bool:
  80         return chan in self.fills
  81
  82     def empty(self) -> bool:
  83         if self.original is None:
  84             return not self.fills
  85         else:
  86             return False
  87
  88     def blender_image(self) -> Optional[bpy.types.Image]:
  89         """If there's an existing Blender image we can use,
  90         returns it. Otherwise (if channels need packing),
  91         returns None.
  92         """
  93         if self.__on_happy_path():
  94             for fill in self.fills.values():
  95                 return fill.image
  96         return None
  97
  98     def __on_happy_path(self) -> bool:
  99         # All src_chans match their dst_chan and come from the same image
 100         return (
 101             all(isinstance(fill, FillImage) for fill in self.fills.values()) and
 102             all(dst_chan == fill.src_chan for dst_chan, fill in self.fills.items()) and
 103             len(set(fill.image.name for fill in self.fills.values())) == 1
 104         )
 105
 106     def encode(self, mime_type: Optional[str]) -> bytes:
 107         self.file_format = {
 108             "image/jpeg": "JPEG",
 109             "image/png": "PNG"
 110         }.get(mime_type, "PNG")
 111
 112         # Happy path = we can just use an existing Blender image
 113         if self.__on_happy_path():
 114             return self.__encode_happy()
 115
 116         # Unhappy path = we need to create the image self.fills describes.
 117         return self.__encode_unhappy()
 118
 119     def __encode_happy(self) -> bytes:
 120         return self.__encode_from_image(self.blender_image())
 121
 122     def __encode_unhappy(self) -> bytes:
 123         # We need to assemble the image out of channels.
 124         # Do it with numpy and image.pixels.
 125
 126         # Find all Blender images used
 127         images = []
 128         for fill in self.fills.values():
 129             if isinstance(fill, FillImage):
 130                 if fill.image not in images:
 131                     images.append(fill.image)
 132
 133         if not images:
 134             # No ImageFills; use a 1x1 white pixel
 135             pixels = np.array([1.0, 1.0, 1.0, 1.0], np.float32)
 136             return self.__encode_from_numpy_array(pixels, (1, 1))
 137
 138         width = max(image.size[0] for image in images)
 139         height = max(image.size[1] for image in images)
 140
 141         out_buf = np.ones(width * height * 4, np.float32)
 142         tmp_buf = np.empty(width * height * 4, np.float32)
 143
 144         for image in images:
 145             if image.size[0] == width and image.size[1] == height:
 146                 image.pixels.foreach_get(tmp_buf)
 147             else:
 148                 # Image is the wrong size; make a temp copy and scale it.
 149                 with TmpImageGuard() as guard:
 150                     _make_temp_image_copy(guard, src_image=image)
 151                     tmp_image = guard.image
 152                     tmp_image.scale(width, height)
 153                     tmp_image.pixels.foreach_get(tmp_buf)
 154
 155             # Copy any channels for this image to the output
 156             for dst_chan, fill in self.fills.items():
 157                 if isinstance(fill, FillImage) and fill.image == image:
 158                     out_buf[int(dst_chan)::4] = tmp_buf[int(fill.src_chan)::4]
 159
 160         tmp_buf = None  # GC this
 161
 162         return self.__encode_from_numpy_array(out_buf, (width, height))
 163
 164     def __encode_from_numpy_array(self, pixels: np.ndarray, dim: Tuple[int, int]) -> bytes:
 165         with TmpImageGuard() as guard:
 166             guard.image = bpy.data.images.new(
 167                 "##gltf-export:tmp-image##",
 168                 width=dim[0],
 169                 height=dim[1],
 170                 alpha=Channel.A in self.fills,
 171             )
 172             tmp_image = guard.image
 173
 174             tmp_image.pixels.foreach_set(pixels)
 175
 176             return _encode_temp_image(tmp_image, self.file_format)
 177
 178     def __encode_from_image(self, image: bpy.types.Image) -> bytes:
 179         # See if there is an existing file we can use.
 180         data = None
 181         if image.source == 'FILE' and not image.is_dirty:
 182             if image.packed_file is not None:
 183                 data = image.packed_file.data
 184             else:
 185                 src_path = bpy.path.abspath(image.filepath_raw)
 186                 if os.path.isfile(src_path):
 187                     with open(src_path, 'rb') as f:
 188                         data = f.read()
 189         # Check magic number is right
 190         if data:
 191             if self.file_format == 'PNG':
 192                 if data.startswith(b'\x89PNG'):
 193                     return data
 194             elif self.file_format == 'JPEG':
 195                 if data.startswith(b'\xff\xd8\xff'):
 196                     return data
 197
 198         # Copy to a temp image and save.
 199         with TmpImageGuard() as guard:
 200             _make_temp_image_copy(guard, src_image=image)
 201             tmp_image = guard.image
 202             return _encode_temp_image(tmp_image, self.file_format)
 203
 204
 205 def _encode_temp_image(tmp_image: bpy.types.Image, file_format: str) -> bytes:
 206     with tempfile.TemporaryDirectory() as tmpdirname:
 207         tmpfilename = tmpdirname + '/img'
 208         tmp_image.filepath_raw = tmpfilename
 209
 210         tmp_image.file_format = file_format
 211
 212         tmp_image.save()
 213
 214         with open(tmpfilename, "rb") as f:
 215             return f.read()
 216
 217
 218 class TmpImageGuard:
 219     """Guard to automatically clean up temp images (use it with `with`)."""
 220     def __init__(self):
 221         self.image = None
 222
 223     def __enter__(self):
 224         return self
 225
 226     def __exit__(self, exc_type, exc_value, traceback):
 227         if self.image is not None:
 228             bpy.data.images.remove(self.image, do_unlink=True)
 229
 230
 231 def _make_temp_image_copy(guard: TmpImageGuard, src_image: bpy.types.Image):
 232     """Makes a temporary copy of src_image. Will be cleaned up with guard."""
 233     guard.image = src_image.copy()
 234     tmp_image = guard.image
 235
 236     tmp_image.update()
 237
 238     if src_image.is_dirty:
 239         # Unsaved changes aren't copied by .copy(), so do them ourselves
 240         tmp_buf = np.empty(src_image.size[0] * src_image.size[1] * 4, np.float32)
 241         src_image.pixels.foreach_get(tmp_buf)
 242         tmp_image.pixels.foreach_set(tmp_buf)