<gh_stars>10-100
from nipype.interfaces.base import BaseInterface, \
    BaseInterfaceInputSpec, traits, File, TraitedSpec, InputMultiPath, Directory
from nipype.utils.filemanip import split_filename
import nibabel as nb
import numpy as np
import os

class ConsensusInputSpec(BaseInterfaceInputSpec):
    in_Files = traits.Either(InputMultiPath(File(exists=True)),
                                Directory(exists=True),
                                traits.Str(),
                                traits.List(),
                                mandatory=True)
    maskfile = File(exists=True, desc='total target mask', mandatory=True)

class ConsensusOutputSpec(TraitedSpec):
    #out_File = File(exists=True, desc="out_File")
    variation_mat = File(exists=True, desc="variation_mat")
    consensus_mat = File(exists=True, desc="consensus_mat")

class Consensus(BaseInterface):
    input_spec = ConsensusInputSpec
    output_spec = ConsensusOutputSpec

    def _get_filelist(self, trait_input):
        if os.path.isdir(trait_input[0]):
            filelist = []
            for directory in trait_input:
                for root, dirnames, fnames in os.walk(directory):
                    for f in fnames:
                        if f.endswith('.nii'):
                            filelist.append(os.path.join(root,f))
            return filelist
        if os.path.isfile(trait_input[0]): 
            return trait_input

    def makeConsensus(self, eachFile, mask):
        clustermap=nb.load(eachFile).get_data()
        maskedmap = clustermap[np.where(np.squeeze(mask))]
        consensus = np.zeros((len(maskedmap),len(maskedmap)))
        for j in range(len(maskedmap)):
            consensus[j] = maskedmap == maskedmap[j]
        return consensus

    def _run_interface(self, runtime):
        mask = nb.load(self.inputs.maskfile).get_data()
        src_paths = self._get_filelist(self.inputs.in_Files)
        _, base, _ = split_filename(self.inputs.in_Files[0])
        
        cons_dim = len(nb.load(src_paths[0]).get_data())
        totalConsensus = np.zeros((cons_dim,cons_dim), dtype=np.float64)
        for src_path in src_paths:
            totalConsensus += self.makeConsensus(src_path)
        ##average across all consensus instances and output##
        totalConsensus = totalConsensus/len(src_paths)
        cImg = nb.Nifti1Image(totalConsensus, None)
        nb.save(cImg, os.path.abspath(base+'_ConsensusMat.nii'))
        ##make consensus into stability measure## remove stability measure for now...
        #likeness = abs(totalConsensus-0.5)
        #stability = np.mean(likeness,axis=0)
        ##make into NiftiImage##
        #nImg = nb.Nifti1Image(stability, None)
        #nb.save(nImg, os.path.abspath(base+'_Stability.nii'))
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        _, base, _ = split_filename(self.inputs.in_Files[0])
        #outputs["out_File"] = os.path.abspath(base+'_Stability.nii')
        outputs["variation_mat"] = os.path.abspath(base+'_VariationMat.nii')
        outputs["consensus_mat"] = os.path.abspath(base+'_ConsensusMat.nii')
        return outputs

from nipype.interfaces.base import BaseInterface, \
    BaseInterfaceInputSpec, traits, File, TraitedSpec, InputMultiPath, Directory
from nipype.utils.filemanip import split_filename
import nibabel as nb
import numpy as np
import os

class ConsensusInputSpec(BaseInterfaceInputSpec):
    in_Files = traits.Either(InputMultiPath(File(exists=True)),
                                Directory(exists=True),
                                traits.Str(),
                                traits.List(),
                                mandatory=True)
    maskfile = File(exists=True, desc='total target mask', mandatory=True)

class ConsensusOutputSpec(TraitedSpec):
    #out_File = File(exists=True, desc="out_File")
    variation_mat = File(exists=True, desc="variation_mat")
    consensus_mat = File(exists=True, desc="consensus_mat")

class Consensus(BaseInterface):
    input_spec = ConsensusInputSpec
    output_spec = ConsensusOutputSpec

    def _get_filelist(self, trait_input):
        if os.path.isdir(trait_input[0]):
            filelist = []
            for directory in trait_input:
                for root, dirnames, fnames in os.walk(directory):
                    for f in fnames:
                        if f.endswith('.nii'):
                            filelist.append(os.path.join(root,f))
            return filelist
        if os.path.isfile(trait_input[0]): 
            return trait_input

    def makeConsensus(self, eachFile, mask):
        clustermap=nb.load(eachFile).get_data()
        maskedmap = clustermap[np.where(np.squeeze(mask))]
        consensus = np.zeros((len(maskedmap),len(maskedmap)))
        for j in range(len(maskedmap)):
            consensus[j] = maskedmap == maskedmap[j]
        return consensus

    def _run_interface(self, runtime):
        mask = nb.load(self.inputs.maskfile).get_data()
        src_paths = self._get_filelist(self.inputs.in_Files)
        _, base, _ = split_filename(self.inputs.in_Files[0])
        
        cons_dim = len(nb.load(src_paths[0]).get_data())
        totalConsensus = np.zeros((cons_dim,cons_dim), dtype=np.float64)
        for src_path in src_paths:
            totalConsensus += self.makeConsensus(src_path)
        ##average across all consensus instances and output##
        totalConsensus = totalConsensus/len(src_paths)
        cImg = nb.Nifti1Image(totalConsensus, None)
        nb.save(cImg, os.path.abspath(base+'_ConsensusMat.nii'))
        ##make consensus into stability measure## remove stability measure for now...
        #likeness = abs(totalConsensus-0.5)
        #stability = np.mean(likeness,axis=0)
        ##make into NiftiImage##
        #nImg = nb.Nifti1Image(stability, None)
        #nb.save(nImg, os.path.abspath(base+'_Stability.nii'))
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        _, base, _ = split_filename(self.inputs.in_Files[0])
        #outputs["out_File"] = os.path.abspath(base+'_Stability.nii')
        outputs["variation_mat"] = os.path.abspath(base+'_VariationMat.nii')
        outputs["consensus_mat"] = os.path.abspath(base+'_ConsensusMat.nii')
        return outputs
#out_File = File(exists=True, desc="out_File") ##average across all consensus instances and output## ##make consensus into stability measure## remove stability measure for now... #likeness = abs(totalConsensus-0.5) #stability = np.mean(likeness,axis=0) ##make into NiftiImage## #nImg = nb.Nifti1Image(stability, None) #nb.save(nImg, os.path.abspath(base+'_Stability.nii')) #outputs["out_File"] = os.path.abspath(base+'_Stability.nii')