Segmentation Mask with MediaPipe in TouchDesigner

The tutorial is an updated version of the MediaPipe Pose using the new segmentation mask function to identify the human body tracked. It used the Script TOP to generate the mask image. Users can further enhance the image with the Threshold TOP for display purpose. Similar to the previous tutorials, it assumes the installation of Python 3.7 and the MediaPipe library through Pip.

The source TouchDesigner project file is available in my TouchDesigner GitHub repository. The Python code is relatively straighforward. The pose tracking results will include an array (segmentation_mask) of the size of the tracked image. Each pixel will have a value between 0.0 to 1.0. Darker value will be the background while brighter value will likely be the tracked body. Here is the full listing.

# me - this DAT
# scriptOp - the OP which is cooking

import numpy as np
import cv2
import mediapipe as mp

mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

pose = mp_pose.Pose(
	min_detection_confidence=0.5,
	min_tracking_confidence=0.5,
	enable_segmentation=True
)

def onSetupParameters(scriptOp):
	return

# called whenever custom pulse parameter is pushed
def onPulse(par):
	return

def onCook(scriptOp):
	input = scriptOp.inputs[0].numpyArray(delayed=True)
	if input is not None:
		image = cv2.cvtColor(input, cv2.COLOR_RGBA2RGB)
		image *= 255
		image = image.astype('uint8')
		results = pose.process(image)
		
		if results.segmentation_mask is not None:
			rgb = cv2.cvtColor(results.segmentation_mask, cv2.COLOR_GRAY2RGB)
			rgb = rgb * 255
			rgb = rgb.astype(np.uint8)
			scriptOp.copyNumpyArray(rgb)
		else:
			black = np.zeros(image.shape, dtype=np.uint8)
			scriptOp.copyNumpyArray(black)
	return