irAEMM/pythonScripts/runSegmentation.py

import os
import json
import re
import subprocess
import nibabel
import shutil
import sys
import pathlib
import SimpleITK
import numpy

#nothing gets done if you do import

def getPatientLabel(row,participantField='PatientId'):
    return row[participantField].replace('/','_') 

def getVisitLabel(row):
    return 'VISIT_'+str(int(row['SequenceNum']))

def getStudyLabel(row,participantField='PatientId'):
    return getPatientLabel(row,participantField)+'-'+getVisitLabel(row)


def updateRow(project,dataset,row,imageResampledField,gzFileNames,\
        participantField='PatientId'):
    row['patientCode']=getPatientLabel(row,participantField)
    row['visitCode']=getVisitLabel(row)
    for im in imageResampledField:
        row[imageResampledField[im]]=gzFileNames[im]
    db.modifyRows('update',project,'study',dataset,[row])
 
def replacePatterns(infile,outfile,replacePatterns):
    of=open(outfile,'w')
    with open(infile,'r') as f:
        data=f.read()
        for p in replacePatterns:
            val=replacePatterns[p]
            data=re.sub(p,val,data)
    of.write(data)
    of.close()
    
def valueSubstitution(pars,val):
    if val.find('__home__')>-1:
        val=re.sub(r'__home__',os.path.expanduser('~'),val)

    return path

def getSuffix(tempFile):
    p=pathlib.Path(tempFile)
    return ''.join(p.suffixes)

def getSegmImagePath(tempFile):
    sfx=getSuffix(tempFile)
    return re.sub(sfx,'_Segm'+sfx,tempFile)

def addVersion(tempFile,version):
    sfx=getSuffix(tempFile)
    return re.sub(sfx,'_'+version+sfx,tempFile)

def normalizeCT(ctFile,maskFile):
    
    im=SimpleITK.ReadImage(ctFile)
    mask=SimpleITK.ReadImage(maskFile)
    nm=SimpleITK.GetArrayFromImage(im)
    nmask=SimpleITK.GetArrayFromImage(mask)
    mu=numpy.mean(nm[nmask>0])
    st=numpy.std(nm[nmask>0])
    nm[nmask>0]=(nm[nmask>0]-mu)/st
    nm[nmask==0]=0
    im1=SimpleITK.GetImageFromArray(nm)
    im1.SetOrigin(im.GetOrigin())
    im1.SetSpacing(im.GetSpacing())
    im1.SetDirection(im.GetDirection())
    SimpleITK.WriteImage(im1,ctFile)

def runDeepMedic(setup,pars):
    args=[]
    args.append(os.path.join(setup['paths']['deepMedicVE'],'bin','python'))
    args.append(setup['paths']['deepMedicRun'])
    args.append('-model')
    args.append(pars['deepmedic']['config']['model']['out'])
    args.append('-test')
    args.append(pars['deepmedic']['config']['test']['out'])
    args.append('-dev')
    args.append('cpu')
    print(args) 
    print(subprocess.run(args,check=True,stdout=subprocess.PIPE).stdout)

def runDeepMedicDocker(setup,pars):
    args=[]
    args.extend(['docker-compose','-f',pars['deepmedic']['segmentationdmYAML'],'up'])
    print(args) 
    print(subprocess.run(args,check=True,stdout=subprocess.PIPE).stdout)


def getSegmentationFile(pars):
    #this is how deep medic stores files
    return getSegmImagePath(\
            os.path.join(pars['tempBase'],'output','predictions','currentSession','predictions',\
            pars['images']['segmentations']['tempFile'])
            )

def runSegmentation(fb,row,pars,setup):
    
     
    #download to temp file (could be a fixed name)
    project=pars['project']
    images=pars['images']
    participantField=pars['participantField']
    baseDir=fb.formatPathURL(project,pars['imageDir']+'/'+\
        getPatientLabel(row,participantField)+'/'+\
        getVisitLabel(row))
    
    #download 
    fullFile={key:os.path.join(pars['tempBase'],images[key]['tempFile']) for key in images}
    for im in images:
        if 'queryField' in images[im]:
            fb.readFileToFile(baseDir+'/'+row[images[im]['queryField']],fullFile[im])
   
    #normalize 

    normalizeCT(fullFile['CT'],fullFile['patientmask'])

    #update templates to know which files to process


    #run deep medic
    runDeepMedicDocker(setup,pars)
    
    #processed file is
    segFile=getSegmentationFile(pars)
    #SimpleITK.WriteImage(outImg,segFile)
    return segFile

        
def main(parameterFile):
    
    fhome=os.path.expanduser('~')


    with open(os.path.join(fhome,".labkey","setup.json")) as f:
        setup=json.load(f)

    sys.path.insert(0,setup["paths"]["labkeyInterface"])
    import labkeyInterface
    import labkeyDatabaseBrowser
    import labkeyFileBrowser

    sys.path.append(setup['paths']['parseConfig'])
    import parseConfig

    with open(parameterFile) as f:
        pars=json.load(f)
    
    pars=parseConfig.convert(pars)
    pars=parseConfig.convertValues(pars)
    print(pars)
    #images=pars['images']
    #ctFile=os.path.join(pars['tempBase'],images['CT']['tempFile'])
    #maskFile=os.path.join(pars['tempBase'],images['patientmask']['tempFile'])
    #normalizeCT(ctFile,maskFile)


    

def doSegmentation(parameterFile):
    fhome=os.path.expanduser('~')


    with open(os.path.join(fhome,".labkey","setup.json")) as f:
        setup=json.load(f)

    sys.path.insert(0,setup["paths"]["labkeyInterface"])
    import labkeyInterface
    import labkeyDatabaseBrowser
    import labkeyFileBrowser

    sys.path.append(setup['paths']['parseConfig'])
    import parseConfig


    fconfig=os.path.join(fhome,'.labkey','network.json')

    net=labkeyInterface.labkeyInterface()
    net.init(fconfig)
    db=labkeyDatabaseBrowser.labkeyDB(net)
    fb=labkeyFileBrowser.labkeyFileBrowser(net)

    with open(parameterFile) as f:
        pars=json.load(f)

    pars=parseConfig.convert(pars)
    pars=parseConfig.convertValues(pars)
    print(pars)
    
    #update the config
    cfg=pars['deepmedic']['config']
    for c in cfg:
        replacePatterns(cfg[c]['template'],\
                cfg[c]['out'],\
                pars['replacePattern'])


    project=pars['project']
    dataset=pars['targetQuery']
    schema=pars['targetSchema']


    tempBase=pars['tempBase']
    if not os.path.isdir(tempBase):
        os.makedirs(tempBase)


    #all images from database
    ds=db.selectRows(project,schema,dataset,[])

    
    #input
    #use webdav to transfer file (even though it is localhost)

 
    i=0
    for row in ds["rows"]:
       

        #check if file is already there
        #dummy tf to get the suffix
        tf=getSegmentationFile(pars,'XX')
        outpath=fb.buildPathURL(pars['project'],[pars['imageDir'],row['patientCode'],row['visitCode']])
        outName=addVersion(\
                getSegmImagePath(\
                    getStudyLabel(row,pars['participantField'])+getSuffix(tf)),\
                pars['version'])

        outFile=outpath+'/'+outName

        #check if file is there
        if not fb.entryExists(outFile):
            segFile=runSegmentation(fb,row,pars,setup)
            #copy file to file
            #normally I would update the targetQuery, but it contains previously set images
            #copy to labkey
            #fb.writeFileToFile(segFile,outFile)

        

        #separate script (set version!)

        #update database
        copyFields=[pars['participantField'],'SequenceNum','patientCode','visitCode']
        filters=[{'variable':v,'value':str(row[v]),'oper':'eq'} for v in copyFields]
        ds1=db.selectRows(pars['project'],pars['segmentationSchema'],pars['segmentationQuery'],filters)
        if len(ds1['rows'])>0:
            mode='update'
            outRow=ds1['rows'][0]
        else:
            mode='insert'
            outRow={v:row[v] for v in copyFields}
        outRow[pars['version']]= outName
        db.modifyRows(mode,pars['project'],pars['segmentationSchema'],pars['segmentationQuery'],[outRow])
        #pull results back to LabKey
        i+=1
        if i==1:
            break
    print("Done")


if __name__ == '__main__':
    main(sys.argv[1])
    #sys.exit()