Blender-TexTools/op_color_convert_texture.py

import bpy
import bmesh
import operator
import math
from mathutils import Vector
from collections import defaultdict

from . import utilities_color
from . import utilities_bake
from . import utilities_ui

material_prefix = "TT_atlas_"
gamma = 2.2

class op(bpy.types.Operator):
	bl_idname = "uv.textools_color_convert_to_texture"
	bl_label = "Pack Texture"
	bl_description = "Pack ID Colors into single texture and UVs"
	bl_options = {'REGISTER', 'UNDO'}
	

	@classmethod
	def poll(cls, context):
		if not bpy.context.active_object:
			return False

		if bpy.context.active_object not in bpy.context.selected_objects:
			return False

		if len(bpy.context.selected_objects) != 1:
			return False

		if bpy.context.active_object.type != 'MESH':
			return False

		#Only in UV editor mode
		if bpy.context.area.type != 'IMAGE_EDITOR':
			return False

		return True
	
	def execute(self, context):
		pack_texture(self, context)
		return {'FINISHED'}



def pack_texture(self, context):
	obj = bpy.context.active_object
	name = material_prefix+obj.name

	if obj.mode != 'OBJECT':
		bpy.ops.object.mode_set(mode='OBJECT')


	# Determine size
	size_pixel = 8
	size_square = math.ceil(math.sqrt( context.scene.texToolsSettings.color_ID_count ))
	size_image = size_square * size_pixel
	size_image_pow = int(math.pow(2, math.ceil(math.log(size_image, 2))))

	# Maximize pixel size
	size_pixel = math.floor(size_image_pow/size_square)

	print("{0} colors = {1} x {1} = ({2}pix)  {3} x {3}  | {4} x {4}".format(
		context.scene.texToolsSettings.color_ID_count, 
		size_square,
		size_pixel,
		size_image,
		size_image_pow
	))

	# Create image
	image = bpy.data.images.new(name, width=size_image_pow, height=size_image_pow)
	pixels = [None] * size_image_pow * size_image_pow

	# Black pixels
	for x in range(size_image_pow):
		for y in range(size_image_pow):
			pixels[(y * size_image_pow) + x] = [0, 0, 0, 1]

	# Pixels
	for c in range(context.scene.texToolsSettings.color_ID_count):
		x = c % size_square
		y = math.floor(c/size_square)
		color = utilities_color.get_color(c).copy()
		for i in range(3):
			color[i] = pow(color[i] , 1.0/gamma)

		for sx in range(size_pixel):
			for sy in range(size_pixel):
				_x = x*size_pixel + sx
				_y = y*size_pixel + sy
				pixels[(_y * size_image_pow) + _x] = [color[0], color[1], color[2], 1]


	# flatten list & assign pixels
	pixels = [chan for px in pixels for chan in px]
	image.pixels = pixels

	# Set background image
	for area in bpy.context.screen.areas:
		if area.type == 'IMAGE_EDITOR':
			area.spaces[0].image = image

	# Edit mesh
	bpy.ops.object.mode_set(mode='EDIT')
	bpy.ops.mesh.select_mode(use_extend=False, use_expand=False, type='FACE')
	bpy.ops.mesh.select_all(action='SELECT')
	# bpy.ops.uv.smart_project(angle_limit=1)
	bpy.ops.uv.unwrap(method='ANGLE_BASED', margin=0.0078)


	bm = bmesh.from_edit_mesh(bpy.context.active_object.data)
	uv_layers = bm.loops.layers.uv.verify();

	for face in bm.faces:
		index = face.material_index

		# Get UV coordinates for index
		x = index%size_square
		y = math.floor(index/size_square)

		x*= (size_pixel / size_image_pow) 
		y*= (size_pixel / size_image_pow)
		x+= size_pixel/size_image_pow/2
		y+= size_pixel/size_image_pow/2

		for loop in face.loops:
			loop[uv_layers].uv = (x, y)

	# Remove Slots & add one
	bpy.ops.object.mode_set(mode='OBJECT')
	bpy.ops.uv.textools_color_clear()
	bpy.ops.object.material_slot_add()
	
	#Create material with image
	obj.material_slots[0].material = utilities_bake.get_image_material(image)

	#Display UVs
	bpy.ops.object.mode_set(mode='EDIT')

	# Switch textured shading
	for area in bpy.context.screen.areas:
		if area.type == 'VIEW_3D':
			for space in area.spaces:
				if space.type == 'VIEW_3D':
					space.shading.type = 'MATERIAL'

	bpy.ops.ui.textools_popup('INVOKE_DEFAULT', message="Packed texture with {} color IDs".format( context.scene.texToolsSettings.color_ID_count ))


bpy.utils.register_class(op)