irAEMM/pythonScripts/generateFigures.py

import nibabel
import os
import json
import sys
import numpy
import matplotlib.pyplot
#import chardet

def buildPath(server,project,imageDir,patientCode,visitCode,imageType):
    path='/'.join([server,'labkey/_webdav',project,'@files',imageDir,patientCode,visitCode])
    tail='_notCropped_2mmVoxel'
    if imageType=='Segm':
        tail='_v5'
    path+='/'+patientCode+'-'+visitCode+'_'+imageType+tail+'.nii.gz'
    return path


def getCOWAxis(seg,val,axis):
#returns center of weight for segmentation image where val is selected
    if axis==0:
        #2,1 or 1,1
        i1=2
        i2=1
    if axis==1:
        #0,1 or 2,0
        i1=2
        i2=0
    if axis==2:
        #0,0 or 1,0
        i1=1
        i2=0
    
    s=numpy.sum(numpy.sum(seg==val,i1),i2)
    x=numpy.arange(len(s))
    s0=0
    for i in x:
        if s[i]==0:
            continue
        s0=i
        break
    s1=len(s)-1
    revIdx=numpy.arange(len(s),0,-1)-1
    for i in revIdx:
        if s[i]==0:
            continue
        s1=i
        break
    s1+=1
    try:
        sm=numpy.average(x,weights=s)
    except ZeroDivisionError:
        print('getCOWaxis - Zero division error')
        raise

    return [s0,numpy.average(x,weights=s),s1]

def getGeometry(seg,val):
#return center of weight of segmentation seg for segment val as a 3D vector
    return [getCOWAxis(seg,val,x) for x in [0,1,2]]

def getCOW(geom):
    return [x[1] for x in geom]

def getRange(geom):
    return [[x[0],x[2]] for x in geom]


def plot(imgs,t,val,tempBase):


    segColors=[[0,0,0],[0.1,0.1,0.1],[0,0.2,0],[1,0,0],[0,0,1],[1,0,1]]

#3-lungs 4-thyroid 5-bowel
    delta=20
    if val==4:
        delta=40
    window=350
    level=40
    geo=getGeometry(imgs['Segm'],val)
    cowF=getCOW(geo)
    rng=getRange(geo)
    #print(rng)
    cowI=[int(x) for x in cowF]
    segment=imgs['Segm']==val

    i0=rng[0][0]-delta
    if i0<0:
        i0=0
    i1=rng[0][1]+delta
    if i1>imgs['CT'].shape[0]:
        i1=imgs['CT'].shape[0]
    k0=rng[2][0]-delta
    if k0<0:
        k0=0
    k1=rng[2][1]+delta
    if k1>imgs['CT'].shape[2]:
        k1=imgs['CT'].shape[2]

    if t=='CT':
        v0=level-0.5*window
        v1=v0+window
        matplotlib.pyplot.imshow(imgs['CT'][i0:i1,cowI[1],k0:k1].transpose(),cmap='gray',
                vmin=v0,vmax=v1)
    if t=='PET':
        matplotlib.pyplot.imshow(imgs['PET'][i0:i1,cowI[1],k0:k1].transpose(),cmap='inferno')
        #blueish

    if t=='CT':
        rgb=segColors[val]
    if t=='PET':
        rgb=[1,1,1]

    colors = [rgb+[c] for c in numpy.linspace(0,1,100)]
                                                                                    
    cmap = matplotlib.colors.LinearSegmentedColormap.from_list('mycmap', colors, N=5)

    matplotlib.pyplot.imshow(segment[i0:i1,cowI[1],k0:k1].transpose(), cmap=cmap, alpha=0.2)
    matplotlib.pyplot.gca().invert_yaxis()
    outfile=os.path.join(tempBase,'slice{}_{}.png'.format(t,val))
    matplotlib.pyplot.savefig(outfile)
    return outfile



def main(parameterFile):
#mask for segmentations

    setupJSON=os.path.join(os.path.expanduser('~'),'.labkey','setup.json')
    with open(setupJSON) as f:
        setup=json.load(f)
            
    sys.path.insert(0,setup["paths"]["nixWrapper"])
    import nixWrapper 
             
    nixWrapper.loadLibrary("labkeyInterface")

    import labkeyInterface
    import labkeyDatabaseBrowser
    import labkeyFileBrowser

    fconfig=os.path.join(os.path.expanduser('~'),'.labkey','network.json')
                 
    net=labkeyInterface.labkeyInterface()
    net.init(fconfig)
    db=labkeyDatabaseBrowser.labkeyDB(net)
    fb=labkeyFileBrowser.labkeyFileBrowser(net)

    tempBase=os.path.join(os.path.expanduser('~'),'temp')
    
    with open(parameterFile) as f:
        pars=json.load(f)

    project=pars['project']
    dataset=pars['targetQuery']
    schema=pars['targetSchema']
    view=pars['viewName']
    
    reportSchema=pars['reportSchema']
    reportQuery=pars['reportQuery']
    participantField=pars['participantField']

    #all images from database
    ds=db.selectRows(project,schema,dataset,[],view)
    
    #input
    imageResampledField={"CT":"ctResampled","PET":"petResampled","patientmask":"ROImask"}

    rows=ds['rows']
    #rows=[r for r in rows if r[participantField]=='8701/08']
    #rows=[ds['rows'][0]]

    for r in rows:
        missingCodes=[r[f]==None for f in ['patientCode','visitCode']]
        if any(missingCodes):
            print('[{}/{}] - Skipping, missing codes'.\
                    format(r[participantField],r['SequenceNum']))
            continue


        #print(r)
        iTypes=['CT','PET','Segm']
        needToCalculate=False
        for t in ['CT','PET']:
            idFilter={'variable':'patientCode','value':r['patientCode'],'oper':'eq'}
            visitFilter={'variable':'visitCode','value':r['visitCode'],'oper':'eq'}
            verFilter={'variable':'version','value':pars['version'],'oper':'eq'}
            typeFilter={'variable':'type','value':t,'oper':'eq'}
            
            ds2=db.selectRows(project,reportSchema,reportQuery,\
                    [idFilter,visitFilter,verFilter,typeFilter])
            
            if len(ds2['rows'])==0:
                #skip if row is present
                #there are in fact multiple rows for multiple organs...
                needToCalculate=True
                break

        if not needToCalculate:
            print('[{}/{}] - done'.format(r[participantField],r['SequenceNum']))
            continue
        
        ds1=db.selectRows(project,pars['segmentationSchema'],pars['segmentationQuery'],\
                [idFilter,visitFilter,verFilter])
              
        #check if CT, PET and Segm images are set
        imagesAvailable=[r[imageResampledField[t]] for t in ['CT','PET']]
        imagesAvailable=[f!=None for f in imagesAvailable]
        try:
            imagesAvailable.append(ds1['rows'][0]['segmentation']!=None)
        except IndexError:
            imagesAvailable.append(False)

        if not all(imagesAvailable):
            print('[{}/{}] Skipping - not all images available :{}'.\
                    format(r[participantField],r['SequenceNum'],imagesAvailable))
            continue



        imgs={} 
        for t in iTypes:
            
            try:
                imagePath=r['_labkeyurl_'+imageResampledField[t]]
            except KeyError:
                imagePath=ds1['rows'][0]['_labkeyurl_segmentation']

            localPath=os.path.join(tempBase,'image'+t+'.nii.gz')
            if os.path.isfile(localPath):
                os.remove(localPath)
            fb.readFileToFile(imagePath,localPath)
            img=nibabel.load(localPath)
            imgs[t]=img.get_fdata()

        print('Loading completed')

        for t in ['CT','PET']:
            for val in [3,4,5]:


                try:
                    outfile=plot(imgs,t,val,tempBase)
                except ZeroDivisionError:
                    print('[{}/{}] - Skipping. Failed to plot for organ [{}]'.\
                            format(r[participantField],r['SequenceNum'],val))
                    continue
                
                remoteDir=fb.buildPathURL(project,
                        [pars['imageDir'],r['patientCode'],r['visitCode']])
                imageFile=r['patientCode']+'-'+r['visitCode']+'_'+t+'_{}'.\
                        format(val)+'_'+pars['version']+'.png'
                remoteFile='/'.join([remoteDir,imageFile])
                fb.writeFileToFile(outfile,remoteFile)
                print('Uploaded {}'.format(remoteFile))
                os.remove(outfile)  
                organFilter={'variable':'organ','value':'{}'.format(val),'oper':'eq'}
                typeFilter['value']=t
                ds3=db.selectRows(project,reportSchema,reportQuery,\
                        [idFilter,visitFilter,verFilter,organFilter,typeFilter])
                if len(ds3['rows'])>0:
                    mode='update'
                    frow=ds3['rows'][0]
                else:
                    mode='insert'
                    frow={}
                    for f in [participantField,'patientCode','visitCode']:
                        frow[f]=r[f]
                frow['organ']='{}'.format(val)
                frow['type']=t
                frow['version']=pars['version']
                frow['file']=imageFile
                db.modifyRows(mode,project,reportSchema,reportQuery,[frow])
        print('Images uploaded')
    print('Done')


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