utils/pose_utils.py

# Pose processing utilities for dwpose-editor

import gradio as gr
from .dwpose_manager import DWPoseManager
from .dwpose_detector import DWPoseDetector
from .error_handler import handle_error, ModelLoadError, PoseDetectionError, safe_execute
from .notifications import notify_success, notify_warning, notify_error, notify_progress, NotificationMessages

# グローバルなDWPoseインスタンス
dwpose_manager = None
dwpose_detector = None

def initialize_dwpose():
    """DWPoseモデルを初期化"""
    global dwpose_manager, dwpose_detector
    
    def _init_process():
        global dwpose_manager, dwpose_detector
        dwpose_manager = DWPoseManager()
        success, message = dwpose_manager.initialize()
        
        if success:
            dwpose_detector = DWPoseDetector(dwpose_manager)
            notify_success(NotificationMessages.MODEL_LOADED)
            return True, "DWPoseモデル初期化完了"
        else:
            raise ModelLoadError(message)
    
    try:
        return safe_execute(
            _init_process,
            "DWPoseモデルの初期化に失敗しました",
            show_error=False
        ) or (False, "初期化処理でエラーが発生しました")
        
    except ModelLoadError as e:
        return False, str(e)
    except Exception as e:
        return False, f"予期しないエラー: {str(e)}"

def safe_detect_pose(image):
    """安全なポーズ検出（Gradio用）"""
    global dwpose_detector
    
    def _detection_process():
        if dwpose_detector is None:
            raise PoseDetectionError("DWPoseモデルが初期化されていません")
        
        if image is None:
            raise PoseDetectionError("画像が選択されていません")
        
        pose_data, error = dwpose_detector.detect(image)
        if error:
            raise PoseDetectionError(error)
        
        return pose_data
    
    try:
        result = safe_execute(
            _detection_process,
            "ポーズ検出に失敗しました",
            show_error=False
        )
        
        if result is not None:
            notify_success(NotificationMessages.POSE_DETECTED)
            return result
        else:
            return None
            
    except PoseDetectionError as e:
        handle_error(e)
        return None
    except Exception as e:
        handle_error(e)
        return None

def safe_detect_pose_with_progress(image):
    """進捗表示付きポーズ検出"""
    global dwpose_detector
    
    try:
        if dwpose_detector is None:
            notify_error("DWPoseモデルが初期化されていません")
            return None
        
        if image is None:
            notify_error("画像が選択されていません")
            return None
        
        # 進捗付きでポーズ検出実行
        notify_progress("画像を処理中...", 0.1)
        
        # 画像前処理
        notify_progress("画像前処理中...", 0.2)
        
        # 人物検出
        notify_progress("人物を検出中...", 0.4)
        
        # ポーズ推定
        notify_progress("ポーズを解析中...", 0.7)
        
        # 実際の検出処理
        pose_data, error = dwpose_detector.detect(image)
        
        if error:
            notify_error(error)
            return None
        
        notify_progress("完了", 1.0)
        notify_success(NotificationMessages.POSE_DETECTED)
        return pose_data
        
    except Exception as e:
        notify_error(f"ポーズ検出中にエラーが発生しました: {str(e)}")
        return None