Adding overlay to preprocess.py

pythonScripts/preprocess.py

@@ -5,6 +5,7 @@ import subprocess
 import nibabel
 import shutil
 import sys
+import matplotlib.pyplot
 
 #nothing gets done if you do import
 
@@ -81,16 +82,30 @@ def getDicom(ofb,row,zipDir,rawDir,im,imageSelector,\
 
     return True    
 
-def updateRow(db,project,dataset,row,imageResampledField,gzFileNames,\
+
+def createOverlay(gzFileNames,overlayName):
+    ct=nibabel.load(gzFileNames['CT'])
+    pet=nibabel.load(gzFileNames['PET'])
+    aCT=ct.get_fdata()
+    aPET=pet.get_fdata()
+    n2=int(aCT.shape[1]/2)
+    matplotlib.pyplot.imshow(aCT[:,n2,:],cmap='gray')
+    matplotlib.pyplot.imshow(aPET[:,n2,:],cmap='viridis',alpha=0.5)
+    matplotlib.pyplot.savefig(overlayName)
+
+def updateRow(db,project,dataset,row,imageResampledField,gzFileNames,overlayName='overlay',\
         participantField='PatientId'):
     row['patientCode']=getPatientLabel(row,participantField)
     row['visitCode']=getVisitLabel(row)
+    
     for im in imageResampledField:
         row[imageResampledField[im]]=gzFileNames[im]
+    row['overlay']=overlayName
     db.modifyRows('update',project,'study',dataset,[row])
  
 
 def main(parameterFile):
+    debug=False
     fhome=os.path.expanduser('~')
     with open(os.path.join(fhome,".labkey","setup.json")) as f:
         setup=json.load(f)
@@ -168,9 +183,11 @@ def main(parameterFile):
         #build/check remote directory structure
         remoteDir=fb.buildPathURL(project,['preprocessedImages',\
             getPatientLabel(row,participantField),getVisitLabel(row)])
+        overlayName=getStudyLabel(row,participantField)+'_overlay.png'
 
         gzRemoteFiles={im:remoteDir+'/'+f\
             for (im,f) in gzFileNames.items()}
+        overlayRemote=remoteDir+'/'+overlayName
 
         remoteFilePresent=[fb.entryExists(f)\
             for f in gzRemoteFiles.values()]
@@ -189,7 +206,7 @@ def main(parameterFile):
         if sampleOK and all(remoteFilePresent):
             print("Entry for row done.")
             updateRow(db,project,dataset,row,imageResampledField,\
-                gzFileNames,participantField)
+                gzFileNames,overlayName,participantField)
             continue
 
     
@@ -218,6 +235,8 @@ def main(parameterFile):
         gzFiles={im:f+".gz"\
             for (im,f) in volumeFiles.items()}
 
+
+
         filesPresent=[os.path.isfile(f) for f in gzFiles.values()]
     
     
@@ -251,6 +270,11 @@ def main(parameterFile):
                 outText=subprocess.check_output([gzip,f])
                 print(outText.decode('utf-8'))
 
+
+        #create overlay
+        overlayLocal=os.path.join(processedDir,overlayName)
+        createOverlay(gzFiles,overlayLocal)
+
         #upload local files to remote
         for im in gzFiles:
         #for local,remote in zip(gzFiles,gzRemoteFiles):
@@ -258,15 +282,18 @@ def main(parameterFile):
             remote=gzRemoteFiles[im]
             print("Uploading {}".format(local))
             fb.writeFileToFile(local,remote)
-
+        #upload overlay:
+        print(f'Uploading {overlayLocal}')
+        fb.writeFileToFile(overlayLocal,overlayRemote)
 
         #update row and let it know where the processed files are
-        updateRow(db,project,dataset,row,imageResampledField,gzFileNames,\
+        updateRow(db,project,dataset,row,imageResampledField,gzFileNames,overlayName,\
             participantField)
    
 
         #cleanup
-        shutil.rmtree(studyDir)
+        if not debug:
+            shutil.rmtree(studyDir)
     
 
         if i==-1: