MetaWhiz/scripts/Main.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Tool - Maya MetaHuman Plugin
Main entry point for the Tool plugin

This module initializes the Tool plugin and provides the main UI and functionality
for working with MetaHuman models in Maya, including DNA editing, calibration,
custom binding, and more.
"""

import os
import sys
import json
import traceback
import platform
import maya.cmds as cmds
import maya.mel as mel
import maya.OpenMayaUI as omui

# Detect Maya version and operating system
MAYA_VERSION = int(cmds.about(version=True))
OS_NAME = platform.system()

# Import configuration
try:
    from config import TOOL_NAME, TOOL_VERSION, TOOL_AUTHOR, SCRIPTS_PATH, ASSETS_PATH, ICONS_PATH, PLUGINS_PATH, DNA_FILE_PATH, DNA_IMG_PATH, PYTHON_VERSION
except ImportError:
    # If unable to import config, use default paths
    TOOL_NAME = "Tool"
    TOOL_VERSION = "1.0.0"
    TOOL_AUTHOR = "Tool Team"
    
    # Get script directory
    METAHUMAN_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    SCRIPTS_PATH = os.path.join(METAHUMAN_PATH, "scripts")
    ASSETS_PATH = os.path.join(METAHUMAN_PATH, "assets")
    ICONS_PATH = os.path.join(METAHUMAN_PATH, "icons")
    PLUGINS_PATH = os.path.join(METAHUMAN_PATH, "plugins")
    DNA_FILE_PATH = os.path.join(ASSETS_PATH, "dna")
    DNA_IMG_PATH = os.path.join(ASSETS_PATH, "images")
    
    # Maya version information
    PYTHON_VERSION = platform.python_version()

# Add scripts directory to Python path
if SCRIPTS_PATH not in sys.path:
    sys.path.append(SCRIPTS_PATH)

# Select appropriate PySide version based on Maya version
QT_VERSION = None

# Maya 2025+ uses PySide6
if MAYA_VERSION >= 2025:
    try:
        from PySide6 import QtCore, QtGui, QtWidgets
        from shiboken6 import wrapInstance
        QT_VERSION = "PySide6"
        print(f"Using PySide6 (Maya {MAYA_VERSION})")
    except ImportError as e:
        print(f"PySide6 import error: {str(e)}")
        try:
            from PySide2 import QtCore, QtGui, QtWidgets
            from shiboken2 import wrapInstance
            QT_VERSION = "PySide2"
            print(f"Using PySide2 (fallback mode, Maya {MAYA_VERSION})")
        except ImportError:
            cmds.error("Unable to import PySide6 or PySide2, please ensure they are installed")

# Maya 2022-2024 uses PySide2
elif MAYA_VERSION >= 2022:
    try:
        from PySide2 import QtCore, QtGui, QtWidgets
        from shiboken2 import wrapInstance
        QT_VERSION = "PySide2"
        print(f"Using PySide2 (Maya {MAYA_VERSION})")
    except ImportError as e:
        print(f"PySide2 import error: {str(e)}")
        try:
            # Try to fallback to PySide6 (some special installations may have PySide6)
            from PySide6 import QtCore, QtGui, QtWidgets
            from shiboken6 import wrapInstance
            QT_VERSION = "PySide6"
            print(f"Using PySide6 (non-standard configuration, Maya {MAYA_VERSION})")
        except ImportError:
            cmds.error("Unable to import PySide2 or PySide6, please ensure they are installed")

# Maya 2020-2021 uses PySide2 or PySide
else:
    try:
        from PySide2 import QtCore, QtGui, QtWidgets
        from shiboken2 import wrapInstance
        QT_VERSION = "PySide2"
        print(f"Using PySide2 (Maya {MAYA_VERSION})")
    except ImportError:
        try:
            from PySide import QtCore, QtGui, QtWidgets
            from shiboken import wrapInstance
            QT_VERSION = "PySide"
            print(f"Using PySide (Maya {MAYA_VERSION})")
        except ImportError:
            cmds.error("Unable to import PySide or PySide2, please ensure they are installed")

# Import UI modules
try:
    from ui.main_window import ToolWindow
    from ui.dna_editor import DNAEditorWidget
    from ui.calibration import CalibrationWidget
    from ui.binding import BindingWidget
    from ui.blendshape import BlendShapeWidget
except ImportError as e:
    cmds.error(f"Error importing UI modules: {str(e)}\n{traceback.format_exc()}")

# Import utility modules
try:
    from utils.dna_utils import DNAUtils
    from utils.maya_utils import MayaUtils
    from utils.mesh_utils import MeshUtils
    from utils.joint_utils import JointUtils
    from utils.blendshape_utils import BlendShapeUtils
except ImportError as e:
    cmds.error(f"Error importing utility modules: {str(e)}\n{traceback.format_exc()}")

class Tool:
    """Main Tool class that initializes and manages the plugin"""
    
    _instance = None
    
    @classmethod
    def instance(cls):
        """Singleton pattern to ensure only one instance of Tool exists"""
        if cls._instance is None:
            cls._instance = Tool()
        return cls._instance
    
    def __init__(self):
        """Initialize the Tool plugin"""
        self.window = None
        self.dna_utils = DNAUtils()
        self.maya_utils = MayaUtils()
        self.mesh_utils = MeshUtils()
        self.joint_utils = JointUtils()
        self.blendshape_utils = BlendShapeUtils()
        
        # Initialize paths
        self._init_paths()
        
        # Load plugin dependencies
        self._load_plugins()
        
        print(f"{TOOL_NAME} {TOOL_VERSION} initialized")
    
    def _init_paths(self):
        """Initialize necessary paths and directories"""
        # Ensure all required directories exist
        paths = [
            SCRIPTS_PATH,
            ASSETS_PATH,
            ICONS_PATH,
            PLUGINS_PATH,
            DNA_FILE_PATH,
            DNA_IMG_PATH
        ]
        
        for path in paths:
            if not os.path.exists(path):
                try:
                    os.makedirs(path)
                    print(f"Created directory: {path}")
                except Exception as e:
                    print(f"Error creating directory: {path}, Error: {str(e)}")
    
    def _load_plugins(self):
        """Load required Maya plugins"""
        required_plugins = ["fbxmaya", "matrixNodes"]
        
        for plugin in required_plugins:
            if not cmds.pluginInfo(plugin, query=True, loaded=True):
                try:
                    cmds.loadPlugin(plugin)
                    print(f"Loaded plugin: {plugin}")
                except Exception as e:
                    print(f"Error loading plugin: {plugin}, Error: {str(e)}")
    
    def show_ui(self):
        """Show the main Tool UI window"""
        if self.window is None or not self.window.isVisible():
            self.window = ToolWindow()
        
        self.window.show()
        self.window.raise_()
        return self.window
    
    def create_metahuman_rig(self, mesh=None):
        """Create a MetaHuman rig from a mesh"""
        if mesh is None:
            # Get the selected mesh
            selection = cmds.ls(selection=True)
            if not selection:
                cmds.warning("No mesh selected for rigging")
                return False
            mesh = selection[0]
        
        # Validate the mesh
        if not self.mesh_utils.validate_metahuman_topology(mesh):
            cmds.warning(f"Selected mesh {mesh} does not have MetaHuman topology")
            return False
        
        # Create the rig
        try:
            result = self.joint_utils.create_metahuman_skeleton(mesh)
            if result:
                print(f"Successfully created MetaHuman rig for {mesh}")
                return True
            else:
                cmds.warning(f"Failed to create MetaHuman rig for {mesh}")
                return False
        except Exception as e:
            cmds.error(f"Error creating MetaHuman rig: {str(e)}")
            return False
    
    def load_dna(self, dna_file_path):
        """Load a DNA file into the scene"""
        if not os.path.exists(dna_file_path):
            cmds.warning(f"DNA file does not exist: {dna_file_path}")
            return False
        
        try:
            result = self.dna_utils.load_dna(dna_file_path)
            if result:
                print(f"Successfully loaded DNA file: {dna_file_path}")
                return True
            else:
                cmds.warning(f"Failed to load DNA file: {dna_file_path}")
                return False
        except Exception as e:
            cmds.error(f"Error loading DNA file: {str(e)}")
            return False
    
    def save_dna(self, output_path):
        """Save the current scene as a DNA file"""
        try:
            result = self.dna_utils.save_dna(output_path)
            if result:
                print(f"Successfully saved DNA file to: {output_path}")
                return True
            else:
                cmds.warning(f"Failed to save DNA file to: {output_path}")
                return False
        except Exception as e:
            cmds.error(f"Error saving DNA file: {str(e)}")
            return False
    
    def export_fbx(self, output_path):
        """Export the current scene as an FBX file"""
        try:
            result = self.maya_utils.export_fbx(output_path)
            if result:
                print(f"Successfully exported FBX to: {output_path}")
                return True
            else:
                cmds.warning(f"Failed to export FBX to: {output_path}")
                return False
        except Exception as e:
            cmds.error(f"Error exporting FBX: {str(e)}")
            return False
    
    def batch_export_blendshapes(self, output_dir):
        """Batch export blendshapes to the specified directory"""
        try:
            # Check if output directory exists, if not create it
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                print(f"Created output directory: {output_dir}")
            
            # Get all blendShape nodes in the scene
            blendshape_nodes = cmds.ls(type="blendShape")
            if not blendshape_nodes:
                cmds.warning("No blendShape nodes found in the scene")
                return False
            
            # Export each blendShape node's targets
            exported_files = []
            for blendshape in blendshape_nodes:
                try:
                    # Create subdirectory for each blendShape
                    blendshape_dir = os.path.join(output_dir, blendshape)
                    if not os.path.exists(blendshape_dir):
                        os.makedirs(blendshape_dir)
                    
                    # Export targets
                    files = self.blendshape_utils.batch_export_targets(blendshape, blendshape_dir)
                    exported_files.extend(files)
                except Exception as e:
                    cmds.warning(f"Error exporting blendShape {blendshape}: {str(e)}")
            
            if exported_files:
                print(f"Successfully exported {len(exported_files)} blendShape targets to: {output_dir}")
                return True
            else:
                cmds.warning(f"No blendShape targets exported to: {output_dir}")
                return False
        except Exception as e:
            error_msg = f"Error exporting blendShape: {str(e)}\n{traceback.format_exc()}"
            cmds.error(error_msg)
            return False

# Main function to initialize and show the Tool UI
def main():
    """Main function to initialize and show the Tool UI"""
    try:
        # Check Maya version compatibility
        version_message = ""
        version_warning = False
        
        if MAYA_VERSION < 2022:
            version_message = f"{TOOL_NAME} recommended to run in Maya 2022 or higher.\nCurrent version: Maya {MAYA_VERSION}\nSome features may be unavailable or perform abnormally."
            version_warning = True
        elif MAYA_VERSION > 2025:
            version_message = f"{TOOL_NAME} has not been tested in Maya {MAYA_VERSION}.\nMay need updates to support new Maya features."
            version_warning = True
        
        if version_warning:
            result = cmds.confirmDialog(
                title="Version Compatibility Warning",
                message=version_message + "\n\nDo you want to continue?",
                button=["Continue", "Cancel"],
                defaultButton="Continue",
                cancelButton="Cancel",
                dismissString="Cancel"
            )
            
            if result != "Continue":
                print(f"{TOOL_NAME} startup cancelled")
                return None
        
        # Initialize and display UI
        tool_instance = Tool.instance()
        ui = tool_instance.show_ui()
        
        # Print version information
        print(f"\n{'-'*50}")
        print(f"{TOOL_NAME} {TOOL_VERSION} started")
        print(f"Maya version: {MAYA_VERSION}")
        print(f"Operating system: {OS_NAME}")
        print(f"Python version: {PYTHON_VERSION}")
        print(f"Qt version: {QT_VERSION}")
        print(f"{'-'*50}\n")
        
        # In status bar display version information
        if hasattr(ui, "status_bar"):
            ui.status_bar.showMessage(f"Tool {TOOL_VERSION} | Maya {MAYA_VERSION} | {QT_VERSION} | {OS_NAME}")
        
        return ui
    except Exception as e:
        error_msg = f"Error starting Tool: {str(e)}\n{traceback.format_exc()}"
        cmds.error(error_msg)
        return None

# If the script is run directly, show the UI
if __name__ == "__main__":
    main()