irAEMM/slicerModules/iraemmBrowser.py

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

#
# 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 ...
    
    fhome=os.path.expanduser('~')
    fsetup=os.path.join(fhome,'.labkey','setup.json')
    try:
        with open(fsetup) as f:
            self.setup=json.load(f)
    except FileNotFoundError:
        self.setup={}

    try:
        pt=self.setup['paths']
    except KeyError:
        self.setup['paths']={}

    try:    
        sys.path.append(self.setup['paths']['labkeyInterface'])
    except KeyError:
        self.setup['paths']['labkeyInterface']=loadLibrary('labkeyInterface')
        with open(fsetup,'w') as f:
            json.dump(self.setup,f,indent='\t')
    
    import labkeyInterface
    import labkeyDatabaseBrowser
    import labkeyFileBrowser


    self.network=labkeyInterface.labkeyInterface() 

    fconfig=os.path.join(os.path.expanduser('~'),'.labkey','network.json')
    self.network.init(fconfig)

    self.db=labkeyDatabaseBrowser.labkeyDB(self.network)
    self.fb=labkeyFileBrowser.labkeyFileBrowser(self.network)
    self.project="iPNUMMretro/Study"
    self.schema='study'
    self.dataset="Imaging1"
    self.reviewDataset="ImageReview"
    self.aeDataset="PET"
    self.segmentList=['liver','bowel','thyroid','lung','kidney','pancreas']

    

    ds=self.db.selectRows(self.project,self.schema,self.dataset,[],"segmentationReview")
    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.participantField=qt.QLabel("PatientId")
    setupFormLayout.addRow("Participant field:",self.participantField)
    
    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)
   

    self.idField=qt.QLabel(self.network.getUserId()['displayName'])
    setupFormLayout.addRow("ID",self.idField)
    self.logic=iraemmBrowserLogic(self)


    #
    # Patients 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.aeResult=qt.QComboBox()
    self.reviewFormLayout.addRow("Is organ suffering from adverse effect?",\
            self.aeResult)
    aeOptions=['Select','Yes','No']
    for opt in aeOptions:
        self.aeResult.addItem(opt)
    #self.aeResult.setCurrentIndex(0)

    self.updateReview=qt.QPushButton("Save")
    self.reviewFormLayout.\
            addRow("Save segmentation and AE decision for current segment",\
            self.updateReview)
    self.updateReview.clicked.connect(self.onUpdateReviewButtonClicked)

    self.reviewBoxLayout.addLayout(self.reviewFormLayout)

    submitFrame=qt.QGroupBox("Submit data")
    
    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.setFlat(1)
    #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.db.selectRows(self.project,self.schema,self.dataset, [idFilter],"segmentationReview")
      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.db.selectRows(self.project,self.schema,self.dataset,[idFilter,sFilter],"segmentationReview")
      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()
      #also bladder,vertebraL1, stomach, heart
      for seg in segmentList:
          if not seg in self.segmentList:
              continue
          #filter to most important ones
          self.reviewSegment.addItem(seg)
      self.logic.loadReview(self.currentRow)
      self.logic.loadAE(self.currentRow)
      for segment in self.segmentList:
          rIdx=self.logic.getReviewResult(segment)
          aIdx=self.logic.getAEResult(segment)
          print("Segment {}: {}/{}".format(segment,rIdx,aIdx))
          try:
              if (rIdx+aIdx)>0:
                  self.updateResult(segment,rIdx,aIdx)
          except TypeError:
              continue
      try:
          self.reviewComment.setPlainText(self.logic.reviewComment)
      except AttributeError:
          pass

      self.onReviewSegmentChanged()
      #except AttributeError:
      #    print("Missing current row")
      #    return
  
  def onReviewSegmentChanged(self):
      segment=self.reviewSegment.currentText
      self.logic.hideSegments()
      self.logic.showSegment(segment)
      #set reviewFlag to stored value
      self.reviewResult.setCurrentIndex(self.logic.getReviewResult(segment))
      self.aeResult.setCurrentIndex(self.logic.getAEResult(segment))

  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)
      self.logic.submitAE(self.currentRow)

  def onUpdateReviewButtonClicked(self):
      print("Save")
      
      segment=self.reviewSegment.currentText
      self.logic.updateReview(segment,\
              self.reviewResult.currentIndex)

      self.logic.updateAE(segment,\
              self.aeResult.currentIndex)
      self.updateResult(segment,self.reviewResult.currentIndex,\
              self.aeResult.currentIndex)

     
  def updateResult(self,segment,reviewResult,aeResult):
      reviewText=self.reviewResult.itemText(reviewResult)
      aeText=self.aeResult.itemText(aeResult)

      idx=self.findCompletedSegment(segment)

      if idx<0:
          qReview=qt.QLabel(reviewText)
          self.submitFormLayout.insertRow(0,segment,qReview)
          qAE=qt.QLabel(aeText)
          self.submitFormLayout.insertRow(1,segment+'AE',qAE)
          try:
              self.segmentsCompleted.append(segment)
              self.segmentsCompleted.append(segment+'AE')
          except AttributeError:
              self.segmentsCompleted=[]
              self.segmentsCompleted.append(segment)
              self.segmentsCompleted.append(segment+'AE')
      else:
          qReview=self.submitFormLayout.itemAt(idx,1).widget()
          qReview.setText(reviewText)
          qAE=self.submitFormLayout.itemAt(idx+1,1).widget()
          qAE.setText(aeText)

      reviewColors=['pink','green','yellow','orange','red']
      qReview.setStyleSheet("background-color: "+reviewColors[reviewResult])
      aeColors=['pink','red','green']
      qAE.setStyleSheet("background-color: "+aeColors[aeResult])



  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()
      self.reviewComment.clear()

      

  



  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.db=parent.db
          self.fb=parent.fb
          self.project=parent.project
          self.schema=parent.schema
          self.participantField=parent.participantField.text
          self.segmentList=parent.segmentList

      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):
      
      tempDir=os.path.join(os.path.expanduser('~'),'temp')
      if not os.path.isdir(tempDir):
          os.mkdir(tempDir)
      
      #fields={'ctResampled':True,'petResampled':False}
      fields={"CT":self.parent.ctField.text,\
              "PET":self.parent.petField.text,\
              "Segmentation":self.parent.segmentationField.text}

      relativePaths={x:['preprocessedImages',row['patientCode'],row['visitCode'],row[y]]\
             for (x,y) in fields.items()}

      self.volumeNode={}
      for f in relativePaths:

          p=relativePaths[f]
          
          localPath=os.path.join(tempDir,p[-1])

          if not os.path.isfile(localPath):
              #download from server
              remotePath=self.fb.formatPathURL(self.project,'/'.join(p))
              if not self.fb.entryExists(remotePath):
                  print("Failed to get {}".format(remotePath))
                  continue

              self.fb.readFileToFile(remotePath,localPath) 
          
          properties={}
          #make sure segmentation gets loaded as a labelmap
          if f=="Segmentation":
              properties["labelmap"]=1
              
          
          self.volumeNode[f]=slicer.util.loadNodeFromFile(localPath,filetype='VolumeFile',properties=properties)
          
          if not keepCached:
              os.remove(localPath)

      #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.1,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={}
      self.aeList={}

  #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 getUniqueRows(self, project, dataset, fields, inputRow):
      filters=[]

      for f in fields:
          filters.append({'variable':f,'value':str(inputRow[f]),'oper':'eq'})
      
      
      ds=self.db.selectRows(project,self.schema,dataset,filters)
      return ds['rows']

  def loadReview(self,currentRow):

      #see if we have already done a review
      currentRow['ModifiedBy']=self.net.getUserId()['id']
      rows=self.getUniqueRows(self.parent.project,self.parent.reviewDataset,\
              [self.participantField,'visitCode','Segmentation','ModifiedBy'],currentRow)

      if len(rows)==0:
          return

      row=rows[0]
      for label in self.segmentList:
          #name=self.segLabel[label]+'Review'
          name=label+'Review'
          try:
              self.updateReview(label,row[name])
          except KeyError:
              continue 
      self.reviewComment=row['reviewComment']

  #submit review to labkey
  def submitReview(self,currentRow,comment):
      currentRow['ModifiedBy']=self.net.getUserId()['id']
      fields=[self.participantField,'visitCode','Segmentation','ModifiedBy']
      rows=self.getUniqueRows(self.parent.project,self.parent.reviewDataset,\
              fields,currentRow)


      mode='update'
      
      if len(rows)==0:
          mode='insert'
          row={}
          for f in fields:
              row[f]=currentRow[f]
          
          frows=self.getUniqueRows(self.parent.project,self.parent.reviewDataset,\
                  [self.participantField,'visitCode'],currentRow)

          row['SequenceNum']=currentRow['SequenceNum']+0.01*len(frows)
      
      else:     
          row=rows[0]
      
      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()
      self.db.modifyRows(mode,self.parent.project,self.parent.schema,\
              self.parent.reviewDataset,[row])
      
      print("review submitted")


#AE management
  def updateAE(self,segment,value):
      try:
          self.aeList[segment]=value
      except AttributeError:
          self.aeList={}
          self.updateAE(segment,value)
  
  def getAEResult(self,segment):
      try:
          return self.aeList[segment]
      except AttributeError:
          #review result not initialized (unknown)
          return 0
      except KeyError:
          #segment not done yet (unknown)
          return 0
  


  def loadAE(self,currentRow):
      currentRow['ModifiedBy']=self.net.getUserId()['id']
      fields=[self.participantField,'petResampled','ModifiedBy']
      rows=self.getUniqueRows(self.parent.project,self.parent.aeDataset,\
              fields,currentRow)
      
      if len(rows)==0:
          return

      print("Found {} rows".format(len(rows)))
      row=rows[0]
     
      for seg in self.segmentList:
          name=seg+'AE'
          try:
              self.updateAE(seg,row[name])
          except AttributeError:
              continue
          except KeyError:
              continue

  def submitAE(self,currentRow):
      currentRow['ModifiedBy']=self.net.getUserId()['id']
      fields=[self.participantField,'petResampled','ModifiedBy']
      rows=self.getUniqueRows(self.parent.project,self.parent.aeDataset,\
              fields,currentRow)
      
      
      if len(rows)==0:
          mode='insert'
          row={}
          for f in fields:
              row[f]=currentRow[f]
          
          frows=self.getUniqueRows(self.parent.project,self.parent.aeDataset,\
                  [self.participantField,'visitCode'],currentRow)

          row['SequenceNum']=currentRow['SequenceNum']+0.01*len(frows)
      
      else:
          mode='update'
          row=rows[0]
     
      for seg in self.segmentList:
          row[seg+'AE']=self.getAEResult(seg) 

      row['Date']=datetime.datetime.now().ctime()
      resp=self.db.modifyRows(mode,self.parent.project,self.parent.schema,\
              self.parent.aeDataset,[row])
      print("Response {}".format(resp))
      print("AE 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
    #