Sample Editor 4 MayaAPI2.0 + QTableView + QAbstractTableModel

SampleEditor4.py

# -*- coding: utf-8 -*-
from maya import  cmds
from maya.app.general.mayaMixin import MayaQWidgetBaseMixin
import maya.api.OpenMaya as om2
import maya.api.OpenMayaAnim as oma2
import time
import imp
try:
    imp.find_module('PySide2')
    from PySide2.QtWidgets import *
    from PySide2.QtGui import *
    from PySide2.QtCore import *
except ImportError:
    from PySide.QtGui import *
    from PySide.QtCore import *
    
#時間計測用関数 
def timer(func):
    def wrapper(*args, **kwargs):
        global WINDOW
        start = time.time()#開始時間
        func(*args, **kwargs)#関数実行
        end = time.time()#終了時間
        msg = func.__name__ + ' : ' + str(end - start)
        print '-'*50
        print msg
        print '-'*50
        WINDOW.time_label.setText(msg)
    return wrapper
    
#ウィンドウ呼び出し
def main():
    global WINDOW
    WINDOW = SampleEditor()
    WINDOW.resize(600, 400)
    WINDOW.show()
    
#サンプルエディタークラス
class SampleEditor(MayaQWidgetBaseMixin, QMainWindow):
    #コンストラクタ
    def __init__(self, parent=None):
        super(self.__class__, self).__init__(parent)
        self.setWindowTitle('Sample Editor API')
        
        #ラッパーウイジェットとレイアウトを設定
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        self.main_layout = QVBoxLayout()
        central_widget.setLayout(self.main_layout)
        
        #ウェイト取得するボタン
        self.get_button = QPushButton('Get Weights')
        self.main_layout.addWidget(self.get_button)
        self.get_button.clicked.connect(self.get_weights)
        
        #ウェイトテーブルの変更を実際のウェイト値に反映するボタン
        self.set_button = QPushButton('Set Weights')
        self.main_layout.addWidget(self.set_button)
        self.set_button.clicked.connect(self.set_weights)
        
        #ウェイトテーブルビュー
        self.table_view = QTableView()
        self.main_layout.addWidget(self.table_view)
        
        #時間計測結果のお知らせ
        self.time_label = QLabel('Execution time :')
        self.main_layout.addWidget(self.time_label)
        
    #選択コンポーネントのウェイトデータを取得してテーブルに反映
    def get_weights(self):
        self.get_skin_data()#まずはウェイトデータの取得
        self.table_setup()#テーブルにデータセット
        
    #QTableWidgetにデータ設定
    @timer
    def table_setup(self):
        #StandardItemModel作成
        self.weight_model = SampleTableModel(self.weights_list, influences=self.influences, vertices=self.vertices)
        self.weight_model.dataChanged.connect(self.update_weights)
        #選択管理モデルを作る
        self.sel_model = QItemSelectionModel(self.weight_model)
        #モデルをビューに設定
        self.table_view.setModel(self.weight_model)
        self.table_view.setSelectionModel(self.sel_model)
    
    #選択頂点、ウェイト、インフルエンス、スキンクラスタを取得
    @timer
    def get_skin_data(self):
        #現在選択しているオブジェクト、コンポーネントをAPIで取得
        sList = om2.MGlobal.getActiveSelectionList()
        #選択リストのイテレータ
        iter = om2.MItSelectionList(sList)
        #イテレーションしてメッシュのDagPathとコンポーネントをセットで取得
        while not iter.isDone():
            meshDag, component = iter.getComponent() 
            iter.next()#nextでイテレーションを進める
        #今回は1メッシュのみサポート（書式をわかりやすくするため）    
        self.meshDag = meshDag
        
        #MFnSkinMeshクラス、頂点アレイ、スキンクラスタ名を取得する関数呼び出し
        self.skinFn, self.skin_cluster = self.om_get_skin_cluster(self.meshDag)
        self.vertexArray= self.convert_to_vertex(self.meshDag, component)
        
        #指定の頂点をコンポーネントとして取得する
        singleIdComp = om2.MFnSingleIndexedComponent()#SingleIDを格納するクラス
        self.vertexComp = singleIdComp.create(om2.MFn.kMeshVertComponent)#頂点コンポーネント（SingleID）を格納する函
        singleIdComp.addElements(self.vertexArray)#頂点アレイを格納
        
        #インフルエンスを取得して情報整理
        infDags = self.skinFn.influenceObjects()#インフルエンスを全取得
        self.infIndices = om2.MIntArray( len(infDags), 0 )#インフルエンスIDを格納するMIntArrayを用意、インフルエンスの数のIntArray
        for x in xrange(len(infDags)):
            self.infIndices[x] = int(self.skinFn.indexForInfluenceObject(infDags[x]))
        #インフルエンス名をフルパスで取得
        self.full_path_influences = [infDags[x].fullPathName() for x in range(len(self.infIndices))]
        #ショート名に変換（列ヘッダー表示用）
        self.influences = [inf.split('|')[-1] for inf in self.full_path_influences]
            
        #ウェイト取得
        self.weights = self.skinFn.getWeights(self.meshDag , self.vertexComp)
        #ウェイトデータ(1次元配列)を頂点ごとの2次元配列にリシェイプ
        self.weights_list = self.conv_weight_shape(len(self.infIndices), self.weights[0])
        
        #頂点番号名を生成（行ヘッダー表示用）
        mesh_name = self.meshDag.fullPathName().split('|')[-1]
        self.vertices = [mesh_name + '.vtx[' + str(vid) + ']'for vid in self.vertexArray]
        
    #ウェイトをバーテックス単位の2次元配列にリシェイプする
    def conv_weight_shape(self, shape, weights):
        return [[weights[i+j*shape] for i in xrange(shape)] for j in xrange(len(weights)/shape)]
        
    #ウェイトテーブルが更新されたらメンバ変数で保有しているウェイト値も更新する
    def update_weights(self, top_left, bottom_right):
        index = top_left
        row = index.row()
        column = index.column() 
        new_value = float(self.weight_model.get_data(index))
        self.weights_list[row][column] = new_value
        
    #コンポーネントやメッシュリストをバーテックスリストに変換する
    def convert_to_vertex(self, meshDag, component):
        cmpType = None
        
        #選択されているコンポーネントタイプを取得
        if component.hasFn(om2.MFn.kMeshVertComponent):
            cmpType = "vtx"
        elif component.hasFn(om2.MFn.kMeshEdgeComponent):
            cmpType = "edge"
        elif component.hasFn(om2.MFn.kMeshPolygonComponent):
            cmpType = "face"
        if cmpType:
            compFn = om2.MFnSingleIndexedComponent(component)
        
        #コンポーネントを頂点情報に変換してIDリストに格納する
        meshFn = om2.MFnMesh(meshDag)
        #頂点選択はそのまま取得
        if "vtx" == cmpType:
            vtxArray = compFn.getElements()
        #エッジを頂点IDに変換する
        elif "edge" == cmpType:
            eid = compFn.getElements()
            eSet = []
            for e in eid:
                evid = meshFn.getEdgeVertices(e)
                eSet.extend(evid)
            vids = list(set(eSet))
            vtxArray = om2.MIntArray()
            [vtxArray.append(id) for id in vids]
        #フェースを頂点IDに変換する
        elif "face" == cmpType:
            fid = compFn.getElements()
            fSet = []
            for f in fid:
                vid = meshFn.getPolygonVertices(f)
                fSet.extend(vid)
            vids = list(set(fSet))
            vtxArray = om2.MIntArray()
            [vtxArray.append(id) for id in vids]
        #メッシュ選択の場合
        else:
            vids = range(meshFn.numVertices)
            vtxArray = om2.MIntArray()
            [vtxArray.append(id) for id in vids]
        return vtxArray
        
    #スキンクラスタ名とMFnSkinClusterクラスを返す
    def om_get_skin_cluster(self, dagPath=None):
        skinCluster = cmds.ls(cmds.listHistory(dagPath.fullPathName()), type='skinCluster')
        clusterName = skinCluster[0]
        sellist = om2.MGlobal.getSelectionListByName(clusterName)
        skinNode = sellist.getDependNode(0)
        skinFn = oma2.MFnSkinCluster( skinNode )
        return skinFn, clusterName
        
    #ウェイト値をスキンメッシュに反映
    @timer
    def set_weights(self):
        #変更されたウェイト値をMDoubleArrayにキャスト
        new_weights = om2.MDoubleArray()
        for weights in self.weights_list:
            new_weights += weights
        #書き込み
        self.skinFn.setWeights(self.meshDag, self.vertexComp, self.infIndices, new_weights, False)
            
class SampleTableModel(QAbstractTableModel):
    def __init__(self, data, parent=None, influences=None,  vertices=None):
        super(self.__class__, self).__init__(parent)
        self._weight_data = data
        self.v_header_list = vertices
        self.h_header_list = influences
    
    #ヘッダー設定をオーバーライド
    def headerData(self, id, orientation, role):
        if orientation == Qt.Horizontal:
            if role == Qt.DisplayRole:
                return self.h_header_list[id]
        if orientation == Qt.Vertical:
            if role == Qt.DisplayRole:
                return self.v_header_list[id]
        return None
        
    #行数を返す、オーバーライド
    def rowCount(self, parent=None):
        return len(self._weight_data)
        
    #列数を返す、オーバーライド
    def columnCount(self, parent=None):
        return len(self._weight_data[0])-1 if self.rowCount() else 0
        
    #データ設定関数をオーバーライド、 ロールに応じて処理
    #ここで表示の値やカラーがすべて扱われます
    def data(self, index, role=Qt.DisplayRole):
        row = index.row()
        column = index.column()
        if role == Qt.DisplayRole:
            weight = self._weight_data[row][column]
            return weight
        elif role == Qt.TextAlignmentRole:
            return Qt.AlignRight#右揃え
                
    #データセットする関数をオーバライド
    def setData(self, index, value, role=Qt.EditRole):
        row = index.row()
        column = index.column()
        if role == Qt.EditRole:
            self._weight_data[row][column] = float(value)
            self.dataChanged.emit(index, index)#データ変更シグナルを送出
            return True
        else:
            return False
            
    # 各セルのインタラクション
    def flags(self, index):
        #インデックスによる選択可否等をつけたければここで条件分岐
        return Qt.ItemIsSelectable | Qt.ItemIsEnabled | Qt.ItemIsEditable
        
    #指定セルのデータを取得
    def get_data(self, index):
        row = index.row()
        column = index.column()
        value  = float(self._weight_data[row][column])
        return value
        
            
#実行
main()