""" COCO Skeleton Graph Definition for ST-GCN This module defines the skeleton graph structure for COCO 17-keypoint format used by YOLOv11-Pose. The graph represents spatial relationships between joints as an adjacency matrix for Spatial-Temporal Graph Convolutional Networks. COCO 17 Keypoints: 0: nose, 1: left_eye, 2: right_eye, 3: left_ear, 4: right_ear 5: left_shoulder, 6: right_shoulder, 7: left_elbow, 8: right_elbow 9: left_wrist, 10: right_wrist, 11: left_hip, 12: right_hip 13: left_knee, 14: right_knee, 15: left_ankle, 16: right_ankle References: - ST-GCN Paper: https://arxiv.org/abs/1801.07455 - COCO Dataset: https://cocodataset.org/#keypoints-2020 """ import numpy as np class Graph: """COCO skeleton graph for ST-GCN.""" def __init__(self, labeling_mode='spatial'): """ Initialize COCO skeleton graph. Args: labeling_mode: Partitioning strategy for skeleton graph - 'spatial': Partition based on spatial distance from center - 'uniform': All edges treated equally (baseline) """ self.num_nodes = 17 # COCO keypoints self.labeling_mode = labeling_mode # Define skeleton connectivity (parent-child relationships) self.edges = self._get_edges() # Create adjacency matrix self.A = self._create_adjacency_matrix() # Get partitioning strategy self.A_with_partitions = self._get_partitioned_adjacency() def _get_edges(self): """ Define COCO skeleton edges (connections between keypoints). Returns: List of tuples representing connected joints """ # COCO skeleton structure (17 keypoints) edges = [ # Head connections (0, 1), (0, 2), # nose to eyes (1, 3), (2, 4), # eyes to ears # Torso connections (5, 6), # shoulders (5, 11), (6, 12), # shoulders to hips (11, 12), # hips # Left arm (5, 7), (7, 9), # shoulder -> elbow -> wrist # Right arm (6, 8), (8, 10), # shoulder -> elbow -> wrist # Left leg (11, 13), (13, 15), # hip -> knee -> ankle # Right leg (12, 14), (14, 16), # hip -> knee -> ankle ] return edges def _create_adjacency_matrix(self): """ Create adjacency matrix from skeleton edges. Returns: A: (V, V) adjacency matrix where V=17 (number of keypoints) """ A = np.zeros((self.num_nodes, self.num_nodes)) # Add edges (bidirectional connections) for i, j in self.edges: A[i, j] = 1 A[j, i] = 1 # Add self-connections A += np.eye(self.num_nodes) return A def _get_partitioned_adjacency(self): """ Partition adjacency matrix based on labeling strategy. For spatial labeling, partitions are: - Partition 0: Self-connections (centripetal group) - Partition 1: Joints closer to skeleton center (centripetal group) - Partition 2: Joints farther from skeleton center (centrifugal group) Returns: A_partitioned: (num_partitions, V, V) stacked adjacency matrices """ if self.labeling_mode == 'uniform': # Uniform labeling: all edges treated equally return self.A[np.newaxis, :, :] elif self.labeling_mode == 'spatial': # Spatial labeling: partition based on distance from center # Center joint is defined as the midpoint between shoulders (joints 5, 6) center_joints = [5, 6] # Left and right shoulders # Initialize partition matrices A_partitions = [] # Partition 0: Self-connections A_self = np.eye(self.num_nodes) A_partitions.append(A_self) # Partition 1: Centripetal (moving toward center) # Partition 2: Centrifugal (moving away from center) A_centripetal = np.zeros((self.num_nodes, self.num_nodes)) A_centrifugal = np.zeros((self.num_nodes, self.num_nodes)) # Compute distances from center for each joint distances = self._compute_center_distances(center_joints) # Classify edges based on distance change (both directions) for i, j in self.edges: if distances[j] < distances[i]: # Moving toward center (j is closer than i) A_centripetal[i, j] = 1 # Reverse direction: moving away from center A_centrifugal[j, i] = 1 elif distances[j] > distances[i]: # Moving away from center (j is farther than i) A_centrifugal[i, j] = 1 # Reverse direction: moving toward center A_centripetal[j, i] = 1 else: # Same distance: treat as centripetal for both directions A_centripetal[i, j] = 1 A_centripetal[j, i] = 1 A_partitions.append(A_centripetal) A_partitions.append(A_centrifugal) # Stack partitions: (3, V, V) A_partitioned = np.stack(A_partitions, axis=0) return A_partitioned else: raise ValueError(f"Unknown labeling mode: {self.labeling_mode}") def _compute_center_distances(self, center_joints): """ Compute graph distance from center joints to all other joints. Uses BFS to compute shortest path distance in graph. Args: center_joints: List of joint indices considered as center Returns: distances: (V,) array of distances from center """ from collections import deque distances = np.full(self.num_nodes, np.inf) queue = deque() # Initialize center joints with distance 0 for joint in center_joints: distances[joint] = 0 queue.append(joint) # BFS to compute distances while queue: current = queue.popleft() current_dist = distances[current] # Check all neighbors for neighbor in range(self.num_nodes): if self.A[current, neighbor] > 0 and neighbor != current: if distances[neighbor] > current_dist + 1: distances[neighbor] = current_dist + 1 queue.append(neighbor) return distances def get_adjacency_matrix(self, normalize=True): """ Get normalized adjacency matrix for ST-GCN. Args: normalize: Whether to apply symmetric normalization (D^-0.5 * A * D^-0.5) Returns: A_normalized: Normalized adjacency matrix """ if self.labeling_mode == 'spatial': # Return partitioned adjacency matrices A = self.A_with_partitions if normalize: # Normalize each partition separately A_normalized = [] for partition in A: A_norm = self._normalize_adjacency(partition) A_normalized.append(A_norm) return np.stack(A_normalized, axis=0) else: return A else: # Return single adjacency matrix A = self.A[np.newaxis, :, :] if normalize: A_norm = self._normalize_adjacency(A[0]) return A_norm[np.newaxis, :, :] else: return A def _normalize_adjacency(self, A): """ Apply symmetric normalization: D^-0.5 * A * D^-0.5 Args: A: (V, V) adjacency matrix Returns: A_normalized: (V, V) normalized adjacency matrix """ # Compute degree matrix D = np.sum(A, axis=1) # Avoid division by zero D[D == 0] = 1 # Compute D^-0.5 D_inv_sqrt = np.power(D, -0.5) # Apply normalization: D^-0.5 * A * D^-0.5 A_normalized = A * D_inv_sqrt[:, np.newaxis] * D_inv_sqrt[np.newaxis, :] return A_normalized def get_coco_graph(labeling_mode='spatial'): """ Convenience function to get COCO skeleton graph. Args: labeling_mode: Partitioning strategy ('spatial' or 'uniform') Returns: Graph object with COCO skeleton structure """ return Graph(labeling_mode=labeling_mode) if __name__ == '__main__': # Test graph construction print("Testing COCO Skeleton Graph...") # Test uniform labeling graph_uniform = Graph(labeling_mode='uniform') print(f"\nUniform labeling:") print(f" Adjacency shape: {graph_uniform.A.shape}") print(f" Partitions shape: {graph_uniform.A_with_partitions.shape}") print(f" Number of edges: {len(graph_uniform.edges)}") # Test spatial labeling graph_spatial = Graph(labeling_mode='spatial') print(f"\nSpatial labeling:") print(f" Adjacency shape: {graph_spatial.A.shape}") print(f" Partitions shape: {graph_spatial.A_with_partitions.shape}") # Get normalized adjacency A_norm = graph_spatial.get_adjacency_matrix(normalize=True) print(f"\nNormalized adjacency shape: {A_norm.shape}") print("\nCOCO skeleton graph construction successful!")