io_mesh_uv_layout/export_uv_png.py

   1 # SPDX-FileCopyrightText: 2011-2023 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy
   6 import gpu
   7 from mathutils import Vector, Matrix
   8 from mathutils.geometry import tessellate_polygon
   9 from gpu_extras.batch import batch_for_shader
  10
  11 # Use OIIO if available, else Blender for writing the image.
  12 try:
  13     import OpenImageIO as oiio
  14 except ImportError:
  15     oiio = None
  16
  17
  18 def export(filepath, face_data, colors, width, height, opacity):
  19     offscreen = gpu.types.GPUOffScreen(width, height)
  20     offscreen.bind()
  21
  22     try:
  23         fb = gpu.state.active_framebuffer_get()
  24         fb.clear(color=(0.0, 0.0, 0.0, 0.0))
  25         draw_image(face_data, opacity)
  26
  27         pixel_data = fb.read_color(0, 0, width, height, 4, 0, 'UBYTE')
  28         pixel_data.dimensions = width * height * 4
  29         save_pixels(filepath, pixel_data, width, height)
  30     finally:
  31         offscreen.unbind()
  32         offscreen.free()
  33
  34
  35 def draw_image(face_data, opacity):
  36     gpu.state.blend_set('ALPHA')
  37
  38     with gpu.matrix.push_pop():
  39         gpu.matrix.load_matrix(get_normalize_uvs_matrix())
  40         gpu.matrix.load_projection_matrix(Matrix.Identity(4))
  41
  42         draw_background_colors(face_data, opacity)
  43         draw_lines(face_data)
  44
  45     gpu.state.blend_set('NONE')
  46
  47
  48 def get_normalize_uvs_matrix():
  49     '''matrix maps x and y coordinates from [0, 1] to [-1, 1]'''
  50     matrix = Matrix.Identity(4)
  51     matrix.col[3][0] = -1
  52     matrix.col[3][1] = -1
  53     matrix[0][0] = 2
  54     matrix[1][1] = 2
  55
  56     # OIIO writes arrays from the left-upper corner.
  57     if oiio:
  58         matrix.col[3][1] *= -1.0
  59         matrix[1][1] *= -1.0
  60
  61     return matrix
  62
  63
  64 def draw_background_colors(face_data, opacity):
  65     coords = [uv for uvs, _ in face_data for uv in uvs]
  66     colors = [(*color, opacity) for uvs, color in face_data for _ in range(len(uvs))]
  67
  68     indices = []
  69     offset = 0
  70     for uvs, _ in face_data:
  71         triangles = tessellate_uvs(uvs)
  72         indices.extend([index + offset for index in triangle] for triangle in triangles)
  73         offset += len(uvs)
  74
  75     shader = gpu.shader.from_builtin('FLAT_COLOR')
  76     batch = batch_for_shader(
  77         shader, 'TRIS',
  78         {"pos": coords, "color": colors},
  79         indices=indices,
  80     )
  81     batch.draw(shader)
  82
  83
  84 def tessellate_uvs(uvs):
  85     return tessellate_polygon([uvs])
  86
  87
  88 def draw_lines(face_data):
  89     coords = []
  90     for uvs, _ in face_data:
  91         for i in range(len(uvs)):
  92             start = uvs[i]
  93             end = uvs[(i + 1) % len(uvs)]
  94             coords.append((start[0], start[1]))
  95             coords.append((end[0], end[1]))
  96
  97     shader = gpu.shader.from_builtin('POLYLINE_UNIFORM_COLOR')
  98     shader.uniform_float("viewportSize", gpu.state.viewport_get()[2:])
  99     shader.uniform_float("lineWidth", 1.0)
 100     shader.uniform_float("color", (0.0, 0.0, 0.0, 1.0))
 101
 102     batch = batch_for_shader(shader, 'LINES', {"pos": coords})
 103     batch.draw(shader)
 104
 105
 106 def save_pixels(filepath, pixel_data, width, height):
 107     if oiio:
 108         spec = oiio.ImageSpec(width, height, 4, "uint8")
 109         image = oiio.ImageOutput.create(filepath)
 110         image.open(filepath, spec)
 111         image.write_image(pixel_data)
 112         image.close()
 113         return
 114
 115     image = bpy.data.images.new("temp", width, height, alpha=True)
 116     image.filepath = filepath
 117     image.pixels = [v / 255 for v in pixel_data]
 118     image.save()
 119     bpy.data.images.remove(image)