utils/coordinate_system.py

# Coordinate system utilities for dwpose-editor

import numpy as np

class CoordinateTransformer:
    def __init__(self, data_resolution=(512, 512), display_resolution=(640, 640)):
        """
        座標変換システム
        
        Args:
            data_resolution: ポーズデータの解像度 (width, height)
            display_resolution: Canvas表示解像度 (width, height)
        """
        self.data_resolution = data_resolution
        self.display_resolution = display_resolution
        self.scale_x = display_resolution[0] / data_resolution[0]
        self.scale_y = display_resolution[1] / data_resolution[1]
    
    def data_to_display(self, x, y):
        """データ座標系から表示座標系に変換"""
        return {
            'x': x * self.scale_x,
            'y': y * self.scale_y
        }
    
    def display_to_data(self, x, y):
        """表示座標系からデータ座標系に変換"""
        return {
            'x': x / self.scale_x,
            'y': y / self.scale_y
        }
    
    def transform_pose_data(self, pose_data):
        """ポーズデータ全体を表示座標系に変換"""
        if not pose_data:
            return pose_data
        
        transformed_data = pose_data.copy()
        
        # ボディキーポイントの変換
        if 'bodies' in transformed_data and 'candidate' in transformed_data['bodies']:
            candidates = []
            for point in transformed_data['bodies']['candidate']:
                if len(point) >= 2:
                    transformed = self.data_to_display(point[0], point[1])
                    new_point = [transformed['x'], transformed['y']]
                    if len(point) > 2:
                        new_point.extend(point[2:])  # 信頼度などを保持
                    candidates.append(new_point)
                else:
                    candidates.append(point)
            transformed_data['bodies']['candidate'] = candidates
        
        # 手キーポイントの変換
        if 'hands' in transformed_data:
            transformed_hands = []
            for hand in transformed_data['hands']:
                if hand and len(hand) > 0:
                    transformed_hand = []
                    for i in range(0, len(hand), 3):
                        if i + 1 < len(hand):
                            transformed = self.data_to_display(hand[i], hand[i + 1])
                            transformed_hand.extend([transformed['x'], transformed['y']])
                            if i + 2 < len(hand):
                                transformed_hand.append(hand[i + 2])  # 信頼度
                        else:
                            transformed_hand.append(hand[i])
                    transformed_hands.append(transformed_hand)
                else:
                    transformed_hands.append(hand)
            transformed_data['hands'] = transformed_hands
        
        # 顔キーポイントの変換
        if 'faces' in transformed_data:
            transformed_faces = []
            for face in transformed_data['faces']:
                if face and len(face) > 0:
                    transformed_face = []
                    for i in range(0, len(face), 3):
                        if i + 1 < len(face):
                            transformed = self.data_to_display(face[i], face[i + 1])
                            transformed_face.extend([transformed['x'], transformed['y']])
                            if i + 2 < len(face):
                                transformed_face.append(face[i + 2])  # 信頼度
                        else:
                            transformed_face.append(face[i])
                    transformed_faces.append(transformed_face)
                else:
                    transformed_faces.append(face)
            transformed_data['faces'] = transformed_faces
        
        return transformed_data
    
    def update_resolution(self, data_resolution=None, display_resolution=None):
        """解像度設定の更新"""
        if data_resolution:
            self.data_resolution = data_resolution
        if display_resolution:
            self.display_resolution = display_resolution
        
        self.scale_x = self.display_resolution[0] / self.data_resolution[0]
        self.scale_y = self.display_resolution[1] / self.data_resolution[1]
    
    def get_scale_factors(self):
        """スケール係数を取得"""
        return {
            'scale_x': self.scale_x,
            'scale_y': self.scale_y,
            'data_resolution': self.data_resolution,
            'display_resolution': self.display_resolution
        }

# グローバル座標変換器
default_transformer = CoordinateTransformer()

def transform_point_to_display(x, y):
    """ポイントを表示座標系に変換"""
    return default_transformer.data_to_display(x, y)

def transform_point_to_data(x, y):
    """ポイントをデータ座標系に変換"""
    return default_transformer.display_to_data(x, y)

def transform_pose_to_display(pose_data):
    """ポーズデータを表示座標系に変換"""
    return default_transformer.transform_pose_data(pose_data)

def update_coordinate_system(data_resolution, display_resolution):
    """座標系設定を更新"""
    default_transformer.update_resolution(data_resolution, display_resolution)

def get_coordinate_info():
    """座標系情報を取得"""
    return default_transformer.get_scale_factors()