From 6bbfbeede514202c91717215e7660dde1ac84aac Mon Sep 17 00:00:00 2001 From: Ashley Stewart Date: Fri, 9 Feb 2024 14:23:37 +1000 Subject: [PATCH] Added --download_weights option and fixed setup.py --- nextqsm/predict_all.py | 10 ++++++++++ setup.py | 1 + 2 files changed, 11 insertions(+) diff --git a/nextqsm/predict_all.py b/nextqsm/predict_all.py index 9e28541..1c4327d 100755 --- a/nextqsm/predict_all.py +++ b/nextqsm/predict_all.py @@ -6,6 +6,8 @@ import argparse import os +import sys + from pathlib import Path def main(ckp_path, ckp_name, source_path, mask_path, out_file, b_vec): @@ -97,6 +99,14 @@ def download_weights(): def cli_main(): + if '--download_weights' in sys.argv: + parser = argparse.ArgumentParser(description="Download weights for NeXtQSM.") + parser.add_argument('--download_weights', action='store_true', help='Only download the weights and exit') + args = parser.parse_args() + + if args.download_weights: + download_weights() + exit() parser = argparse.ArgumentParser( description="NeXtQSM: Deep Learning QSM Algorithm", diff --git a/setup.py b/setup.py index 0cf535c..a59f631 100644 --- a/setup.py +++ b/setup.py @@ -3,6 +3,7 @@ setup( name="nextqsm", + long_description="NeXtQSM is a complete deep learning pipeline for data-consistent quantitative susceptibility mapping trained with hybrid data", version="1.0.2", packages=find_packages(), package_dir={'nextqsm': 'nextqsm'},