ep_surface_fit.py

#!/usr/bin/env python
import itertools
import os
import subprocess as sp
import sys
from argparse import ArgumentParser, Namespace, ArgumentDefaultsHelpFormatter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Optional

from chris_plugin import chris_plugin, PathMapper
from loguru import logger

__version__ = '0.5.1'


DISPLAY_TITLE = r"""
                                  __                  __ _ _   
                                 / _|                / _(_) |  
  ___ _ __ ______ ___ _   _ _ __| |_ __ _  ___ ___  | |_ _| |_ 
 / _ \ '_ \______/ __| | | | '__|  _/ _` |/ __/ _ \ |  _| | __|
|  __/ |_) |     \__ \ |_| | |  | || (_| | (_|  __/ | | | | |_ 
 \___| .__/      |___/\__,_|_|  |_| \__,_|\___\___| |_| |_|\__|
     | |                                        ______         
     |_|                                       |______|        

             Parameterized batch experiment

"""

parser = ArgumentParser(description='surface_fit wrapper',
                        formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('--no-fail', dest='no_fail', action='store_true',
                    help='Produce exit code 0 even if any subprocesses do not.')
parser.add_argument('-V', '--version', action='version',
                    version=f'%(prog)s {__version__}')
parser.add_argument('-t', '--threads', type=int, default=0,
                    help='Number of threads to use for parallel jobs. '
                         'Pass 0 to use number of visible CPUs.')

parser.add_argument('--size', type=str, default='81920', help='number of polygons')
parser.add_argument('--stretch-weight', type=str, default='100', help='stretch weight')
parser.add_argument('--laplacian-weight', type=str, default='5e-6', help='laplacian weight')
parser.add_argument('--iter-outer', type=str, default='1000', help='total number of iterations per stage')
parser.add_argument('--iter-inner', type=str, default='50', help='save every few iterations')
parser.add_argument('--iso-value', type=str, default='10',
                    help='Chamfer value of laplacian map indicating mask boundary (i.e. target value)')
parser.add_argument('--step-size', type=str, default='0.10', help='Step size per iteration')
parser.add_argument('--oversample', type=str, default='0', help='subsampling (0=none, n=#points extra along edge)')
parser.add_argument('--self-dist', type=str, default='0.01', help='distance to check for self-intersection')
parser.add_argument('--self-weight', type=str, default='1.0', help='weight for self-intersection constraint')
parser.add_argument('--taubin', type=str, default='0',
                    help='iterations of taubin smoothing to perform between cycles of surface_fit')


@chris_plugin(
    parser=parser,
    title='surface_fit experiment',
    category='Experiment',
    min_memory_limit='1Gi',
    min_cpu_limit='1000m',
)
def main(options: Namespace, inputdir: Path, outputdir: Path):
    print(DISPLAY_TITLE, file=sys.stderr, flush=True)

    params = [
        '-size',
        options.size,
        '-sw',
        options.stretch_weight,
        '-lw',
        options.laplacian_weight,
        '-iter-outer',
        options.iter_outer,
        '-iter-inner',
        options.iter_inner,
        '-iso-value',
        options.iso_value,
        '-step-size',
        options.step_size,
        '-oversample',
        options.oversample,
        '-self-dist',
        options.self_dist,
        '-self-weight',
        options.self_weight,
        '-taubin',
        options.taubin
    ]

    if options.threads > 0:
        nproc = options.threads
    else:
        nproc = len(os.sched_getaffinity(0))
    logger.info('Using {} threads.', nproc)

    mapper = PathMapper.file_mapper(inputdir, outputdir, glob='**/*.mnc', suffix='.obj')
    with ThreadPoolExecutor(max_workers=nproc) as pool:
        results = pool.map(lambda t, p: run_surface_fit(*t, p), mapper, itertools.repeat(params))

    if not options.no_fail and not all(results):
        sys.exit(1)


def run_surface_fit(grid: Path, output_surf: Path, params: list[str]) -> bool:
    """
    :return: True if successful
    """
    starting_surface = locate_surface_for(grid)
    if starting_surface is None:
        logger.error('No starting surface found for {}', grid)
        return False

    extra_args = [
        '-disterr', output_surf.with_suffix('.disterr.txt'),
        '-disterr-abs', output_surf.with_suffix('.disterr.abs.txt')
    ]

    cmd = ['surface_fit_script.pl', *params, *extra_args, grid, starting_surface, output_surf]
    log_file = output_surf.with_name(output_surf.name + '.log')
    logger.info('Starting: {}', ' '.join(map(str, cmd)))
    with log_file.open('wb') as log_handle:
        job = sp.run(cmd, stdout=log_handle, stderr=log_handle)
    rc_file = log_file.with_suffix('.rc')
    rc_file.write_text(str(job.returncode))

    if job.returncode == 0:
        logger.info('Finished: {} -> {}', starting_surface, output_surf)
        return True

    logger.error('FAILED -- check log file for details: {}', log_file)
    return False


def locate_surface_for(mask: Path) -> Optional[Path]:
    if any(mask.name.startswith(p) for p in ('lh.', 'rh.')):
        prefix = mask.name[:3]
    else:
        prefix = ''

    glob = mask.parent.glob(f'{prefix}*.obj')
    first = next(glob, None)
    second = next(glob, None)
    if second is not None:
        return None
    return first


if __name__ == '__main__':
    main()