diff --git a/nodes_utility.py b/nodes_utility.py index 76fc512..14f9090 100644 --- a/nodes_utility.py +++ b/nodes_utility.py @@ -500,6 +500,12 @@ class FaceMaskFromPoseKeypoints: for i, pose_frame in enumerate(pose_frames): selected_idx, prev_center = self.select_closest_person(pose_frame, person_index if i == 0 else prev_center) np_frames.append(self.draw_kps(pose_frame, selected_idx)) + + if not np_frames: + # Handle case where no frames were processed + log.warning("No valid pose frames found, returning empty mask") + return (torch.zeros((1, 64, 64), dtype=torch.float32),) + np_frames = np.stack(np_frames, axis=0) tensor = torch.from_numpy(np_frames).float() / 255. print("tensor.shape:", tensor.shape) @@ -510,39 +516,149 @@ class FaceMaskFromPoseKeypoints: people = pose_frame["people"] if not people: return -1, None + centers = [] - for person in people: + valid_people_indices = [] + + for idx, person in enumerate(people): + # Check if face keypoints exist and are valid + if "face_keypoints_2d" not in person or not person["face_keypoints_2d"]: + continue + kps = np.array(person["face_keypoints_2d"]) + if len(kps) == 0: + continue + n = len(kps) // 3 + if n == 0: + continue + facial_kps = rearrange(kps, "(n c) -> n c", n=n, c=3)[:, :2] + + # Check if we have valid coordinates (not all zeros) + if np.all(facial_kps == 0): + continue + center = facial_kps.mean(axis=0) + + # Check if center is valid (not NaN or infinite) + if np.isnan(center).any() or np.isinf(center).any(): + continue + centers.append(center) + valid_people_indices.append(idx) + + if not centers: + return -1, None + if isinstance(prev_center_or_index, (int, np.integer)): - # First frame: use person_index - idx = prev_center_or_index if 0 <= prev_center_or_index < len(people) else 0 - return idx, centers[idx] - else: + # First frame: use person_index, but map to valid people + if 0 <= prev_center_or_index < len(valid_people_indices): + idx = valid_people_indices[prev_center_or_index] + return idx, centers[prev_center_or_index] + elif valid_people_indices: + # Fallback to first valid person + idx = valid_people_indices[0] + return idx, centers[0] + else: + return -1, None + elif prev_center_or_index is not None: # Find closest to previous center prev_center = np.array(prev_center_or_index) dists = [np.linalg.norm(center - prev_center) for center in centers] - idx = int(np.argmin(dists)) - return idx, centers[idx] + min_idx = int(np.argmin(dists)) + actual_idx = valid_people_indices[min_idx] + return actual_idx, centers[min_idx] + else: + # prev_center_or_index is None, fallback to first valid person + if valid_people_indices: + idx = valid_people_indices[0] + return idx, centers[0] + else: + return -1, None def draw_kps(self, pose_frame, person_index): import cv2 width, height = pose_frame["canvas_width"], pose_frame["canvas_height"] canvas = np.zeros((height, width, 3), dtype=np.uint8) people = pose_frame["people"] + if person_index < 0 or person_index >= len(people): return canvas # Out of bounds, return blank + person = people[person_index] - n = len(person["face_keypoints_2d"]) // 3 - facial_kps = rearrange(np.array(person["face_keypoints_2d"]), "(n c) -> n c", n=n, c=3)[:, :2] + + # Check if face keypoints exist and are valid + if "face_keypoints_2d" not in person or not person["face_keypoints_2d"]: + return canvas # No face keypoints, return blank + + face_kps_data = person["face_keypoints_2d"] + if len(face_kps_data) == 0: + return canvas # Empty keypoints, return blank + + n = len(face_kps_data) // 3 + if n < 17: # Need at least 17 points for outer contour + return canvas # Not enough keypoints, return blank + + facial_kps = rearrange(np.array(face_kps_data), "(n c) -> n c", n=n, c=3)[:, :2] + + # Check if we have valid coordinates (not all zeros) + if np.all(facial_kps == 0): + return canvas # All keypoints are zero, return blank + + # Check for NaN or infinite values + if np.isnan(facial_kps).any() or np.isinf(facial_kps).any(): + return canvas # Invalid coordinates, return blank + + # Check for negative coordinates or coordinates that would create streaks + if np.any(facial_kps < 0): + return canvas # Negative coordinates, likely bad detection + + # Check if coordinates are reasonable (not too close to edges which might indicate bad detection) + min_margin = 5 # Minimum distance from edges + if (np.any(facial_kps[:, 0] < min_margin) or + np.any(facial_kps[:, 1] < min_margin) or + np.any(facial_kps[:, 0] > width - min_margin) or + np.any(facial_kps[:, 1] > height - min_margin)): + # Check if this looks like a streak to corner (many points near 0,0) + corner_points = np.sum((facial_kps[:, 0] < min_margin) & (facial_kps[:, 1] < min_margin)) + if corner_points > 3: # Too many points near corner, likely bad detection + return canvas + facial_kps = facial_kps.astype(np.int32) + + # Ensure coordinates are within canvas bounds + facial_kps[:, 0] = np.clip(facial_kps[:, 0], 0, width - 1) + facial_kps[:, 1] = np.clip(facial_kps[:, 1], 0, height - 1) + part_color = (255, 255, 255) - outer_contour = facial_kps[:17] - cv2.fillPoly(canvas, pts=[outer_contour], color=part_color) + + # Additional validation for the contour before drawing + # Check if contour points are too spread out (indicating bad detection) + if len(outer_contour) >= 3: + # Calculate bounding box of the contour + min_x, min_y = np.min(outer_contour, axis=0) + max_x, max_y = np.max(outer_contour, axis=0) + contour_width = max_x - min_x + contour_height = max_y - min_y + + # If contour spans more than 80% of canvas, likely bad detection + if (contour_width > 0.8 * width or contour_height > 0.8 * height): + return canvas + + # Check if we have a valid contour (at least 3 unique points) + unique_points = np.unique(outer_contour, axis=0) + if len(unique_points) >= 3: + # Final check: ensure the contour is reasonable + # Calculate area to see if it's too large or too small + contour_area = cv2.contourArea(outer_contour) + canvas_area = width * height + + # If contour is less than 0.1% or more than 50% of canvas, skip + if 0.001 * canvas_area <= contour_area <= 0.5 * canvas_area: + cv2.fillPoly(canvas, pts=[outer_contour], color=part_color) + return canvas NODE_CLASS_MAPPINGS = {