extractApplyPose.py

#!/usr/bin/env python

import lx
import lxu
import lxifc
import lxu.command

class ListUVMaps (lxifc.Visitor):
	def __init__ (self, meshmap):
		self.meshmap = meshmap
		self.meshmap_IDs = []
		self.meshmap_names = []

	def vis_Evaluate (self):
		self.meshmap_names.append (self.meshmap.Name ())
		self.meshmap_IDs.append (self.meshmap.ID ())

class DuplicateDeformedMesh_Cmd(lxu.command.BasicCommand):

#______________________________________________________________________________________________ SETUP AND INITIALISATION

	def __init__(self):
		lxu.command.BasicCommand.__init__(self)

	def cmd_UserName(self):
 		return 'Extract Deformation to Mesh'

 	def cmd_Desc(self):
 		return 'Create a new mesh with the deformations from the current mesh baked in.'

 	def cmd_Tooltip(self):
 		return 'Create a new mesh with the deformations from the current mesh baked in.'

 	def cmd_Help(self):
 		return 'http://www.farfarer.com/'

	def basic_ButtonName(self):
		return 'Extract Deformation to Mesh'

	def cmd_Flags(self):
		return lx.symbol.fCMD_MODEL | lx.symbol.fCMD_UNDO

	def basic_Enable(self, msg):
		return True

	def cmd_Interact(self):
		pass

#_______________________________________________________________________________________________ MAIN FUNCTION EXECUTION

	def basic_Execute(self, msg, flags):
		msg = lx.object.Message (msg)
		sel_svc = lx.service.Selection ()
		scn_svc = lx.service.Scene ()
		tag = lx.object.StringTag ()

		time = sel_svc.GetTime ()
		scene = lx.object.Scene (lxu.select.SceneSelection ().current ())
		chan_read = lx.object.ChannelRead (scene.Channels (None, time))
		chan_write_setup = lx.object.ChannelWrite (scene.Channels (lx.symbol.s_ACTIONLAYER_SETUP, 0.0))
		chan_write_edit = lx.object.ChannelWrite (scene.Channels (lx.symbol.s_ACTIONLAYER_EDIT, time))

		mesh_item = None

		layer_svc = lx.service.Layer ()
		layer_scan = lx.object.LayerScan (layer_svc.ScanAllocate (lx.symbol.f_LAYERSCAN_PRIMARY))
		if layer_scan.test ():
			layer_scan_count = layer_scan.Count ()
			if layer_scan_count > 0:
				mesh_item = layer_scan.MeshItem (0)
			layer_scan.Apply ()
		del layer_scan

		if mesh_item is None:
			lx.out('Cannot find primary mesh item.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		# Create new mesh item.
		itype_mesh = scn_svc.ItemTypeLookup (lx.symbol.sITYPE_MESH)
		new_mesh_item = scene.ItemAdd (itype_mesh)
		new_mesh_item.SetName (mesh_item.UniqueName() + ' Deformed')
		new_mesh_ident = new_mesh_item.Ident ()

		item_pkt_trans = lx.object.ItemPacketTranslation (sel_svc.Allocate (lx.symbol.sSELTYP_ITEM))
		sel_type_item = sel_svc.LookupType (lx.symbol.sSELTYP_ITEM)
		pkt = item_pkt_trans.Packet (new_mesh_item)
		sel_svc.Select (sel_type_item, pkt)

		mesh = None
		new_mesh = None

		mesh_layer = None
		new_mesh_layer = None

		layer_scan = lx.object.LayerScan (layer_svc.ScanAllocate (lx.symbol.f_LAYERSCAN_EDIT))
		if layer_scan.test ():
			layer_scan_count = layer_scan.Count ()
			for layer_idx in xrange (layer_scan_count):
				itm = layer_scan.MeshItem (layer_idx)
				if itm.Ident () == new_mesh_ident:
					# Grab the new mesh.
					new_mesh = layer_scan.MeshEdit (layer_idx)
					if not new_mesh.test ():
						lx.out ('Can\'t edit new mesh.')
						msg.SetCode (lx.symbol.e_FAILED)
						return

					# Evaluate the deformed mesh.
					index = mesh_item.ChannelLookup (lx.symbol.sICHAN_MESH_MESH)
					mfilt = lx.object.MeshFilter (chan_read.ValueObj (mesh_item, index))
					if mfilt.test ():
						mesh = lx.object.Mesh (mfilt.Generate ())
						if mesh.test ():
							point = lx.object.Point(mesh.PointAccessor())
							new_point = lx.object.Point(new_mesh.PointAccessor())

							polygon = lx.object.Polygon(mesh.PolygonAccessor())
							new_polygon = lx.object.Polygon(new_mesh.PolygonAccessor())

							meshmap = lx.object.MeshMap(mesh.MeshMapAccessor())
							new_meshmap = lx.object.MeshMap(new_mesh.MeshMapAccessor())

							visitor = ListUVMaps (meshmap)
							meshmap.FilterByType (lx.symbol.i_VMAP_TEXTUREUV)
							meshmap.Enumerate (lx.symbol.iMARK_ANY, visitor, 0)
							meshmap_IDs = tuple(visitor.meshmap_IDs)
							meshmap_names = tuple(visitor.meshmap_names)
							new_meshmap_IDs = []
							meshmap.FilterByType (0)

							tag.set(new_mesh_item)
							tagID = lxu.lxID4('MCMT')
							tag.Set(tagID, mesh_item.Ident() + ':' + str(time))

							if point.test () and  new_point.test () and polygon.test () and  new_polygon.test () and meshmap.test () and  new_meshmap.test ():
								new_points = []
								new_polygons = []

								for meshmap_name in meshmap_names:
									new_meshmap_IDs.append (new_meshmap.New(lx.symbol.i_VMAP_TEXTUREUV, meshmap_name))

								for p in xrange (mesh.PointCount()):
									point.SelectByIndex (p)
									new_point_ID = new_point.New (point.Pos ())
									new_points.append (new_point_ID)
									new_point.Select (new_point_ID)
									for meshmap_ID, new_meshmap_ID in zip(meshmap_IDs, new_meshmap_IDs):
										uv = lx.object.storage ('f', 2)
										try:
											res = point.MapValue (meshmap_ID, uv)
										except:
											pass
										else:
											if res:
												new_point.SetMapValue (new_meshmap_ID, uv)

								for p in xrange (mesh.PolygonCount ()):
									polygon_points = []
									polygon.SelectByIndex (p)

									tag.set(polygon)
									matr = None
									part = None
									pick = None
									try:
										matr = tag.Get(lx.symbol.i_PTAG_MATR)
									except LookupError:
										pass

									try:
										part = tag.Get(lx.symbol.i_PTAG_PART)
									except LookupError:
										pass

									try:
										pick = tag.Get(lx.symbol.i_PTAG_PICK)
									except:
										pass

									for v in xrange (polygon.VertexCount ()):
										point.Select (polygon.VertexByIndex (v))
										polygon_points.append (point.Index ())
									point_storage = lx.object.storage ('p', len(polygon_points))
									point_storage.set ([new_points[x] for x in polygon_points])
									new_polygon_id = new_polygon.New (polygon.Type (), point_storage, len(polygon_points), 0)
									new_polygons.append (new_polygon_id)
									new_polygon.Select (new_polygon_id)

									for v in xrange (polygon.VertexCount ()):
										point_ID = polygon.VertexByIndex (v)
										point.Select (point_ID)
										point_index = point.Index ()
										for meshmap_ID, new_meshmap_ID in zip(meshmap_IDs, new_meshmap_IDs):
											uv = lx.object.storage ('f', 2)
											try:
												res = polygon.MapValue (meshmap_ID, point_ID, uv)
											except:
												pass
											else:
												if res:
													new_polygon.SetMapValue (new_points[point_index], new_meshmap_ID, uv)

									tag.set(new_polygon)
									if matr is not None:
										tag.Set(lx.symbol.i_PTAG_MATR, matr)
									if part is not None:
										tag.Set(lx.symbol.i_PTAG_PART, part)
									if pick is not None:
										tag.Set(lx.symbol.i_PTAG_PICK, pick)
						layer_scan.SetMeshChange (layer_idx, lx.symbol.f_MESHEDIT_POSITION | lx.symbol.f_MESHEDIT_GEOMETRY | lx.symbol.f_MESHEDIT_MAP_OTHER | lx.symbol.f_MESHEDIT_MAP_UV | lx.symbol.f_MESHEDIT_POL_TAGS | lx.symbol.f_MESHEDIT_POL_TYPE)

					# Deselect the primary layer and hide it.
					pkt = item_pkt_trans.Packet (mesh_item)
					sel_svc.Deselect (sel_type_item, pkt)
					visible_channel_idx = mesh_item.ChannelLookup (lx.symbol.sICHAN_LOCATOR_VISIBLE)
					chan_write_edit.Integer (mesh_item, visible_channel_idx, 2)

					break
		layer_scan.Apply ()

#________________________________________________________________________________________ BLESS FUNCTION AS MODO COMMAND

lx.bless (DuplicateDeformedMesh_Cmd, 'ffr.extractPose')


class RestoreDeformedMesh_Cmd(lxu.command.BasicCommand):

#______________________________________________________________________________________________ SETUP AND INITIALISATION

	def __init__(self):
		lxu.command.BasicCommand.__init__(self)
		self.dyna_Add ('morph', lx.symbol.sTYPE_STRING)

	def cmd_UserName(self):
 		return 'Apply Altered Pose Mesh as Morph'

 	def cmd_Desc(self):
 		return 'Take the vertex delta from this mesh and apply it back to the source mesh as a morph map.'

 	def cmd_Tooltip(self):
 		return 'Take the vertex delta from this mesh and apply it back to the source mesh as a morph map.'

 	def cmd_Help(self):
 		return 'http://www.farfarer.com/'

	def basic_ButtonName(self):
		return 'Apply Altered Pose Mesh as Morph'

	def cmd_Flags(self):
		return lx.symbol.fCMD_MODEL | lx.symbol.fCMD_UNDO | lx.symbol.fCMD_SELECT

	def basic_Enable(self, msg):
		return True

	def arg_UIHints (self, index, hints):
		if index == 0:
			hints.Label ('Target Morph Map')

	def cmd_Interact(self):
		pass

#_______________________________________________________________________________________________ MAIN FUNCTION EXECUTION

	def basic_Execute(self, msg, flags):

		morph_name = ''
		if self.dyna_IsSet (0):
			morph_name = self.dyna_String(0)

		if len(morph_name) < 1:
			lx.out('Invalid Morph Name')
			return

		sel_svc = lx.service.Selection ()
		scn_svc = lx.service.Scene ()
		tag = lx.object.StringTag ()

		time = sel_svc.GetTime ()
		scene = lx.object.Scene (lxu.select.SceneSelection ().current ())
		chan_read_eval = lx.object.ChannelRead (scene.Channels (None, time))
		chan_write_edit = lx.object.ChannelWrite (scene.Channels (lx.symbol.s_ACTIONLAYER_EDIT, time))

		mesh_deformed = None
		mesh_corrected = None
		mesh_base = None

		mesh_item_base = None
		mesh_ident_base = None

		# Grab the primary layer, this should be our corrected mesh.
		layer_svc = lx.service.Layer ()
		layer_scan = lx.object.LayerScan (layer_svc.ScanAllocate (lx.symbol.f_LAYERSCAN_PRIMARY))
		if not layer_scan.test ():
			lx.out ('Layer scan failed.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		layer_scan_count = layer_scan.Count ()
		if layer_scan_count == 0:
			lx.out('No primary layer.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		# Grab it's item.
		mesh_item = layer_scan.MeshItem (0)
		mesh_ident_corrected = mesh_item.Ident ()
		layer_scan.Apply ()
		del layer_scan

		# Try and grab the corrected mesh item from the scene.
		try:
			mesh_item_corrected = scene.ItemLookupIdent (mesh_ident_corrected)
		except:
			lx.out('Can\'t find corrected mesh.')
			msg.SetCode (lx.symbol.e_FAILED)
			return
		else:
			# Read the tags on it that tell us what mesh it was created from and at what time.
			tag.set(mesh_item_corrected)
			tagID = lxu.lxID4('MCMT')
			mesh_info = tag.Get(tagID)
			mesh_info = mesh_info.split (':')
			if len(mesh_info) != 2:
				lx.out('Tag info is bad.')
				msg.SetCode (lx.symbol.e_FAILED)
				return

			mesh_ident_base = mesh_info[0]
			mesh_time = float(mesh_info[1])

		# Try and find the mesh item it was created from.
		try:
			mesh_item_base = scene.ItemLookupIdent (mesh_ident_base)
		except:
			lx.out('Can\'t find original mesh.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		# Select the base mesh item and get it's edit mesh via layerscan (i.e. get the base undeformed mesh, writeable).
		item_pkt_trans = lx.object.ItemPacketTranslation (sel_svc.Allocate (lx.symbol.sSELTYP_ITEM))
		sel_type_item = sel_svc.LookupType (lx.symbol.sSELTYP_ITEM)
		pkt = item_pkt_trans.Packet (mesh_item_base)
		sel_svc.Select (sel_type_item, pkt)

		layer_scan = lx.object.LayerScan (layer_svc.ScanAllocate (lx.symbol.f_LAYERSCAN_ACTIVE | lx.symbol.f_LAYERSCAN_WRITEMESH))
		if not layer_scan.test ():
			lx.out ('Layer scan 2 failed.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		layer_scan_count = layer_scan.Count ()
		mesh_base_idx = None
		for layer_idx in xrange (layer_scan_count):
			itm = layer_scan.MeshItem (layer_idx)
			if itm.Ident () == mesh_ident_base:
				mesh_base = layer_scan.MeshEdit (layer_idx)

				if not mesh_base.test():
					mesh_base = None

				mesh_base_idx = layer_idx

				break
		else:
			lx.out ('Didn\'t find base mesh in layer scan.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		chan_read_eval1 = lx.object.ChannelRead (scene.Channels (None, mesh_time))
		chan_read_eval2 = lx.object.ChannelRead (scene.Channels (None, mesh_time))

		# Grab the base mesh item and evaluate it's mesh channel at that time (i.e. get the deformed mesh).
		mesh_chan_deformed_index = mesh_item_base.ChannelLookup (lx.symbol.sICHAN_MESH_MESH)
		mfilt_deformed = lx.object.MeshFilter (chan_read_eval1.ValueObj (mesh_item_base, mesh_chan_deformed_index))
		if mfilt_deformed.test ():
			mesh_deformed = lx.object.Mesh (mfilt_deformed.Generate ())
			if not mesh_deformed.test ():
				mesh_deformed = None

		# Grab the corrected mesh item and evaluate it's mesh channel at that time (i.e. get the corrected mesh).
		mesh_chan_corrected_index = mesh_item_corrected.ChannelLookup (lx.symbol.sICHAN_MESH_MESH)
		mfilt_corrected = lx.object.MeshFilter (chan_read_eval2.ValueObj (mesh_item_corrected, mesh_chan_corrected_index))
		if mfilt_corrected.test ():
			mesh_corrected = lx.object.Mesh (mfilt_corrected.Generate ())
			if not mesh_corrected.test ():
				mesh_corrected = None

		if mesh_deformed is None or mesh_corrected is None or mesh_base is None:
			lx.out ('Couldn\'t find all the meshes.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		pc_base = mesh_base.PointCount ()
		pc_deformed = mesh_deformed.PointCount ()
		pc_corrected = mesh_corrected.PointCount ()

		if (pc_base != pc_deformed) or (pc_base != pc_corrected) or (pc_deformed != pc_corrected):
			lx.out ('Point count differs between meshes.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		# Whew!
		point_base = lx.object.Point(mesh_base.PointAccessor())
		point_deformed = lx.object.Point(mesh_deformed.PointAccessor())
		point_corrected = lx.object.Point(mesh_corrected.PointAccessor())
		meshmap_base = lx.object.MeshMap(mesh_base.MeshMapAccessor())

		if not point_base.test () or not point_deformed.test () or not point_corrected.test () or not meshmap_base.test ():
			lx.out('Can\'t get the accessors.')
			msg.SetCode (lx.symbol.e_FAILED)
			return

		try:
			meshmap_base.SelectByName (lx.symbol.i_VMAP_MORPH, morph_name)
		except:
			meshmapID = meshmap_base.New (lx.symbol.i_VMAP_MORPH, morph_name)
		else:
			meshmap_base.Clear ()
			meshmapID = meshmap_base.ID ()

		morph_storage = lx.object.storage('f', 3)

		for p in xrange(pc_base):
			point_base.SelectByIndex(p)
			point_deformed.SelectByIndex(p)
			point_corrected.SelectByIndex(p)

			pos_base = point_base.Pos ()
			pos_deformed = point_deformed.Pos ()
			pos_corrected = point_corrected.Pos ()

			morph1 = [pos_base[x] - pos_deformed[x] for x in xrange(3)]
			morph2 = [pos_base[x] - pos_corrected[x] for x in xrange(3)]

			morph3 = [((pos_base[x] - morph1[x]) + morph2[x]) for x in xrange(3)]

			morph4 = [pos_base[x] - morph3[x] for x in xrange(3)]

			if morph4[0] == 0.0 and morph4[1] == 0.0 and morph4[2] == 0.0:
				continue

			morph_storage.set(morph4)

			point_base.SetMapValue (meshmapID, morph_storage)

		if mesh_base_idx is not None:
			layer_scan.SetMeshChange (mesh_base_idx, lx.symbol.f_MESHEDIT_MAP_OTHER)

		layer_scan.Apply ()

		pkt = item_pkt_trans.Packet (mesh_item_corrected)
		sel_svc.Deselect (sel_type_item, pkt)
		visible_channel_idx = mesh_item_corrected.ChannelLookup (lx.symbol.sICHAN_LOCATOR_VISIBLE)
		chan_write_edit.Integer (mesh_item_corrected, visible_channel_idx, 2)

#________________________________________________________________________________________ BLESS FUNCTION AS MODO COMMAND

lx.bless (RestoreDeformedMesh_Cmd, 'ffr.applyPoseAsMorph')