MediaPipe를 사용한 포즈 추정

import cv2
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

# MediaPipe Pose 초기화 (MediaPipe 0.10.x 버전)
# MediaPipe 0.10.x에서는 새로운 tasks API를 사용합니다

import urllib.request
import os

# 모델 파일 저장 디렉토리 설정
model_dir = '../../.data'
model_path = os.path.join(model_dir, 'pose_landmarker.task')
model_url = 'https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/1/pose_landmarker_heavy.task'

# 디렉토리가 없으면 생성
os.makedirs(model_dir, exist_ok=True)

# 모델 파일 다운로드 (없는 경우에만)
if not os.path.exists(model_path):
    print(f"모델 파일을 다운로드 중입니다: {model_url}")
    urllib.request.urlretrieve(model_url, model_path)
    print(f"다운로드 완료! 저장 위치: {model_path}")
else:
    print(f"모델 파일이 이미 존재합니다: {model_path}")

# PoseLandmarker 옵션 설정
base_options = python.BaseOptions(
    model_asset_path=model_path  # 다운로드한 모델 파일 경로
)

pose_landmarker_options = vision.PoseLandmarkerOptions(
    base_options=base_options,
    output_segmentation_masks=False,  # 배경 분할 사용 여부
    min_pose_detection_confidence=0.5,  # 최소 감지 신뢰도
    min_pose_presence_confidence=0.5,  # 최소 포즈 존재 신뢰도
    min_tracking_confidence=0.5,  # 최소 추적 신뢰도
    num_poses=1  # 감지할 최대 포즈 수
)

# PoseLandmarker 생성
pose_landmarker = vision.PoseLandmarker.create_from_options(pose_landmarker_options)
print("PoseLandmarker 초기화 완료!")

def detect_pose(image_path):
    """
    이미지에서 포즈를 추정하는 함수 (MediaPipe 0.10.x 버전)
    
    Args:
        image_path: 이미지 파일 경로
        
    Returns:
        image: 원본 이미지 (BGR)
        mp_image: MediaPipe Image 객체
        detection_result: MediaPipe 포즈 추정 결과
    """
    # 이미지 읽기
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"이미지를 읽을 수 없습니다: {image_path}")
    
    # BGR을 RGB로 변환
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # MediaPipe Image 객체 생성
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image_rgb)
    
    # 포즈 추정 수행
    detection_result = pose_landmarker.detect(mp_image)
    
    return image, mp_image, detection_result

# 사용 예시 (이미지 경로를 실제 경로로 변경하세요)
# image, mp_image, detection_result = detect_pose('path/to/your/image.jpg')

def draw_pose_landmarks(image, detection_result):
    """
    포즈 랜드마크를 이미지에 그리는 함수 (MediaPipe 0.10.x 버전)
    
    Args:
        image: 원본 이미지 (BGR)
        detection_result: MediaPipe 포즈 추정 결과
        
    Returns:
        annotated_image: 랜드마크가 그려진 이미지
    """
    # 이미지 복사
    annotated_image = image.copy()
    
    # RGB로 변환
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    
    # 포즈 랜드마크 그리기
    if detection_result.pose_landmarks:
        for pose_landmarks in detection_result.pose_landmarks:
            # 랜드마크 연결선 그리기
            connections = vision.PoseLandmarksConnections.POSE_LANDMARKS
            for connection in connections:
                start_idx = connection.start
                end_idx = connection.end
                if start_idx < len(pose_landmarks) and end_idx < len(pose_landmarks):
                    start_point = pose_landmarks[start_idx]
                    end_point = pose_landmarks[end_idx]
                    start_x = int(start_point.x * annotated_image_rgb.shape[1])
                    start_y = int(start_point.y * annotated_image_rgb.shape[0])
                    end_x = int(end_point.x * annotated_image_rgb.shape[1])
                    end_y = int(end_point.y * annotated_image_rgb.shape[0])
                    cv2.line(annotated_image_rgb, (start_x, start_y), (end_x, end_y), (0, 255, 0), 2)
            
            # 랜드마크 포인트 그리기
            for landmark in pose_landmarks:
                x = int(landmark.x * annotated_image_rgb.shape[1])
                y = int(landmark.y * annotated_image_rgb.shape[0])
                cv2.circle(annotated_image_rgb, (x, y), 5, (0, 0, 255), -1)
    
    # RGB를 BGR로 다시 변환
    annotated_image = cv2.cvtColor(annotated_image_rgb, cv2.COLOR_RGB2BGR)
    
    return annotated_image

def get_landmark_coordinates(detection_result, image_shape):
    """
    랜드마크 좌표를 픽셀 좌표로 변환하는 함수 (MediaPipe 0.10.x 버전)
    
    Args:
        detection_result: MediaPipe 포즈 추정 결과
        image_shape: 이미지 크기 (height, width)
        
    Returns:
        landmarks: 랜드마크 좌표 딕셔너리 {landmark_name: (x, y, z)}
    """
    landmarks = {}
    
    if detection_result.pose_landmarks:
        height, width = image_shape[:2]
        
        # MediaPipe 랜드마크 인덱스 정의 (33개 포인트)
        landmark_names = {
            0: 'NOSE',
            2: 'LEFT_EYE_INNER', 4: 'LEFT_EYE', 5: 'LEFT_EYE_OUTER',
            7: 'RIGHT_EYE_INNER', 9: 'RIGHT_EYE', 10: 'RIGHT_EYE_OUTER',
            11: 'LEFT_EAR', 12: 'RIGHT_EAR',
            13: 'MOUTH_LEFT', 14: 'MOUTH_RIGHT',
            15: 'LEFT_SHOULDER', 16: 'RIGHT_SHOULDER',
            17: 'LEFT_ELBOW', 18: 'RIGHT_ELBOW',
            19: 'LEFT_WRIST', 20: 'RIGHT_WRIST',
            21: 'LEFT_PINKY', 22: 'RIGHT_PINKY',
            23: 'LEFT_INDEX', 24: 'RIGHT_INDEX',
            25: 'LEFT_THUMB', 26: 'RIGHT_THUMB',
            27: 'LEFT_HIP', 28: 'RIGHT_HIP',
            29: 'LEFT_KNEE', 30: 'RIGHT_KNEE',
            31: 'LEFT_ANKLE', 32: 'RIGHT_ANKLE',
            33: 'LEFT_HEEL', 34: 'RIGHT_HEEL',
            35: 'LEFT_FOOT_INDEX', 36: 'RIGHT_FOOT_INDEX'
        }
        
        # 첫 번째 포즈의 랜드마크 사용
        pose_landmarks = detection_result.pose_landmarks[0]
        
        for idx, landmark in enumerate(pose_landmarks):
            # 정규화된 좌표를 픽셀 좌표로 변환
            x = int(landmark.x * width)
            y = int(landmark.y * height)
            z = landmark.z  # 깊이 정보 (상대적)
            
            if idx in landmark_names:
                landmarks[landmark_names[idx]] = (x, y, z)
            else:
                landmarks[f'LANDMARK_{idx}'] = (x, y, z)
    
    return landmarks

def process_image(image_path):
    """
    이미지에서 포즈를 추정하고 결과를 시각화하는 완전한 함수 (MediaPipe 0.10.x 버전)
    
    Args:
        image_path: 입력 이미지 경로
        
    Returns:
        annotated_image: 랜드마크가 그려진 이미지
        landmarks: 랜드마크 좌표 딕셔너리
    """
    # 이미지 읽기 및 포즈 추정
    image, mp_image, detection_result = detect_pose(image_path)
    
    # 포즈가 감지되었는지 확인
    if detection_result.pose_landmarks:
        print("포즈가 성공적으로 감지되었습니다!")
        
        # 랜드마크 그리기
        annotated_image = draw_pose_landmarks(image, detection_result)
        
        # 랜드마크 좌표 추출
        landmarks = get_landmark_coordinates(detection_result, image.shape)
        
        # 주요 관절 좌표 출력
        key_points = ['LEFT_SHOULDER', 'RIGHT_SHOULDER', 'LEFT_ELBOW', 
                     'RIGHT_ELBOW', 'LEFT_WRIST', 'RIGHT_WRIST',
                     'LEFT_HIP', 'RIGHT_HIP', 'LEFT_KNEE', 'RIGHT_KNEE',
                     'LEFT_ANKLE', 'RIGHT_ANKLE']
        
        print("\n주요 관절 좌표:")
        for point in key_points:
            if point in landmarks:
                x, y, z = landmarks[point]
                print(f"{point}: ({x}, {y}, z={z:.3f})")
        
        return annotated_image, landmarks
    else:
        print("포즈를 감지할 수 없습니다.")
        return image, None

# 사용 예시
# annotated_image, landmarks = process_image('input.jpg', 'output.jpg')

def visualize_pose_comparison(image_path):
    """
    원본 이미지와 포즈 추정 결과를 나란히 비교하는 함수 (MediaPipe 0.10.x 버전)
    
    Args:
        image_path: 이미지 파일 경로
    """
    # 이미지 읽기 및 포즈 추정
    image, mp_image, detection_result = detect_pose(image_path)
    
    # 포즈 랜드마크 그리기
    annotated_image = draw_pose_landmarks(image, detection_result)
    
    # Matplotlib으로 표시
    fig, axes = plt.subplots(1, 2, figsize=(15, 7))
    
    # 원본 이미지
    axes[0].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    axes[0].set_title('원본 이미지', fontsize=14)
    axes[0].axis('off')
    
    # 포즈 추정 결과
    axes[1].imshow(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
    axes[1].set_title('포즈 추정 결과', fontsize=14)
    axes[1].axis('off')
    
    plt.tight_layout()
    plt.show()
    
    return annotated_image, detection_result

# 사용 예시
# visualize_pose_comparison('path/to/your/image.jpg')

# process_image 함수를 사용한 간단한 방법
image_path = '../../.data/man.jpg'
annotated_image, landmarks = process_image(image_path)

# 결과 시각화
if landmarks:
    fig, axes = plt.subplots(1, 2, figsize=(16, 8))
    
    # 원본 이미지 읽기
    original = cv2.imread(image_path)
    axes[0].imshow(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))
    axes[0].set_title('Original Image', fontsize=16, fontweight='bold')
    axes[0].axis('off')
    
    # 포즈 추정 결과
    axes[1].imshow(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
    axes[1].set_title('Pose Estimation Result (Skeleton)', fontsize=16, fontweight='bold')
    axes[1].axis('off')
    
    plt.tight_layout()
    plt.show()

# 리소스 정리 (MediaPipe 0.10.x 버전)
pose_landmarker.close()