139 lines
5.5 KiB
Python
139 lines
5.5 KiB
Python
class DirectionNode:
|
|
def __init__(self):
|
|
self.model_path = None # Initialize to None
|
|
self.model = None
|
|
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
# Force recalculation on every change
|
|
return float("NaN")
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pose_kps": ("POSE_KEYPOINT",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "INT")
|
|
RETURN_NAMES = ("direction", "direction_code")
|
|
FUNCTION = "main"
|
|
|
|
CATEGORY = "OpenPose"
|
|
|
|
def main(self, pose_kps):
|
|
# Get the first person's keypoints (assuming there's at least one person)
|
|
if not pose_kps or len(pose_kps) == 0:
|
|
return ("missing keypoints", -1)
|
|
|
|
person = pose_kps[0]
|
|
|
|
# Extract face keypoints
|
|
if 'people' not in person or not person['people'] or 'face_keypoints_2d' not in person['people'][0]:
|
|
return ("missing keypoints", -1)
|
|
|
|
face_kps = person['people'][0]['face_keypoints_2d']
|
|
if len(face_kps) < 12: # Minimum required face keypoints
|
|
return ("missing keypoints", -1)
|
|
|
|
# Extract key face points (assuming 68-point model)
|
|
if len(face_kps) >= 68 * 3: # 68-point face model
|
|
# Get left and right eye points
|
|
left_eye_x = sum(face_kps[i*3] for i in range(36, 42)) / 6
|
|
right_eye_x = sum(face_kps[i*3] for i in range(42, 48)) / 6
|
|
# Get nose tip
|
|
nose_tip_x = face_kps[30*3] # Nose tip is point 30
|
|
# Get mouth corners
|
|
left_mouth_x = face_kps[48*3]
|
|
right_mouth_x = face_kps[54*3]
|
|
# Get jaw points
|
|
left_jaw_x = face_kps[0*3] # Left jaw point
|
|
right_jaw_x = face_kps[16*3] # Right jaw point
|
|
# Get eyebrow points
|
|
left_eyebrow_x = sum(face_kps[i*3] for i in range(17, 22)) / 5
|
|
right_eyebrow_x = sum(face_kps[i*3] for i in range(22, 27)) / 5
|
|
else: # Fallback for smaller face models
|
|
# Use first two pairs as eyes, next two as mouth corners
|
|
left_eye_x = face_kps[0]
|
|
right_eye_x = face_kps[3]
|
|
nose_tip_x = face_kps[6] # Assuming nose is point 6
|
|
left_mouth_x = face_kps[9]
|
|
right_mouth_x = face_kps[12]
|
|
left_jaw_x = face_kps[0] # Use eye point as jaw point
|
|
right_jaw_x = face_kps[3] # Use eye point as jaw point
|
|
left_eyebrow_x = face_kps[0] # Use eye point as eyebrow point
|
|
right_eyebrow_x = face_kps[3] # Use eye point as eyebrow point
|
|
|
|
# Calculate face symmetry metrics
|
|
eye_mid_x = (left_eye_x + right_eye_x) / 2
|
|
mouth_mid_x = (left_mouth_x + right_mouth_x) / 2
|
|
jaw_mid_x = (left_jaw_x + right_jaw_x) / 2
|
|
eyebrow_mid_x = (left_eyebrow_x + right_eyebrow_x) / 2
|
|
|
|
# Calculate face center using all features
|
|
face_center_x = (eye_mid_x + mouth_mid_x + jaw_mid_x + eyebrow_mid_x) / 4
|
|
|
|
# Calculate symmetry metrics for each feature
|
|
eye_symmetry = abs(nose_tip_x - eye_mid_x)
|
|
mouth_symmetry = abs(nose_tip_x - mouth_mid_x)
|
|
jaw_symmetry = abs(nose_tip_x - jaw_mid_x)
|
|
eyebrow_symmetry = abs(nose_tip_x - eyebrow_mid_x)
|
|
|
|
# Calculate feature distances
|
|
eye_distance = abs(left_eye_x - right_eye_x)
|
|
mouth_distance = abs(left_mouth_x - right_mouth_x)
|
|
jaw_distance = abs(left_jaw_x - right_jaw_x)
|
|
eyebrow_distance = abs(left_eyebrow_x - right_eyebrow_x)
|
|
|
|
# Calculate average feature distance for normalization
|
|
avg_feature_distance = (eye_distance + mouth_distance + jaw_distance + eyebrow_distance) / 4
|
|
|
|
# Calculate nose offset from face center
|
|
nose_offset = nose_tip_x - face_center_x
|
|
nose_offset_ratio = abs(nose_offset) / avg_feature_distance
|
|
|
|
# Calculate symmetry scores (0 to 1, where 1 is perfectly symmetric)
|
|
symmetry_scores = [
|
|
eye_symmetry / eye_distance,
|
|
mouth_symmetry / mouth_distance,
|
|
jaw_symmetry / jaw_distance,
|
|
eyebrow_symmetry / eyebrow_distance
|
|
]
|
|
avg_symmetry = sum(symmetry_scores) / len(symmetry_scores)
|
|
|
|
# Calculate feature alignment scores
|
|
feature_alignments = [
|
|
abs(eye_mid_x - face_center_x) / eye_distance,
|
|
abs(mouth_mid_x - face_center_x) / mouth_distance,
|
|
abs(jaw_mid_x - face_center_x) / jaw_distance,
|
|
abs(eyebrow_mid_x - face_center_x) / eyebrow_distance
|
|
]
|
|
avg_alignment = sum(feature_alignments) / len(feature_alignments)
|
|
|
|
# Thresholds based on example data analysis
|
|
symmetry_threshold = 0.15 # 15% deviation allowed for forward-facing
|
|
alignment_threshold = 0.12 # 12% deviation allowed for feature alignment
|
|
offset_threshold = 0.18 # 18% of average feature distance
|
|
|
|
# Check if face is symmetric and well-aligned (forward facing)
|
|
if (avg_symmetry < symmetry_threshold and
|
|
avg_alignment < alignment_threshold and
|
|
nose_offset_ratio < offset_threshold):
|
|
return ("forward", 0)
|
|
|
|
# Determine left/right based on nose position relative to face center
|
|
if nose_offset > 0:
|
|
return ("right", 2)
|
|
else:
|
|
return ("left", 1)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OpenPose - Get direction": DirectionNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DirectionNode": "OpenPose - Get direction"
|
|
}
|