stgcn/graph.py

"""
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!")