diff --git a/modules/gui/pymod/init_context_menu.py b/modules/gui/pymod/init_context_menu.py
index 2807a10bd5915a7720574375231afce46cbd21f2..7a6145ccdbccbaea4b193c84905f42b02d044fe4 100644
--- a/modules/gui/pymod/init_context_menu.py
+++ b/modules/gui/pymod/init_context_menu.py
@@ -3,7 +3,98 @@ from PyQt4 import QtCore, QtGui
from ost import geom, gfx, gui
from ost import settings
from ost.bindings import tmtools
-
+
+class SelectRefDialog(QtGui.QDialog):
+ def __init__(self, ent_list, parent=None):
+ QtGui.QDialog.__init__(self, parent)
+ self.ent_list_ = ent_list
+ vb = QtGui.QVBoxLayout()
+ self.setLayout(vb)
+ self.setWindowTitle("Select Reference Object")
+ self.label = QtGui.QLabel("Please Select the Reference Object")
+ self.list = QtGui.QTableWidget(self)
+ self.list.setColumnCount(2)
+ self.list.setColumnWidth(0, 30)
+ self.list.setColumnWidth(1,220)
+ self.list.verticalHeader().setVisible(False)
+ self.list.horizontalHeader().setVisible(False)
+ self.list.setSelectionBehavior(QtGui.QAbstractItemView.SelectRows)
+ self.list.setSelectionMode(QtGui.QAbstractItemView.SingleSelection)
+ vb.addWidget(self.label)
+ vb.addWidget(self.list)
+ hb = QtGui.QHBoxLayout()
+ hb.setDirection(QtGui.QBoxLayout.LeftToRight)
+ cancel_btn = QtGui.QPushButton("Cancel", self)
+ load_btn = QtGui.QPushButton("Select", self)
+ hb.addStretch(1)
+ hb.addWidget(cancel_btn)
+ hb.addWidget(load_btn)
+ vb.addItem(hb)
+ load_btn.setDefault(True)
+ QtCore.QObject.connect(load_btn, QtCore.SIGNAL("clicked()"), self.Select)
+ QtCore.QObject.connect(cancel_btn, QtCore.SIGNAL("clicked()"), self.reject)
+
+ row = 0
+ for ent in self.ent_list_:
+ variant = QtCore.QVariant(ent)
+ self.list.insertRow(row)
+ new_item = QtGui.QTableWidgetItem("%i"%(row+1))
+ new_item.setFlags(QtCore.Qt.ItemIsSelectable|QtCore.Qt.ItemIsEnabled)
+ new_item.setData(QtCore.Qt.UserRole,variant)
+ self.list.setItem(row, 0, new_item)
+ new_item = QtGui.QTableWidgetItem(ent.GetName())
+ new_item.setFlags(QtCore.Qt.ItemIsSelectable|QtCore.Qt.ItemIsEnabled)
+ self.list.setItem(row, 1, new_item)
+ row += 1
+
+ def Select(self):
+ items = self.list.selectedItems()
+ for item in items:
+ if(item.column()==0):
+ ent = item.data(QtCore.Qt.UserRole).toPyObject()
+ self.ent_list_.remove(ent)
+ self.ent_list_.insert(0,ent)
+ self.accept()
+
+ def GetEntities(self):
+ return self.ent_list_
+
+class ShowResultDialog(QtGui.QDialog):
+ def __init__(self, ent_list, res_list, parent=None):
+ QtGui.QDialog.__init__(self, parent)
+ self.ent_list_ = ent_list
+ vb = QtGui.QVBoxLayout()
+ self.setLayout(vb)
+ self.setWindowTitle("Alignment result")
+ self.label = QtGui.QLabel("Alignment results with %s as reference"%ent_list[0].GetName())
+ self.list = QtGui.QTableWidget(self)
+ self.list.setColumnCount(3)
+ self.list.verticalHeader().setVisible(False)
+ self.list.setHorizontalHeaderLabels (["Name","RMSD","TMScore"])
+ self.list.setSelectionBehavior(QtGui.QAbstractItemView.SelectRows)
+ vb.addWidget(self.label)
+ vb.addWidget(self.list)
+ hb = QtGui.QHBoxLayout()
+ hb.setDirection(QtGui.QBoxLayout.LeftToRight)
+ ok_btn = QtGui.QPushButton("OK", self)
+ hb.addStretch(1)
+ hb.addWidget(ok_btn)
+ vb.addItem(hb)
+ ok_btn.setDefault(True)
+ QtCore.QObject.connect(ok_btn, QtCore.SIGNAL("clicked()"), self.accept)
+
+ for i in range(0, len(res_list)):
+ self.list.insertRow(i)
+ new_item = QtGui.QTableWidgetItem(ent_list[i+1].GetName())
+ new_item.setFlags(QtCore.Qt.ItemIsEnabled)
+ self.list.setItem(i, 0, new_item)
+ new_item = QtGui.QTableWidgetItem("%.2f"%res_list[i].rmsd)
+ new_item.setFlags(QtCore.Qt.ItemIsEnabled)
+ self.list.setItem(i, 1, new_item)
+ new_item = QtGui.QTableWidgetItem("%i"%res_list[i].tm_score)
+ new_item.setFlags(QtCore.Qt.ItemIsEnabled)
+ self.list.setItem(i, 2, new_item)
+
class AlignmentContextMenu(QtCore.QObject):
def __init__(self, context_menu):
try:
@@ -12,27 +103,43 @@ class AlignmentContextMenu(QtCore.QObject):
self.action = QtGui.QAction("Align", self)
QtCore.QObject.connect(self.action,QtCore.SIGNAL("triggered()"), self.Align)
- print gui.ContextActionType.ENTITY | gui.ContextActionType.MULTI
context_menu.AddAction(self.action, gui.ContextActionType.ENTITY | gui.ContextActionType.MULTI)
except settings.FileNotFound:
return
def Align(self):
scene_selection = gui.SceneSelection.Instance()
- if scene_selection.GetActiveNodeCount() >= 2:
- node = scene_selection.GetActiveNode(0)
- if isinstance(node, gfx.Entity):
- ref = node.view.handle
- for i in range(1,scene_selection.GetActiveNodeCount()):
- node = scene_selection.GetActiveNode(i)
- if isinstance(node, gfx.Entity):
- print tmtools.TMAlign(node.view.handle, ref)
- node.UpdatePositions()
-
-
-
+ ent_list = list()
+ for i in range(0,scene_selection.GetActiveNodeCount()):
+ ent_list.append(scene_selection.GetActiveNode(i))
+ if len(ent_list) == 2:
+ self.__Align(ent_list)
+ elif len(ent_list) > 2:
+ sd = SelectRefDialog(ent_list)
+ if(sd.exec_()):
+ self.__Align(sd.GetEntities())
+
+ def __Align(self, ent_list):
+ node = ent_list[0]
+ res_list = list()
+ if isinstance(node, gfx.Entity):
+ ref = node.view.handle
+ for i in range(1,len(ent_list)):
+ node = ent_list[i]
+ if isinstance(node, gfx.Entity):
+ res_list.append(tmtools.TMAlign(node.view.handle, ref))
+ node.UpdatePositions()
+ self.__ShowScore(ent_list, res_list)
+
+ def __ShowScore(self, ent_list, res_list):
+ if(len(res_list)==1):
+ res = res_list[0]
+ string = "RMSD: %.2f, TMScore: %i"%(res.rmsd, res.tm_score)
+ gui.GostyApp.Instance().perspective.StatusMessage(string)
+ elif(len(res_list)>1):
+ ShowResultDialog(ent_list, res_list).exec_()
-
+
def _InitContextMenu():
cm=gui.GostyApp.Instance().GetSceneWin().GetContextMenu()
AlignmentContextMenu(cm)
\ No newline at end of file