blenderkit/image_utils.py

   1 import bpy
   2 import os
   3 import time
   4
   5 def get_orig_render_settings():
   6     rs = bpy.context.scene.render
   7     ims = rs.image_settings
   8
   9     vs = bpy.context.scene.view_settings
  10
  11     orig_settings = {
  12         'file_format': ims.file_format,
  13         'quality': ims.quality,
  14         'color_mode': ims.color_mode,
  15         'compression': ims.compression,
  16         'exr_codec': ims.exr_codec,
  17         'view_transform': vs.view_transform
  18     }
  19     return orig_settings
  20
  21
  22 def set_orig_render_settings(orig_settings):
  23     rs = bpy.context.scene.render
  24     ims = rs.image_settings
  25     vs = bpy.context.scene.view_settings
  26
  27     ims.file_format = orig_settings['file_format']
  28     ims.quality = orig_settings['quality']
  29     ims.color_mode = orig_settings['color_mode']
  30     ims.compression = orig_settings['compression']
  31     ims.exr_codec = orig_settings['exr_codec']
  32
  33     vs.view_transform = orig_settings['view_transform']
  34
  35
  36 def img_save_as(img, filepath='//', file_format='JPEG', quality=90, color_mode='RGB', compression=15, view_transform = 'Raw', exr_codec = 'DWAA'):
  37     '''Uses Blender 'save render' to save images - BLender isn't really able so save images with other methods correctly.'''
  38
  39     ors = get_orig_render_settings()
  40
  41     rs = bpy.context.scene.render
  42     vs = bpy.context.scene.view_settings
  43
  44     ims = rs.image_settings
  45     ims.file_format = file_format
  46     ims.quality = quality
  47     ims.color_mode = color_mode
  48     ims.compression = compression
  49     ims.exr_codec = exr_codec
  50     vs.view_transform = view_transform
  51
  52
  53     img.save_render(filepath=bpy.path.abspath(filepath), scene=bpy.context.scene)
  54
  55     set_orig_render_settings(ors)
  56
  57 def set_colorspace(img, colorspace):
  58     '''sets image colorspace, but does so in a try statement, because some people might actually replace the default
  59     colorspace settings, and it literally can't be guessed what these people use, even if it will mostly be the filmic addon.
  60     '''
  61     try:
  62         if colorspace == 'Non-Color':
  63             img.colorspace_settings.is_data = True
  64         else:
  65             img.colorspace_settings.name = colorspace
  66     except:
  67         print(f'Colorspace {colorspace} not found.')
  68
  69 def generate_hdr_thumbnail():
  70     import numpy
  71     scene = bpy.context.scene
  72     ui_props = scene.blenderkitUI
  73     hdr_image = ui_props.hdr_upload_image#bpy.data.images.get(ui_props.hdr_upload_image)
  74
  75     base, ext = os.path.splitext(hdr_image.filepath)
  76     thumb_path = base + '.jpg'
  77     thumb_name = os.path.basename(thumb_path)
  78
  79     max_thumbnail_size = 2048
  80     size = hdr_image.size
  81     ratio = size[0] / size[1]
  82
  83     imageWidth = size[0]
  84     imageHeight = size[1]
  85     thumbnailWidth = min(size[0], max_thumbnail_size)
  86     thumbnailHeight = min(size[1], int(max_thumbnail_size / ratio))
  87
  88     tempBuffer = numpy.empty(imageWidth * imageHeight * 4, dtype=numpy.float32)
  89     inew = bpy.data.images.new(thumb_name, imageWidth, imageHeight, alpha=False, float_buffer=False)
  90
  91     hdr_image.pixels.foreach_get(tempBuffer)
  92
  93     inew.filepath = thumb_path
  94     set_colorspace(inew, 'Linear')
  95     inew.pixels.foreach_set(tempBuffer)
  96
  97     bpy.context.view_layer.update()
  98     if thumbnailWidth < imageWidth:
  99         inew.scale(thumbnailWidth, thumbnailHeight)
 100
 101     img_save_as(inew, filepath=inew.filepath)
 102
 103
 104 def find_color_mode(image):
 105     if not isinstance(image, bpy.types.Image):
 106         raise(TypeError)
 107     else:
 108         depth_mapping = {
 109             8: 'BW',
 110             24: 'RGB',
 111             32: 'RGBA',#can also be bw.. but image.channels doesn't work.
 112             96: 'RGB',
 113             128: 'RGBA',
 114         }
 115         return depth_mapping.get(image.depth,'RGB')
 116
 117 def find_image_depth(image):
 118     if not isinstance(image, bpy.types.Image):
 119         raise(TypeError)
 120     else:
 121         depth_mapping = {
 122             8: '8',
 123             24: '8',
 124             32: '8',#can also be bw.. but image.channels doesn't work.
 125             96: '16',
 126             128: '16',
 127         }
 128         return depth_mapping.get(image.depth,'8')
 129
 130 def can_erase_alpha(na):
 131     alpha = na[3::4]
 132     alpha_sum = alpha.sum()
 133     if alpha_sum == alpha.size:
 134         print('image can have alpha erased')
 135     # print(alpha_sum, alpha.size)
 136     return alpha_sum == alpha.size
 137
 138
 139 def is_image_black(na):
 140     r = na[::4]
 141     g = na[1::4]
 142     b = na[2::4]
 143
 144     rgbsum = r.sum() + g.sum() + b.sum()
 145
 146     # print('rgb sum', rgbsum, r.sum(), g.sum(), b.sum())
 147     if rgbsum == 0:
 148         print('image can have alpha channel dropped')
 149     return rgbsum == 0
 150
 151 def is_image_bw(na):
 152     r = na[::4]
 153     g = na[1::4]
 154     b = na[2::4]
 155
 156     rg_equal = r == g
 157     gb_equal = g == b
 158     rgbequal = rg_equal.all() and gb_equal.all()
 159     if rgbequal:
 160         print('image is black and white, can have channels reduced')
 161
 162     return rgbequal
 163
 164
 165 def numpytoimage(a, iname, width=0, height=0, channels=3):
 166     t = time.time()
 167     foundimage = False
 168
 169     for image in bpy.data.images:
 170
 171         if image.name[:len(iname)] == iname and image.size[0] == a.shape[0] and image.size[1] == a.shape[1]:
 172             i = image
 173             foundimage = True
 174     if not foundimage:
 175         if channels == 4:
 176             bpy.ops.image.new(name=iname, width=width, height=height, color=(0, 0, 0, 1), alpha=True,
 177                               generated_type='BLANK', float=True)
 178         if channels == 3:
 179             bpy.ops.image.new(name=iname, width=width, height=height, color=(0, 0, 0), alpha=False,
 180                               generated_type='BLANK', float=True)
 181
 182     i = None
 183
 184     for image in bpy.data.images:
 185         # print(image.name[:len(iname)],iname, image.size[0],a.shape[0],image.size[1],a.shape[1])
 186         if image.name[:len(iname)] == iname and image.size[0] == width and image.size[1] == height:
 187             i = image
 188     if i is None:
 189         i = bpy.data.images.new(iname, width, height, alpha=False, float_buffer=False, stereo3d=False, is_data=False, tiled=False)
 190
 191     # dropping this re-shaping code -  just doing flat array for speed and simplicity
 192     #    d = a.shape[0] * a.shape[1]
 193     #    a = a.swapaxes(0, 1)
 194     #    a = a.reshape(d)
 195     #    a = a.repeat(channels)
 196     #    a[3::4] = 1
 197     i.pixels.foreach_set(a)  # this gives big speedup!
 198     print('\ntime ' + str(time.time() - t))
 199     return i
 200
 201
 202 def imagetonumpy_flat(i):
 203     t = time.time()
 204
 205     import numpy
 206
 207     width = i.size[0]
 208     height = i.size[1]
 209     # print(i.channels)
 210
 211     size = width * height * i.channels
 212     na = numpy.empty(size, numpy.float32)
 213     i.pixels.foreach_get(na)
 214
 215     # dropping this re-shaping code -  just doing flat array for speed and simplicity
 216     #    na = na[::4]
 217     #    na = na.reshape(height, width, i.channels)
 218     #    na = na.swapaxnes(0, 1)
 219
 220     # print('\ntime of image to numpy ' + str(time.time() - t))
 221     return na
 222
 223 def imagetonumpy(i):
 224     t = time.time()
 225
 226     import numpy as np
 227
 228     width = i.size[0]
 229     height = i.size[1]
 230     # print(i.channels)
 231
 232     size = width * height * i.channels
 233     na = np.empty(size, np.float32)
 234     i.pixels.foreach_get(na)
 235
 236     # dropping this re-shaping code -  just doing flat array for speed and simplicity
 237     # na = na[::4]
 238     na = na.reshape(height, width, i.channels)
 239     na = na.swapaxes(0, 1)
 240
 241     # print('\ntime of image to numpy ' + str(time.time() - t))
 242     return na
 243
 244
 245 def downscale(i):
 246     minsize = 128
 247
 248     sx, sy = i.size[:]
 249     sx = round(sx / 2)
 250     sy = round(sy / 2)
 251     if sx > minsize and sy > minsize:
 252         i.scale(sx, sy)
 253
 254
 255 def get_rgb_mean(i):
 256     '''checks if normal map values are ok.'''
 257     import numpy
 258
 259     na = imagetonumpy_flat(i)
 260
 261     r = na[::4]
 262     g = na[1::4]
 263     b = na[2::4]
 264
 265     rmean = r.mean()
 266     gmean = g.mean()
 267     bmean = b.mean()
 268
 269     rmedian = numpy.median(r)
 270     gmedian = numpy.median(g)
 271     bmedian = numpy.median(b)
 272
 273     #    return(rmedian,gmedian, bmedian)
 274     return (rmean, gmean, bmean)
 275
 276 def check_nmap_mean_ok(i):
 277     '''checks if normal map values are in standard range.'''
 278
 279     rmean,gmean,bmean = get_rgb_mean(i)
 280
 281     #we could/should also check blue, but some ogl substance exports have 0-1, while 90% nmaps have 0.5 - 1.
 282     nmap_ok = 0.45< rmean < 0.55 and .45 < gmean < .55
 283
 284     return nmap_ok
 285
 286
 287 def check_nmap_ogl_vs_dx(i, mask = None, generated_test_images = False):
 288     '''
 289     checks if normal map is directX or OpenGL.
 290     Returns - String value - DirectX and OpenGL
 291     '''
 292     import numpy
 293     width = i.size[0]
 294     height = i.size[1]
 295
 296
 297
 298     rmean, gmean, bmean = get_rgb_mean(i)
 299
 300     na = imagetonumpy(i)
 301
 302     if mask:
 303         mask = imagetonumpy(mask)
 304
 305     red_x_comparison = numpy.zeros((width, height), numpy.float32)
 306     green_y_comparison = numpy.zeros((width, height), numpy.float32)
 307
 308     if generated_test_images:
 309         red_x_comparison_img = numpy.empty((width, height, 4), numpy.float32) #images for debugging purposes
 310         green_y_comparison_img = numpy.empty((width, height, 4), numpy.float32)#images for debugging purposes
 311
 312     ogl = numpy.zeros((width, height), numpy.float32)
 313     dx = numpy.zeros((width, height), numpy.float32)
 314
 315     if generated_test_images:
 316         ogl_img = numpy.empty((width, height, 4), numpy.float32)  # images for debugging purposes
 317         dx_img = numpy.empty((width, height, 4), numpy.float32)  # images for debugging purposes
 318
 319     for y in range(0, height):
 320         for x in range(0, width):
 321             #try to mask with UV mask image
 322             if mask is None or mask[x,y,3]>0:
 323
 324                 last_height_x = ogl[max(x - 1, 0), min(y, height - 1)]
 325                 last_height_y = ogl[max(x,0), min(y - 1,height-1)]
 326
 327                 diff_x = ((na[x, y, 0] - rmean) / ((na[x, y, 2] - 0.5)))
 328                 diff_y = ((na[x, y, 1] - gmean) / ((na[x, y, 2] - 0.5)))
 329                 calc_height = (last_height_x + last_height_y) \
 330                          -  diff_x - diff_y
 331                 calc_height = calc_height /2
 332                 ogl[x, y] = calc_height
 333                 if generated_test_images:
 334                     rgb = calc_height *.1 +.5
 335                     ogl_img[x,y] = [rgb,rgb,rgb,1]
 336
 337                 # green channel
 338                 last_height_x = dx[max(x - 1, 0), min(y, height - 1)]
 339                 last_height_y = dx[max(x, 0), min(y - 1, height - 1)]
 340
 341                 diff_x = ((na[x, y, 0] - rmean) / ((na[x, y, 2] - 0.5)))
 342                 diff_y = ((na[x, y, 1] - gmean) / ((na[x, y, 2] - 0.5)))
 343                 calc_height = (last_height_x + last_height_y) \
 344                               - diff_x + diff_y
 345                 calc_height = calc_height / 2
 346                 dx[x, y] = calc_height
 347                 if generated_test_images:
 348                     rgb = calc_height * .1 + .5
 349                     dx_img[x, y] = [rgb, rgb, rgb, 1]
 350
 351
 352     ogl_std = ogl.std()
 353     dx_std = dx.std()
 354
 355     # print(mean_ogl, mean_dx)
 356     # print(max_ogl, max_dx)
 357     print(ogl_std, dx_std)
 358     print(i.name)
 359     #    if abs(mean_ogl) > abs(mean_dx):
 360     if abs(ogl_std) > abs(dx_std):
 361         print('this is probably a DirectX texture')
 362     else:
 363         print('this is probably an OpenGL texture')
 364
 365
 366     if generated_test_images:
 367         # red_x_comparison_img = red_x_comparison_img.swapaxes(0,1)
 368         # red_x_comparison_img = red_x_comparison_img.flatten()
 369         #
 370         # green_y_comparison_img = green_y_comparison_img.swapaxes(0,1)
 371         # green_y_comparison_img = green_y_comparison_img.flatten()
 372         #
 373         # numpytoimage(red_x_comparison_img, 'red_' + i.name, width=width, height=height, channels=1)
 374         # numpytoimage(green_y_comparison_img, 'green_' + i.name, width=width, height=height, channels=1)
 375
 376         ogl_img = ogl_img.swapaxes(0, 1)
 377         ogl_img = ogl_img.flatten()
 378
 379         dx_img = dx_img.swapaxes(0, 1)
 380         dx_img = dx_img.flatten()
 381
 382         numpytoimage(ogl_img, 'OpenGL', width=width, height=height, channels=1)
 383         numpytoimage(dx_img, 'DirectX', width=width, height=height, channels=1)
 384
 385     if abs(ogl_std) > abs(dx_std):
 386             return 'DirectX'
 387     return 'OpenGL'
 388
 389 def make_possible_reductions_on_image(teximage, input_filepath, do_reductions=False, do_downscale=False):
 390     '''checks the image and saves it to drive with possibly reduced channels.
 391     Also can remove the image from the asset if the image is pure black
 392     - it finds it's usages and replaces the inputs where the image is used
 393     with zero/black color.
 394     currently implemented file type conversions:
 395     PNG->JPG
 396     '''
 397     colorspace = teximage.colorspace_settings.name
 398     teximage.colorspace_settings.name = 'Non-Color'
 399     #teximage.colorspace_settings.name = 'sRGB' color correction mambo jambo.
 400
 401     JPEG_QUALITY = 90
 402     # is_image_black(na)
 403     # is_image_bw(na)
 404
 405     rs = bpy.context.scene.render
 406     ims = rs.image_settings
 407
 408     orig_file_format = ims.file_format
 409     orig_quality = ims.quality
 410     orig_color_mode = ims.color_mode
 411     orig_compression = ims.compression
 412     orig_depth = ims.color_depth
 413
 414     # if is_image_black(na):
 415     #     # just erase the image from the asset here, no need to store black images.
 416     #     pass;
 417
 418     # fp = teximage.filepath
 419
 420     # setup  image depth, 8 or 16 bit.
 421     # this should normally divide depth with number of channels, but blender always states that number of channels is 4, even if there are only 3
 422
 423     print(teximage.name)
 424     print(teximage.depth)
 425     print(teximage.channels)
 426
 427     bpy.context.scene.display_settings.display_device = 'None'
 428
 429     image_depth = find_image_depth(teximage)
 430
 431     ims.color_mode = find_color_mode(teximage)
 432     #image_depth = str(max(min(int(teximage.depth / 3), 16), 8))
 433     print('resulting depth set to:', image_depth)
 434
 435     fp = input_filepath
 436     if do_reductions:
 437         na = imagetonumpy_flat(teximage)
 438
 439         if can_erase_alpha(na):
 440             print(teximage.file_format)
 441             if teximage.file_format == 'PNG':
 442                 print('changing type of image to JPG')
 443                 base, ext = os.path.splitext(fp)
 444                 teximage['original_extension'] = ext
 445
 446                 fp = fp.replace('.png', '.jpg')
 447                 fp = fp.replace('.PNG', '.jpg')
 448
 449                 teximage.name = teximage.name.replace('.png', '.jpg')
 450                 teximage.name = teximage.name.replace('.PNG', '.jpg')
 451
 452                 teximage.file_format = 'JPEG'
 453                 ims.quality = JPEG_QUALITY
 454                 ims.color_mode = 'RGB'
 455
 456             if is_image_bw(na):
 457                 ims.color_mode = 'BW'
 458
 459     ims.file_format = teximage.file_format
 460     ims.color_depth = image_depth
 461
 462     # all pngs with max compression
 463     if ims.file_format == 'PNG':
 464         ims.compression = 100
 465     # all jpgs brought to reasonable quality
 466     if ims.file_format == 'JPG':
 467         ims.quality = JPEG_QUALITY
 468
 469     if do_downscale:
 470         downscale(teximage)
 471
 472
 473
 474     # it's actually very important not to try to change the image filepath and packed file filepath before saving,
 475     # blender tries to re-pack the image after writing to image.packed_image.filepath and reverts any changes.
 476     teximage.save_render(filepath=bpy.path.abspath(fp), scene=bpy.context.scene)
 477     if len(teximage.packed_files) > 0:
 478         teximage.unpack(method='REMOVE')
 479     teximage.filepath = fp
 480     teximage.filepath_raw = fp
 481     teximage.reload()
 482
 483     teximage.colorspace_settings.name = colorspace
 484
 485     ims.file_format = orig_file_format
 486     ims.quality = orig_quality
 487     ims.color_mode = orig_color_mode
 488     ims.compression = orig_compression
 489     ims.color_depth = orig_depth