op_texel_checker_map.py

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

from . import utilities_texel


texture_modes = ['UV_GRID','COLOR_GRID','GRAVITY','NONE']



class op(bpy.types.Operator):
	bl_idname = "uv.textools_texel_checker_map"
	bl_label = "Checker Map"
	bl_description = "Add a checker map to the selected model and UV view"
	bl_options = {'REGISTER', 'UNDO'}
	
	@classmethod
	def poll(cls, context):
		if len(get_valid_objects()) == 0:
			return False

		return True

	def execute(self, context):
		assign_checker_map(
			bpy.context.scene.texToolsSettings.size[0], 
			bpy.context.scene.texToolsSettings.size[1]
		)
		return {'FINISHED'}




def assign_checker_map(size_x, size_y):
	# Force Object mode
	if bpy.context.view_layer.objects.active != None and bpy.context.object.mode != 'OBJECT':
		bpy.ops.object.mode_set(mode='OBJECT')

	# Collect Objects
	objects = get_valid_objects()
	
	if len(objects) == 0:
		self.report({'ERROR_INVALID_INPUT'}, "No UV mapped objects selected" )

	#Change View mode to TEXTURED
	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'


	if len(objects) > 0:

		# Detect current Checker modes
		mode_count = {}
		for mode in texture_modes:
			mode_count[mode] = 0

		# Image sizes
		image_sizes_x = []
		image_sizes_y = []

		# Collect current modes in selected objects
		for obj in objects:
			image = utilities_texel.get_object_texture_image(obj)
			mode = 'NONE'
			if image:
				if "GRAVITY" in image.name.upper():
					mode = 'GRAVITY'

				elif image.generated_type in texture_modes:
					# Generated checker maps
					mode = image.generated_type

					# Track image sizes
					if image.size[0] not in image_sizes_x:
						image_sizes_x.append(image.size[0])
					if image.size[1] not in image_sizes_y:
						image_sizes_y.append(image.size[1])

			mode_count[mode]+=1


		# Sort by count (returns tuple list of key,value)
		mode_max_count = sorted(mode_count.items(), key=operator.itemgetter(1))
		mode_max_count.reverse()

		for key,val in mode_max_count:
			print("{} = {}".format(key, val))


		# Determine next mode
		mode = 'NONE'
		if mode_max_count[0][1] == 0:
			# There are no maps
			mode = texture_modes[0]

		elif mode_max_count[0][0] in texture_modes:
			if mode_max_count[-1][1] > 0:
				# There is more than 0 of another mode, complete existing mode first
				mode = mode_max_count[0][0]

			else:
				# Switch to next checker mode
				index = texture_modes.index(mode_max_count[0][0])
				
				if texture_modes[ index ] != 'NONE' and len(image_sizes_x) > 1 or len(image_sizes_y) > 1:
					# There are multiple resolutions on selected objects
					mode = texture_modes[ index ]
				elif texture_modes[ index ] != 'NONE' and (len(image_sizes_x) > 0 and image_sizes_x[0] != size_x) and (len(image_sizes_y) > 0 and image_sizes_y[0] != size_y):
					# The selected objects have a different resolution
					mode = texture_modes[ index ]
				else:
					# Next mode
					mode = texture_modes[ (index+1)%len(texture_modes) ]


		print("Mode: "+mode)

		if mode == 'NONE':
			for obj in objects:
				remove_material(obj)

		elif mode == 'GRAVITY':
			image = load_image("checker_map_gravity")
			for obj in objects:
				apply_image(obj, image)

		else:
			name = utilities_texel.get_checker_name(mode, size_x, size_y)
			image = get_image(name, mode, size_x, size_y)
			for obj in objects:
				apply_image(obj, image)
	
	# Restore object selection
	bpy.ops.object.mode_set(mode='OBJECT')
	bpy.ops.object.select_all(action='DESELECT')
	for obj in objects:				
		obj.select_set( state = True, view_layer = None)

	# Clean up images and materials
	utilities_texel.checker_images_cleanup()

	# Force redraw of viewport to update texture
	# bpy.context.scene.update()
	bpy.context.view_layer.update()




def load_image(name):
	pathTexture = icons_dir = os.path.join(os.path.dirname(__file__), "resources/{}.png".format(name))
	image = bpy.ops.image.open(filepath=pathTexture, relative_path=False)
	if "{}.png".format(name) in bpy.data.images:
		bpy.data.images["{}.png".format(name)].name = name #remove extension in name
	return bpy.data.images[name];



def get_valid_objects():
	# Collect Objects
	objects = []
	for obj in bpy.context.selected_objects:
		if obj.type == 'MESH' and obj.data.uv_layers:
				objects.append(obj)

	return objects







def remove_material(obj):
	bpy.ops.object.mode_set(mode='OBJECT')
	bpy.ops.object.select_all(action='DESELECT')
	obj.select_set( state = True, view_layer = None)
	bpy.context.view_layer.objects.active = obj
	count = len(obj.material_slots)
	for i in range(count):
		bpy.ops.object.material_slot_remove()



def apply_image(obj, image):

	bpy.ops.object.mode_set(mode='OBJECT')
	bpy.ops.object.select_all(action='DESELECT')
	obj.select_set( state = True, view_layer = None)
	bpy.context.view_layer.objects.active = obj

	# Assign Cycles material with image

	# Get Material
	material = None
	if image.name in bpy.data.materials:
		material = bpy.data.materials[image.name]
	else:
		material = bpy.data.materials.new(image.name)
		material.use_nodes = True

	# Assign material
	if len(obj.data.materials) > 0:
		obj.data.materials[0] = material
	else:
		obj.data.materials.append(material)

	# Setup Node
	tree = material.node_tree
	node = None
	if "checker" in tree.nodes:
		node = tree.nodes["checker"]
	else:
		node = tree.nodes.new("ShaderNodeTexImage")
	node.name = "checker"
	node.select = True
	tree.nodes.active = node
	node.image = image

	# LINKANDO:
	tree = obj.data.materials[0].node_tree
	links = tree.links
	nodo1 = tree.nodes['checker']
	nodo2 = tree.nodes['Principled BSDF']
	links.new(nodo1.outputs['Color'], nodo2.inputs['Base Color'])


def get_image(name, mode, size_x, size_y):
	# Image already exists?
	if name in bpy.data.images:
		# Update texture UV checker mode
		bpy.data.images[name].generated_type = mode
		return bpy.data.images[name];

	# Create new image instead
	image = bpy.data.images.new(name, width=size_x, height=size_y)
	image.generated_type = mode #UV_GRID or COLOR_GRID
	image.generated_width = int(size_x)
	image.generated_height = int(size_y)

	return image

bpy.utils.register_class(op)