generated from FNNDSC/python-chrisapp-template
-
Notifications
You must be signed in to change notification settings - Fork 0
/
spleendata.py
111 lines (93 loc) · 3.85 KB
/
spleendata.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
#!/usr/bin/env python
from pathlib import Path
import shutil
from argparse import ArgumentParser, Namespace, ArgumentDefaultsHelpFormatter
from chris_plugin import chris_plugin, PathMapper
from monai.apps.utils import download_and_extract
__version__ = "1.0.10"
DISPLAY_TITLE = r"""
_ _ _ _
| | | | | | | |
_ __ | |______ ___ _ __ | | ___ ___ _ __ __| | __ _| |_ __ _
| '_ \| |______/ __| '_ \| |/ _ \/ _ \ '_ \ / _` |/ _` | __/ _` |
| |_) | | \__ \ |_) | | __/ __/ | | | (_| | (_| | || (_| |
| .__/|_| |___/ .__/|_|\___|\___|_| |_|\__,_|\__,_|\__\__,_|
| | | |
|_| |_|
"""
parser = ArgumentParser(
description="""
A ChRIS FS plugin that downloads a spleen data set for training
and inference. Based off a MONAI exemplar:
https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb
""",
formatter_class=ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"--skipDownload",
default=False,
action="store_true",
help="""If specified, skip the download. Only useful for debugging/testing really.""",
)
parser.add_argument(
"--trainingOnly",
default=False,
action="store_true",
help="""If specified, only preserve the training data (saving about 462Mb)""",
)
parser.add_argument(
"--testingOnly",
default=False,
action="store_true",
help="""If specified, only preserve the testing data (saving about 1.2Gb)""",
)
parser.add_argument(
"-V", "--version", action="version", version=f"%(prog)s {__version__}"
)
def dir_findAndDelete(startdir: Path, target: str):
for item in startdir.iterdir():
if item.is_dir():
if item.name == target:
shutil.rmtree(str(item))
# item.rmdir()
print(f"Deleted directory: {target}")
break
else:
dir_findAndDelete(item, target)
# The main function of this *ChRIS* plugin is denoted by this ``@chris_plugin`` "decorator."
# Some metadata about the plugin is specified here. There is more metadata specified in setup.py.
#
# documentation: https://fnndsc.github.io/chris_plugin/chris_plugin.html#chris_plugin
@chris_plugin(
parser=parser,
title="Spleen data downloader",
category="", # ref. https://chrisstore.co/plugins
min_memory_limit="100Mi", # supported units: Mi, Gi
min_cpu_limit="1000m", # millicores, e.g. "1000m" = 1 CPU core
min_gpu_limit=0, # set min_gpu_limit=1 to enable GPU
)
def main(options: Namespace, outputdir: Path):
"""
*ChRIS* plugins usually have two positional arguments: an **input directory** containing
input files and an **output directory** where to write output files. Command-line arguments
are passed to this main method implicitly when ``main()`` is called below without parameters.
:param options: non-positional arguments parsed by the parser given to @chris_plugin
:param inputdir: directory containing (read-only) input files
:param outputdir: directory where to write output files
"""
print(DISPLAY_TITLE)
resource: str = "https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar"
md5: str = "410d4a301da4e5b2f6f86ec3ddba524e"
compressed_file: Path = outputdir / "Task09_Spleen.tar"
data_dir: Path = outputdir / "Task09_Spleen"
if not data_dir.exists() or options.skipDownload:
download_and_extract(resource, str(compressed_file), str(outputdir), md5)
if compressed_file.exists():
compressed_file.unlink()
if options.trainingOnly:
dir_findAndDelete(outputdir, "imagesTs")
if options.testingOnly:
dir_findAndDelete(outputdir, "imagesTr")
dir_findAndDelete(outputdir, "labelsTr")
if __name__ == "__main__":
main()