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)
    
    reviewFormLayout = qt.QFormLayout(reviewCollapsibleButton)

   
    self.reviewResult=qt.QComboBox()
    reviewFormLayout.addRow("What do you think about the segmentation:",\
            self.reviewResult)
    self.reviewResult.addItem("Excellent")
    self.reviewResult.addItem("Minor deficiencies")
    self.reviewResult.addItem("Major deficiencies")
    self.reviewResult.addItem("Unusable")
    
    self.reviewComment=qt.QLineEdit("this is a test")
    reviewFormLayout.addRow("Comments (optional)",\
            self.reviewComment)
    
    self.submitReviewButton=qt.QPushButton("Submit")
    reviewFormLayout.addRow("Submit to database",\
            self.submitReviewButton)
    self.submitReviewButton.clicked.connect(self.onSubmitReviewButtonClicked)


  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())
      self.logic.compileSegmentation()
      #except AttributeError:
      #    print("Missing current row")
      #    return

  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.reviewResult.currentIndex,\
              self.reviewComment.text)

  def onPatientClearButtonClicked(self):
      self.logic.clearVolumesAndSegmentations()



  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

  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 rp.code==200:
              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)

  def compileSegmentation(self):
      try:
          labelmapVolumeNode = self.volumeNode['Segmentation']
      except KeyError:
          print("No segmentaion volumeNode available")
          return
     
      seg = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLSegmentationNode')
      slicer.modules.segmentations.logic().\
              ImportLabelmapToSegmentationNode(labelmapVolumeNode, seg)
      
      segLabel={'1':'liver','2':'spleen','3':'lung','4':'thyroid',\
              '5':'kidney','6':'pancreas','7':'gallbladder','8':'bladder',\
              '9':'aorta','10':'trachea','11':'sternum','12':'vertebra L1',\
              '13':'adrenal','14':'psoas major','15':'rectus',\
              '16':'bowel','17':'stomach','18':'heart'}
      
      for i in range(seg.GetSegmentation().GetNumberOfSegments()):
          segment=seg.GetSegmentation().GetNthSegment(i)
          segment.SetName(segLabel[segment.GetName()])

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

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

  def submitReview(self,currentRow,idx,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]
      
      row['reviewResult']=idx+1 #labkey has 1-based arrays
      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
    #