irAEMM/slicerModule/iraemmBrowser.py

import os
import unittest
from __main__ import vtk, qt, ctk, slicer
from slicer.ScriptedLoadableModule import *
import slicerNetwork
import loadDicom
import json
import datetime

#
# labkeySlicerPythonExtension
#

class iraemmBrowser(ScriptedLoadableModule):
  """Uses ScriptedLoadableModule base class, available at:
  https://github.com/Slicer/Slicer/blob/master/Base/Python/slicer/ScriptedLoadableModule.py
  """

  def __init__(self, parent):
    ScriptedLoadableModule.__init__(self, parent)
    self.parent.title = "irAEMM Browser" # TODO make this more human readable by adding spaces
    self.parent.categories = ["LabKey"]
    self.parent.dependencies = []
    self.parent.contributors = ["Andrej Studen (UL/FMF)"] # replace with "Firstname Lastname (Organization)"
    self.parent.helpText = """
    Interface to irAEMM files in LabKey
    """
    self.parent.acknowledgementText = """
    Developed within the medical physics research programme of the Slovenian research agency.
    """ # replace with organization, grant and thanks.

#
# labkeySlicerPythonExtensionWidget
#

class iraemmBrowserWidget(ScriptedLoadableModuleWidget):
  """Uses ScriptedLoadableModuleWidget base class, available at:
  https://github.com/Slicer/Slicer/blob/master/Base/Python/slicer/ScriptedLoadableModule.py
  """

  def setup(self):
    print("Setting up iraemmBrowserWidget")
    ScriptedLoadableModuleWidget.setup(self)
    # Instantiate and connect widgets ...
    self.network=slicerNetwork.labkeyURIHandler()

    fconfig=os.path.join(os.path.expanduser('~'),'.labkey','network.json')
    self.network.parseConfig(fconfig)
    self.network.initRemote()
    self.project="iPNUMMretro/Study"
    self.dataset="Imaging1"
    self.reviewDataset="ImageReview"


    self.logic=iraemmBrowserLogic(self)
    

    ds=self.network.filterDataset(self.project,self.dataset,[])
    ids=[row['PatientId'] for row in ds['rows']]
    ids=list(set(ids))

    
    #
    # Setup Area
    #
    setupCollapsibleButton = ctk.ctkCollapsibleButton()
    setupCollapsibleButton.text = "Setup"
    self.layout.addWidget(setupCollapsibleButton)
    
    setupFormLayout = qt.QFormLayout(setupCollapsibleButton)

   
    self.ctField=qt.QLabel("ctResampled")
    setupFormLayout.addRow("Data field (CT):",self.ctField)

    self.petField=qt.QLabel("petResampled")
    setupFormLayout.addRow("Data field (PET):",self.petField)

    self.segmentationField=qt.QLabel("Segmentation")
    setupFormLayout.addRow("Data field (Segmentation):",self.segmentationField)


    #
    # Patienrs Area
    #
    patientsCollapsibleButton = ctk.ctkCollapsibleButton()
    patientsCollapsibleButton.text = "Patients"
    self.layout.addWidget(patientsCollapsibleButton)

    patientsFormLayout = qt.QFormLayout(patientsCollapsibleButton)


    self.patientList=qt.QComboBox()
    for id in ids:
        self.patientList.addItem(id)
    self.patientList.currentIndexChanged.connect(self.onPatientListChanged)
    patientsFormLayout.addRow("Patient:",self.patientList)

    self.visitList=qt.QComboBox()
    self.visitList.currentIndexChanged.connect(self.onVisitListChanged)
    patientsFormLayout.addRow("Visit:",self.visitList)


    self.ctCode=qt.QLabel("ctCode")
    patientsFormLayout.addRow("CT:",self.ctCode)
    
    self.petCode=qt.QLabel("petCode")
    patientsFormLayout.addRow("PET:",self.petCode)

    self.segmentationCode=qt.QLabel("segmentationCode")
    patientsFormLayout.addRow("Segmentation",self.segmentationCode)

    self.patientLoad=qt.QPushButton("Load")
    self.patientLoad.clicked.connect(self.onPatientLoadButtonClicked)
    patientsFormLayout.addRow("Load patient",self.patientLoad)

    self.patientClear=qt.QPushButton("Clear")
    self.patientClear.clicked.connect(self.onPatientClearButtonClicked)
    patientsFormLayout.addRow("Clear patient",self.patientClear)
    
    self.keepCached=qt.QCheckBox("keep Cached")
    self.keepCached.setChecked(1)
    patientsFormLayout.addRow("Keep cached",self.keepCached)
    
    #set to a defined state
    self.onPatientListChanged(0)


    #
    # Review Area
    #
    reviewCollapsibleButton = ctk.ctkCollapsibleButton()
    reviewCollapsibleButton.text = "Review"
    self.layout.addWidget(reviewCollapsibleButton)
    
    self.reviewBoxLayout = qt.QVBoxLayout(reviewCollapsibleButton)

    self.reviewFormLayout = qt.QFormLayout()


    self.reviewSegment=qt.QComboBox()
    self.reviewSegment.currentIndexChanged.connect(self.onReviewSegmentChanged)
    self.reviewFormLayout.addRow("Selected region:",self.reviewSegment)
    
   

    self.reviewResult=qt.QComboBox()
    self.reviewFormLayout.addRow("What do you think about the segmentation:",\
            self.reviewResult)
    reviewOptions=['Select','Excellent','Minor deficiencies','Major deficiencies','Unusable']
    for opt in reviewOptions:
        self.reviewResult.addItem(opt)
    
    self.updateReview=qt.QPushButton("Save")
    self.reviewFormLayout.addRow("Save segmentation decision for current segment",\
            self.updateReview)
    self.updateReview.clicked.connect(self.onUpdateReviewButtonClicked)

    self.reviewBoxLayout.addLayout(self.reviewFormLayout)

    submitFrame=qt.QFrame()
    
    self.submitFormLayout=qt.QFormLayout()

    self.reviewComment=qt.QTextEdit("this is a test")
    self.submitFormLayout.addRow("Comments (optional)",\
            self.reviewComment)
    
    self.submitReviewButton=qt.QPushButton("Submit")
    self.submitFormLayout.addRow("Submit to database",\
            self.submitReviewButton)
    self.submitReviewButton.clicked.connect(self.onSubmitReviewButtonClicked)
    submitFrame.setLayout(self.submitFormLayout)
    submitFrame.setFrameShape(qt.QFrame.StyledPanel)
    submitFrame.setFrameShadow(qt.QFrame.Sunken)
    submitFrame.setStyleSheet("background-color:rgba(220,215,180,45)")
    self.reviewBoxLayout.addWidget(submitFrame)

  def onPatientListChanged(self,i):
      idFilter={'variable':'PatientId','value':self.patientList.currentText,'oper':'eq'}
      ds=self.network.filterDataset(self.project,self.dataset, [idFilter])
      seq=[int(row['SequenceNum']) for row in ds['rows']]
      self.visitList.clear()  
            
      for s in seq:
          self.visitList.addItem("Visit "+str(s))
      self.onVisitListChanged(0)

  def onVisitListChanged(self,i):
      try:
        s=self.visitList.currentText.split(' ')[1]
      except IndexError:
        return
      print("Visit: Selected item: {}->{}".format(i,s))
      idFilter={'variable':'PatientId',\
              'value':self.patientList.currentText,'oper':'eq'}
      sFilter={'variable':'SequenceNum','value':s,'oper':'eq'}
      ds=self.network.filterDataset(self.project,self.dataset,[idFilter,sFilter])
      if not len(ds['rows'])==1:
          print("Found incorrect number {} of matches for [{}]/[{}]".\
                  format(len(ds['rows']),\
                  self.patientList.currentText,s))
      row=ds['rows'][0]

      #copy row properties for data access
      self.currentRow=row
      self.petCode.setText(row[self.petField.text])
      self.ctCode.setText(row[self.ctField.text])
      self.segmentationCode.setText(row[self.segmentationField.text])

  def onPatientLoadButtonClicked(self):
      print("Load")
      #delegate loading to logic
      #try:
      self.logic.loadImage(self.currentRow,self.keepCached.isChecked())
      segmentList=self.logic.compileSegmentation()
      importantSegmentList=['liver','bowel','thyroid','bladder']
      for seg in segmentList:
          if not seg in importantSegmentList:
              continue
          #filter to most important ones
          self.reviewSegment.addItem(seg)
      self.logic.loadReview(self.currentRow)
      self.onReviewSegmentChanged()
      #except AttributeError:
      #    print("Missing current row")
      #    return
  
  def onReviewSegmentChanged(self):
      self.logic.hideSegments()
      self.logic.showSegment(self.reviewSegment.currentText)
      #set reviewFlag to stored value
      self.reviewResult.setCurrentIndex(self.logic.getReviewResult(self.reviewSegment.currentText)) 

  def onSubmitReviewButtonClicked(self):
      print("Submit")
      print("Selected review:{}/{}".format(self.reviewResult.currentIndex,
          self.reviewResult.currentText))
      print("Comment:{}".format(self.reviewComment))
      self.logic.submitReview(self.currentRow,\
              self.reviewComment.plainText)

  def onUpdateReviewButtonClicked(self):
      print("Save")
      
      self.logic.updateReview(self.reviewSegment.currentText,\
              self.reviewResult.currentIndex)
      
      idx=self.findCompletedSegment(self.reviewSegment.currentText)
      if idx<0:
          qReview=qt.QLabel(self.reviewResult.currentText)
          self.submitFormLayout.insertRow(0,self.reviewSegment.currentText,qReview)
          try:
              self.segmentsCompleted.append(self.reviewSegment.currentText)
          except AttributeError:
              self.segmentsCompleted=[]
              self.segmentsCompleted.append(self.reviewSegment.currentText)
      else:
          qReview=self.submitFormLayout.itemAt(idx,1).widget()
          qReview.setText(self.reviewResult.currentText)

      colors=['pink','green','yellow','orange','red']
      qReview.setStyleSheet("background-color: "+colors[self.reviewResult.currentIndex])
  
  def findCompletedSegment(self,segment):

      for i in range(self.submitFormLayout.rowCount()):
          if self.submitFormLayout.itemAt(i,0).widget().text==segment:
              return i
      return -1

  def removeCompletedSegments(self):
      
      try:
          segments=self.segmentsCompleted
      except AttributeError:
          return 

      for seg in segments:
          idx=self.findCompletedSegment(seg)
          if idx>-1:
              self.submitFormLayout.removeRow(idx)
      
      self.segmentsCompleted=[]


  def onPatientClearButtonClicked(self):
      self.logic.clearVolumesAndSegmentations()
      self.reviewSegment.clear()
      self.removeCompletedSegments()

      

  



  def cleanup(self):
    pass

#
# irAEMMBrowserLogic
#

class iraemmBrowserLogic(ScriptedLoadableModuleLogic):
  """This class should implement all the actual
  computation done by your module.  The interface
  should be such that other python code can import
  this class and make use of the functionality without
  requiring an instance of the Widget.
  Uses ScriptedLoadableModuleLogic base class, available at:
  https://github.com/Slicer/Slicer/blob/master/Base/Python/slicer/ScriptedLoadableModule.py
  """
  def __init__(self,parent=None):
      ScriptedLoadableModuleLogic.__init__(self, parent)
      if not parent==None:
          #assume parent has the network set up
          self.parent=parent
          self.net=parent.network
          self.project=parent.project

      self.segLabel={'1':'liver','2':'spleen','3':'lung','4':'thyroid',\
              '5':'kidney','6':'pancreas','7':'gallbladder','8':'bladder',\
              '9':'aorta','10':'trachea','11':'sternum','12':'vertebraL1',\
              '13':'adrenal','14':'psoasMajor','15':'rectus',\
              '16':'bowel','17':'stomach','18':'heart'}

  def setLabkeyInterface(self,net):
      #additional way of setting the labkey network interface 
      #if no parent was provided in logic initialization (stand-alone mode)
      self.net=net

  def setLabkeyProject(self,project):
      self.project=project

  def loadImage(self,row,keepCached):
      
      
      #fields={'ctResampled':True,'petResampled':False}
      fields={"CT":self.parent.ctField.text,\
              "PET":self.parent.petField.text,\
              "Segmentation":self.parent.segmentationField.text}

      relativePaths={x:self.project+'/@files/preprocessedImages/'\
             +row['patientCode']+'/'+row['visitCode']+'/'+row[y]\
             for (x,y) in fields.items()}

      self.volumeNode={}
      for f in relativePaths:
          p=relativePaths[f]
          labkeyPath=self.net.GetLabkeyPathFromRelativePath(p)
          rp=self.net.head(labkeyPath)
          if not slicerNetwork.labkeyURIHandler.HTTPStatus(rp):
              print("Failed to get {}".format(labkeyPath))
              continue

          #pushes it to background
          properties={}
          #make sure segmentation gets loaded as a labelmap
          if f=="Segmentation":
              properties["labelmap"]=1

          self.volumeNode[f]=self.net.loadNode(p,'VolumeFile',\
                  properties=properties,returnNode=True,keepCached=keepCached)

      #mimic abdominalCT standardized window setting
      self.volumeNode['CT'].GetScalarVolumeDisplayNode().\
              SetWindowLevel(1400, -500)
      #set colormap for PET to PET-Heat (this is a verbatim setting from
      #the Volumes->Display->Lookup Table colormap identifier)
      self.volumeNode['PET'].GetScalarVolumeDisplayNode().\
              SetAndObserveColorNodeID(\
              slicer.util.getNode('PET-Heat').GetID())
      slicer.util.setSliceViewerLayers(background=self.volumeNode['CT'],\
          foreground=self.volumeNode['PET'],foregroundOpacity=0.5,fit=True)


  #segmentations

  def compileSegmentation(self):
      try:
          labelmapVolumeNode = self.volumeNode['Segmentation']
      except KeyError:
          print("No segmentaion volumeNode available")
          return
     
      self.segmentationNode = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLSegmentationNode')
      slicer.modules.segmentations.logic().\
              ImportLabelmapToSegmentationNode(labelmapVolumeNode, self.segmentationNode)
      
      
      segmentList=[]
      
      seg=self.segmentationNode.GetSegmentation()

      for i in range(seg.GetNumberOfSegments()):
          segment=seg.GetNthSegment(i)
          segment.SetName(self.segLabel[segment.GetName()])
          segmentList.append(segment.GetName())

      #seg.CreateClosedSurfaceRepresentation()
      slicer.mrmlScene.RemoveNode(labelmapVolumeNode)
      self.volumeNode.pop('Segmentation',None)

      #return list of segment names
      return segmentList
  
  def hideSegments(self):

      try:
          displayNode=self.segmentationNode.GetDisplayNode()
      except AttributeError:
          return

      seg=self.segmentationNode.GetSegmentation()
      for i in range(seg.GetNumberOfSegments()):
          #segment=self.segmentationNode.GetSegmentation().GetNthSegment(i)
          segmentID=seg.GetNthSegmentID(i)
          displayNode.SetSegmentVisibility(segmentID, False)
      #print("Done")
  
  def showSegment(self,name):
      try:
          displayNode=self.segmentationNode.GetDisplayNode()
      except AttributeError:
          return
      
      seg=self.segmentationNode.GetSegmentation()
      for i in range(seg.GetNumberOfSegments()):
          segment=seg.GetNthSegment(i)
          if not segment.GetName()==name:
              continue
          segmentID=seg.GetNthSegmentID(i)
          displayNode.SetSegmentVisibility(segmentID, True)
          break
      #print("Done")
  

  #clear 

  def clearVolumesAndSegmentations(self):
      nodes=slicer.util.getNodesByClass("vtkMRMLVolumeNode")
      nodes.extend(slicer.util.getNodesByClass("vtkMRMLSegmentationNode"))
      res=[slicer.mrmlScene.RemoveNode(f) for f in nodes] 
      self.segmentationNode=None
      self.reviewResult={}

  #reviews by segment

  def updateReview(self,segment,value):
      try:
          self.reviewResult[segment]=value
      except AttributeError:
          self.reviewResult={}
          self.updateReview(segment,value)
  
  def getReviewResult(self,segment):
      try:
          return self.reviewResult[segment]
      except AttributeError:
          #review result not initialized
          return 0
      except KeyError:
          #segment not done yet
          return 0
  
  #load review from labkey
  def loadReview(self,currentRow):

      #see if we have already done a review
      filters=[]
      fields=['PatientId','SequenceNum']


      for f in fields:
          filters.append({'variable':f,'value':str(currentRow[f]),'oper':'eq'})
      
      ds=self.net.filterDataset(self.parent.project,\
              self.parent.reviewDataset,filters)

      if len(ds['rows'])==0:
          return

      row=ds['rows'][0]
      for label in self.segLabel:
          name=self.segLabel[label]+'Review'
          try:
              self.updateReview(self.segLabel[label],row[name])
          except KeyError:
              continue 
   
  #submit review to labkey
  def submitReview(self,currentRow,comment):
      row={}

      fields=['PatientId','SequenceNum']

      #see if we have to update or insert
      filters=[]

      for f in fields:
          filters.append({'variable':f,'value':str(currentRow[f]),'oper':'eq'})
      
      ds=self.net.filterDataset(self.parent.project,\
              self.parent.reviewDataset,filters)

      mode='insert'

      if len(ds['rows'])>0:
          row=ds['rows'][0]
          mode='update'
      else:
          for f in fields:
              row[f]=currentRow[f]
      
      seg=self.segmentationNode.GetSegmentation()
 
      for i in range(seg.GetNumberOfSegments()):
          segment=seg.GetNthSegment(i)
          fieldName=segment.GetName()+'Review'
          value=self.getReviewResult(segment.GetName())
          row[fieldName]=value

      row['reviewComment']=comment
      row['Date']=datetime.datetime.now().ctime()
      ds=self.net.modifyDataset(mode,self.parent.project,\
              self.parent.reviewDataset,[row])
      
      print("review submitted")


class irAEMMBrowserTest(ScriptedLoadableModuleTest):

  """
  This is the test case for your scripted module.
  Uses ScriptedLoadableModuleTest base class, available at:
  https://github.com/Slicer/Slicer/blob/master/Base/Python/slicer/ScriptedLoadableModule.py
  """

  def setUp(self):
    """ Do whatever is needed to reset the state - typically a scene clear will be enough.
    """
    slicer.mrmlScene.Clear(0)

  def runTest(self):
    """Run as few or as many tests as needed here.
    """
    self.setUp()
    self.test_irAEMMBrowser()

  def test_irAEMMBrowser(self):
    """ Ideally you should have several levels of tests.  At the lowest level
    tests sould exercise the functionality of the logic with different inputs
    (both valid and invalid).  At higher levels your tests should emulate the
    way the user would interact with your code and confirm that it still works
    the way you intended.
    One of the most important features of the tests is that it should alert other
    developers when their changes will have an impact on the behavior of your
    module.  For example, if a developer removes a feature that you depend on,
    your test should break so they know that the feature is needed.
    """

    self.delayDisplay("Starting the test")
    #
    # first, get some data
    #