irAEMM/pythonScripts/runSegmentationnnUNet.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 addnnUNetCode(tempFile,fileNumber=0):
    sfx=getSuffix(tempFile)
    return re.sub(sfx,'_'+'{:04d}'.format(fileNumber)+sfx,tempFile)

def runnnUNet(setup,pars):
    args=[]
    #set the environment
    args.append(setup['paths']['nnUNetRunInference'])
    #location of input images
    args.extend(['-i',os.path.join(pars['tempBase'],'CT')])
    #output path is segmentations
    args.extend(['-o',os.path.join(pars['tempBase'],'segmentations')])
    #modelid, nnUNet internal rules.
    args.extend(['-t',pars['nnUNet']['ModelId']])
    #specify configuration (3d_fullres)
    args.extend(['-m',pars['nnUNet']['configuration']])
    print(args) 
    my_env = os.environ
    for key in pars['nnUNet']['env']:
        my_env[key]=pars['nnUNet']['env'][key]
      
    print(subprocess.run(args,env=my_env,check=True,stdout=subprocess.PIPE).stdout)

def getSegmentationFile(pars):
    #this is how deep medic stores files
    return os.path.join(pars['tempBase'],'segmentations',\
            pars['images']['CT']['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 CT
    ctDir=os.path.join(pars['tempBase'],'CT')
    if not os.path.isdir(ctDir):
        os.mkdir(ctDir)
    fullFile=os.path.join(ctDir,images['CT']['tempFile']) 

    fullFile=addnnUNetCode(fullFile)
    fb.readFileToFile(baseDir+'/'+row[images['CT']['queryField']],fullFile)
    
    #debug

    #run deep medic
    runnnUNet(setup,pars)

    #processed file is
    segFile=getSegmentationFile(pars)
    #SimpleITK.WriteImage(outImg,segFile)
    return segFile

        
def test(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"]["nixWrapper"])
    
    import nixWrapper
    
    nixWrapper.loadLibrary("labkeyInterface")#force reload
    import labkeyInterface
    import labkeyDatabaseBrowser
    import labkeyFileBrowser

    nixWrapper.loadLibrary("parseConfig")
    import parseConfig

    with open(parameterFile) as f:
        pars=json.load(f)
    
    pars=parseConfig.convert(pars)
    pars=parseConfig.convertValues(pars)
    #print(pars)


    

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"]["nixWrapper"])
    
    import nixWrapper
    
    nixWrapper.loadLibrary("labkeyInterface")#force reload
    import labkeyInterface
    import labkeyDatabaseBrowser
    import labkeyFileBrowser

    nixWrapper.loadLibrary("parseConfig")
    import parseConfig

    with open(parameterFile) as f:
        pars=json.load(f)
    
    pars=parseConfig.convert(pars)
    pars=parseConfig.convertValues(pars)
    
    project=pars['project']
    dataset=pars['targetQuery']
    schema=pars['targetSchema']
    view=pars['viewName']


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

    #start the database interface
    fconfig=os.path.join(fhome,'.labkey','network.json')
    net=labkeyInterface.labkeyInterface()
    net.init(fconfig)
    db=labkeyDatabaseBrowser.labkeyDB(net)
    fb=labkeyFileBrowser.labkeyFileBrowser(net)


    #all images from database
    ds=db.selectRows(project,schema,dataset,[],view)
    
    #input
    #use webdav to transfer file (even though it is localhost)

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

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

        outFile=outpath+'/'+outName

        #check if file is there
        if not fb.entryExists(outFile):

            
            segFile=getSegmentationFile(pars)
            #remove existing file
            if os.path.isfile(segFile):
                os.remove(segFile)
            
            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)
            print(segFile)
        #debug 

        #update database
        copyFields=[pars['participantField'],'SequenceNum','patientCode','visitCode']
        row['SequenceNum']+=0.001*float(pars['versionNumber'])
        filters=[{'variable':v,'value':str(row[v]),'oper':'eq'} for v in copyFields]
        filters.append({'variable':'Version','value':pars['version'],'oper':'eq'})

        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['Version']= pars['version']
        outRow['Segmentation']= outName
        print(db.modifyRows(mode,pars['project'],pars['segmentationSchema'],pars['segmentationQuery'],[outRow]))
        #pull results back to LabKey
    print("Done")


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