added options to weightpaint normalize script

- armature only - so when using weight groups for other things only armature groups are affected,
- Active only - so you can normalize all weight groups

1 files changed, 88 insertions, 38 deletions

diff --git a/release/scripts/weightpaint_normalize.py b/release/scripts/weightpaint_normalize.py
index 204868a79fc..1f6c3d201aa 100644
--- a/release/scripts/weightpaint_normalize.py
+++ b/release/scripts/weightpaint_normalize.py
@@ -40,54 +40,100 @@ proportion of the veighting is unchanged.
 # ***** END GPL LICENCE BLOCK *****
 # --------------------------------------------------------------------------
 
-from Blender import Scene, Draw
+from Blender import Scene, Draw, Object, Modifier
 import BPyMesh
 SMALL_NUM= 0.000001
-def actWeightNormalize(me, PREF_PEAKWEIGHT, PREF_KEEP_PROPORTION):
+
+def getArmatureGroups(ob, me):
+	
+	arm_obs = []
+	
+	arm = ob.parent
+	if arm and arm.type == 'Armature' and ob.parentType == Object.ParentTypes.ARMATURE:
+		arm_obs.append(arm)
+	
+	for m in ob.modifiers:
+		if m.type== Modifier.Types.ARMATURE:
+			arm = m[Modifier.Settings.OBJECT]
+			if arm:
+				arm_obs.append(arm)
+	
+	# convert to a dict and back, should be a set! :/ - python 2.3 dosnt like.
+	return dict([ (bonename, None) for arm in arm_obs for bonename in arm.data.bones.keys() ]).keys()
+
+
+
+def actWeightNormalize(me, ob, PREF_PEAKWEIGHT, PREF_ACTIVE_ONLY, PREF_ARMATURE_ONLY, PREF_KEEP_PROPORTION):
 	
 	groupNames, vWeightDict= BPyMesh.meshWeight2Dict(me)
 	new_weight= max_weight= -1.0
 	act_group= me.activeGroup
 	
-	vWeightDictUsed=[False] * len(vWeightDict)
-	
-	for i, wd in enumerate(vWeightDict):
-		try:
-			new_weight= wd[act_group]
-			if new_weight > max_weight:
-				max_weight= new_weight
-			vWeightDictUsed[i]=wd
-		except:
-			pass
-	
-	if max_weight < SMALL_NUM or new_weight == -1:
-		Draw.PupMenu('No verts to normalize. exiting.')
-		return
+	if PREF_ACTIVE_ONLY:
+		normalizeGroups = [act_group]
+	else:
+		normalizeGroups  = groupNames[:]
 	
-	if abs(max_weight-PREF_PEAKWEIGHT) < SMALL_NUM:
-		Draw.PupMenu('Vert Weights are alredy normalized.')
-		return
+	if PREF_ARMATURE_ONLY:
+		
+		armature_groups = getArmatureGroups(ob, me)
+		
+		i = len(normalizeGroups)
+		while i:
+			i-=1
+			if not normalizeGroups[i] in armature_groups:
+				del normalizeGroups[i]
 	
-	max_weight= max_weight/PREF_PEAKWEIGHT
 	
-	if PREF_KEEP_PROPORTION:
-		# TODO, PROPORTIONAL WEIGHT SCALING.
-		for wd in vWeightDictUsed:
-			if wd: # not false.
-				if len(wd) == 1:
-					# Only 1 group for thsi vert. Simple
-					wd[act_group] /= max_weight
-				else:
-					# More then 1 group. will need to scale all users evenly.
-					local_maxweight= max(wd.itervalues()) / PREF_PEAKWEIGHT
-					for weight in wd.iterkeys():
-						wd[weight] /= local_maxweight
-				
+	for act_group in normalizeGroups:
+		vWeightDictUsed=[False] * len(vWeightDict)
+		
+		for i, wd in enumerate(vWeightDict):
+			try:
+				new_weight= wd[act_group]
+				if new_weight > max_weight:
+					max_weight= new_weight
+				vWeightDictUsed[i]=wd
+			except:
+				pass
 				
-	else: # Simple, just scale the weights up.
-		for wd in vWeightDictUsed:
-			if wd: # not false.
-				wd[act_group] /= max_weight
+		# These can be skipped for now, they complicate things when using multiple vgroups,
+		'''
+		if max_weight < SMALL_NUM or new_weight == -1:
+			Draw.PupMenu('No verts to normalize. exiting.')
+			#return
+		
+		if abs(max_weight-PREF_PEAKWEIGHT) < SMALL_NUM:
+			Draw.PupMenu('Vert Weights are alredy normalized.')
+			#return
+		'''
+		max_weight= max_weight/PREF_PEAKWEIGHT
+		
+		if PREF_KEEP_PROPORTION:
+			# TODO, PROPORTIONAL WEIGHT SCALING.
+			for wd in vWeightDictUsed:
+				if wd: # not false.
+					if len(wd) == 1:
+						# Only 1 group for thsi vert. Simple
+						wd[act_group] /= max_weight
+					else:
+						# More then 1 group. will need to scale all users evenly.
+						if PREF_ARMATURE_ONLY:
+							local_maxweight= max([v for k, v in wd.iteritems() if k in armature_groups]) / PREF_PEAKWEIGHT
+							if local_maxweight > 0.0:
+								# So groups that are not used in any bones are ignored.
+								for weight in wd.iterkeys():
+									if weight in armature_groups:
+										wd[weight] /= local_maxweight
+						else:
+							local_maxweight= max(wd.itervalues()) / PREF_PEAKWEIGHT
+							for weight in wd.iterkeys():
+								wd[weight] /= local_maxweight
+		
+		else: # Simple, just scale the weights up. we alredy know this is in an armature group (if needed)
+			for wd in vWeightDictUsed:
+				if wd: # not false.
+					wd[act_group] /= max_weight
 		
 	# Copy weights back to the mesh.
 	BPyMesh.dict2MeshWeight(me, groupNames, vWeightDict)
@@ -104,17 +150,21 @@ def main():
 	me= ob.getData(mesh=1)
 	
 	PREF_PEAKWEIGHT= Draw.Create(1.0)
+	PREF_ACTIVE_ONLY= Draw.Create(1)
 	PREF_KEEP_PROPORTION= Draw.Create(1)
+	PREF_ARMATURE_ONLY= Draw.Create(0)
 	
 	pup_block= [\
 	('Peak Weight:', PREF_PEAKWEIGHT, 0.01, 1.0, 'Upper weight for normalizing.'),\
+	('Active Only', PREF_ACTIVE_ONLY, 'Only Normalize groups that have matching bones in an armature (when an armature is used).'),\
 	('Proportional', PREF_KEEP_PROPORTION, 'Scale other weights so verts (Keep weights with other groups in proportion).'),\
+	('Armature Only', PREF_ARMATURE_ONLY, 'Only Normalize groups that have matching bones in an armature (when an armature is used).'),\
 	]
 	
 	if not Draw.PupBlock('Clean Selected Meshes...', pup_block):
 		return
 	
-	actWeightNormalize(me, PREF_PEAKWEIGHT.val, PREF_KEEP_PROPORTION.val)
+	actWeightNormalize(me, ob, PREF_PEAKWEIGHT.val, PREF_ACTIVE_ONLY.val, PREF_ARMATURE_ONLY.val, PREF_KEEP_PROPORTION.val)
 	
 if __name__=='__main__':
 	main()
\ No newline at end of file