#!/usr/bin/env python
#
# SHBASECOPYRIGHT
#
# FSL installer script.
#
"""This is the FSL installation script, which can be used to install FSL.

This script must:

 - be able to be executed with Python 2.7 or newer.

 - be able to be executed in a "vanilla" Python environment, with no third
   party dependencies.

 - be self-contained, with no dependencies on any other modules (apart from
   the Python standard library).

 - be importable as a Python module - this script contains functions and
   classes that may be used by other scripts.
"""


from __future__ import print_function, division, unicode_literals

import functools      as ft
import os.path        as op
import subprocess     as sp
import textwrap       as tw
import                   argparse
import                   contextlib
import                   datetime
import                   fnmatch
import                   getpass
import                   hashlib
import                   json
import                   logging
import                   os
import                   platform
import                   pwd
import                   readline
import                   shlex
import                   shutil
import                   ssl
import                   sys
import                   tempfile
import                   threading
import                   time
import                   traceback

try:                import urllib.request as urlrequest
except ImportError: import urllib         as urlrequest


try:                import urllib.parse as urlparse
except ImportError: import                 urlparse

try:                import queue
except ImportError: import Queue as queue


PYVER = sys.version_info[:2]


log = logging.getLogger(__name__)


# this sometimes gets set to fslinstaller.pyc, so rstrip c
__absfile__ = op.abspath(__file__).rstrip('c')


__version__ = '3.5.3'
"""Installer script version number. This must be updated
whenever a new version of the installer script is released.
"""


DEFAULT_INSTALLATION_DIRECTORY = op.join(op.expanduser('~'), 'fsl')
"""Default FSL installation directory. """


DEFAULT_ROOT_INSTALLATION_DIRECTORY = '/usr/local/fsl/'
"""Default FSL installation directory when the installer is run as root. """


FSL_RELEASE_MANIFEST = 'https://fsl.fmrib.ox.ac.uk/fsldownloads/' \
                       'fslconda/releases/manifest.json'
"""URL to download the FSL installer manifest file from. The installer
manifest file is a JSON file which contains information about available FSL
versions.

See the download_manifest function, and an example manifest file
in test/data/manifest.json, for more details.

A custom manifest URL can be specified with the -a/--manifest command-line
option.
"""


FSL_DEV_RELEASES = 'https://fsl.fmrib.ox.ac.uk/fsldownloads/' \
                   'fslconda/releases/devreleases.txt'
"""URL to the devreleases.txt file, which contains a list of available
internal/development FSL releases. See the download_dev_releases function
for more details.
"""


# List of modifiers which can be used to change how
# a message is printed by the printmsg function.
INFO      = 1
IMPORTANT = 2
QUESTION  = 3
PROMPT    = 4
WARNING   = 5
ERROR     = 6
EMPH      = 7
EMPHASIS  = 7
UNDERLINE = 8
RESET     = 9
ANSICODES = {
    INFO      : '\033[37m',         # Light grey
    IMPORTANT : '\033[92m',         # Green
    QUESTION  : '\033[36m\033[4m',  # Blue+underline
    PROMPT    : '\033[36m\033[1m',  # Bright blue+bold
    WARNING   : '\033[93m',         # Yellow
    ERROR     : '\033[91m',         # Red
    EMPHASIS  : '\033[1m',          # White+bold
    UNDERLINE : '\033[4m',          # Underline
    RESET     : '\033[0m',          # Used internally
}


def printmsg(*args, **kwargs):
    """Prints a sequence of strings formatted with ANSI codes. Expects
    positional arguments to be of the form::

        printable, ANSICODE, printable, ANSICODE, ...

    :arg log: Must be specified as a keyword argument. If True (default),
              the message is logged.

    All other keyword arguments are passed through to the print function.
    """

    args     = list(args)
    blockids = [i for i in range(len(args)) if (args[i] not in ANSICODES)]
    logmsg   = kwargs.pop('log', True)
    uncoded  = ''

    for i, idx in enumerate(blockids):
        if i == len(blockids) - 1:
            slc = slice(idx + 1, None)
        else:
            slc = slice(idx + 1, blockids[i + 1])

        msg      = args[idx]
        msgcodes = args[slc]
        msgcodes = [ANSICODES[c] for c in msgcodes]
        msgcodes = ''.join(msgcodes)
        uncoded += msg

        print('{}{}{}'.format(msgcodes, msg, ANSICODES[RESET]), end='')

    if len(blockids) > 0:
        print(**kwargs)
        if logmsg:
            log.debug(uncoded)

    sys.stdout.flush()


def prompt(promptmsg, *msgtypes, **kwargs):
    """Prompts the user for some input. msgtypes and kwargs are passed
    through to the printmsg function.
    """
    printmsg(promptmsg, *msgtypes, end='', log=False, **kwargs)

    if PYVER[0] == 2: response = raw_input(' ').strip()
    else:             response = input(    ' ').strip()

    log.debug('%s: %s', promptmsg, response)

    return response


def identify_platform():
    """Figures out what platform we are running on. Returns a platform
    identifier string - one of:

      - "linux-64" (Linux, x86_64)
      - "macos-64" (macOS, x86_64)
      - "macos-M1" (macOS, M1)

    Note that these identifiers are for FSL releases, and are not the
    same as the platform identifiers used by conda.
    """

    platforms = {
        ('linux',  'x86_64') : 'linux-64',
        ('darwin', 'x86_64') : 'macos-64',
        ('darwin', 'arm64')  : 'macos-M1',
    }

    system = platform.system().lower()
    cpu    = platform.machine()
    key    = (system, cpu)

    if key not in platforms:
        supported = ', '.join(['[{}, {}]' for s, c in platforms])
        raise Exception('This platform [{}, {}] is unrecognised or '
                        'unsupported! Supported platforms: {}'.format(
                            system, cpu, supported))

    return platforms[key]


def check_need_admin(dirname):
    """Returns True if dirname needs administrator privileges to write to,
    False otherwise.
    """
    # os.supports_effective_ids added in
    # python 3.3, so can't be used here
    return not os.access(dirname, os.W_OK | os.X_OK)


def get_admin_password(action='install FSL'):
    """Prompt the user for their administrator password. An Exception is raised
    if an incorrect password is entered three times.a

    :arg action: String which describes what the password is needed for, i.e.:
                 "Your administrator password is needed to {action}"
    :returns:    the validated administrator password
    """

    def validate_admin_password(password):
        proc = Process.sudo_popen(['true'], password, stdin=sp.PIPE)
        proc.communicate()
        return proc.returncode == 0

    msg = 'Your administrator password is needed to {}'.format(action)

    for attempt in range(3):
        if attempt == 0: msg = '{}:'.format(msg)
        else:            msg = '{} [attempt {} of 3]:'.format(msg, attempt + 1)
        printmsg(msg, IMPORTANT, end='')
        password = getpass.getpass('')
        valid    = validate_admin_password(password)

        if valid:
            printmsg('Password accepted', INFO)
            break
        else:
            printmsg('Incorrect password', WARNING)

    if not valid:
        raise Exception('Incorrect password')

    return password


def isstr(s):
    """Returns True if s is a string, False otherwise. Works on python 2.7
    and >=3.3.
    """
    try:              return isinstance(s, basestring)
    except Exception: return isinstance(s, str)


def match_any(s, patterns):
    """Test if the string s matches any of the fnmatch-style patterns.
    Returns the matched pattern, or None.
    """
    for pat in patterns:
        if fnmatch.fnmatch(s, pat):
            return pat
    return None


@contextlib.contextmanager
def tempdir(override_dir=None, change_into=True, delete=True):
    """Returns a context manager which creates, changes into, and returns a
    temporary directory, and then deletes it on exit (unless delete is False).

    If override_dir is not None, instead of creating and changing into a
    temporary directory, this function just changes into override_dir.
    """

    if override_dir is None: tmpdir = tempfile.mkdtemp()
    else:                    tmpdir = override_dir

    prevdir = os.getcwd()

    try:
        if change_into:
            os.chdir(tmpdir)
        yield tmpdir

    finally:
        if change_into:
            os.chdir(prevdir)
        if delete and override_dir is None:
            shutil.rmtree(tmpdir)


@contextlib.contextmanager
def tempfilename(permissions=None, delete=True):
    """Returns a context manager which creates a temporary file, yields its
    name, then deletes the file on exit.
    """

    fname = None

    try:
        tmpf  = tempfile.NamedTemporaryFile(delete=False)
        fname = tmpf.name

        tmpf.close()

        if permissions:
            os.chmod(fname, permissions)

        yield fname

    finally:
        if delete and fname and op.exists(fname):
            os.remove(fname)


def sha256(filename, check_against=None, blocksize=1048576):
    """Calculate the SHA256 checksum of the given file. If check_against
    is provided, it is compared against the calculated checksum, and an
    error is raised if they are not the same.
    """

    hashobj = hashlib.sha256()

    with open(filename, 'rb') as f:
        while True:
            block = f.read(blocksize)
            if len(block) == 0:
                break
            hashobj.update(block)

    checksum = hashobj.hexdigest()

    if check_against is not None:
        if checksum != check_against:
            raise Exception('File {} does not match expected checksum '
                            '({})'.format(filename, check_against))

    return checksum


def clean_environ():
    """Return a dict containing a set of sanitised environment variables.

    All FSL and conda related variables are removed.
    """
    env = os.environ.copy()
    for v in list(env.keys()):
        if any(('FSL' in v, 'CONDA' in v, 'PYTHON' in v)):
            env.pop(v)
    return env


def install_environ(fsldir, username=None, password=None):
    """Returns a dict containing some environment variables that should
    be added to the shell environment when the FSL conda environment is
    being installed.
    """
    env = {}
    # post-link scripts call $FSLDIR/share/fsl/sbin/createFSLWrapper
    # (part of fsl/base), which will only do its thing if the following
    # env vars are set
    env['FSL_CREATE_WRAPPER_SCRIPTS'] = '1'
    env['FSLDIR']                     = fsldir

    # Make sure HTTP proxy variables, if set,
    # are available to the conda env command
    for v in ['http_proxy', 'https_proxy',
              'HTTP_PROXY', 'HTTPS_PROXY']:
        if v in os.environ:
            env[v] = os.environ[v]

    # FSL environments which source packages from the internal
    # FSL conda channel will refer to the channel as:
    #
    # http://${FSLCONDA_USERNAME}:${FSLCONDA_PASSWORD}/abc.com/
    #
    # so we need to set those variables
    if username: env['FSLCONDA_USERNAME'] = username
    if password: env['FSLCONDA_PASSWORD'] = password

    return env


def download_file(url,
                  destination,
                  progress=None,
                  blocksize=131072,
                  ssl_verify=True):
    """Download a file from url, saving it to destination. """

    def default_progress(downloaded, total):
        pass

    if progress is None:
        progress = default_progress

    log.debug('Downloading %s ...', url)

    # Path to local file
    if op.exists(url):
        url = 'file:' + urlrequest.pathname2url(op.abspath(url))

    # We create and use an unconfigured SSL
    # context to disable SSL verification.
    # Otherwise pass None causes urlopen to
    # use default behaviour. The context
    # argument is not available in py3.3
    kwargs = {}
    if (not ssl_verify) and (PYVER != (3, 3)):
        printmsg('Skipping SSL verification - this '
                 'is not recommended!', WARNING)
        kwargs['context'] = ssl.SSLContext(ssl.PROTOCOL_TLS)

    req = None

    try:
        # py2: urlopen result cannot be
        # used as a context manager
        req = urlrequest.urlopen(url, **kwargs)
        with open(destination, 'wb') as outf:

            try:             total = int(req.headers['content-length'])
            except KeyError: total = None

            downloaded = 0

            progress(downloaded, total)
            while True:
                block = req.read(blocksize)
                if len(block) == 0:
                    break
                downloaded += len(block)
                outf.write(block)
                progress(downloaded, total)

    finally:
        if req:
            req.close()


def download_manifest(url, workdir=None):
    """Downloads the installer manifest file, which contains information
    about available FSL versions, and the most recent version number of the
    installer (this script).

    The manifest file is a JSON file. Lines beginning
    with a double-forward-slash are ignored. See test/data/manifest.json
    for an example.

    This function modifies the manifest structure by adding a 'version'
    attribute to all FSL build entries.
    """

    log.debug('Downloading FSL installer manifest from %s', url)

    with tempdir(workdir):

        try:
            download_file(url, 'manifest.json')
        except Exception as e:
            log.debug('Error downloading FSL release manifest from %s',
                      url, exc_info=True)
            raise Exception('Unable to download FSL release manifest '
                            'from {} [{}]!'.format(url, str(e)))

        with open('manifest.json') as f:
            lines = f.readlines()

    # Drop comments
    lines = [l for l in lines if not l.lstrip().startswith('//')]

    manifest = json.loads('\n'.join(lines))

    # Add "version" to every build
    for version, builds in manifest['versions'].items():
        if version == 'latest':
            continue
        for build in builds:
            build['version'] = version

    return manifest


def download_dev_releases(url, workdir=None):
    """Downloads the FSL_DEV_RELEASES file. This file contains a list of
    available development manifest URLS. Returns a list of tuples, one
    for each development release, with each tuple containing:

      - URL to the manifest file
      - Most recent release tag
      - Date of the release
      - Commit hash (on the fsl/conda/manifest repository)
      - Branch name (on the fsl/conda/manifest repository)

    The list is sorted by date, newest first.
    """

    # parse a dev manifest file name, returning
    # a sequence containing the tage, date, commit
    # hash, and branch name. Dev manifest files
    # are named like so:
    #
    #   manifest-<tag>.<date>.<commit>.<branch>.json
    #
    # where <tag> is the tag of the most recent
    # public FSL release, and everything else is
    # self-explanatory.
    def parse_devrelease_name(url):
        name = urlparse.urlparse(url).path
        name = op.basename(name)
        name = name.lstrip('manifest-').rstrip('.json')
        # Awkward - the tag may have periods in it
        name = name.rsplit('.', 3)
        return name

    # list of (url, tag, date, commit, branch)
    devreleases = []

    with tempdir(workdir):

        try:
            download_file(url, 'devreleases.txt')
        except Exception as e:
            log.debug('Error downloading devreleases.txt from %s',
                      url, exc_info=True)
            raise Exception('Unable to download development manifest '
                            'list from {}!'.format(url))

        with open('devreleases.txt', 'rt') as f:
            urls = f.read().strip().split('\n')
            urls = [l.strip() for l in urls]

        for url in urls:
            devreleases.append([url] + parse_devrelease_name(url))

    # sort by date, newest first
    return sorted(devreleases, key=lambda r: r[2], reverse=True)


class Progress(object):
    """Simple progress reporter. Displays one of the following:

       - If both a value and total are provided, a progress bar is shown
       - If only a value is provided, a cumulative count is shown
       - If nothing is provided, a spinner is shown.

    Use as a context manager, and call the update method to report progress,
    e,g:

        with Progress('%') as p:
            for i in range(100):
                p.update(i + 1, 100)
    """

    def __init__(self,
                 label='',
                 transform=None,
                 fmt='{:.1f}',
                 total=None,
                 width=None):
        """Create a Progress reporter.

        :arg label:     Units (e.g. "MB", "%",)

        :arg transform: Function to transform values (see e.g.
                        Progress.bytes_to_mb)

        :arg fmt:       Template string used to format value / total.

        :arg total:     Maximum value - overrides the total value passed to
                        the update method.

        :arg width:     Maximum width, if a progress bar is displayed. Default
                        is to automatically infer the terminal width (see
                        Progress.get_terminal_width).
        """

        if transform is None:
            transform = Progress.default_transform

        self.width     = width
        self.fmt       = fmt.format
        self.total     = total
        self.label     = label
        self.transform = transform

        # used by the spin function
        self.__last_spin = None

    @staticmethod
    def default_transform(val, total):
        return val, total

    @staticmethod
    def bytes_to_mb(val, total):
        if val   is not None: val   = val   / 1048576
        if total is not None: total = total / 1048576
        return val, total

    @staticmethod
    def percent(val, total):
        if val is None or total is None:
            return val, total
        return 100 * (val / total), 100

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        printmsg('', log=False)

    def update(self, value=None, total=None):

        if total is None:
            total = self.total

        value, total = self.transform(value, total)

        if value is None and total is None:
            self.spin()
        elif value is not None and total is None:
            self.count(value)
        elif value is not None and total is not None:
            self.progress(value, total)

    def spin(self):

        symbols = ['|', '/', '-',  '\\']

        if self.__last_spin is not None: last = self.__last_spin
        else:                            last = symbols[-1]

        idx  = symbols.index(last)
        idx  = (idx + 1) % len(symbols)
        this = symbols[idx]

        printmsg(this, end='\r', log=False)
        self.__last_spin = this

    def count(self, value):

        value = self.fmt(value)

        if self.label is None: line = '{} ...'.format(value)
        else:                  line = '{}{} ...'.format(value, self.label)

        printmsg(line, end='\r', log=False)

    def progress(self, value, total):

        value = min(value, total)

        # arbitrary fallback of 50 columns if
        # terminal width cannot be determined
        if self.width is None: width = Progress.get_terminal_width(50)
        else:                  width = self.width

        fvalue = self.fmt(value)
        ftotal = self.fmt(total)
        suffix = '{} / {} {}'.format(fvalue, ftotal, self.label).rstrip()

        # +5: - square brackets around bar
        #     - space between bar and tally
        #     - space+spin at the end
        width     = width - (len(suffix) + 5)
        completed = int(round(width * (value  / total)))
        remaining = width - completed
        progress  = '[{}{}] {}'.format('#' * completed,
                                       ' ' * remaining,
                                       suffix)

        printmsg(progress, end='', log=False)
        printmsg(' ', end='', log=False)
        self.spin()
        printmsg(end='\r', log=False)


    @staticmethod
    def get_terminal_width(fallback=None):
        """Return the number of columns in the current terminal, or fallback
        if it cannot be determined.
        """
        # os.get_terminal_size added in python
        # 3.3, so we try it but fall back to
        # COLUMNS, or tput as a last resort.
        try:
            return shutil.get_terminal_size()[0]
        except Exception:
            pass

        try:
            return int(os.environ['COLUMNS'])
        except Exception:
            pass

        try:
            result = Process.check_output('tput cols', log_output=False)
            return int(result.strip())
        except Exception:
            return fallback


class Process(object):
    """Container for a subprocess.Popen object, allowing non-blocking
    line-based access to its standard output and error streams via separate
    queues, while logging all outputs.

    Don't create a Process directly - use one of the following static methods:
     - Process.check_output
     - Process.check_call
     - Process.monitor_progress
    """


    def __init__(self,
                 cmd,
                 admin=False,
                 password=None,
                 log_output=True,
                 print_output=False,
                 append_env=None,
                 **kwargs):
        """Run the specified command. Starts threads to capture stdout and
        stderr.

        :arg cmd:          Command to run - passed through shlex.split, then
                           passed to subprocess.Popen

        :arg admin:        Run the command with administrative privileges

        :arg password:     Administrator password - can be None if admin is
                           False.

        :arg log_output:   If True, the command and all of its stdout/stderr
                           are logged.

        :arg print_output: If True, the command and all of its stdout/stderr
                           are logged.

        :arg append_env:   Dictionary of additional environment to be set when
                           the command is run.

        :arg kwargs:       Passed to subprocess.Popen
        """

        self.cmd     = cmd
        self.stdoutq = queue.Queue()
        self.stderrq = queue.Queue()

        if log_output:
            log.debug('Running %s [as admin: %s]', cmd, admin)

        self.popen = Process.popen(cmd, admin, password,
                                   append_env=append_env, **kwargs)

        # threads for consuming stdout/stderr
        self.stdout_thread = threading.Thread(
            target=Process.forward_stream,
            args=(self.popen.stdout, self.stdoutq, cmd,
                  'stdout', log_output, print_output))
        self.stderr_thread = threading.Thread(
            target=Process.forward_stream,
            args=(self.popen.stderr, self.stderrq, cmd,
                  'stderr', log_output, print_output))

        self.stdout_thread.daemon = True
        self.stderr_thread.daemon = True
        self.stdout_thread.start()
        self.stderr_thread.start()


    def wait(self):
        """Waits for the process to terminate, then waits for the stdout
        and stderr consumer threads to finish.
        """
        self.popen.wait()
        self.stdout_thread.join()
        self.stderr_thread.join()


    @property
    def returncode(self):
        """Process return code. Returns None until the process has terminated,
        and the stdout/stderr consumer threads have finished.
        """
        if self.popen.returncode is None: return None
        if self.stdout_thread.is_alive(): return None
        if self.stderr_thread.is_alive(): return None
        return self.popen.returncode


    @staticmethod
    def check_output(cmd, *args, **kwargs):
        """Behaves like subprocess.check_output. Runs the given command, then
        waits until it finishes, and return its standard output. An error
        is raised if the process returns a non-zero exit code.

        :arg cmd: The command to run, as a string
        """

        proc = Process(cmd, *args, **kwargs)
        proc.wait()

        if proc.returncode != 0:
            raise RuntimeError('This command returned an error: ' + cmd)

        stdout = ''
        while True:
            try:
                stdout += proc.stdoutq.get_nowait()
            except queue.Empty:
                break

        return stdout


    @staticmethod
    def check_call(cmd, *args, **kwargs):
        """Behaves like subprocess.check_call. Runs the given command, then
        waits until it finishes. An error is raised if the process returns a
        non-zero exit code.

        :arg cmd: The command to run, as a string
        """
        proc = Process(cmd, *args, **kwargs)
        proc.wait()
        if proc.returncode != 0:
            raise RuntimeError('This command returned an error: ' + cmd)


    @staticmethod
    def monitor_progress(cmd, total=None, *args, **kwargs):
        """Runs the given command(s), and shows a progress bar under the
        assumption that cmd will produce "total" number of lines of output.

        :arg cmd:      The commmand to run as a string, or a sequence of
                       multiple commands.

        :arg total:    Total number of lines of standard output to expect.

        :arg timeout:  Refresh rate in seconds. Must be passed as a keyword
                       argument.

        :arg progfunc: Function which returns a number indicating how far
                       the process has progressed.  If provided, this
                       function is called, instead of standard output
                       lines being monitored. The function is passed a
                       reference to the Process object. Must be passed as a
                       keyword argument.
        """

        timeout  = kwargs.pop('timeout',  0.5)
        progfunc = kwargs.pop('progfunc', None)

        if total is None: label = None
        else:             label = '%'

        if progfunc is None:
            nlines = [0]
            def progfunc(proc):
                try:
                    _         = proc.stdoutq.get_nowait()
                    nlines[0] = nlines[0] + 1
                except queue.Empty:
                    pass
                return nlines[0]

        if isstr(cmd): cmds = [cmd]
        else:          cmds =  cmd

        with Progress(label=label,
                      fmt='{:.0f}',
                      transform=Progress.percent) as prog:

            progcount = 0 if total else None

            for cmd in cmds:

                proc = Process(cmd, *args, **kwargs)
                prog.update(progcount, total)

                while proc.returncode is None:
                    time.sleep(timeout)
                    progcount = progfunc(proc) if total else None
                    prog.update(progcount, total)
                    proc.popen.poll()

                # force progress bar to 100% when finished
                if proc.returncode == 0:
                    prog.update(total, total)
                else:
                    raise RuntimeError('This command returned '
                                       'an error: ' + cmd)


    @staticmethod
    def forward_stream(stream,
                       queue,
                       cmd,
                       streamname,
                       log_output,
                       print_output):
        """Reads lines from stream and pushes them onto queue until popen
        is finished. Logs every line.

        :arg stream:       stream to forward
        :arg queue:        queue.Queue to push lines onto
        :arg cmd:          string - the command that is running
        :arg streamname:   string - 'stdout' or 'stderr'
        :arg log_output:   If True, log all stdout/stderr.
        :arg print_output: If True, print all stdout/stderr.
        """

        while True:
            line = stream.readline().decode('utf-8')
            if line == '':
                break
            queue.put(line)
            if log_output:
                log.debug(' [%s]: %s', streamname, line.rstrip())
            if print_output:
                print(' [{}]: {}'.format(streamname, line.rstrip()))


    @staticmethod
    def popen(cmd, admin=False, password=None, append_env=None, **kwargs):
        """Runs the given command via subprocess.Popen, as administrator if
        requested.

        :arg cmd:        The command to run, as a string

        :arg admin:      Whether to run with administrative privileges

        :arg pssword:    Administrator password. Only required if admin is
                         True.

        :arg append_env: Dictionary of additional environment to be set when
                         the command is run.

        :arg kwargs:     Passed to subprocess.Popen. stdin, stdout, and stderr
                         will be silently clobbered

        :returns:        The subprocess.Popen object.
        """

        admin = admin and os.getuid() != 0

        cmd              = shlex.split(cmd)
        kwargs['stdin']  = sp.PIPE
        kwargs['stdout'] = sp.PIPE
        kwargs['stderr'] = sp.PIPE

        if admin:
            proc = Process.sudo_popen(cmd, password, append_env, **kwargs)
        else:
            # if append_env has been specified,
            # add it to the normal env option.
            if append_env is not None:
                env = kwargs.get('env', os.environ.copy())
                env.update(append_env)
                kwargs['env'] = env

            proc = sp.Popen(cmd, **kwargs)

        return proc


    @staticmethod
    def sudo_popen(cmd, password, append_env=None, **kwargs):
        """Runs "sudo cmd" using subprocess.Popen. Used by Process.popen.
        Assumes that kwargs contains stdin=sp.PIPE
        """

        # sudo will not necessarily propagate environment
        # variables, and there is no guarantee that the
        # sudo -E option will work. So here we create a
        # wrapper shell script with "export VAR=VALUE"
        # statements for all environment variables that
        # are set.
        if append_env is None:
            append_env = {}

        # Make the wrapper script delete itself
        # after the command has been executed.
        with tempfilename(0o755, delete=False) as wrapper:
            with open(wrapper, 'wt') as f:
                f.write('#!/usr/bin/env sh\n')
                f.write('set -e\n')
                f.write('thisfile=$0\n')
                f.write('thisdir=$(cd $(dirname $0) && pwd)\n')
                for k, v in append_env.items():
                    f.write('export {}="{}"\n'.format(k, v))
                # shlex.join not available in py27
                f.write(' '.join(cmd) + '\n')
                f.write('cd ${thisdir} && rm ${thisfile}\n')

        cmd  = ['sudo', '-S', '-k', wrapper]
        proc = sp.Popen(cmd, **kwargs)
        proc.stdin.write('{}\n'.format(password).encode())
        proc.stdin.flush()
        return proc


@ft.total_ordering
class Version(object):
    """Class to represent and compare version strings.  Accepted version
    strings are of the form W.X.Y.Z, where W, X, Y, and Z are all integers.
    """
    def __init__(self, verstr):
        # Version identifiers for official FSL
        # releases will have up to four
        # components (X.Y.Z.W), but we accept
        # any number of (integer) components,
        # as internal releases may have more.
        components = []

        # ignore a leading "v", e.g. v1.2.3
        verstr = verstr.lower()
        if verstr.startswith('v'):
            verstr = verstr[1:]

        for comp in verstr.split('.'):
            try:              components.append(int(comp))
            except Exception: break

        self.components = components
        self.verstr     = verstr

    def __str__(self):
        return self.verstr

    def __eq__(self, other):
        for sn, on in zip(self.components, other.components):
            if sn != on:
                return False
        return len(self.components) == len(other.components)

    def __lt__(self, other):
        for p1, p2 in zip(self.components, other.components):
            if p1 < p2: return True
            if p1 > p2: return False
        return len(self.components) < len(other.components)


class Context(object):
    """Bag of information and settings created in the main function, and passed
    around this script.

    Several settings are lazily evaluated on first access, but once evaluated,
    their values are immutable.
    """

    def __init__(self, args, destdir=None, action='install FSL'):
        """Create the context with the argparse.Namespace object containing
        parsed command-line arguments.

        :arg args:    argparse.Namespace containing command-line arguments
        :arg destdir: Installation directory. If not provided, read from
                      args.dest, or read from the user,
        :arg action:  Passed to get_admin_password as a prompt.
        """

        self.args  = args
        self.shell = op.basename(os.environ.get('SHELL', 'sh')).lower()

        # These attributes are updated on-demand via
        # the property accessors defined below, or are
        # all updated via the finalise_settings method.
        self.__platform         = None
        self.__manifest         = None
        self.__devmanifest      = None
        self.__candidate_builds = None
        self.__build            = None
        self.__destdir          = destdir
        self.__need_admin       = None
        self.__admin_password   = None
        self.__action           = action

        # If the destination directory already exists,
        # and the user chooses to overwrite it, it is
        # moved so that, if the installation fails, it
        # can be restored. The new path is stored
        # here - refer to overwrite_destdir.
        self.old_destdir = None

        # The download_fsl_environment function stores
        # the path to the FSL conda environment file
        # and list of conda channels
        self.environment_file     = None
        self.environment_channels = None

        # The config_logging function stores the path
        # to the fslinstaller log file here.
        self.logfile = None


    def finalise_settings(self):
        """Finalise values for all information and settings in the Context.
        """
        self.manifest
        self.candidate_builds
        self.platform
        self.build
        self.destdir
        self.need_admin
        self.admin_password


    @property
    def platform(self):
        """The platform we are running on, e.g. "linux-64", "macos-64",
        "macos-M1". This identifier is used to determine which FSL build to
        install.

        Note that this function may report a different platform identifier than
        the actual platform - specifically, if running on a M1 mac, and there
        is no M1 FSL build for the requested FSL version, this function will
        report "macos-64". This is because some older FSL releases do not have
        M1 builds available.
        """
        if self.__platform is None:
            plat = identify_platform()

            # if M1, check that we have a suitable
            # FSL build, falling back to x86 if not.
            if plat == 'macos-M1':
                candidates = self.candidate_builds
                if not any(c['platform'] == 'macos-M1' for c in candidates):
                    plat = 'macos-64'

            self.__platform = plat

        return self.__platform


    @property
    def candidate_builds(self):
        """Query the manifest and return a list of available builds for the
        requested FSL release, for all platforms.
        """
        if self.__candidate_builds is not None:
            return self.__candidate_builds

        # defaults to "latest" if
        # not specified by the user
        fslversion = self.args.fslversion
        if fslversion is None:
            fslversion = 'latest'

        if fslversion not in self.manifest['versions']:
            available = ', '.join(self.manifest['versions'].keys())
            raise Exception(
                'FSL version "{}" is not available - available '
                'versions: {}'.format(fslversion, available))

        if fslversion == 'latest':
            fslversion = self.manifest['versions']['latest']

        self.__candidate_builds = list(self.manifest['versions'][fslversion])

        return self.__candidate_builds


    @property
    def build(self):
        """Returns a suitable FSL build (a dictionary entry from the FSL
        installer manifest) for the target platform.

        The returned dictionary has the following elements:
          - 'version'      FSL version.
          - 'platform':    Platform identifier (e.g. 'linux-64')
          - 'environment': Environment file to download
          - 'sha256':      Checksum of environment file
          - 'output':      Number of lines of expected output, for reporting
                           progress
        """

        if self.__build is not None:
            return self.__build

        # Find refs to a suitable build for this
        # platform. We assume that there is only
        # one default build for each platform.
        # in the list of builds for a given FSL
        # version.
        candidates = self.candidate_builds
        build      = None

        for candidate in candidates:
            if candidate['platform'] == self.platform:
                build = candidate
                break
        else:
            raise Exception(
                'Cannot find a version of FSL matching '
                'platform {}'.format(self.platform))

        printmsg('FSL {} selected for installation'.format(build['version']))

        self.__build = build
        return build


    @property
    def destdir(self):
        """Installation directory. If not specified at the command line, the
        user is prompted to enter a directory.
        """

        if self.__destdir is not None:
            return self.__destdir

        if os.getuid() != 0: defdestdir = DEFAULT_INSTALLATION_DIRECTORY
        else:                defdestdir = DEFAULT_ROOT_INSTALLATION_DIRECTORY

        # The loop below validates the destination directory
        # both when specified at commmand line or
        # interactively.  In either case, if invalid, the
        # user is re-prompted to enter a new destination.
        destdir = None
        if self.args.dest is not None: response = self.args.dest
        else:                          response = None

        while destdir is None:

            if response is None:
                printmsg('\nWhere do you want to install FSL?',
                         IMPORTANT, EMPHASIS)
                printmsg('Press enter to install to the default location '
                         '[{}]\n'.format(defdestdir), INFO)
                response = prompt('FSL installation directory [{}]:'.format(
                    defdestdir), QUESTION, EMPHASIS)
                response = response.rstrip(op.sep)

                if response == '':
                    response = defdestdir

            response  = op.expanduser(op.expandvars(response))
            response  = op.abspath(response)
            parentdir = op.dirname(response)
            if op.exists(parentdir):
                destdir = response
            else:
                printmsg('Destination directory {} does not '
                         'exist!'.format(parentdir), ERROR)
                response = None

        self.__destdir = destdir
        return self.__destdir


    @property
    def conda(self):
        """Return a path to the ``conda`` or ``mamba`` executable. """

        condabin = op.join(self.destdir, 'bin', 'conda')
        mambabin = op.join(self.destdir, 'bin', 'mamba')

        # If mamba is present, prefer it over conda
        if not self.args.conda: candidates = [mambabin, condabin]
        else:                   candidates = [condabin, mambabin]

        for c in candidates:
            if op.exists(c):
                return c


    @property
    def need_admin(self):
        """Returns True if administrator privileges will be needed to install
        FSL.
        """
        if self.__need_admin is not None:
            return self.__need_admin
        parentdir = op.dirname(self.destdir)
        self.__need_admin = check_need_admin(parentdir)
        return self.__need_admin


    @property
    def admin_password(self):
        """Returns the user's administrator password, prompting them if needed.
        """
        if self.__admin_password is not None:
            return self.__admin_password
        # need_admin may be None or False,
        # so don't rely on truthiness.
        if self.__need_admin == False:
            return None
        self.__admin_password = get_admin_password(self.__action)


    @property
    def manifest(self):
        """Returns the FSL installer manifest as a dictionary. """

        if self.__manifest is None:
            if self.devmanifest is not None:
                self.args.manifest = self.devmanifest

            self.__manifest = download_manifest(self.args.manifest,
                                                self.args.workdir)
        return self.__manifest


    @property
    def devmanifest(self):
        """Returns a URL to a development manifest to use for installation.
        This will only return a value if the --devrelease or --devlatest
        options are active.

        If this is the case, the FSL_DEV_RELEASES file is downloaded - this
        file contains a list of available development manifest URLS. The
        user is then prompted to choose which development manifest to use
        for the installation, unless --devlatest is active, in which case
        the newest manifest is selected.
        """
        if not self.args.devrelease:
            return None
        if self.__devmanifest == 'na':
            return None
        elif self.__devmanifest is not None:
            return self.__devmanifest

        devreleases = download_dev_releases(FSL_DEV_RELEASES,
                                            self.args.workdir)

        if len(devreleases) == 0:
            self.__devmanifest = 'na'
            return None

        self.__devmanifest = prompt_dev_release(devreleases,
                                                self.args.devlatest)

        return self.__devmanifest


    def run(self, process_func, *args, **kwargs):
        """Run a command via a static Process method.  Handles sudo/
        administrator authentication, and ensures that the shell
        environment in which the command is executed is sanitised.

        Can be used with Process.check_call, Process.check_output, and
        Process.monitor_progress. For example:

            ctx = Context(...)
            ctx.run(Process.check_call, 'my_command')
            ctx.run(Process.monitor_progress, 'my_command', total=100)
        """

        process_func = ft.partial(process_func, *args, **kwargs)

        # Clear any environment variables that refer to
        # existing FSL or conda installations, and ensure
        # that some specific FSL environment variables
        # are set while the command is running.  See
        # clean_environ and install_environ for more
        # details, and see Process.sudo_popen regarding
        # append_env.
        env        = clean_environ()
        append_env = install_environ(self.destdir,
                                     self.args.username,
                                     self.args.password)
        return process_func(admin=self.need_admin,
                            password=self.admin_password,
                            env=env,
                            append_env=append_env)


def check_rosetta_status(ctx):
    """Called from the main routine, before installation is attempted.  If
    running on a M1 macos machine, and a x86 version of FSL has been selected
    for installation, checks whether rosetta emulation is enabled. If so,
    does nothing further. Otherwise, prints a message and exits.
    """

    if not all((identify_platform() == 'macos-M1',
                ctx.platform        == 'macos-64')):
        return

    # Using the strategy discussed at
    # https://forum.latenightsw.com/t/possible-for-a-script-\
    #   to-test-whether-rosetta-2-is-installed/3207/6
    #
    # The pkgutil command should return 0 if
    # rosetta is installed, non-0 otherwise.
    try:
        Process.check_output('pkgutil --files com.apple.pkg.RosettaUpdateAuto')
    except RuntimeError:
        printmsg('Rosetta emulation does not appear to be enabled!\n', ERROR)
        printmsg('Enable rosetta emulation, and then run this installer '
                 'again. You can enable rosetta emulation by running this '
                 'command:\n', INFO)
        printmsg('  /usr/sbin/softwareupdate --install-rosetta '
                 '--agree-to-license\n', IMPORTANT)
        printmsg('Aborting installation', ERROR)
        sys.exit(1)
    # pkgutil command not found - should
    # never happen, but print a warning
    # just in case
    except Exception as e:
        printmsg('An error occurred calling the pkgutil command - this '
                 'may not be a problem, so I\'ll attempt to proceed '
                 'with the installation. ({}'.format(e), WARNING)


def list_available_versions(manifest):
    """Lists available FSL versions. """
    printmsg('Available FSL versions:', EMPHASIS)
    for version in manifest['versions']:
        if version == 'latest':
            continue
        printmsg(version, IMPORTANT, EMPHASIS)
        for build in manifest['versions'][version]:
            printmsg('  {}'.format(build['platform']), EMPHASIS, end=' ')
            printmsg(build['environment'], INFO)


def prompt_dev_release(devreleases, latest):
    """Prompts the user to select a development release.

    :arg devreleases: List of development releases, as returned by
                      download_dev_releases.
    :arg latest:      If True, the latest develeopment release is returned.
    """

    if len(devreleases) == 0:
        return None

    # automatically choose latest dev manifest?
    if latest:
        return devreleases[0][0]

    # show the user a list, ask them which one they want
    printmsg('Available development releases:', EMPHASIS)
    for i, (url, tag, date, commit, branch) in enumerate(devreleases):
        printmsg('  [{}]: {} [{} commit {}]'.format(
            i + 1, date, branch, commit), IMPORTANT)
    while True:
        selection = prompt('Which release would you like to '
                           'install? [1]:', PROMPT)
        if selection == '':
            selection = '1'
        try:
            selection = int(selection) - 1
        except Exception:
            continue
        if selection >= 0 and selection < len(devreleases):
            break
    return devreleases[selection][0]


def download_fsl_environment(ctx):
    """Downloads the environment specification file for the selected FSL
    version.

    Internal/development FSL versions may source packages from the internal
    FSL conda channel, which requires a username+password to authenticate.

    These are referred to in the environment file as ${FSLCONDA_USERNAME}
    and ${FSLCONDA_PASSWORD}.

    If the user has not provided a username+password on the command-line, they
    are prompted for them.

    The downloaded environment file may be modified - if the (hidden)
    --exclude_package option has been used.
    """

    build        = ctx.build
    url          = build['environment']
    checksum     = build.get('sha256', None)

    printmsg('Downloading FSL environment specification '
             'from {}...'.format(url))
    fname = url.split('/')[-1]
    download_file(url, fname)
    ctx.environment_file = op.abspath(fname)
    if (checksum is not None) and (not ctx.args.no_checksum):
        sha256(fname, checksum)

    # Environment files for internal/dev FSL versions
    # will list the internal FSL conda channel with
    # ${FSLCONDA_USERNAME} and ${FSLCONDA_PASSWORD}
    # as placeholders for the username/password.
    with open(fname, 'rt') as f:
        need_auth = '${FSLCONDA_USERNAME}' in f.read()

    # We need a username/password to access the internal
    # FSL conda channel. Prompt the user if they haven't
    # provided credentials.
    if need_auth and (ctx.args.username is None):
        printmsg('A username and password are required to install '
                 'this version of FSL.', WARNING, EMPHASIS)
        ctx.args.username = prompt('Username:').strip()
        ctx.args.password = getpass.getpass('Password: ').strip()

    # Conda expands environment variables within a
    # .condarc file, but *not* within an environment.yml
    # file. So to authenticate to our internal channel
    # without storing credentials anywhere in plain text,
    # we *move* the channel list from the environment.yml
    # file into $FSLDIR/.condarc.
    #
    # Here we extract the channels from the environment
    # file, and save them to ctx.environment_channels.
    # The install_miniconda function will then add the
    # channels to $FSLDIR/.condarc.
    #
    # We also remove any packages that the user has
    # requested to exclude from the installation.
    copy     = '.' + op.basename(ctx.environment_file)
    channels = []

    shutil.move(ctx.environment_file, copy)
    with open(copy,                 'rt') as inf, \
         open(ctx.environment_file, 'wt') as outf:

        in_channels_section = False

        for line in inf:

            # start of channels list
            if line.strip() == 'channels:':
                in_channels_section = True
                continue

            if in_channels_section:
                # end of channels list
                if not line.strip().startswith('-'):
                    in_channels_section = False
                else:
                    channels.append(line.split()[-1])
                    continue

            # Include/exclude packages upon user request
            pkgname = line.strip(' -').split()[0]
            exclude = match_any(pkgname, ctx.args.exclude_package)
            if exclude:
                log.debug('Excluding package %s (matched '
                          '--exclude_package %s)', line, exclude)
            else:
                outf.write(line)

    ctx.environment_channels = channels


def download_miniconda(ctx):
    """Downloads the miniconda/miniforge installer and saves it as
    "miniconda.sh".

    This function assumes that it is run within a temporary/scratch directory.
    """

    if ctx.args.miniconda is None:
        metadata = ctx.manifest['miniconda'][ctx.platform]
        url      = metadata['url']
        checksum = metadata['sha256']
    else:
        url      = ctx.args.miniconda
        checksum = None

    # Download
    printmsg('Downloading miniconda from {}...'.format(url))
    with Progress('MB', transform=Progress.bytes_to_mb) as prog:
        download_file(url, 'miniconda.sh', prog.update)
    if (not ctx.args.no_checksum) and (checksum is not None):
        sha256('miniconda.sh', checksum)


def install_miniconda(ctx):
    """Downloads the miniconda/miniforge installer, and installs it to the
    destination directory.

    This function assumes that it is run within a temporary/scratch directory.
    """

    metadata = ctx.manifest['miniconda'][ctx.platform]
    output   = metadata.get('output', '').strip()

    if output == '': output = None
    else:            output = int(output)

    # Install
    printmsg('Installing miniconda at {}...'.format(ctx.destdir))
    cmd = 'bash miniconda.sh -b -p {}'.format(ctx.destdir)
    ctx.run(Process.monitor_progress, cmd, total=output)

    # Avoid WSL filesystem issue
    # https://github.com/conda/conda/issues/9948
    cmd = 'find {} -type f -exec touch {{}} +'.format(ctx.destdir)
    ctx.run(Process.check_call, cmd)

    # Generate $FSLDIR/.condarc which contains
    # some default/fixed conda settings
    condarc = generate_condarc(ctx.destdir,
                               ctx.environment_channels,
                               ctx.args.skip_ssl_verify)
    with open('.condarc', 'wt') as f:
        f.write(condarc)

    ctx.run(Process.check_call, 'cp -f .condarc {}'.format(ctx.destdir))


def generate_condarc(fsldir, channels, skip_ssl_verify=False):
    """Called by install_miniconda. Generates content for a .condarc file to
    be saved in $FSLDIR/.condarc. This file contains some default values, and
    also enforces some settings so that they cannot be overridden by the
    user. For example. the list of conda channels is configured so that it
    cannot be overridden by a user's ~/.condarc file.

    See the following web pages for more details:
     - https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html
     - https://www.anaconda.com/blog/conda-configuration-engine-power-users
     - https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-channels.html
    """

    # Create .condarc config file
    pkgsdir = op.join(fsldir, 'pkgs')
    condarc = tw.dedent("""
    # FSL conda configuration file, auto-generated by the fslinstaller script.
    #
    # WARNING: This file may be automatically re-generated
    # without warning, so it is recommended that any custom
    # conda settings are stored elsewhere. Refer to the conda
    # documentation for more guidance:
    #
    # https://conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html
    # https://www.anaconda.com/blog/conda-configuration-engine-power-users
    # https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-channels.html

    # Try and make package downloads more robust
    remote_read_timeout_secs:    240
    remote_connect_timeout_secs: 20
    remote_max_retries:          10
    remote_backoff_factor:       5
    safety_checks:               warn

    # Disable caching of remote channel repodata.
    # This is a hack which is combined with the
    # WSL1 filesystem hack in the fslinstaller
    # script - because we have modified file
    # timestamps, conda will assume that its
    # channel repodata cache is up to date, and
    # will not bother refreshing it in the
    # commands that we run in the install_fsl
    # function. When we remove the WSL1 hack from
    # the fslinstaller, we can remove this config
    # setting.
    # https://github.com/conda/conda/issues/9948
    local_repodata_ttl: 0

    # Channel priority is important. In older versions
    # of FSL we placed the FSL conda channel at the
    # bottom (lowest priority) for legacy reasons (to
    # ensure that conda-forge versions of e.g. VTK were
    # preferred over legacy FSL conda versions).
    #
    # Use final/top/bottom marks to prevent the channel
    # priority order being modified by user ~/.condarc
    # configuration files.
    channel_priority: strict #!final

    # Fix the package cache at $FSLDIR/pkgs/
    pkgs_dirs: #!final
     - {} #!top #!bottom
    """.format(pkgsdir))

    if skip_ssl_verify:
        printmsg('Configuring conda to skip SSL verification '
                 '- this is not recommended!', WARNING)
        condarc += tw.dedent("""
        # Disable SSL verification when accessing
        # conda channels over https:// (the
        # fslinstaller script was called with
        # --skip_ssl_verify). NOT RECOMMENDED.
        ssl_verify: false
        """)

    channels = list(channels)
    if len(channels) > 0:
        channels[0]  += ' #!top'
        channels[-1] += ' #!bottom'
    condarc      += '\nchannels: #!final\n'
    for channel in channels:
        condarc += ' - {}\n'.format(channel)

    return condarc


def get_install_fsl_progress_reporting_method(ctx):
    """Figure out which reporting mechansim to use for reporting progress
    whilst FSL is being installed. The mechanism that is used has changed
    a few times.

    Returns a tuple containing values to pass to the Progress.monitor_progress
    function, either of which may be None if progress cannot be reported:

      - an integer value to pass as the total
      - a function to pass as the progfunc.
    """

    # We calculate installation progress in
    # one of a few ways, as we have changed
    # the mechanism a few times.  The
    # 'output/install' field in the manifest
    # gives us information about how to
    # report installation progress.
    fslver     = ctx.build['version']
    progparams = ctx.build.get('output', {}).get('install', None)
    pkgdir     = op.join(ctx.destdir, 'pkgs')

    # The first method (version 1) involves
    # progress reporting by monitoring number of
    # lines of standard output produced by
    # "conda env install". This is set to None,
    # as it is the default behaviour of the
    # Progress.monitor_progress function.
    progress_v1 = None

    # The second method involves progress
    # reporting by monitoring the number of
    # package files created in $FSLDIR/pkgs/
    def progress_v2(_):
        pkgs = os.listdir(pkgdir)
        pkgs = [p for p in pkgs if p.endswith('.conda') or p.endswith('.bz2')]
        return len(pkgs)

    progresses      = {}
    progresses['1'] = None
    progresses['2'] = progress_v2

    progval  = None
    progfunc = None

    # The output field may be either a
    # string, in which case we assume
    # version 2, or a dictionary containing
    # the progress reporting version, and
    # an integer value.
    if isstr(progparams):
        progval  = int(progparams)
        progfunc = progresses['2']

    # output field is a dict - versioned
    # progress reporting
    elif isinstance(progparams, dict):
        progval  = int(progparams['value'])
        progfunc = progresses[progparams['version']]

    return progval, progfunc


def install_fsl(ctx):
    """Install FSL into ctx.destdir (which is assumed to be a miniconda
    installation.

    This function assumes that it is run within a temporary/scratch directory.
    """

    progval, progfunc = get_install_fsl_progress_reporting_method(ctx)

    # We install FSL simply by running conda env
    # update -f env.yml.
    cmd = ctx.conda + ' env update -n base -f ' + ctx.environment_file

    # Make conda/mamba super verbose if the
    # hidden --debug option was specified.
    if ctx.args.debug:
        cmd += ' -v -v -v'

    printmsg('Installing FSL into {}...'.format(ctx.destdir))
    ctx.run(Process.monitor_progress, cmd,
            timeout=1, total=progval, progfunc=progfunc)


def finalise_installation(ctx):
    """Performs some finalisation tasks. Includes:
      - Saving the installed version to $FSLDIR/etc/fslversion
      - Saving this installer script and the environment file to
        $FSLDIR/etc/
    """
    with open('fslversion', 'wt') as f:
        f.write(ctx.build['version'])

    etcdir = op.join(ctx.destdir, 'etc')
    cmds   = [
        'cp fslversion {}'.format(etcdir),
        'cp {} {}'        .format(ctx.environment_file, etcdir)]

    for cmd in cmds:
        ctx.run(Process.check_call, cmd)


def post_install_cleanup(ctx, tmpdir):
    """Cleans up the FSL directory after installation. """

    cmds = [ctx.conda + ' clean -y --all']

    if tmpdir is not None:
        cmds.append('rm -rf ' + tmpdir)

    for cmd in cmds:
        ctx.run(Process.check_call, cmd)


def patch_file(filename, searchline, numlines, content):
    """Used by configure_shell and configure_matlab. Adds to, modifies,
    or creates the specified file.

    If a line matching searchline is found in the file, numlines (starting
    from searchline) are replaced with content.

    Otherwise, content is appended to the end of the file.
    """

    content = content.split('\n')

    if op.isfile(filename):
        with open(filename) as f:
            lines = [l.strip() for l in f.readlines()]
    else:
        lines = []

    # replace block
    try:
        idx   = lines.index(searchline)
        lines = lines[:idx] + content + lines[idx + numlines:]

    # append to end
    except ValueError:
        lines = lines + [''] + content

    with open(filename, 'wt') as f:
        f.write('\n'.join(lines))


def configure_shell(shell, homedir, fsldir):
    """Configures the user's shell environment (e.g. ~/.bash_profile).

    :arg shell:   User's shell (taken from the $SHELL environment variable
    :arg homedir: User's home directory, presumed to contain shell profile
                  file(s).
    :arg fsldir:  FSL installation directory
    """

    bourne_shells  = ['sh', 'bash', 'zsh', 'dash']
    csh_shells     = ['csh', 'tcsh']

    # we edit the first file that exists in
    # the list of candidate profile files.
    # They are attached as an attribute of
    # this function just for testing purposes
    # (see after function definition)
    shell_profiles = configure_shell.shell_profiles

    # DO NOT CHANGE the format of these configurations -
    # they are kept exactly as-is for compatibility with
    # legacy FSL installations, i.e. so we can modify
    # profiles with an existing configuration from older
    # FSL versions
    bourne_cfg = tw.dedent("""
    # FSL Setup
    FSLDIR={fsldir}
    PATH=${{FSLDIR}}/share/fsl/bin:${{PATH}}
    export FSLDIR PATH
    . ${{FSLDIR}}/etc/fslconf/fsl.sh
    """).format(fsldir=fsldir).strip()

    csh_cfg = tw.dedent("""
    # FSL Setup
    setenv FSLDIR {fsldir}
    setenv PATH ${{FSLDIR}}/share/fsl/bin:${{PATH}}
    source ${{FSLDIR}}/etc/fslconf/fsl.csh
    """).format(fsldir=fsldir).strip()

    if shell not in bourne_shells + csh_shells:
        printmsg('Shell {} not recognised - skipping environment '
                 'setup'.format(shell), WARNING, EMPHASIS)
        return

    if shell in bourne_shells: cfg = bourne_cfg
    else:                      cfg = csh_cfg

    # find the profile file to edit
    profile    = None
    candidates = [op.join(homedir, p)
                  for p in shell_profiles[shell]]
    for candidate in candidates:
        if op.isfile(candidate):
            profile = candidate
            break

    # if no candidate profile files
    # exist, fall back to the first one
    if profile is None:
        profile = candidates[0]

    printmsg('Adding FSL configuration to {}'.format(profile))

    patch_file(profile, '# FSL Setup', len(cfg.split('\n')), cfg)
configure_shell.shell_profiles = {'sh'   : ['.profile'],
                                  'bash' : ['.bash_profile', '.profile'],
                                  'dash' : ['.bash_profile', '.profile'],
                                  'zsh'  : ['.zprofile'],
                                  'csh'  : ['.cshrc'],
                                  'tcsh' : ['.tcshrc']}


def configure_matlab(homedir, fsldir):
    """Creates/appends FSL configuration code to ~/Documents/MATLAB/startup.m.
    """

    # DO NOT CHANGE the format of this configuration -
    # see in-line comments in configure_shell.
    cfg = tw.dedent("""
    % FSL Setup
    setenv( 'FSLDIR', '{fsldir}' );
    setenv('FSLOUTPUTTYPE', 'NIFTI_GZ');
    fsldir = getenv('FSLDIR');
    fsldirmpath = sprintf('%s/etc/matlab',fsldir);
    path(path, fsldirmpath);
    clear fsldir fsldirmpath;
    """).format(fsldir=fsldir).strip()

    matlab_dir = op.expanduser(op.join(homedir, 'Documents', 'MATLAB'))
    startup_m  = op.join(matlab_dir, 'startup.m')

    if not op.exists(matlab_dir):
        os.makedirs(matlab_dir)

    printmsg('Adding FSL configuration to {}'.format(startup_m))

    patch_file(startup_m, '% FSL Setup', len(cfg.split('\n')), cfg)


def self_update(manifest, workdir, checksum):
    """Checks to see if a newer version of the installer (this script) is
    available and if so, downloads it to a temporary file, and runs it in
    place of this script.
    """

    thisver   = Version(__version__)
    latestver = Version(manifest['installer']['version'])

    if latestver <= thisver:
        log.debug('Installer is up to date (this version: %s, '
                  'latest version: %s)', thisver, latestver)
        return

    log.debug('New version of installer is available '
              '(%s) - self-updating', latestver)

    tmpf = tempfile.NamedTemporaryFile(
        prefix='new_fslinstaller', delete=False, dir=workdir)
    tmpf.close()
    tmpf = tmpf.name

    download_file(manifest['installer']['url'], tmpf)

    if checksum:
        try:
            sha256(tmpf, manifest['installer']['sha256'])
        except Exception as e:
            printmsg('New installer file does not match expected '
                     'checksum! Skipping update.', WARNING)
            return

    # Don't try and update again - if for some
    # reason the online manifest reports a newer
    # version than what is available, we would
    # otherwise enter into an infinite loop.
    cmd = [sys.executable, tmpf] + sys.argv[1:] + ['--no_self_update']
    log.debug('Running new installer: %s', cmd)
    os.execv(sys.executable, cmd)


def overwrite_destdir(ctx):
    """Called by main if the destination directory already exists. Asks the
    user if they want to overwrite it. If they do, or if the --overwrite
    option was specified, the directory is moved, and then deleted after
    the installation succeeds.

    This function assumes that it is run within a temporary/scratch directory.
    """

    if not ctx.args.overwrite:
        printmsg()
        printmsg('Destination directory [{}] already exists!'
                 .format(ctx.destdir), WARNING, EMPHASIS)
        response = prompt('Do you want to overwrite it [y/N]?',
                          QUESTION, EMPHASIS)
        if response.lower() not in ('y', 'yes'):
            printmsg('Aborting installation', ERROR, EMPHASIS)
            sys.exit(1)

    # generate a unique name for the old
    # destination directory (to avoid
    # collisions if using the same workdir
    # repeatedly)
    i = 0
    while True:
        ctx.old_destdir = op.abspath('old_destdir{}'.format(i))
        i              += 1
        if not op.exists(ctx.old_destdir):
            break

    printmsg('Deleting directory {}'.format(ctx.destdir), IMPORTANT)
    ctx.run(Process.check_call,
            'mv {} {}'.format(ctx.destdir, ctx.old_destdir))


def parse_args(argv=None, include=None):
    """Parse command-line arguments, returns an argparse.Namespace object.

    :arg argv:    Command-line arguments.

    :arg include: List of arguments to parse. May be used by other scripts
                  which re-use some of the routines defined in this script.
                  The resulting argparse.Namespace object will contain values
                  of None for all arguments that are not included.
    """

    uid = os.getuid()

    if uid != 0: destdir = DEFAULT_INSTALLATION_DIRECTORY
    else:        destdir = DEFAULT_ROOT_INSTALLATION_DIRECTORY

    # on macOS, when Python is run with sudo,
    # op.expanduser('~') will return the
    # calling user's home directory, and not
    # the root home directory. This doesn't
    # really matter, as homedir is only used
    # for modifying the shell/matlab profile,
    # and this is automatically disabled via
    # the --no_env option when run as root. But
    # in case the user wants the root shell
    # profile modified (via the hidden
    # --root_env option), we use getpwuid to
    # determine the appropriate home directory.
    homedir = pwd.getpwuid(uid).pw_dir

    username = os.environ.get('FSLCONDA_USERNAME', None)
    password = os.environ.get('FSLCONDA_PASSWORD', None)

    options = {
        # regular options
        'version'      : ('-v', {'action'  : 'version',
                                 'version' : __version__}),
        'dest'         : ('-d', {'metavar' : 'DESTDIR'}),
        'overwrite'    : ('-o', {'action'  : 'store_true'}),
        'listversions' : ('-l', {'action'  : 'store_true'}),
        'no_env'       : ('-n', {'action'  : 'store_true'}),
        'no_shell'     : ('-s', {'action'  : 'store_true'}),
        'no_matlab'    : ('-m', {'action'  : 'store_true'}),
        'fslversion'   : ('-V', {'default' : 'latest'}),

        # hidden options
        'debug'           : (None, {'action'  : 'store_true'}),
        'username'        : (None, {'default' : username}),
        'password'        : (None, {'default' : password}),
        'no_checksum'     : (None, {'action'  : 'store_true'}),
        'skip_ssl_verify' : (None, {'action'  : 'store_true'}),
        'workdir'         : (None, {}),
        'homedir'         : (None, {'default' : homedir}),
        'devrelease'      : (None, {'action'  : 'store_true'}),
        'devlatest'       : (None, {'action'  : 'store_true'}),
        'manifest'        : (None, {}),
        'miniconda'       : (None, {}),
        'conda'           : (None, {'action'  : 'store_true'}),
        'no_self_update'  : (None, {'action'  : 'store_true'}),
        'exclude_package' : (None, {'action'  : 'append'}),
        'root_env'        : (None, {'action'  : 'store_true'}),
    }

    if include is None:
        include = list(options.keys())

    helps = {
        'version'      : 'Print installer version number and exit',
        'listversions' : 'List available FSL versions and exit',
        'dest'         : 'Install FSL into this folder (default: '
                         '{})'.format(destdir),
        'overwrite'    : 'Delete existing destination directory if it exists, '
                         'without asking',
        'no_env'       : 'Do not modify your shell or MATLAB configuration '
                         '(implies --no_shell and --no_matlab). When running '
                         'the installer script as the root user, the root '
                         'shell profile is never modified.',
        'no_shell'     : 'Do not modify your shell configuration',
        'no_matlab'    : 'Do not modify your MATLAB configuration',
        'fslversion'   : 'Install this specific version of FSL',

        # Enable verbose output when calling
        # mamba/conda.
        'debug'           : argparse.SUPPRESS,

        # Username / password for accessing
        # internal FSL conda channel, if an
        # internal/development release is being
        # installed. If not set, will be read from
        # FSLCONDA_USERNAME/FSLCONDA_PASSWORD
        # environment variables.
        'username'        : argparse.SUPPRESS,
        'password'        : argparse.SUPPRESS,

        # Do not automatically update the installer script,
        'no_self_update'  : argparse.SUPPRESS,

        # Install a development release. This
        # option will cause the installer to
        # download the devrelreases.txt file,
        # which contains a list of available
        # internal/development manifests. The
        # user will be prompted to choose one,
        # which will be propagated on to the
        # --manifest option. If --devlatest
        # is used, the most recent developmet
        # release is automatically selected.
        'devrelease'      : argparse.SUPPRESS,
        'devlatest'       : argparse.SUPPRESS,

        # Path/URL to alternative FSL release
        # manifest.
        'manifest'        : argparse.SUPPRESS,

        # Install miniconda from this path/URL,
        # instead of the one specified in the
        # FSL release manifest
        'miniconda'       : argparse.SUPPRESS,

        # Use conda and not mamba
        'conda'           : argparse.SUPPRESS,

        # Print debugging messages
        'debug'           : argparse.SUPPRESS,

        # Disable SHA256 checksum validation of downloaded files
        'no_checksum'     : argparse.SUPPRESS,

        # Store temp files in this directory
        # rather than in a temporary directory
        'workdir'         : argparse.SUPPRESS,

        # Treat this directory as user's home
        # directory, for the purposes of shell
        # configuration. Must already exist.
        'homedir'         : argparse.SUPPRESS,

        # Configure conda to skip SSL verification.
        # Not recommended.
        'skip_ssl_verify' : argparse.SUPPRESS,

        # Do not install packages matching this
        # fnmatch-style wildcard pattern. Can
        # be used multiple times.
        'exclude_package' : argparse.SUPPRESS,

        # If the installer is run as root, the
        # --no_env flag is automatically enabled
        # UNLESS this flag is also provided.
        'root_env'        : argparse.SUPPRESS,
    }

    # parse args
    parser = argparse.ArgumentParser()
    for option in include:
        shortflag, kwargs = options[option]
        flags             = ['--{}'.format(option)]
        if shortflag is not None:
            flags.insert(0, shortflag)
        parser.add_argument(*flags, help=helps[option], **kwargs)

    # parse_known_args so that newly added
    # args are ignored by older versions,
    # but will be parsed after self_update
    args = parser.parse_known_args(argv)[0]

    if getattr(args, 'fslversion', 'latest') != 'latest':
        if Version(args.fslversion) < Version('6.0.6'):
            printmsg(
                'This script can only be used to install FSL 6.0.6 or newer. '
                'Visit https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FslInstallation '
                'for information on installing older versions.', ERROR,
                EMPHASIS)
            sys.exit(1)


    # add placeholder values for excluded args
    for option in options.keys():
        if option not in include:
            setattr(args, option, None)

    # alternate home dir (for debugging)
    if args.homedir is not None:
        args.homedir = op.abspath(args.homedir)
        if not op.isdir(args.homedir):
            printmsg('Home directory {} does not exist!'.format(args.homedir),
                     ERROR, EMPHASIS)
            sys.exit(1)

    # --no-env is automatically enabled
    #  when installer is run as root
    if os.getuid() == 0 and not (args.root_env):
        args.no_env = True

    # don't modify shell profile
    if args.no_env:
        args.no_shell  = True
        args.no_matlab = True

    # use workdir rather than a tempdir
    if args.workdir is not None:
        args.workdir = op.abspath(args.workdir)
        if not op.exists(args.workdir):
            os.mkdir(args.workdir)

    # manifest takes priority over devrelease/devlatest
    if args.manifest is not None:
        args.devrelease = False
        args.devlatest  = False

    if args.manifest is None:
        args.manifest = FSL_RELEASE_MANIFEST

    if args.devlatest:
        args.devrelease = True

    if args.exclude_package is None:
        args.exclude_package = []

    # accept local path for manifest and environment
    if args.manifest is not None and op.exists(args.manifest):
        args.manifest = op.abspath(args.manifest)

    # accept local path for miniconda installer
    if args.miniconda is not None and op.exists(args.miniconda):
        args.miniconda = op.abspath(args.miniconda)

    return args


def config_logging(prefix='fslinstaller_', logdir=None):
    """Configures logging. Log messages are directed to
    $TMPDIR/fslinstaller_<unique_token>.log, or
    logdir/fslinstaller_<unique_token>.log
    """
    if logdir is None:
        logdir = tempfile.gettempdir()

    # Use a unique name for the log file
    # (important for multi-user systems)
    logfilef, logfile = tempfile.mkstemp(prefix=prefix,
                                         suffix='.log',
                                         dir=logdir)
    os.close(logfilef)

    handler   = logging.FileHandler(logfile)
    formatter = logging.Formatter(
        '%(asctime)s %(filename)s:%(lineno)4d: %(message)s', '%H:%M:%S')
    handler.setFormatter(formatter)
    log.addHandler(handler)
    log.setLevel(logging.DEBUG)

    return logfile


@contextlib.contextmanager
def handle_error(ctx):
    """Used by main as a context manager around the main installation steps.
    If an error occurs, prints some messages, performs some clean-up/
    restoration tasks, and exits.
    """

    try:
        yield

    except Exception as e:
        printmsg('\nERROR occurred during installation!', ERROR, EMPHASIS)
        printmsg('    {}\n'.format(e), INFO)

        # send traceback to log file
        tb = traceback.format_tb(sys.exc_info()[2])
        log.debug(''.join(tb))

        if op.exists(ctx.destdir):
            printmsg('Removing failed installation directory '
                     '{}'.format(ctx.destdir), WARNING)
            ctx.run(Process.check_call, 'rm -r ' + ctx.destdir)

        # overwrite_destdir moves the existing
        # destdir to a temp location, so we can
        # restore it if the installation fails
        if not op.exists(ctx.destdir) and (ctx.old_destdir is not None):
            printmsg('Restoring contents of {}'.format(ctx.destdir),
                     WARNING)
            ctx.run(Process.check_call,
                    'mv {} {}'.format(ctx.old_destdir, ctx.destdir))

        # copy log file to ~/ so it is
        # easier for the user to access
        date    = datetime.datetime.today().strftime('%Y%m%d%H%M%S')
        logfile = 'fsl_installation_{}.log'.format(date)
        logfile = op.join(op.expanduser('~'), logfile)
        shutil.copy(ctx.logfile, logfile)

        printmsg('\nFSL installation failed!\n', ERROR, EMPHASIS)
        printmsg('Please check the log file - it may contain some '
                 'more information to help you diagnose the problem: '
                 '{}\n'.format(logfile), WARNING, EMPHASIS)
        sys.exit(1)


def main(argv=None):
    """Installer entry point. Downloads and installs miniconda and FSL, and
    configures the user's environment.
    """

    printmsg('FSL installer version:', EMPHASIS, UNDERLINE, end='')
    printmsg(' {}'.format(__version__))
    printmsg('Press CTRL+C at any time to cancel installation', INFO)

    if os.getuid() == 0:
        printmsg('Running the installer script as root user is discouraged! '
                 'You should run this script as a regular user - you will be '
                 'asked for your administrator password if required.',
                 WARNING, EMPHASIS)

    args    = parse_args(argv)
    logfile = config_logging(logdir=args.workdir)

    log.debug(' '.join(sys.argv))
    printmsg('Installation log file: {}\n'.format(logfile), INFO)

    ctx         = Context(args)
    ctx.logfile = logfile

    if not args.no_self_update:
        self_update(ctx.manifest, args.workdir, not args.no_checksum)

    if args.listversions:
        list_available_versions(ctx.manifest)
        sys.exit(0)

    try:
        ctx.finalise_settings()
    except Exception as e:
        printmsg('An error has occurred: {}'.format(e), ERROR)
        sys.exit(1)

    check_rosetta_status(ctx)

    # Do everything in a temporary directory,
    # but don't delete it, as some operations
    # may be run as root. The tempdir is
    # deleted within the post_install_cleanup
    # function.
    with tempdir(args.workdir, delete=False) as tmpdir:

        if args.workdir is not None:
            tmpdir = None

        # Ask the user if they want to overwrite
        # an existing installation
        if op.exists(ctx.destdir):
            overwrite_destdir(ctx)

        download_fsl_environment(ctx)

        printmsg('\nInstalling FSL in {}\n'.format(ctx.destdir), EMPHASIS)
        with handle_error(ctx):
            download_miniconda(ctx)
            install_miniconda(ctx)
            install_fsl(ctx)
            finalise_installation(ctx)
            post_install_cleanup(ctx, tmpdir)

    if not args.no_shell:
        configure_shell(ctx.shell, args.homedir, ctx.destdir)
    if not args.no_matlab:
        configure_matlab(args.homedir, ctx.destdir)

    printmsg('\nFSL successfully installed\n', IMPORTANT)
    if not args.no_shell:
        printmsg('Open a new terminal, or log out and log back in, '
                 'for the environment changes to take effect.', INFO)


if __name__ == '__main__':
    sys.exit(main())

fslinstaller.py
RicardoZiTseng/SACNet
from . import *
import warnings
warnings.filterwarnings("ignore", category=UserWarning)

print("\n\nPlease cite the following paper when using SACNet. Oh shit, we don't have paper yet.")
print("If you have any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).\n")

try:
    from art import *
    tprint("SACNet")
except:
    pass

sacnet/__init__.py
RicardoZiTseng/SACNet
from __future__ import absolute_import
from . import *
from .network import *
from .transform import *

sacnet/architectures/__init__.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import torch
import torch.nn as nn
import torch.nn.functional as F

def predict_flow(in_planes, num_of_flow, dist):
    flow = nn.Conv3d(in_planes, num_of_flow, kernel_size=3, padding=1)
    flow.weight = nn.Parameter(dist.sample(flow.weight.shape))
    flow.bias = nn.Parameter(torch.zeros(flow.bias.shape))
    return flow

class conv_block(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        """
        Instiatiate the conv block
            :param in_channels: number of input channels
            :param out_channels: number of output channels
            :param stride: stride of the convolution
        """
        super(conv_block, self).__init__()
        self.main = nn.Conv3d(in_channels, out_channels, 3, stride, 1)
        self.activation = nn.LeakyReLU(0.2)

    def forward(self, x):
        """
        Pass the input through the conv_block
        """
        out = self.main(x)
        out = self.activation(out)
        return out

class SingleNormalization(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, img):
        img_min = torch.min(img)
        img_max = torch.max(img)
        norm_img = torch.sub(img, img_min) / torch.sub(img_max, img_min)
        return norm_img

class MultiNormalization(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, img_1, img_2):
        img_1_min = torch.min(img_1)
        img_2_min = torch.min(img_2)
        img_min = torch.min(img_1_min, img_2_min)

        img_1_max = torch.max(img_1)
        img_2_max = torch.max(img_2)
        img_max = torch.max(img_1_max, img_2_max)

        norm_img_1 = torch.sub(img_1, img_min) / torch.sub(img_max, img_min)
        norm_img_2 = torch.sub(img_2, img_min) / torch.sub(img_max, img_min)

        return norm_img_1, norm_img_2

class ResizeImage(nn.Module):
    def __init__(self, downsample_factor=4):
        super().__init__()
        if downsample_factor == 1:
            self.sampler = None
        else:
            self.sampler = nn.Sequential(
                *[nn.AvgPool3d(kernel_size=3, stride=2, padding=1)] * int((downsample_factor // 2))
            )
    
    def forward(self, img):
        if self.sampler is not None:
            return self.sampler(img)
        else:
            return img

class ResizeTransform(nn.Module):
    """
    [ResizeTransform] Resize a transform, which involves resizing the vector field *and* rescaling it.
    """

    def __init__(self, vel_resize, ndims):
        super(ResizeTransform, self).__init__()
        self.factor = 1.0 / vel_resize
        self.mode = 'linear'
        if ndims == 2:
            self.mode = 'bi' + self.mode
        elif ndims == 3:
            self.mode = 'tri' + self.mode

    def forward(self, x):
        if self.factor < 1:
            # resize first to save memory
            x = F.interpolate(x, align_corners=True, scale_factor=self.factor, mode=self.mode)
            x = self.factor * x

        elif self.factor > 1:
            # multiply first to save memory
            x = self.factor * x
            x = F.interpolate(x, align_corners=True, scale_factor=self.factor, mode=self.mode)

        # don't do anything if resize is 1
        return x

class PreActBlock(nn.Module):
    '''Pre-activation version of the BasicBlock.'''
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, bias=False):
        super(PreActBlock, self).__init__()
        self.conv1 = nn.Conv3d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=bias)
        self.conv2 = nn.Conv3d(planes, planes, kernel_size=3, stride=1, padding=1, bias=bias)

        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv3d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=bias))

    def forward(self, x):
        out = F.leaky_relu(x, negative_slope=0.2)
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv1(out)
        out = self.conv2(F.leaky_relu(out, negative_slope=0.2))
        out += shortcut
        return out



sacnet/architectures/layers.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from sacnet.architectures.network import ResUNet
import torch
import torch.nn as nn
from sacnet.architectures.layers import (MultiNormalization, ResizeImage,
                                          ResizeTransform, SingleNormalization)
from sacnet.architectures.transform import EPIWarp


class AbstractModel(nn.Module):
    """The Abstract class of EPI distortion correction network.
    
    Parameters
    ----------
    in_dim : int
        The number of input images which will be fed into network.
    problem_type : str "MultiPE" or "SinglePE". Default is "MultiPE".
        ...
    """
    def __init__(self, in_dim, problem_type="MultiPE", **kargs):
        super().__init__()
        self.problem_type = problem_type
        if self.problem_type == 'SinglePE':
            self.norm = SingleNormalization()
        elif self.problem_type == 'MultiPE':
            self.norm_1 = MultiNormalization()
            self.norm_2 = SingleNormalization()

    def forward(self, input_list, train=True):
        raise NotImplementedError("Do not use this abstract class.")

    def normalize(self, input_list):
        """The instance specific normalization layer.
        For single PE problem, each image is normalized by min-max normalization.
        For multiple PE problem, the two b0 images are co-normalized.
        """
        normalize_input_list = None
        if self.problem_type == 'SinglePE':
            normalize_input_list = [self.norm(img) for img in input_list]
        elif self.problem_type == 'MultiPE':
            if len(input_list) == 2:
                normalize_input_list = self.norm_1(*input_list)
            elif len(input_list) == 3:
                norm_0, norm_1 = self.norm_1(input_list[0], input_list[1])
                norm_2 = self.norm_2(input_list[2])
                normalize_input_list = [norm_0, norm_1, norm_2]
        else:
            raise ValueError("Length of input_list must be 2 or 3, but got {}".format(len(input_list)))
        return normalize_input_list

class MultiScaleModel(AbstractModel):
    """The multi-scale multi-stage recursive training network class.

    Parameters
    ----------
    in_dim : int
        The number of input images which will be fed into network.
    problem_type : str "MultiPE" or "SinglePE". Default is "MultiPE".
    img_shape : tuple or list. Default is [160, 176, 160].
        The padded image size of input images.
    downsample_factor : int. Default is 2.
        The downsample factor in this stage.
    direction : str
        The phase encoding direction. 'x' means the LR/RL direction, 'y' means the AP/PA direction.
    previous_model : MultiScaleModel. Default is None.
        The model in the previous stage.
    load_previous_weights : bool True or False. Default is True.
        Whether or not to use the previous model's parameters as initialization.
    """
    def __init__(self, in_dim, problem_type="MultiPE", img_shape=[160, 176, 160], downsample_factor=2, direction='x',
                previous_model=None, load_previous_weights=True, **kargs):
        super().__init__(in_dim, problem_type=problem_type, **kargs)
        self.img_shape = img_shape
        if downsample_factor == 1:
            self.down_sampler = nn.Identity()
        else:
            self.down_sampler = ResizeImage(downsample_factor=downsample_factor)
        
        self.resize = ResizeTransform(vel_resize=0.5, ndims=3)
        self.downsample_factor = downsample_factor
        self.unet = ResUNet(in_dim=in_dim)
        self.epi_warp = EPIWarp(size=[d//self.downsample_factor for d in self.img_shape], direction=direction)
        self.load_previous_weights = load_previous_weights
        self.set_up_previous_model(previous_model)

    def forward(self, input_list, train=True):        
        """
        Parameters
        ----------
        input_list : list
        train : bool True or False. Default is True.

        Returns
        ----------
        1. if param train is set to True, then return estimated inhomogeneity field B and normalized images.
        2. if param train is set to False, then only return estimated inhomogeneity field B.
        """
        if self.previous_model:
            constant_flow, _ = self.previous_model(input_list)
            constant_flow = self.resize(constant_flow)
        else:
            batch_size = input_list[0].shape[0]
            constant_flow = torch.zeros(size=[batch_size, 1] + [d//self.downsample_factor for d in self.img_shape], device=input_list[0].device)

        normalize_input_list = self.normalize(input_list=input_list)
        downsample_input_list = [self.down_sampler(img) for img in normalize_input_list]

        warped_downsample_input_list = []
        if self.problem_type == 'SinglePE':
            warped_downsample_I = self.epi_warp(downsample_input_list[0], constant_flow)
            warped_downsample_input_list = [warped_downsample_I, downsample_input_list[1]]
        elif self.problem_type == 'MultiPE':
            warped_downsample_I_1 = self.epi_warp(downsample_input_list[0], constant_flow)
            warped_downsample_I_2 = self.epi_warp(downsample_input_list[1], -constant_flow)
            warped_downsample_input_list = [warped_downsample_I_1, warped_downsample_I_2]
            if len(downsample_input_list) == 3:
                warped_downsample_input_list.append(downsample_input_list[2])
        else:
            raise ValueError("param probelem_type must be `SinglePE` or `MultiPE`, but got {}".format(self.problem_type))

        flow = self.unet(warped_downsample_input_list)
        flow = flow + constant_flow

        if train:
            return flow, downsample_input_list
        else:
            return flow

    def set_up_previous_model(self, previous_model):
        self.previous_model = previous_model
        if self.previous_model:
            if self.load_previous_weights:
                self.unet.load_state_dict(self.previous_model.unet.state_dict())
            for param in self.previous_model.parameters():
                param.requires_grad = False

    def unfreeze_previous_model_parameters(self):
        if self.previous_model:
            for param in self.previous_model.parameters():
                param.requires_grad = True


sacnet/architectures/models.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import torch
import torch.nn as nn
from sacnet.architectures.layers import PreActBlock, predict_flow, conv_block
from torch.distributions.normal import Normal

class ResUNet(nn.Module):
    def __init__(self, in_dim=3, block_num=2):
        super().__init__()
        self.feature_extraction = conv_block(in_channels=in_dim, out_channels=16)
        self.res_layer_1 = self.res_layer(in_channels=16, out_channels=32, stride=2, block_num=block_num)
        self.res_layer_2 = self.res_layer(in_channels=32, out_channels=32, stride=2, block_num=block_num)
        self.res_layer_3 = self.res_layer(in_channels=64, out_channels=32, stride=1, block_num=block_num)
        self.res_layer_4 = self.res_layer(in_channels=48, out_channels=16, stride=1, block_num=block_num)
        nd = Normal(0, 1e-5)
        self.flow = predict_flow(16, 1, nd)
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')

    def res_layer(self, in_channels, out_channels, stride=1, bias=False, block_num=2):
        if block_num == 1:
            return PreActBlock(in_planes=in_channels, planes=out_channels, stride=stride, bias=bias)
        else:
            return nn.Sequential(
                PreActBlock(in_planes=in_channels, planes=out_channels, stride=stride, bias=bias),
                *[PreActBlock(in_planes=out_channels, planes=out_channels, stride=1, bias=bias)]*(block_num - 1)
            )

    def forward(self, input_list):
        feed = torch.cat(input_list, dim=1)
        feature = self.feature_extraction(feed)
        res_feature_1 = self.res_layer_1(feature)
        res_feature_2 = self.res_layer_2(res_feature_1)
        
        res_feature_3 = self.upsample(res_feature_2)
        res_feature_3 = torch.cat([res_feature_1, res_feature_3], dim=1)
        res_feature_3 = self.res_layer_3(res_feature_3)

        res_feature_4 = self.upsample(res_feature_3)
        res_feature_4 = torch.cat([feature, res_feature_4], dim=1)
        res_feature_4 = self.res_layer_4(res_feature_4)
        res_feature_4 = nn.LeakyReLU(0.2)(res_feature_4)

        flow = self.flow(res_feature_4)
        return flow


sacnet/architectures/network.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import torch
import torch.nn as nn
import torch.nn.functional as F

class TransformTo3D(nn.Module):
    """
    [TransformTo3D] represents a block that fullfill zeros in other 
    dimension accoarding to the param `direction`. 
    """
    def __init__(self, direction='x'):
        """
        Instiatiate the block
            :param direction: direction of phase encoding
        """
        super(TransformTo3D, self).__init__()
        self.direction = direction
    
    def forward(self, flow):
        zeros = torch.zeros(flow.shape).float().to(flow.device)
        if self.direction == 'x':
            flow_3d = torch.cat([flow, zeros, zeros], dim=1)
        elif self.direction == 'y':
            flow_3d = torch.cat([zeros, flow, zeros], dim=1)
        else:
            raise ValueError("Expected `direction` to be `x` or `y`, but got `{}`".format(self.direction))
        return flow_3d

class SpatialTransformer(nn.Module):
    """
    [SpatialTransformer] represesents a spatial transformation block
    that uses the output from the UNet to preform an grid_sample
    https://pytorch.org/docs/stable/nn.functional.html#grid-sample
    """
    def __init__(self, size, mode='bilinear'):
        """
        Instiatiate the block
            :param size: size of input to the spatial transformer block
            :param mode: method of interpolation for grid_sampler
        """
        super(SpatialTransformer, self).__init__()

        # Create sampling grid
        vectors = [ torch.arange(0, s) for s in size ] 
        grids = torch.meshgrid(vectors) 
        grid  = torch.stack(grids) # y, x, z
        grid  = torch.unsqueeze(grid, 0)  #add batch
        grid = grid.type(torch.FloatTensor)
        self.register_buffer('grid', grid)

        self.mode = mode

    def forward(self, src, flow):   
        """
        Push the src and flow through the spatial transform block
            :param src: the original moving image
            :param flow: the output from the U-Net
        """
        new_locs = self.grid + flow 

        shape = flow.shape[2:]

        # Need to normalize grid values to [-1, 1] for resampler
        for i in range(len(shape)):
            new_locs[:,i,...] = 2*(new_locs[:,i,...]/(shape[i]-1) - 0.5)

        if len(shape) == 2:
            new_locs = new_locs.permute(0, 2, 3, 1) 
            new_locs = new_locs[..., [1,0]]
        elif len(shape) == 3:
            new_locs = new_locs.permute(0, 2, 3, 4, 1) 
            new_locs = new_locs[..., [2,1,0]]

        return F.grid_sample(src, new_locs, mode=self.mode)

class FlowGradient(nn.Module):
    """
    [FlowGradient] represents a FlowGradient block that calculate the 
    gradient of a given flow along the spcific direction. 
    """
    def __init__(self, direction='x'):
        """
        Instiatiate the block
            :param direction: direction of phase encoding
        """
        super(FlowGradient, self).__init__()
        self.direction = direction
    
    def forward(self, flow):
        """
        The calculation method is the same as np.gradient 
        https://numpy.org/doc/stable/reference/generated/numpy.gradient.html.
        """
        if self.direction == 'x':
            if flow.shape[1] == 3:
                flow = flow[:,0:1,...]
        elif self.direction == 'y':
            if flow.shape[1] == 3:
                flow = flow[:,1:2,...]
        else:
            raise ValueError("Wrong parameter `direction`.")

        grad_flow = self._gradient(flow)
        return grad_flow
    
    def _gradient(self, flow):
        if self.direction == 'x':
            left_pad_edge  = 2 * flow[:,:,:1,:,:]  - flow[:,:,1:2,:,:]
            right_pad_edge = 2 * flow[:,:,-1:,:,:] - flow[:,:,-2:-1,:,:]
            left_sample_point = torch.cat([left_pad_edge, flow[:,:,:-1,:,:]], dim=2)
            right_sample_point = torch.cat([flow[:,:,1:,:,:], right_pad_edge], dim=2)
            grad_flow = (right_sample_point - left_sample_point) / 2.0

        elif self.direction == 'y':
            left_pad_edge  = 2 * flow[:,:,:,:1,:]  - flow[:,:,:,1:2,:]
            right_pad_edge = 2 * flow[:,:,:,-1:,:] - flow[:,:,:,-2:-1,:]
            left_sample_point = torch.cat([left_pad_edge, flow[:,:,:,:-1,:]], dim=3)
            right_sample_point = torch.cat([flow[:,:,:,1:,:], right_pad_edge], dim=3)
            grad_flow = (right_sample_point - left_sample_point) / 2.0
        else:
            raise ValueError("Wrong parameter `direction`.")
        
        return grad_flow

class JacobianDet(nn.Module):
    def __init__(self, direction='x'):
        super(JacobianDet, self).__init__()
        self.flow_grad = FlowGradient(direction=direction)
    
    def forward(self, flow):
        return 1 + self.flow_grad(flow)

class EPIWarp(nn.Module):
    """
    [EPIWarp] represents a EPIWarp block that uses the output 
    inhomogeneity field B from UNet to perform geometric correction (grid sample) 
    and intensity correction (multiply by Jacobian determinant of B).
    """
    def __init__(self, size, direction='x', mode='bilinear'):
        """
        Instiatiate the block
            :param size: size of input to the spatial transformer block
            :param direction: direction of phase encoding
            :param mode: method of interpolation for grid_sampler in spatial transformer block
        """
        super(EPIWarp, self).__init__()
        self.to3d_op = TransformTo3D(direction=direction)
        self.warp_op = SpatialTransformer(size=size, mode=mode)
        self.jaco = JacobianDet(direction=direction)
    
    def forward(self, src, flow, cal_Jaco=True):
        """
        Push the src and flow through the spatial transform block
            :param src: the original moving image
            :param flow: the output from the U-Net
            :param cal_Jaco: whether or not to do intensity correction
        """
        if flow.shape[1] == 1:
            flow = self.to3d_op(flow)
        
        warp_src = self.warp_op(src, flow)

        if cal_Jaco:
            jaco_det = self.jaco(flow)
            warp_src = warp_src * torch.clamp(jaco_det, min=0)
        
        return warp_src

class DWIWarp(nn.Module):
    def __init__(self, size, pos_or_neg_list, fieldmap, direction='x', mode='bilinear', cal_Jaco=True):
        """
        Instiatiate the block
            :param size: size of input to the spatial transformer block
            :param direction: direction of phase encoding
            :param mode: method of interpolation for grid_sampler in spatial transformer block
        """
        super(DWIWarp, self).__init__()
        self.to3d_op = TransformTo3D(direction=direction)
        self.warp_op = SpatialTransformer(size=size, mode=mode)
        self.jaco = JacobianDet(direction=direction)
        self.pos_or_neg = pos_or_neg_list

        if fieldmap.shape[1] == 1:
            fieldmap = self.to3d_op(fieldmap)

        self.pos_fieldmap = fieldmap
        self.neg_fieldmap = fieldmap * -1

        if cal_Jaco:
            self.pos_jaco_det = torch.clamp(self.jaco(self.pos_fieldmap), min=0)
            self.neg_jaco_det = torch.clamp(self.jaco(self.neg_fieldmap), min=0)
        else:
            self.pos_jaco_det = torch.ones(size=(1, 1, *size)).to(fieldmap.device)
            self.neg_jaco_det = torch.ones(size=(1, 1, *size)).to(fieldmap.device)
    
    def forward(self, dwi):
        corr_dwi_arr = []
        num_of_vols = dwi.shape[0]
        for i in range(num_of_vols):
            per_img = dwi[i:i+1, ...]
            (fmap, jaco_det) = (self.pos_fieldmap, self.pos_jaco_det) if self.pos_or_neg[i] == '+' else (self.neg_fieldmap, self.neg_jaco_det)
            per_warp_img = self.warp_op(per_img, fmap)
            per_warp_img = per_warp_img * jaco_det
            corr_dwi_arr.append(per_warp_img)
        corr_dwi_arr = torch.cat(corr_dwi_arr, dim=0)
        return corr_dwi_arr

sacnet/architectures/transform.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import argparse
import torch
import numpy as np
import nibabel as nib
from sacnet.architectures.transform import DWIWarp
from sacnet.utilities.tools import analyze_acqparams

def convert_array(numpy_array):
    if len(numpy_array.shape) == 3:
        numpy_array = numpy_array[np.newaxis, np.newaxis, ...]
    if len(numpy_array.shape) == 4:
        numpy_array = numpy_array[np.newaxis, ...]
        numpy_array = np.transpose(numpy_array, axes=(4, 0, 1, 2, 3))
    tensor_array = torch.from_numpy(numpy_array).float()
    return tensor_array

def read_index(index_file):
    f = open(index_file, 'r')
    indexes = []
    for line in f.readlines():
        indexes.append(int(line))
    return indexes

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--imain", required=True, help="File containing all the images to estimate distortions.")
    parser.add_argument("--fieldmap", required=True, help="The path to fieldmap file")
    parser.add_argument("--acqp", required=True, help="name of text file with PE directions")
    parser.add_argument("--index", required=True, help="File containing indices for all volumes in --imain into --acqp")
    parser.add_argument("--output", required=True, help="The path to the output corrected image path")
    parser.add_argument("--no_intensity_correction", required=False, default=False, action="store_true")

    args = parser.parse_args()

    imain = args.imain
    fieldmap = args.fieldmap
    acqparams = analyze_acqparams(args.acqp)
    indexes = read_index(args.index)
    direction = acqparams[0][-1]
    signs = [acqparams[index-1][0] for index in indexes]
    output = args.output
    no_intensity_correction = args.no_intensity_correction
    intensity_correction = not no_intensity_correction
    device = torch.device("cpu")

    dwi_img = nib.load(imain)
    dwi_arr = dwi_img.get_fdata()
    affine = dwi_img.get_affine()
    header = dwi_img.get_header()
    image_size = dwi_arr.shape[:3]
    fieldmap = nib.load(fieldmap).get_fdata()

    dwi_arr = convert_array(dwi_arr).to(device)
    fieldmap = convert_array(fieldmap).to(device)
    warp_op = DWIWarp(size=image_size, pos_or_neg_list=signs, fieldmap=fieldmap, direction=direction, cal_Jaco=intensity_correction).to(device)
    corr_dwi_arr = warp_op(dwi_arr)
    corr_dwi_arr = np.transpose(corr_dwi_arr.cpu().numpy(), axes=(1,2,3,4,0))[0]
    corr_dwi_img = nib.Nifti1Image(corr_dwi_arr, affine, header)
    nib.save(corr_dwi_img, output)

    print("DWI volume unwarp completed.")

if __name__ == '__main__':
    main()

sacnet/run/apply_fieldmap.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from batchgenerators.utilities.file_and_folder_operations import join, isfile

def get_available_models():
    available_models = {
        "HCP_MultiPE_LRRL": {
            'description': "Trained with the inverse PEs(along the LR-RL direction) b0 images on HCP dataset.",
            'resolution': "1.25x1.25x1.25 mm^3",
            'image size': "145x174x145",
            'direction': "LR-RL",
        },
        "HCP_MultiPEWithT1w_LRRL": {
            'description': "Trained with the inverse PEs(along the LR-RL direction) b0 images and T1w image on HCP dataset",
            'resolution': "1.25x1.25x1.25 mm^3",
            'image size': "145x174x145",
            'direction': "LR-RL",
        },
        "HCP_MultiPEWithT2w_LRRL": {
            'description': "Trained with the inverse PEs(along the LR-RL direction) b0 images and T2w image on HCP dataset",
            'resolution': "1.25x1.25x1.25 mm^3",
            'image size': "145x174x145",
            'direction': "LR-RL",
        },
        "HCPD_MultiPE_APPA": {
            'description': "Trained with the inverse PEs(along the AP-PA direction) b0 images on HCP-D dataset.",
            'resolution': "1.5x1.5x1.5 mm^3",
            'image size': "140x140x92",
            'direction': "AP-PA",
        },
        "HCPD_MultiPEWithT1w_APPA": {
            'description': "Trained with the inverse PEs(along the AP-PA direction) b0 images and T1w image on HCP-D dataset",
            'resolution': "1.5x1.5x1.5 mm^3",
            'image size': "140x140x92",
            'direction': "AP-PA",
        },
        "HCPD_MultiPEWithT2w_APPA": {
            'description': "Trained with the inverse PEs(along the AP-PA direction) b0 images and T2w image on HCP-D dataset",
            'resolution': "1.5x1.5x1.5 mm^3",
            'image size': "140x140x92",
            'direction': "AP-PA",
        },
        "dHCP_MultiPE_APPA": {
            'description': "Trained with the inverse PEs(along the AP-PA direction) b0 images on dHCP dataset.",
            'resolution': "1.17x1.17x1.5 mm^3",
            'image size': "128x128x64",
            'direction': "AP-PA",
        },
        "dHCP_MultiPEWithT2w_APPA": {
            'description': "Trained with the inverse PEs(along the AP-PA direction) b0 images and T2w image on dHCP dataset",
            'resolution': "1.17x1.17x1.5 mm^3",
            'image size': "128x128x64",
        },
        "dHCP_MultiPE_LRRL": {
            'description': "Trained with the inverse PEs(along the LR-RL direction) b0 images on dHCP dataset.",
            'resolution': "1.17x1.17x1.5 mm^3",
            'image size': "128x128x64",
        },
        "dHCP_MultiPEWithT2w_LRRL": {
            'description': "Trained with the inverse PEs(along the LR-RL direction) b0 images and T2w image on dHCP dataset",
            'resolution': "1.17x1.17x1.5 mm^3",
            'image size': "128x128x64",
        },
        "CBD_SinglePEWithT1w_AP": {
            'description': "Trained with the single PE(along the AP direction) b0 image and T1w image on CBD dataset",
            'resolution': "2x2x2 mm^3",
            'image size': "112x112x70",
        },
        "CBD_SinglePEWithT2w_AP": {
            'description': "Trained with the single PE(along the AP direction) b0 image and T2w image on CBD dataset",
            'resolution': "2x2x2 mm^3",
            'image size': "112x112x70",
        }
    }
    return available_models

def print_available_pretrained_models():
    print('The following pretrained models are available:\n')
    av_models = get_available_models()
    for m in av_models.keys():
        print('')
        print(m)
        print(av_models[m]['description'])

def print_pretrained_model_info():
    import argparse
    parser = argparse.ArgumentParser(description="Use this to see the properties of a pretrained model.")
    parser.add_argument("--name", required=True, type=str, help='Name of the pretrained model. To see '
                                                                   'available task names, run sacnet_show_avaliable_model_info')
    args = parser.parse_args()
    name = args.name
    av = get_available_models()
    if name not in av.keys():
        raise RuntimeError("Invalid name. This pretrained model does not exist. To see available task names, "
                           "run sacnet_show_avaliable_model_info.")
    for key in av[name].keys():
        print(key, av[name][key])

sacnet/run/display_pretrained_model.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import argparse
import os
import numpy as np
import torch
import nibabel as nib
from sacnet.architectures.models import MultiScaleModel
from sacnet.architectures.transform import EPIWarp
from sacnet.utilities.misc import calculate_pad_size, pad_imgshape, pad, pad_to_specific_shape, unpad
from sacnet.utilities.tools import convert_array
from batchgenerators.utilities.file_and_folder_operations import join

def preprocess_data(data, pad_size, pad_imgshape, device):
    data = pad(data, pad_size)
    data = pad_to_specific_shape(data, pad_imgshape)
    data = np.expand_dims(np.expand_dims(data, axis=0), axis=0)
    data = torch.from_numpy(data).float().to(device)
    return data

def create_model(in_dim, direction, image_shape):
    network_1 = MultiScaleModel(in_dim=in_dim, problem_type="MultiPE", img_shape=image_shape, 
                                downsample_factor=4, direction=direction, previous_model=None,
                                load_previous_weights=False)
    network_2 = MultiScaleModel(in_dim=in_dim, problem_type="MultiPE", img_shape=image_shape, 
                                downsample_factor=2, direction=direction, previous_model=network_1,
                                load_previous_weights=False)
    network_3 = MultiScaleModel(in_dim=in_dim, problem_type="MultiPE", img_shape=image_shape, 
                                downsample_factor=1, direction=direction, previous_model=network_2,
                                load_previous_weights=False)
    return network_3

def predict(pos_b0, neg_b0, ts, direction, model_weight, intensity_correction, device, base_name):
    pos_vol = nib.load(pos_b0)
    affine = pos_vol.get_affine()
    header = pos_vol.get_header()
    pos_dat = pos_vol.get_fdata()
    neg_dat = nib.load(neg_b0).get_fdata()

    if model_weight is None:
        raise ValueError("The path to the pretrained model is not provided.")

    I_1 = pos_dat
    I_2 = neg_dat
    if ts is not None:
        I_a = nib.load(ts).get_fdata()
        in_dim = 3
    else:
        I_a = None
        in_dim = 2

    image_shape = pos_dat.shape[:3]
    pad_size = calculate_pad_size(image_shape, 16)
    pad_image_shape = pad_imgshape(image_shape, pad_size)

    network = create_model(in_dim, direction, pad_image_shape).to(device)
    if str(device) == 'cpu':
        checkpoint = torch.load(model_weight, map_location='cpu')['model']
    else:
        checkpoint = torch.load(model_weight)['model']
    network.load_state_dict(checkpoint)

    I_1 = preprocess_data(I_1, pad_size, pad_image_shape, device)
    I_2 = preprocess_data(I_2, pad_size, pad_image_shape, device)

    if I_a is not None:
        I_a = preprocess_data(I_a, pad_size, pad_image_shape, device)
        flow = network([I_1, I_2, I_a], False)
    else:
        flow = network([I_1, I_2], False)

    if type(flow) == tuple:
        flow = flow[0]

    flow = unpad(flow.data.cpu().numpy()[0,0], pad_size)
    flow = convert_array(flow)

    warp_op = EPIWarp(size=image_shape, direction=direction)
    pos_dat = convert_array(pos_dat)
    neg_dat = convert_array(neg_dat)
    corr_pos_dat = warp_op(pos_dat, flow, cal_Jaco=intensity_correction).cpu().numpy()[0, 0]
    corr_neg_dat = warp_op(neg_dat, -flow, cal_Jaco=intensity_correction).cpu().numpy()[0, 0]
    corr_dat = 2 * (corr_pos_dat * corr_neg_dat) / (corr_pos_dat + corr_neg_dat + 1e-5)
    fieldmap = flow.cpu().numpy()[0, 0]

    corr_dat_vol = nib.Nifti1Image(corr_dat, affine, header)
    corr_pos_vol = nib.Nifti1Image(corr_pos_dat, affine, header)
    corr_neg_vol = nib.Nifti1Image(corr_neg_dat, affine, header)
    fieldmap_vol = nib.Nifti1Image(fieldmap, affine, header)

    nib.save(corr_dat_vol, join("{}_corrected.nii.gz".format(base_name)))
    nib.save(corr_pos_vol, join("{}_corrected_pos.nii.gz".format(base_name)))
    nib.save(corr_neg_vol, join("{}_corrected_neg.nii.gz".format(base_name)))
    nib.save(fieldmap_vol, join("{}_fieldmap.nii.gz".format(base_name)))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--pos_b0", required=True, help="name of b0 images along the positive phase encoding direction")
    parser.add_argument("--neg_b0", required=True, help="name of b0 images along the negative phase encoding direction")
    parser.add_argument("--t1", required=False, help="name of input T1w image.")
    parser.add_argument("--t2", required=False, help="name of input T2w image.")
    parser.add_argument("--direction", required=True, type=str, help="Set to `x` when phase encoding direction is LR-RL, set to `y` when phase encoding direction is AP-PA.")
    parser.add_argument("--no_intensity_correction", required=False, default=False, action="store_true")
    parser.add_argument("--pretrained_model_name_by_developers", required=False, 
                        help="The name of pretrained model and configuration file provided by developers. \
                              Use `sacnet_show_avaliable_model_info` to see detailed information about \
                              avaliable models provided by developers.")
    parser.add_argument("--pretrained_model_name_by_users", required=False, 
                        help="The path to user-custmized model folder.")
    parser.add_argument("--gpu_id", required=False, default="0", type=str, help="The index of GPU to be used. \
                        `-1` denotes to use CPU. Default value is 0.")
    parser.add_argument("--out", required=True, help="base-name of output files")
    args = parser.parse_args()

    pos_b0 = args.pos_b0
    neg_b0 = args.neg_b0
    t1 = args.t1
    t2 = args.t2
    direction = args.direction
    no_intensity_correction = args.no_intensity_correction
    intensity_correction = not no_intensity_correction
    pretrained_model_name_by_developers = args.pretrained_model_name_by_developers
    pretrained_model_name_by_users = args.pretrained_model_name_by_users
    gpu_id = args.gpu_id
    base_name = args.out

    if direction not in ['x', 'y']:
        raise ValueError("The valid value of --direction shoule be `x` or `y`, but got `{}`.".format(direction))

    ts = None
    StrType = None
    if t1 is not None and t2 is not None:
        raise ValueError("`--t1` and `--t2` cannot be assigned at the same time.")
    else:
        if t1 is not None:
            ts = t1
            StrType = "T1w"
        elif t2 is not None:
            ts = t2
            StrType = "T2w"

    if int(gpu_id) >= 0:
        os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
        device = torch.device("cuda:0")
    else:
        device = torch.device("cpu")

    if pretrained_model_name_by_developers == None and pretrained_model_name_by_users == None:
        raise ValueError("You did not specify the value of `--pretrained_model_name_by_developers` or `--pretrained_model_folder_by_users`.")
    elif pretrained_model_name_by_developers is not None and pretrained_model_name_by_users is not None:
        raise ValueError("Both `--pretrained_model_name_by_developers` and `--pretrained_model_folder_by_users` have been specified. \
                          Please Select only one option to assign.")
    else:
        model_weight = None
        if pretrained_model_name_by_developers is not None:
            if StrType is not None and StrType not in pretrained_model_name_by_developers:
                raise ValueError("Invalid modality of input structural image with the developer pretrained model.")
            elif StrType is None:
                if "T1w" in pretrained_model_name_by_developers or "T2w" in pretrained_model_name_by_developers:
                    raise ValueError("You did not specify the value of `--t1` or `--t2`.")
            SCRIPT_PATH = os.path.realpath(__file__)
            SACNet_RUN_FOLDER = os.path.split(SCRIPT_PATH)[0]
            model_weight = join(os.path.split(SACNet_RUN_FOLDER)[0], "pretrained_weights", pretrained_model_name_by_developers, "weights.pth")
        elif pretrained_model_name_by_users is not None:
            user_folder=join(os.environ['SACNet_RESULTS_FOLDER'], pretrained_model_name_by_users)
            if not os.path.exists(user_folder):
                raise ValueError("The folder of user-custmized model does not exist.")
            model_weight = join(user_folder, "level_3", "best.pth")
        predict(pos_b0, neg_b0, ts, direction, model_weight, intensity_correction, device, base_name)

if __name__ == '__main__':
    main()

sacnet/run/predict_multipe.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import argparse
import os
import numpy as np
import torch
import nibabel as nib
from sacnet.architectures.models import MultiScaleModel
from sacnet.architectures.transform import EPIWarp
from sacnet.utilities.misc import calculate_pad_size, pad_imgshape, pad, pad_to_specific_shape, unpad
from sacnet.utilities.tools import convert_array
from batchgenerators.utilities.file_and_folder_operations import join

def preprocess_data(data, pad_size, pad_imgshape, device):
    data = pad(data, pad_size)
    data = pad_to_specific_shape(data, pad_imgshape)
    data = np.expand_dims(np.expand_dims(data, axis=0), axis=0)
    data = torch.from_numpy(data).float().to(device)
    return data

def create_model(in_dim, direction, image_shape):
    network_1 = MultiScaleModel(in_dim=in_dim, problem_type="SinglePE", img_shape=image_shape, 
                                downsample_factor=4, direction=direction, previous_model=None,
                                load_previous_weights=False)
    network_2 = MultiScaleModel(in_dim=in_dim, problem_type="SinglePE", img_shape=image_shape, 
                                downsample_factor=2, direction=direction, previous_model=network_1,
                                load_previous_weights=False)
    network_3 = MultiScaleModel(in_dim=in_dim, problem_type="SinglePE", img_shape=image_shape, 
                                downsample_factor=1, direction=direction, previous_model=network_2,
                                load_previous_weights=False)
    return network_3

def predict(b0, ts, direction, model_weight, intensity_correction, device, base_name):
    b0_vol = nib.load(b0)
    b0_dat = b0_vol.get_fdata()
    affine = b0_vol.get_affine()
    header = b0_vol.get_header()

    if model_weight is None:
        raise ValueError("The path to the pretrained model is not provided.")

    I = b0_dat
    if ts is not None:
        I_a = nib.load(ts).get_fdata()
    else:
        raise ValueError("No structural image has been provided. Please check `--t1` or `--t2`.")

    image_shape = b0_dat.shape[:3]
    pad_size = calculate_pad_size(image_shape, 16)
    pad_image_shape = pad_imgshape(image_shape, pad_size)
    in_dim = 2
    
    network = create_model(in_dim, direction, pad_image_shape).to(device)
    if str(device) == 'cpu':
        checkpoint = torch.load(model_weight, map_location='cpu')['model']
    else:
        checkpoint = torch.load(model_weight)['model']    
    network.load_state_dict(checkpoint)

    I = preprocess_data(I, pad_size, pad_image_shape, device)
    I_a = preprocess_data(I_a, pad_size, pad_image_shape, device)
    flow = network([I, I_a], False)

    if type(flow) == tuple:
        flow = flow[0]

    flow = unpad(flow.data.cpu().numpy()[0,0], pad_size)

    flow = convert_array(flow)
    imain_dat = convert_array(b0_dat)
    warp_op = EPIWarp(size=image_shape, direction=direction)
    corr_imain_dat = warp_op(imain_dat, flow, cal_Jaco=intensity_correction).cpu().numpy()[0,0]
    fieldmap = flow.cpu().numpy()[0, 0]
    corr_final_vol = nib.Nifti1Image(corr_imain_dat, affine, header)
    fieldmap_vol = nib.Nifti1Image(fieldmap, affine, header)

    nib.save(corr_final_vol, join("{}_corrected.nii.gz".format(base_name)))
    nib.save(fieldmap_vol, join("{}_fieldmap.nii.gz".format(base_name)))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--b0", required=True, help="name of 3D b0 image")
    parser.add_argument("--t1", required=False, help="name of input T1w image.")
    parser.add_argument("--t2", required=False, help="name of input T2w image.")
    parser.add_argument("--direction", required=True, help="Set to `x` when phase encoding direction is LR-RL, set to `y` when phase encoding direction is AP-PA.")
    parser.add_argument("--no_intensity_correction", required=False, default=False, action="store_true")
    parser.add_argument("--pretrained_model_name_by_developers", required=False, 
                        help="The name of pretrained model and configuration file provided by developers. \
                              Use `sacnet_show_avaliable_model_info` to see detailed information about \
                              avaliable models provided by developers.")
    parser.add_argument("--pretrained_model_name_by_users", required=False, 
                        help="The path to user-custmized model folder.")
    parser.add_argument("--gpu_id", required=False, default="0", type=str, help="The index of GPU to be used. \
                        `-1` denotes to use CPU. Default value is 0.")
    parser.add_argument("--out", required=True, help="base-name of output files")

    args = parser.parse_args()

    b0 = args.b0
    t1 = args.t1
    t2 = args.t2
    direction = args.direction
    no_intensity_correction = args.no_intensity_correction
    intensity_correction = not no_intensity_correction
    pretrained_model_name_by_developers = args.pretrained_model_name_by_developers
    pretrained_model_name_by_users = args.pretrained_model_name_by_users
    gpu_id = args.gpu_id
    base_name = args.out

    if direction not in ['x', 'y']:
        raise ValueError("The valid value of --direction shoule be `x` or `y`, but got `{}`.".format(direction))

    ts = None
    StrType = None
    if t1 is not None and t2 is not None:
        raise ValueError("`--t1` and `--t2` cannot be assigned at the same time.")
    else:
        if t1 is not None:
            ts = t1
            StrType = "T1w"
        elif t2 is not None:
            ts = t2
            StrType = "T2w"
        else:
            raise ValueError("The value of `--t1` or `--t2` must be asigned.")

    if int(gpu_id) >= 0:
        os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
        device = torch.device("cuda:0")
    else:
        device = torch.device("cpu")

    if pretrained_model_name_by_developers == None and pretrained_model_name_by_users == None:
        raise ValueError("You did not specify the value of `--pretrained_model_name_by_developers` or `--pretrained_model_folder_by_users`.")
    elif pretrained_model_name_by_developers is not None and pretrained_model_name_by_users is not None:
        raise ValueError("Both `--pretrained_model_name_by_developers` and `--pretrained_model_folder_by_users` have been specified. \
                          Please Select only one option to assign.")
    else:
        model_weight = None
        if pretrained_model_name_by_developers is not None:
            if StrType is not None and StrType not in pretrained_model_name_by_developers:
                raise ValueError("Invalid modality of input structural image with the pretrained model.")
            SCRIPT_PATH = os.path.realpath(__file__)
            SACNet_RUN_FOLDER = os.path.split(SCRIPT_PATH)[0]
            model_weight = join(os.path.split(SACNet_RUN_FOLDER)[0], "pretrained_weights", pretrained_model_name_by_developers, "weights.pth")
        elif pretrained_model_name_by_users is not None:
            user_folder=join(os.environ['SACNet_RESULTS_FOLDER'], pretrained_model_name_by_users)
            if not os.path.exists(user_folder):
                raise ValueError("The folder of user-custmized model does not exist.")
            model_weight = join(user_folder, "level_3", "best.pth")
        predict(b0, ts, direction, model_weight, intensity_correction, device, base_name)

if __name__ == '__main__':
    main()

sacnet/run/predict_singlepe.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import argparse
import numpy as np

def read_matrix(matrix_file):
    with open(matrix_file) as f:
        lines = f.readlines()
    lines = lines[:3]
    elems = []
    for line in lines:
        line = [float(num) for num in line.split( )][:3]
        elems += line
    m11, m12, m13, m21, m22, m23, m31, m32, m33 = elems
    return m11, m12, m13, m21, m22, m23, m31, m32, m33

def read_bvecs(bvec_file):
    with open(bvec_file) as f:
        lines = f.readlines()
    bvec_vals = []
    for line in lines:
        line = [float(num) for num in line.split( )]
        bvec_vals.append(line)
    bvec_vals = np.array(bvec_vals)
    bvec_X = bvec_vals[0, :]
    bvec_Y = bvec_vals[1, :]
    bvec_Z = bvec_vals[2, :]
    return bvec_X, bvec_Y, bvec_Z

def rotate_bvecs(bvec_file, matrix_file):
    m11, m12, m13, m21, m22, m23, m31, m32, m33 = read_matrix(matrix_file)
    bvec_X, bvec_Y, bvec_Z = read_bvecs(bvec_file)
    bvec_X_rot = m11 * bvec_X + m12 * bvec_Y + m13 * bvec_Z
    bvec_Y_rot = m21 * bvec_X + m22 * bvec_Y + m23 * bvec_Z
    bvec_Z_rot = m31 * bvec_X + m32 * bvec_Y + m33 * bvec_Z
    return bvec_X_rot, bvec_Y_rot, bvec_Z_rot

def main():
    parse = argparse.ArgumentParser()
    parse.add_argument("--input", required=True, help="Path to the input bvec file.")
    parse.add_argument("--matrix", required=True, help="Path to the FSL-fomat rotation matrix.")
    parse.add_argument("--output", required=True, help="Path to the output bvec file.")

    args = parse.parse_args()
    input = args.input
    matrix = args.matrix
    output = args.output

    bvec_X_rot, bvec_Y_rot, bvec_Z_rot = rotate_bvecs(input, matrix)

    bvec_X_rot = list(bvec_X_rot)
    bvec_X_rot = ["{:.6f}".format(b) for b in bvec_X_rot]
    bvec_X_rot = " ".join(bvec_X_rot)

    bvec_Y_rot = list(bvec_Y_rot)
    bvec_Y_rot = ["{:.6f}".format(b) for b in bvec_Y_rot]
    bvec_Y_rot = " ".join(bvec_Y_rot)

    bvec_Z_rot = list(bvec_Z_rot)
    bvec_Z_rot = ["{:.6f}".format(b) for b in bvec_Z_rot]
    bvec_Z_rot = " ".join(bvec_Z_rot)

    with open(output, mode='a') as f:
        f.write(bvec_X_rot)
        f.write("\n")
        f.write(bvec_Y_rot)
        f.write("\n")
        f.write(bvec_Z_rot)
        f.write("\n")
    
if __name__ == "__main__":
    main()

sacnet/run/rotate_bvecs.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import json
import argparse
from sacnet.utilities.misc import Params
from sacnet.training.network_trainer.network_trainer import SinglePENetworkTrainer, MultiPENetworkTrainer

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--train_config", required=True, help="Path to the configuration file for training.")
    parser.add_argument("-d", "--data_config", required=True, help="Path to the configuration file for dataset.")
    args = parser.parse_args()

    with open(args.train_config, 'r') as f:
        train_params = Params(json.load(f))

    if train_params.problem_type == 'SinglePE':
        dataloader_name = "SinglePEDataLoader"
        criterion_name = "SinglePECriterion"
        trainer_class = SinglePENetworkTrainer
    elif train_params.problem_type == 'MultiPE':
        dataloader_name = "MultiPEDataLoader"
        criterion_name = "MultiPECriterion"
        trainer_class = MultiPENetworkTrainer
    else:
        raise ValueError("param `problem_type` must be \"SinglePE\" or \"MultiPE\" but got \"{}\"".format(train_params.problem_type))
    trainer = trainer_class(params=train_params, data_json=args.data_config, \
                            _dataloader=dataloader_name, _criterion=criterion_name)
    trainer.run()

if __name__ == '__main__':
    main()

sacnet/run/train.py
RicardoZiTseng/SACNet
from __future__ import absolute_import
from . import *
from .dataset_loading import *
from .generic_loading import *
sacnet/training/dataloading/__init__.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import os
import torch
import nibabel as nib
from sacnet.training.dataloading.generic_loading import GenericDataLoader

class MultiPEDataLoader(GenericDataLoader):
    def __init__(self, data_json, train_or_valid='train', anat_type="T2w", **kargs):
        super(MultiPEDataLoader, self).__init__(data_json=data_json, train_or_valid=train_or_valid, **kargs)
        assert anat_type in ['T1w', 'T2w', None], "anat_type must be `T1w`, `T2w` or None, but got {}".format(anat_type)
        self.anat_type = anat_type
        if self.anat_type is None:
            self.input_dimension = 2
        else:
            self.input_dimension = 3
    
    def __getitem__(self, index):
        """Geneterate dataset dict.
        """
        subject_dir = os.path.join(self.data_dir, self.identifier[index])
        I_1 = self._pad(nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.file_key['I_1']))).get_fdata(), self.pad_size)
        I_2 = self._pad(nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.file_key['I_2']))).get_fdata(), self.pad_size)

        I_1 = self._pad_to_specific_shape(I_1, self.pad_imgshape)
        I_2 = self._pad_to_specific_shape(I_2, self.pad_imgshape)

        I_1, I_2 = self.to_tensor(I_1), self.to_tensor(I_2)

        feed_dict = {
            'I_1': I_1,
            'I_2': I_2
        }

        if self.anat_type is None:
            if self.train_or_valid == 'test':
                affine = nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.file_key['I_1']))).get_affine()
                return feed_dict, affine
            else:
                return feed_dict
        else:
            I_a = self._pad(nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.anat_type))).get_fdata(), self.pad_size)
            I_a = self._pad_to_specific_shape(I_a, self.pad_imgshape)
            mask = self._pad(nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.file_key['mask']))).get_fdata(), self.pad_size)
            mask = self._pad_to_specific_shape(mask, self.pad_imgshape)
            I_a = I_a * mask

            I_a = self.to_tensor(I_a)
            mask = self.to_tensor(mask).type(torch.bool)

            feed_dict['I_a'] = I_a
            feed_dict['mask'] = mask

            return feed_dict

class SinglePEDataLoader(GenericDataLoader):
    def __init__(self, data_json, train_or_valid='train', anat_type="T2w", **kargs):
        super(SinglePEDataLoader, self).__init__(data_json=data_json, train_or_valid=train_or_valid, **kargs)
        assert anat_type in ['T1w', 'T2w'], "anat_type must be `T1w` or `T2w`, but got {}".format(anat_type)
        self.anat_type = anat_type
        self.input_dimension = 2
    
    def __getitem__(self, index):
        """Geneterate dataset dict.
        """
        subject_dir = os.path.join(self.data_dir, self.identifier[index])
        I = self._pad(nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.file_key['I']))).get_fdata(), self.pad_size)
        I = self._pad_to_specific_shape(I, self.pad_imgshape)
        I_a = self._pad(nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.anat_type))).get_fdata(), self.pad_size)
        I_a = self._pad_to_specific_shape(I_a, self.pad_imgshape)
        mask = self._pad(nib.load(os.path.join(subject_dir, "{}.nii.gz".format(self.file_key['mask']))).get_fdata(), self.pad_size)
        mask = self._pad_to_specific_shape(mask, self.pad_imgshape)
        I_a = I_a * mask

        I = self.to_tensor(I)
        I_a = self.to_tensor(I_a)
        mask = self.to_tensor(mask).type(torch.bool)

        feed_dict = {
            'I': I,
            'I_a': I_a,
            'mask': mask
        }

        return feed_dict

sacnet/training/dataloading/dataset_loading.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import json
import numpy as np
import torch
from torch.utils.data import Dataset
from torch.utils.data._utils.collate import default_collate
from sacnet.utilities.misc import calculate_pad_size, pad_imgshape

def collate_fn(feed_dict_list):
    output_dict = {}

    for feed_dict in feed_dict_list:
        for key in feed_dict.keys():
            if key not in output_dict.keys():
                output_dict[key] = [feed_dict[key]]
            else:
                output_dict[key].append(feed_dict[key])
    
    for key in output_dict.keys():
        output_dict[key] = default_collate(output_dict[key])

    return output_dict

class GenericDataLoader(Dataset):
    """The Generic class of data loader.

    Parameters
    ----------
    data_json : str
        The path to the data json file.
    train_or_valid : str 'train' or 'valid' or 'test'. Default is 'train'.
        Determine which subset to be choosen.
    anat_type : str 'T1w' or 'T2w', or NoneType. Default is None.
        Determine which type of structural image to be loaded. `None` means that no structural image will be loaded.

    
    Attributes
    ----------
    direction : str
        The phase encoding direction. 'x' means the LR/RL direction, 'y' means the AP/PA direction.
    data_dir : str
        The path to the data folder.
    file_key : str
        The string determine the b0 image to be loaded. 'I_1' and 'I_2' for the multiple PE data. 'I' for the single PE data.
    train_or_valid : str
        ...
    pad_size : tuple or list
        The pad width
    input_dimension : int
        The number of images to be fed into the network.
    anat_type : str or NoneType
        ...
    """
    def __init__(self, data_json=None, train_or_valid='train', anat_type=None, **kargs):
        super().__init__()

        with open(data_json, 'r') as f:
            self.data_profile = json.load(f)

        self.direction = self.data_profile['direction']
        self.data_dir = self.data_profile['data_dir']
        self.file_key = self.data_profile['file_key']
        self.train_or_valid = train_or_valid
        self.pad_size = calculate_pad_size(self.data_profile['image_size'], 16)
        self.pad_imgshape = pad_imgshape(self.data_profile['image_size'], self.pad_size)
        self.input_dimension = None
        self.anat_type = anat_type
        
        if train_or_valid == 'train':
            self.identifier = self.data_profile['training']
        elif train_or_valid == 'valid':
            self.identifier = self.data_profile['validation']
        else:
            raise ValueError("...")

    def _pad(self, x, pad=None):
        if pad is None:
            return x
        else:
            return np.pad(x, pad, mode='constant', constant_values=0.0)

    def _pad_to_specific_shape(self, x, target_shape):
        x_shape = list(x.shape)
        if x_shape == target_shape:
            return x
        else:
            pads = []
            for i in range(len(x_shape)):
                x_length_i = x_shape[i]
                target_length_i = target_shape[i]
                pad_left = int((target_length_i - x_length_i) // 2)
                pad_right = int(target_length_i - x_length_i - pad_left)
                pads.append([pad_left, pad_right])
            return self._pad(x, pads)

    def to_tensor(self, x):
        x = np.expand_dims(x, axis=0)
        x = torch.from_numpy(x).float()
        return x

    def __getitem__(self, index):
        raise NotImplementedError
    
    def __len__(self):
        return len(self.identifier)


sacnet/training/dataloading/generic_loading.py
RicardoZiTseng/SACNet
from __future__ import absolute_import

from . import *
from .criterion import *

sacnet/training/loss_functions/__init__.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from sacnet.training.loss_functions.loss import *
from sacnet.architectures.transform import EPIWarp
import torch.nn as nn

class GenericCriterion(nn.Module):
    """The generic class of criterion.
    """
    def __init__(self, **kargs):
        super(GenericCriterion, self).__init__()
        self.warp = None

    def forward(self, x):
        """
        return dict, which contains all types of loss terms, which must include the 'total_loss' term.
        """
        raise NotImplementedError

class SinglePECriterion(GenericCriterion):
    """The optimization criterion for single PE data.

    Parameters
    ----------
    anat : float
        The parameter for the anatomical regularization.
    smooth : float
        The parameter for the smoothness regularization.
    jaco : float
        The parameter for the diffeomorphism preservation function.
    inshape : tuple or list
        The image size of the input image.
    direction : str
        The phase encoding direction. 'x' means the LR/RL direction, 'y' means the AP/PA direction.
    anat_loss : str 'NCC' or 'MIND' or 'NGF'. Default is 'NGF'.
        The choice of anatomical regularization function.
    cal_Jaco : bool True or False. Default is True.
        Whether or not to use the jacobian determinant for intensity correction.

    Attributes
    ----------
    anat_loss : nn.Module
        The anatomical regularization loss function.
    grad_loss : nn.Module
        The smoothness regularization loss function.
    jaco_loss : nn.Module
        The diffeomorphism preservation loss function.
    warp : nn.Module
        The EPI warp module.
    """
    def __init__(self, anat, smooth, jaco, inshape=[112,112,80], direction='y', cal_Jaco=True, **kargs):
        super(SinglePECriterion, self).__init__(**kargs)
        self.anat = anat
        self.smooth = smooth
        self.jaco = jaco
        self.cal_Jaco = cal_Jaco
        self.anat_loss = NGF()
        self.grad_loss = Gradient()
        self.jaco_loss = OneSideJacoLoss(direction=direction)
        self.warp = EPIWarp(size=inshape, direction=direction)

    def forward(self, I, I_a, flow, mask):
        """
        Parameters
        ----------
        I : torch.Tensor
            The tobe corrected image tensor.
        I_a : torch.Tensor
            The tobe aligned structural image tensor.
        flow : torch.Tensor
            The estimated inhomogeneity field B.
        mask : torch.Tensor
            The brain mask tensor of I_a.

        Returns
        ----------
        loss : dict
        """
        loss = {}
        warp_I = self.warp(I, flow, self.cal_Jaco)
        loss['anat_loss'] = self.anat_loss(warp_I, I_a, mask) * self.anat
        loss['smooth_loss'] = self.grad_loss(flow) * self.smooth
        loss['jaco_loss'] = self.jaco_loss(flow) * self.jaco

        loss['total_loss'] = sum([loss[k] for k in loss.keys() if k.endswith('loss')])

        return loss

class MultiPECriterion(GenericCriterion):
    """The optimization criterion for multiple PE data.

    Parameters
    ----------
    pair : float
        The parameter for the pair-wise loss.
    anat_1 : float
        The parameter for the anatomical regularization between structral image and corrected b0 image.
    anat_2 : float
        The parameter for the anatomical regularization between structral image and corrected b0 image in one PE direction.
    smooth : float
        The parameter for the smoothness regularization.
    jaco : float
        The parameter for the diffeomorphism preservation function.
    inshape : tuple or list
        The image size of the input image.
    direction : str
        The phase encoding direction. 'x' means the LR/RL direction, 'y' means the AP/PA direction.
    anat_loss : str 'NCC' or 'MIND' or 'NGF'. Default is 'NGF'.
        The choice of anatomical regularization function.
    cal_Jaco : bool True or False. Default is True.
        Whether or not to use the jacobian determinant for intensity correction.

    Attributes
    ----------
    pair_loss : nn.Module
        The pair-wise loss function.
    anat_loss : nn.Module
        The anatomical regularization loss function.
    grad_loss : nn.Module
        The smoothness regularization loss function.
    jaco_loss : nn.Module
        The diffeomorphism preservation loss function.
    warp : nn.Module
        The EPI warp module.
    """
    def __init__(self, pair, anat_1, anat_2, smooth, jaco, inshape=[160, 176, 160], direction='x', cal_Jaco=True, **kargs):
        super(MultiPECriterion, self).__init__(**kargs)
        self.pair = pair
        self.anat_1 = anat_1
        self.anat_2 = anat_2
        self.smooth = smooth
        self.jaco = jaco
        self.cal_Jaco = cal_Jaco
        self.pair_loss = MSELoss()
        self.anat_loss = NGF()
        self.grad_loss = Gradient()
        self.jaco_loss = PotentialWellJacoLoss(direction=direction)
        self.warp = EPIWarp(size=inshape, direction=direction)
    
    def forward(self, I_1, I_2, flow, I_a=None, mask=None):
        """
        Parameters
        ----------
        I_1 : torch.Tensor
            The tobe corrected image tensor in one PE direction.
        I_2 : torch.Tensor
            The tobe corrected image tensor in other PE direction.
        flow : torch.Tensor
            The estimated inhomogeneity field B.
        I_a : torch.Tensor or None.
            The tobe aligned structural image tensor. If I_a is set to None, then the anatomical regularization function will not be calculated.
        mask : torch.Tensor
            The brain mask tensor of I_a.

        Returns
        ----------
        loss : dict
        """
        loss = {}

        warp_I_1, warp_I_2 = self.warp(I_1, flow, cal_Jaco=self.cal_Jaco), self.warp(I_2, -flow, cal_Jaco=self.cal_Jaco)
        loss['sim_loss'] = self.pair_loss(warp_I_1, warp_I_2) * self.pair
        loss['smooth_loss'] = self.grad_loss(flow) * self.smooth
        loss['jaco_loss'] = self.jaco_loss(flow) * self.jaco

        if I_a is not None:
            if self.anat_1 > 0:
                K = 2 * (warp_I_1 * warp_I_2) / (warp_I_1 + warp_I_2 + 1e-5)
                loss['anat_1_loss'] = self.anat_loss(K, I_a, mask) * self.anat_1
            
            if self.anat_2 > 0:
                loss['anat_2_loss'] = (self.anat_loss(warp_I_1, I_a, mask) + self.anat_loss(warp_I_2, I_a, mask)) * self.anat_2

        loss['total_loss'] = sum([loss[k] for k in loss.keys() if k.endswith('loss')])

        return loss

sacnet/training/loss_functions/criterion.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

def one_order_diff(flow, direction='x'):
    if direction == 'x':
        if flow.shape[1] == 3:
            flow = flow[:,0:1,...]
        grad_flow = flow[:,:,1:,:,:] - flow[:,:,:-1,:,:]
    elif direction == 'y':
        if flow.shape[1] == 3:
            flow = flow[:,1:2,...]
        grad_flow = flow[:,:,:,1:,:] - flow[:,:,:,:-1,:]
    else:
        raise ValueError("Wrong parameter `direction`.")
    return grad_flow


class OneSideJacoLoss(nn.Module):
    """The Diffeomorphism preservation loss for the single PE data.
    Parameters
    ----------
    direction : str
        The phase encoding direction. 'x' means the LR/RL direction, 'y' means the AP/PA direction.
    R0 : floar. Default is -1.
        ...
    sigma : float. Default is 0.05.
        ...
    """
    def __init__(self, direction='x', R0=-1, sigma=0.05):
        super(OneSideJacoLoss, self).__init__()
        self.direction = direction
        self.R0 = R0
        self.sigma = sigma
    
    def forward(self, flow):
        grad_flow = one_order_diff(flow, self.direction)
        jaco_loss = self.woods_saxon_potential_one_side(grad_flow, self.R0, self.sigma)
        return torch.mean(jaco_loss)
    
    def woods_saxon_potential_one_side(self, x, R0, sigma):
        return (-1 / (1 + torch.exp((R0 - x)/sigma)) + 1) * torch.pow(x, 2)


class PotentialWellJacoLoss(nn.Module):
    """The Diffeomorphism preservation loss for the multiple PE data.
    Parameters
    ----------
    direction : str
        The phase encoding direction. 'x' means the LR/RL direction, 'y' means the AP/PA direction.
    R0 : floar. Default is -1.
        ...
    sigma : float. Default is 0.05.
        ...
    """
    def __init__(self, direction='x', R0=1, sigma=0.05):
        super(PotentialWellJacoLoss, self).__init__()
        self.direction = direction
        self.R0 = R0
        self.sigma = sigma
    
    def forward(self, flow):
        grad_flow = one_order_diff(flow, self.direction)
        jaco_loss = self.woods_saxon_potential(grad_flow, self.R0, self.sigma)
        # assert torch.isnan(jaco_loss).sum() == 0, print('jaco_loss!!!')
        return torch.mean(jaco_loss)

    def woods_saxon_potential(self, x, R0, sigma):
        return (-1 / (1 + torch.exp((torch.abs(x) - R0)/sigma)) + 1) * torch.pow(x, 2)


class MSELoss(nn.Module):
    def __init__(self):
        super(MSELoss, self).__init__()
    
    def forward(self, x, y):
        return torch.mean( (x - y) ** 2 )


class Gradient(nn.Module):
    """The N-D gradient loss.
    """
    def __init__(self, penalty='l2', loss_mult=None):
        super(Gradient, self).__init__()
        self.penalty = penalty
        self.loss_mult = loss_mult

    def forward(self, y_pred):
        dy = torch.abs(y_pred[:, :, 1:, :, :] - y_pred[:, :, :-1, :, :])
        dx = torch.abs(y_pred[:, :, :, 1:, :] - y_pred[:, :, :, :-1, :])
        dz = torch.abs(y_pred[:, :, :, :, 1:] - y_pred[:, :, :, :, :-1])

        if self.penalty == 'l2':
            dy = dy * dy
            dx = dx * dx
            dz = dz * dz

        d = torch.mean(dx) + torch.mean(dy) + torch.mean(dz)
        grad = d / 3.0

        if self.loss_mult is not None:
            grad *= self.loss_mult

        # assert torch.isnan(grad).sum() == 0, print('grad!!!')
        return grad

class NGF(nn.Module):
    """The Normalized Gradient Fields Image Loss.
    """
    def __init__(self, eps=1e-5):
        super(NGF, self).__init__()
        self.eps = eps
    
    def forward(self, x, y, mask=None):
        #print("Using NGF Loss!!!")
        ngf_x = self._ngf_loss(x) # [batch_size, 3, x, y, z]
        ngf_y = self._ngf_loss(y) # [batch_size, 3, x, y, z]

        value = 0
        for dim in range(3):
            value = value + ngf_x[:, dim:dim+1, ...] * ngf_y[:, dim:dim+1, ...]
        
        if mask is not None:
            value = value[mask]
        
        ngf_loss = 0.5 * (1 - torch.pow(value, 2))
        
        return torch.mean(ngf_loss)

    def _ngf_loss(self, img):
        img_dx = self._calculate_gradient(img, 2)
        img_dy = self._calculate_gradient(img, 3)
        img_dz = self._calculate_gradient(img, 4)
        norm = torch.sqrt(torch.pow(img_dx, 2) + torch.pow(img_dy, 2) + torch.pow(img_dz, 2) + self.eps ** 2)
        return torch.cat([img_dx, img_dy, img_dz], dim=1) / norm

    def _calculate_gradient(self, img, dim=2):
        img_temp = img
        img_temp = torch.transpose(img_temp, 0, dim)
        left_pad_edge = 2 * img_temp[:1, ...] - img_temp[1:2, ...]
        right_pad_edge = 2 * img_temp[-1:, ...] - img_temp[-2:-1, ...]
        left_sample_point = torch.cat([left_pad_edge, img_temp[:-1, ...]], dim=0)
        right_sample_point = torch.cat([img_temp[1:, ...], right_pad_edge], dim=0)
        grad_img = (right_sample_point - left_sample_point) / 2.0
        grad_img = torch.transpose(grad_img, dim, 0)
        return grad_img

sacnet/training/loss_functions/loss.py
RicardoZiTseng/SACNet
from __future__ import absolute_import
from . import *
from .network_trainer import *

sacnet/training/network_trainer/__init__.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import os
import sys
from abc import abstractmethod
import logging
import matplotlib.pyplot as plt
import sacnet.training.dataloading as dataloading
import sacnet.training.loss_functions as criterion
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
from torch.optim import Adam
import torch.utils.data as Data
from batchgenerators.utilities.file_and_folder_operations import (
    join, maybe_mkdir_p)
from sacnet.architectures.layers import ResizeTransform
from sacnet.architectures.models import MultiScaleModel
from sacnet.architectures.transform import EPIWarp
from sacnet.training.dataloading.generic_loading import collate_fn
from sacnet.utilities.misc import Clock, Params

class GenericNetworkTrainer(object):
    """
    The derived class which inherit the GenericNetworkTrainer need to implement 4 abstract method:
    1. loss_per_step(self, feed_dict: dict)
        This function is used to define how to calculate the loss of model per step.

    2. valid(self, epoch, level)
        This function is used to calculate and analyze the validation metrics per, and plot the validation figures per epoch per level.

    3. valid_per_step(self, feed_dict: dict)
        This function is used to define the metrics which are chosen to evaluate the model per step.

    Parameters
    ----------
    params : Params.
        The training setting file.
    _dataloader : str "SinglePEDataLoader" or "MultiPEDataLoader". Default is "MultiPEDataLoader".
    _criterion : str "SinglePECriterion" or "MultiPECriterion". Default is "MultiPECriterion".
    """
    def __init__(self, params: Params, data_json: str, _dataloader="MultiPEDataLoader", _criterion="MultiPECriterion"):
        self.params = params
        self.data_json = data_json
        self.work_path = os.environ['SACNet_RESULTS_FOLDER']
        self.save_path = join(self.work_path, params.name)
        maybe_mkdir_p(join(self.save_path, 'level_1'))
        maybe_mkdir_p(join(self.save_path, "level_2"))
        maybe_mkdir_p(join(self.save_path, "level_3"))
        log_path = join(self.save_path, 'logging.log')
        self.logger = self.create_log(params.name, log_path)
        self.logger.info("Starting multi-scale and multi-stage training.")
        self.params.save(join(self.save_path, "training.json"))

        self.clock = Clock()
        self.train_loss_list = {}
        self.valid_metric_list = {}

        os.environ['CUDA_VISIBLE_DEVICES'] = params.gpu
        if not torch.cuda.is_available():
            self.logger.info("No gpu device available.")
            sys.exit(1)

        np.random.seed(params.seed)
        # accelerate the computational speed
        cudnn.benchmark = True
        torch.manual_seed(params.seed)
        cudnn.enabled = True
        torch.cuda.manual_seed(params.seed)

        self.data_gen_class = getattr(dataloading, _dataloader)
        self.criterion_class = getattr(criterion, _criterion)
        self.network_class = MultiScaleModel
        self.optim_class = Adam
        self.inference_json_path = join(self.save_path, "inference.json")
        
        if isinstance(self.params.train_params['epochs'], list) and len(self.params.train_params['epochs']) == 3:
            self.epochs = self.params.train_params['epochs']
        else:
            raise ValueError("You need to set the number of training epochs in each training stage.")

        self.hyper_parameter_in_each_level = self.transform_hyper_parameter(self.params.hyper_parameter)

        self.logger.info('GPU device = %s.' % params.gpu)
        self.logger.info('Data Generator Class : {}.'.format(self.data_gen_class.__name__))
        self.logger.info('Criterion Class : {}.'.format(self.criterion_class.__name__))
        self.logger.info('Network Class: {}.'.format(self.network_class.__name__))
        self.logger.info('Optimizer Class: {}.'.format(self.optim_class.__name__))
        self.logger.info('Hyper-parameter in level-1 is {}.'.format(self.hyper_parameter_in_each_level[0]))
        self.logger.info('Hyper-parameter in level-2 is {}.'.format(self.hyper_parameter_in_each_level[1]))
        self.logger.info('Hyper-parameter in level-3 is {}.'.format(self.hyper_parameter_in_each_level[2]))

    def run_init(self):
        """Initialize some of the necessary components for the network training.

        Attributes
        ----------
        train_loader : Data.DataLoader
        valid_loader : Data.DataLoader
        in_dim : int
            The number of images fed into the network.
        pad_imgshape : tuple
        warp_op : EPIWarp
            The EPI warp module which used to warp EPI images in the validation phase.
        start_epoch : int
            The start epoch.
        """

        train_gen = self.data_gen_class(data_json=self.data_json, train_or_valid='train', anat_type=self.params.anat_type)
        valid_gen = self.data_gen_class(data_json=self.data_json, train_or_valid='valid', anat_type=self.params.anat_type)
        self.train_loader = Data.DataLoader(train_gen, batch_size=self.params.train_params['batch_size'], shuffle=True, num_workers=4, collate_fn=collate_fn)
        self.valid_loader = Data.DataLoader(valid_gen, batch_size=1, shuffle=False, num_workers=4, collate_fn=collate_fn)
        self.pad_size = train_gen.pad_size
        self.in_dim = train_gen.input_dimension
        self.pad_imgshape = train_gen.pad_imgshape
        self.direction = train_gen.direction
        self.warp_op = EPIWarp(size=self.pad_imgshape, direction=self.direction)

    def run(self):
        """Training the network in 3 resolutions.
        """
        self.run_init()
        self.start_epoch = 0

        if self.params.train_params['checkpoint'] == None:
            self.logger.info("No Loading files.")
            # level-1 training
            network_1 = self.network_class(in_dim=self.in_dim, downsample_factor=4, problem_type=self.params.problem_type, direction=self.direction,
                                            previous_model=None, img_shape=self.pad_imgshape, load_previous_weights=False)
            criterion_1 = self.criterion_class(inshape=[s//4 for s in self.pad_imgshape], **self.hyper_parameter_in_each_level[0])
            optimizer_1 = self.optim_class(network_1.parameters(), lr=self.params.train_params['lr'])
            self.switch_network_criterion_optimizer(network_1, criterion_1, optimizer_1)
            self.logger.info("Training the model at level-1 ...")
            self.train(level=1)

            # level-2 training
            network_1.load_state_dict(torch.load(join(self.save_path, 'level_1', '{}.pth'.format(self.epochs[0])))['model'])
            network_2 = self.network_class(in_dim=self.in_dim, downsample_factor=2, problem_type=self.params.problem_type, direction=self.direction,
                                            previous_model=network_1, img_shape=self.pad_imgshape, load_previous_weights=True)
            self.logger.info("Initialize network_2 with the weight of network_1 from model file of {}.".format(join(self.save_path, 'level_1', '{}.pth'.format(self.epochs[0]))))
            criterion_2 = self.criterion_class(inshape=[s//2 for s in self.pad_imgshape], **self.hyper_parameter_in_each_level[1])
            optimizer_2 = self.optim_class(network_2.parameters(), lr=self.params.train_params['lr'])
            self.switch_network_criterion_optimizer(network_2, criterion_2, optimizer_2)
            self.logger.info("Training the model at level-2 ...")
            self.train(level=2)

            # level-3 training
            network_2.load_state_dict(torch.load(join(self.save_path, 'level_2', '{}.pth'.format(self.epochs[0]+self.epochs[1])))['model'])
            network_3 = self.network_class(in_dim=self.in_dim, downsample_factor=1, problem_type=self.params.problem_type, direction=self.direction,
                                            previous_model=network_2, img_shape=self.pad_imgshape, load_previous_weights=True)
            self.logger.info("Initialize network_3 with the weight of network_2 from model file of {}.".format(join(self.save_path, 'level_2', '{}.pth'.format(self.epochs[0]+self.epochs[1]))))
            criterion_3 = self.criterion_class(inshape=self.pad_imgshape, **self.hyper_parameter_in_each_level[2])
            optimizer_3 = self.optim_class(network_3.parameters(), lr=self.params.train_params['lr'])
            self.switch_network_criterion_optimizer(network_3, criterion_3, optimizer_3)
            self.logger.info("Training the model at level-3 ...")
            self.train(level=3)

        else:
            self.logger.info('Loading pretrain model from {}.'.format(self.params.train_params['checkpoint']))
            model_checkpoint = torch.load(self.params.train_params['checkpoint'])["model"]
            pretrain_input_size = list(model_checkpoint["epi_warp.warp_op.grid"].shape[2:])

            temp_network_1 = self.network_class(in_dim=self.in_dim, downsample_factor=4, problem_type=self.params.problem_type, direction=self.direction,
                                            previous_model=None, img_shape=pretrain_input_size, load_previous_weights=False)
            temp_network_2 = self.network_class(in_dim=self.in_dim, downsample_factor=2, problem_type=self.params.problem_type, direction=self.direction,
                                            previous_model=temp_network_1, img_shape=pretrain_input_size, load_previous_weights=False)
            temp_network_3 = self.network_class(in_dim=self.in_dim, downsample_factor=1, problem_type=self.params.problem_type, direction=self.direction,
                                            previous_model=temp_network_2, img_shape=pretrain_input_size, load_previous_weights=False)
            
            temp_network_3.load_state_dict(model_checkpoint)

            network_1 = self.network_class(in_dim=self.in_dim, downsample_factor=4, problem_type=self.params.problem_type, direction=self.direction,
                                            previous_model=None, img_shape=self.pad_imgshape, load_previous_weights=False)
            network_1.unet.load_state_dict(temp_network_3.previous_model.previous_model.unet.state_dict())
            self.logger.info("Initialized model at level-1 with pretrained weight.")
            criterion_1 = self.criterion_class(inshape=[s//4 for s in self.pad_imgshape], **self.hyper_parameter_in_each_level[0])
            optimizer_1 = self.optim_class(network_1.parameters(), lr=self.params.train_params['lr'])
            self.switch_network_criterion_optimizer(network_1, criterion_1, optimizer_1)
            self.logger.info("Training the model at level-1 ...")
            self.train(level=1)

            network_1.load_state_dict(torch.load(join(self.save_path, 'level_1', '{}.pth'.format(self.epochs[0])))['model'])
            network_2 = self.network_class(in_dim=self.in_dim, downsample_factor=2, problem_type=self.params.problem_type, direction=self.direction,
                                previous_model=network_1, img_shape=self.pad_imgshape, load_previous_weights=True)
            network_2.unet.load_state_dict(temp_network_3.previous_model.unet.state_dict())
            self.logger.info("Initialized model at level-2 with pretrained weight.")
            criterion_2 = self.criterion_class(inshape=[s//2 for s in self.pad_imgshape], **self.hyper_parameter_in_each_level[1])
            optimizer_2 = self.optim_class(network_2.parameters(), lr=self.params.train_params['lr'])
            self.switch_network_criterion_optimizer(network_2, criterion_2, optimizer_2)
            self.logger.info("Training the model at level-2 ...")
            self.train(level=2)

            network_2.load_state_dict(torch.load(join(self.save_path, 'level_2', '{}.pth'.format(self.epochs[0]+self.epochs[1])))['model'])
            network_3 = self.network_class(in_dim=self.in_dim, downsample_factor=1, problem_type=self.params.problem_type, direction=self.direction,
                                previous_model=network_2, img_shape=self.pad_imgshape, load_previous_weights=True)
            network_3.unet.load_state_dict(temp_network_3.unet.state_dict())
            self.logger.info("Initialized model at level-3 with pretrained weight.")
            criterion_3 = self.criterion_class(inshape=self.pad_imgshape, **self.hyper_parameter_in_each_level[2])
            optimizer_3 = self.optim_class(network_3.parameters(), lr=self.params.train_params['lr'])
            self.switch_network_criterion_optimizer(network_3, criterion_3, optimizer_3)
            self.logger.info("Training the model at level-3 ...")
            self.train(level=3)

    def switch_to_train_mode(self):
        self.network.train()

    def switch_to_eval_mode(self):
        self.network.eval()

    def train(self, level=3):
        self.cuda()

        for _e in range(self.start_epoch+1, self.epochs[level-1]+self.start_epoch+1):
            self.clock.tic()
            all_loss_per_epoch = self.train_per_epoch(_e)
            self.clock.toc()
            self.logger.info('Epoch {}, training time = {} seconds.'.format(_e, self.clock.diff))

            if level >= 1:
                if _e % self.params.train_params['validate_per_epoch'] == 0:
                    self.clock.tic()
                    self.valid(_e, level)
                    self.clock.toc()
                    self.logger.info('Epoch {}, validation time = {} seconds.'.format(_e, self.clock.diff))
                self.create_valid_figure()

            if _e % self.params.train_params['save_per_epoch'] == 0:
                save_file_name = join(self.save_path, 'level_{}'.format(level), '{}.pth'.format(_e))
                self.save_checkpoint(epoch=_e, save_path=save_file_name)
            
            self.save_history()

            if level >= 2:
                if _e == (self.start_epoch + int(self.epochs[level-1] / 3) + 1):
                    self.logger.info('Unfreezing previous model\'s parameters for level {}.'.format(level))
                    self.network.unfreeze_previous_model_parameters()

        self.start_epoch = self.start_epoch + self.epochs[level-1]

    @abstractmethod
    def loss_per_step(self, feed_dict: dict):
        """This function is used to define how to calculate the loss of model per step.
        
        Parameters
        ----------
        feed_dict : dict
        
        Returns
        ----------
        loss : dict
            A dictory variable which mapping the training loss name to the training loss tensor.
            NOTICE: It must contain the key value of "total_loss".
        """

    def train_per_step(self, feed_dict: dict):
        """Get loss dictory and optimize the network in each step.
        """
        self.optimizer.zero_grad()
        loss = self.loss_per_step(feed_dict)
        total_loss = loss['total_loss']
        total_loss.backward()
        self.optimizer.step()
        return loss

    def train_per_epoch(self, epoch):
        """Get loss dictory and optimize the network in each epoch.
        """
        self.switch_to_train_mode()
        all_loss = {}
        
        for step, data in enumerate(self.train_loader):
            for k in data.keys():
                data[k] = data[k].cuda()
            loss = self.train_per_step(data)
            

            iteration = int((epoch - 1) * len(self.train_loader) + step)
            log = "Iteration: {}".format(iteration)
            for k in loss.keys():
                log += ", {}: {}".format(k, loss[k].item())
                if k not in all_loss.keys():
                    all_loss[k] = [loss[k].item()]
                else:
                    all_loss[k].append(loss[k].item())
            
            self.logger.info(log)
            print(log)

        for k in all_loss.keys():
            if k not in self.train_loss_list.keys():
                self.train_loss_list[k] = [np.mean(all_loss[k])]
            else:
                self.train_loss_list[k].append(np.mean(all_loss[k]))
        
        return all_loss

    @abstractmethod
    def valid(self, epoch, level):
        raise NotImplementedError("Need to implement this function.")
    
    @abstractmethod
    def valid_per_step(self, feed_dict: dict):
        """This function is used to define the metrics which are chosen to evaluate the model per step.
        """

    def valid_per_epoch(self):
        """This function is used to calculate the validation metrics per epoch.
        """
        self.switch_to_eval_mode()
        all_metrics = {}

        for step, data in enumerate(self.valid_loader):
            for k in data.keys():
                data[k] = data[k].cuda()
            metric = self.valid_per_step(data)

            for k in metric.keys():
                if k not in all_metrics.keys():
                    all_metrics[k] = [metric[k]]
                else:
                    all_metrics[k].append(metric[k])

        for k in all_metrics.keys():
            if k not in self.valid_metric_list.keys():
                self.valid_metric_list[k] = [np.mean(all_metrics[k])] 
            else:
                self.valid_metric_list[k].append(np.mean(all_metrics[k]))
        
        return all_metrics
    
    @property
    def warp(self):
        return self.warp_op

    def cuda(self):
        self.network.cuda()
        self.criterion.cuda()
        self.upsample_flow.cuda()
        self.warp_op.cuda()

    def save_checkpoint(self, epoch=None, save_path=None):
        checkpoint = {
            'model': self.network.state_dict(),
            'optimizer': self.optimizer.state_dict()
        }
        if epoch is not None:
            checkpoint['epoch'] = epoch
        torch.save(checkpoint, save_path)

    def load_checkpoint(self, load_path, return_start_epoch=True):
        checkpoint = torch.load(load_path)
        self.network.load_state_dict(checkpoint['model'])
        self.optimizer.load_state_dict(checkpoint['optimizer'])
        for state in self.optimizer.state.values():
            for k, v in state.items():
                if torch.is_tensor(v):
                    state[k] = v.cuda()
        if return_start_epoch:
            return checkpoint['epoch']
        else:
            return 0

    def switch_network_criterion_optimizer(self, network, criterion, optimizer):
        self.network = network
        self.criterion = criterion
        self.optimizer = optimizer

        if self.network.downsample_factor > 1:
            self.upsample_flow = ResizeTransform(1 / self.network.downsample_factor, ndims=3)
        elif self.network.downsample_factor == 1:
            self.upsample_flow = nn.Identity()

    def create_log(self, name, log_path):
        logger = logging.getLogger(name)
        logger.setLevel(logging.DEBUG)
        hfile = logging.FileHandler(log_path)
        hfile.setLevel(logging.INFO)
        formmatter = logging.Formatter('%(asctime)s %(lineno)d %(message)s', datefmt='%a, %d %b %Y %H:%M:%S')
        hfile.setFormatter(formmatter)
        logger.addHandler(hfile)
        return logger

    def create_valid_figure(self):
        for k in self.valid_metric_list.keys():
            y = self.valid_metric_list[k]
            x = list(range(1, len(y)+1))
            plt.clf()
            plt.plot(x, y)
            plt.savefig(join(self.save_path, "valid_{}".format(k)))

    def save_history(self):
        np.save(join(self.save_path, "train_loss.npy"), self.train_loss_list)
        np.save(join(self.save_path, "valid_metric.npy"), self.valid_metric_list)

    @abstractmethod
    def transform_hyper_parameter(self, origin_hyper_parameter):
        raise NotImplementedError("Need to implement this function.")


sacnet/training/network_trainer/generic_trainer.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import os
import numpy as np
import torch
from sacnet.training.network_trainer.generic_trainer import GenericNetworkTrainer
from sacnet.utilities.metrics import calMSE, calNMI, folding_calculate
from sacnet.utilities.misc import unpad, num2list
from sacnet.utilities.misc import Params

class SinglePENetworkTrainer(GenericNetworkTrainer):
    def __init__(self, params: Params, data_json: str, _dataloader, _criterion):
        super(SinglePENetworkTrainer, self).__init__(params, data_json=data_json, _dataloader=_dataloader, _criterion=_criterion)
        self.max_sim = -np.inf

    def loss_per_step(self, feed_dict: dict):
        I, I_a, mask = feed_dict['I'], feed_dict['I_a'], feed_dict['mask']
        flow, downsample_input_list = self.network([I, I_a], True)
        downsample_I, downsample_I_a = downsample_input_list[0], downsample_input_list[1]
        downsample_mask = (downsample_I_a > 0).type(torch.bool)
        loss = self.criterion(downsample_I, downsample_I_a, flow, downsample_mask)
        return loss

    def valid(self, epoch, level):
        all_metrics_per_epoch = self.valid_per_epoch()
        mean_nmi = np.mean(all_metrics_per_epoch['nmi'])

        if mean_nmi > self.max_sim:
            save_path = os.path.join(self.save_path, 'level_{}'.format(level), 'best.pth')
            self.save_checkpoint(epoch=epoch, save_path=save_path)
            self.max_sim = mean_nmi

    def valid_per_step(self, feed_dict: dict):
        I, I_a, mask = feed_dict['I'], feed_dict['I_a'], feed_dict['mask']
        flow = self.network([I, I_a], False)
        flow = self.upsample_flow(flow)

        warp_I = self.warp(I, flow, cal_Jaco=self.params.hyper_parameter['cal_Jaco'])
        I_a = unpad(I_a.data.cpu().numpy()[0,0], self.pad_size)
        warp_I = unpad(warp_I.data.cpu().numpy()[0,0], self.pad_size)
        flow = unpad(flow.data.cpu().numpy()[0,0], self.pad_size)
        mask = unpad(mask.data.cpu().numpy()[0,0], self.pad_size)
        nmi = calNMI(I_a, warp_I, mask)
        fold_count = folding_calculate(flow, -1, None, self.direction)
        return {
            'nmi': nmi,
            'fold_count': fold_count
        }

    def transform_hyper_parameter(self, origin_hyper_parameter):
        hyper_parameter_in_each_level = []

        anat_coef = num2list(origin_hyper_parameter["anat"], 3)
        smooth_coef = num2list(origin_hyper_parameter["smooth"], 3)
        jaco_coef = num2list(origin_hyper_parameter["jaco"], 3)

        for i in range(3):
            hyper_parameter_in_each_level.append({
                "anat": None,
                "smooth": None,
                "jaco": None,
                "cal_Jaco": origin_hyper_parameter["cal_Jaco"]
            })
            hyper_parameter_in_each_level[i]["anat"] = anat_coef[i]
            hyper_parameter_in_each_level[i]["smooth"] = smooth_coef[i]
            hyper_parameter_in_each_level[i]["jaco"] = jaco_coef[i]

        return hyper_parameter_in_each_level


class MultiPENetworkTrainer(GenericNetworkTrainer):
    def __init__(self, params: Params, data_json: str, _dataloader, _criterion):
        super(MultiPENetworkTrainer, self).__init__(params, data_json=data_json, _dataloader=_dataloader, _criterion=_criterion)
        self.min_dissim = np.inf

    def loss_per_step(self, feed_dict: dict):
        I_1, I_2 = feed_dict['I_1'], feed_dict['I_2']
        if 'I_a' in feed_dict:
            I_a, mask = feed_dict['I_a'], feed_dict['mask']
            flow, downsample_input_list = self.network([I_1, I_2, I_a], True)
            downsample_I_1, downsample_I_2, downsample_I_a = downsample_input_list[0], downsample_input_list[1], downsample_input_list[2]
            downsample_mask = (downsample_I_a > 0).type(torch.bool)
            loss = self.criterion(downsample_I_1, downsample_I_2, flow, downsample_I_a, downsample_mask)
        else:
            flow, downsample_input_list = self.network([I_1, I_2], True)
            downsample_I_1, downsample_I_2 = downsample_input_list[0], downsample_input_list[1]
            loss = self.criterion(downsample_I_1, downsample_I_2, flow)
        return loss

    def valid(self, epoch, level):
        all_metrics_per_epoch = self.valid_per_epoch()
        mean_mse = np.mean(all_metrics_per_epoch['mse'])
        
        if mean_mse < self.min_dissim:
            save_path = os.path.join(self.save_path, 'level_{}'.format(level), 'best.pth')
            self.save_checkpoint(epoch=epoch, save_path=save_path)
            self.min_dissim = mean_mse
    
    def valid_per_step(self, feed_dict: dict):
        I_1, I_2 = feed_dict['I_1'], feed_dict['I_2']

        if 'I_a' in feed_dict:
            I_a = feed_dict['I_a']
            flow = self.network([I_1, I_2, I_a], False)
        else:
            flow = self.network([I_1, I_2], False)
        
        flow = self.upsample_flow(flow)
        warp_I_1, warp_I_2 = self.warp(I_1, flow, cal_Jaco=self.params.hyper_parameter['cal_Jaco']), self.warp(I_2, -flow, cal_Jaco=self.params.hyper_parameter['cal_Jaco'])
        warp_I_1 = unpad(warp_I_1.data.cpu().numpy()[0,0], self.pad_size)
        warp_I_2 = unpad(warp_I_2.data.cpu().numpy()[0,0], self.pad_size)
        flow = unpad(flow.data.cpu().numpy()[0,0], self.pad_size)
        mse = calMSE(warp_I_1, warp_I_2)
        fold_count = folding_calculate(flow, -1, 1, self.direction)
        return {
            'mse': mse,
            'fold_count': fold_count
        }

    def transform_hyper_parameter(self, origin_hyper_parameter):
        hyper_parameter_in_each_level = []

        pair_coef = num2list(origin_hyper_parameter["pair"], 3)
        anat_1_coef = num2list(origin_hyper_parameter["anat_1"], 3)
        anat_2_coef = num2list(origin_hyper_parameter["anat_2"], 3)
        smooth_coef = num2list(origin_hyper_parameter["smooth"], 3)
        jaco_coef = num2list(origin_hyper_parameter["jaco"], 3)

        for i in range(3):
            hyper_parameter_in_each_level.append({
                "pair": None,
                "anat_1": None,
                "anat_2": None,
                "smooth": None,
                "jaco": None,
                "cal_Jaco": origin_hyper_parameter["cal_Jaco"]
            })
            hyper_parameter_in_each_level[i]["pair"] = pair_coef[i]
            hyper_parameter_in_each_level[i]["anat_1"] = anat_1_coef[i]
            hyper_parameter_in_each_level[i]["anat_2"] = anat_2_coef[i]
            hyper_parameter_in_each_level[i]["smooth"] = smooth_coef[i]
            hyper_parameter_in_each_level[i]["jaco"] = jaco_coef[i]

        return hyper_parameter_in_each_level


sacnet/training/network_trainer/network_trainer.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import numpy as np
import torch
import torch.nn.functional as F

def gradient(x, direction='x'):
    """
    x's shape is [w, h, d]
    """
    assert len(x.shape) == 3, "x's shape length must be equal to 3."
    if direction == 'x':
        left_pad_edge  = 2 * x[:1,:,:]  - x[1:2,:,:]
        right_pad_edge = 2 * x[-1:,:,:] - x[-2:-1,:,:]
        left_sample_point = np.concatenate([left_pad_edge, x[:-1,:,:]], axis=0)
        right_sample_point = np.concatenate([x[1:,:,:], right_pad_edge], axis=0)
        grad_x = (right_sample_point - left_sample_point) / 2.0
    elif direction == 'y':
        left_pad_edge  = 2 * x[:,:1,:]  - x[:,1:2,:]
        right_pad_edge = 2 * x[:,-1:,:] - x[:,-2:-1,:]
        left_sample_point = np.concatenate([left_pad_edge, x[:,:-1,:]], axis=1)
        right_sample_point = np.concatenate([x[:,1:,:], right_pad_edge], axis=1)
        grad_x = (right_sample_point - left_sample_point) / 2.0
    else:
        raise ValueError("Wrong parameter `direction`.")
    
    return grad_x

def folding_calculate(x, min=-1, max=None, direction='x'):
    assert len(x.shape) == 3, "x's shape length must be equal to 3."
    grad_x = gradient(x, direction)
    if max == None:
        count = np.count_nonzero(grad_x <= min)
    else:
        count = np.count_nonzero(grad_x <= min) + np.count_nonzero(grad_x >= max)
    return count

def calNMI(tgt, src, mask=None, bins=256):
    if mask is not None:
        mask = mask.astype(np.bool)
        tgt = tgt[mask]
        src = src[mask]
    else:
        tgt = tgt.ravel()
        src = src.ravel()

    tgt = (tgt - np.min(tgt)) / (np.max(tgt) - np.min(tgt))
    src = (src - np.min(src)) / (np.max(src) - np.min(src))
    hist_1, bin_1 = np.histogram(tgt, bins=bins)
    hist_2, bin_2 = np.histogram(src, bins=bins)
    hist_1_2, _3, _4 = np.histogram2d(tgt, src, bins=[bin_1, bin_2])
    hist_1 = hist_1 / np.sum(hist_1)
    hist_2 = hist_2 / np.sum(hist_2)
    hist_1_2 = hist_1_2 / np.sum(hist_1_2)

    MI = 0.0
    for i in range(len(hist_1)):
        for j in range(len(hist_2)):
            if hist_1_2[i,j] !=0 and hist_1[i] * hist_2[j] != 0:
                MI = MI + hist_1_2[i,j] * np.log2(hist_1_2[i, j]/(hist_1[i]*hist_2[j]))
    H_1 = 0.0
    for i in range(len(hist_1)):
        if hist_1[i] != 0:
            H_1 = H_1 - hist_1[i] * np.log2(hist_1[i])
    
    H_2 = 0.0
    for i in range(len(hist_2)):
        if hist_2[i] != 0:
            H_2 = H_2 - hist_2[i] * np.log2(hist_2[i])
    
    NMI = 2 * MI / (H_1 + H_2)
    return NMI
    # return MI

def calMSE(tgt, src, mask=None):
    if mask is not None:
        tgt = tgt * mask
        src = src * mask
    diff = tgt - src
    diff = np.abs(diff)
    if mask is not None:
        mask = mask.astype(np.bool)
        diff = diff[mask]
    mse = np.mean(diff**2)
    return mse


sacnet/utilities/metrics.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import json
import os
import time
import datetime
import numpy as np

def print_shape(tensor):
    print(tensor.shape)

def unpad(x, pad_width):
    slices = []
    for c in pad_width:
        e = None if c[1] == 0 else -c[1]
        slices.append(slice(c[0], e))
    return x[tuple(slices)]

def calculate_pad_size(image_size, down_size):
    pad = []
    for s in image_size:
        pad_s_length = np.ceil(s / down_size) * down_size
        pad_left = int((pad_s_length - s) // 2)
        pad_right = int(pad_s_length - s - pad_left)
        pad.append([pad_left, pad_right])
    return pad

def pad(x, pad=None):
    if pad is None:
        return x
    else:
        return np.pad(x, pad, mode='constant', constant_values=0.0)
    
def pad_imgshape(imgshape, pad_width):
    pad_img_shape = [imgshape[i] + pad_width[i][0]+ pad_width[i][1] for i in range(len(imgshape))]
    return pad_img_shape

def pad_to_specific_shape(x, target_shape):
    x_shape = list(x.shape)
    if x_shape == target_shape:
        return x
    else:
        pads = []
        for i in range(len(x_shape)):
            x_length_i = x_shape[i]
            target_length_i = target_shape[i]
            pad_left = int((target_length_i - x_length_i) // 2)
            pad_right = int(target_length_i - x_length_i - pad_left)
            pads.append([pad_left, pad_right])
        return pad(x, pads)



def num2list(value, length=3):
    if isinstance(value, list):
        if len(value) == length:
            return value
        else:
            raise ValueError("The length of value must be equal to {}.".format(length))
    else:
        value = [value] * length
        return value

def change_path(subjects, path):
    subj_list = []
    for subject in subjects:
        _id = subject.split("/")[-1]
        subj_list.append(os.path.join(path, _id))
    return subj_list

def save_subjs(subj_list, path):
    subj_list = sorted(subj_list)
    with open(path, 'a') as f:
        for subj in subj_list:
            f.write(subj)
            f.write("\n")

def read_subjs(path):
    subj_list = []
    for line in open(path, 'r'):
        line = line.strip("\n")
        subj_list.append(line)
    return subj_list

def change_path(subjects, path):
    subj_list = []
    for subject in subjects:
        _id = subject.split("/")[-1]
        subj_list.append(os.path.join(path, _id))
    return subj_list

class Params(object):
    def __init__(self, param):
        if not isinstance(param, dict):
            raise ValueError("Wrong value type, expected `dict`, but got {}".format(type(param)))
        self.param = param
    
    def __getattr__(self, name):
        return self.param[name]

    def save(self, path):
        with open(path, 'w') as f:
            json.dump(self.param, f, indent=4, sort_keys=True)

class Clock(object):
    '''
    A simple timer.
    '''

    def __init__(self):
        self.init_time = time.time()
        self.total_time = 0.
        self.calls = 0
        self.start_time = 0.
        self.diff = 0.
        self.average_time = 0.
        self.std_time = 0.
        self.remain_time = 0.
        self.time_pool = []

    def tic(self):
        # using time.time instead of time.clock because time time.clock
        # does not normalize for multithreading
        self.start_time = time.time()

    def toc(self, average=False):
        self.diff = time.time() - self.start_time
        self.total_time += self.diff
        self.calls += 1
        self.time_pool.append(self.diff)
        self.average_time = self.total_time / self.calls
        self.std_time = np.std(self.time_pool)
        if average:
            return self.average_time
        else:
            return self.diff

    def remain(self, iters, max_iters):
        if iters == 0:
            self.remain_time = 0
        else:
            self.remain_time = (time.time() - self.init_time) * \
                (max_iters - iters) / iters
        return str(datetime.timedelta(seconds=int(self.remain_time)))


sacnet/utilities/misc.py
RicardoZiTseng/SACNet
# If you use this code, please cite our paper.
#
# Copyright (C) 2023 Zilong Zeng
# For any questions, please contact Dr.Zeng (zilongzeng@mail.bnu.edu.cn) or Dr.Zhao (tengdazhao@bnu.edu.cn).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import numpy as np
import torch

def analyze_acqparams(acq_file):
    f = open(acq_file, 'r')
    filenames = []
    for line in f.readlines():
        line = line.split(" ")[:3]
        line = np.array([int(val) for val in line])
        sum_line = np.sum(line)
        arg_ind = np.argmax(np.abs(line))
        if sum_line == 1:
            if arg_ind == 0:
                filenames.append("+x")
            elif arg_ind == 1:
                filenames.append("+y")
        elif sum_line == -1:
            if arg_ind == 0:
                filenames.append("-x")
            elif arg_ind == 1:
                filenames.append("-y")
        else:
            raise ValueError("...")
    return filenames

def convert_array(numpy_array):
    if len(numpy_array.shape) == 3:
        numpy_array = numpy_array[np.newaxis, np.newaxis, ...]
    if len(numpy_array.shape) == 4:
        numpy_array = numpy_array[np.newaxis, ...]
        numpy_array = np.transpose(numpy_array, axes=(4, 0, 1, 2, 3))
    tensor_array = torch.from_numpy(numpy_array).float()
    return tensor_array

sacnet/utilities/tools.py
RicardoZiTseng/SACNet
from setuptools import setup, find_namespace_packages

setup(
    name='sacnet',
    packages=find_namespace_packages(include=["sacnet", "sacnet.*"]),
    version='1.0.0',
    description='SACNet: A Multiscale Diffeomorphic Convolutional Registration Network with Prior Neuroanatomical Constraints for Flexible Susceptibility Artifact Correction in Echo Planar Imaging.',
    url='https://github.com/RicardoZiTseng/SACNet',
    author='Zilong Zeng',
    author_email='zilongzeng@mail.bnu.edu.cn',
    license='GNU AFFERO GENERAL PUBLIC LICENSE Version 3, 19 November 2007',
    install_requires=[],
    entry_points={
        'console_scripts': [
            'sacnet_train = sacnet.run.train:main',
            'sacnet_predict_multipe = sacnet.run.predict_multipe:main',
            'sacnet_predict_singlepe = sacnet.run.predict_singlepe:main',
            'sacnet_rotate_bvecs = sacnet.run.rotate_bvecs:main',
            'sacnet_apply_fieldmap = sacnet.run.apply_fieldmap:main',
            'sacnet_show_avaliable_model_info = sacnet.run.display_pretrained_model:print_available_pretrained_models',
            'sacnet_print_pretrained_model_info = sacnet.run.display_pretrained_model:print_pretrained_model_info',
        ]
    },
    # python_requires="==3.6.*",
    # install_requires=[
    #     "art==5.3",
    #     "torch==1.9.1",
    #     "nibabel==3.2.1",
    #     "matplotlib==3.3.4",
    #     "batchgenerators==0.23"
    # ],
    keywords=['susceptibility artifact correction', 'echo planar imaging', 
            'diffeomorphic image registration', 'deep learning']
)

setup.py
RicardoZiTseng/SACNet