mirror of
https://github.com/sbrl/research-rainfallradar
synced 2024-11-04 17:13:02 +00:00
ai: start creating initial scaffolding
This commit is contained in:
parent
1ec502daea
commit
dac6919fcd
2 changed files with 74 additions and 0 deletions
|
@ -0,0 +1,33 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
|
||||
from loguru import logger
|
||||
|
||||
def init_logging():
|
||||
pass
|
||||
|
||||
from parse_args import parse_args
|
||||
|
||||
def main():
|
||||
subcommand, args = parse_args()
|
||||
if args == None:
|
||||
return
|
||||
|
||||
imported_module = importlib.import_module(f"subcommands.{subcommand}")
|
||||
# TODO: Support multiple subcommands here
|
||||
match subcommand:
|
||||
case "pretrain":
|
||||
imported_module.pretrain(args)
|
||||
case _:
|
||||
sys.stderr.write(f"Error: The subcommand '{subcommand}' hasn't been registered in index.py yet.\nThis is a bug.\n")
|
||||
exit(1)
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
else:
|
||||
print("This script must be run directly. It cannot be imported.")
|
||||
exit(1)
|
41
aimodel/src/parse_args.py
Normal file
41
aimodel/src/parse_args.py
Normal file
|
@ -0,0 +1,41 @@
|
|||
|
||||
import argparse
|
||||
import sys
|
||||
import re
|
||||
import importlib
|
||||
|
||||
# import pysnooper
|
||||
|
||||
# @pysnooper.snoop()
|
||||
def parse_args():
|
||||
"""Defines and parses the CLI arguments."""
|
||||
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
sys.stderr.write("""
|
||||
This program trains, manipulates, visualises, and manages a contrastive learning based rainfall radar → water depth prediction model.
|
||||
It functions by first finding relationships between the rainfall radar data and the water depth + heightmap data (the 'pretrain' subcommand). After this, a decoder model to predict water depth (modelled as an image segmentation task), can then be trained.
|
||||
|
||||
Available subcommands:
|
||||
pretrain Pretrain a contrastive learning model as an encoder.
|
||||
|
||||
For more information, do src/index.py <subcommand> --help.
|
||||
""")
|
||||
exit(0)
|
||||
|
||||
subcommand = re.sub(r'[^a-z0-9-]', '', sys.argv[1])
|
||||
|
||||
subcommand_argparser = importlib.import_module(f"subcommands.{subcommand}").parse_args
|
||||
|
||||
parser = subcommand_argparser()
|
||||
# sys.stderr.write(f"Error: Unknown subcommand '{subcommand} (try --help).\n")
|
||||
# exit(1)
|
||||
if parser == None:
|
||||
sys.stderr.write(f"Error: The subcommand '{subcommand}' did not return an argument parser. This is a bug.\n")
|
||||
exit(1)
|
||||
|
||||
parser.add_argument("--only-gpu",
|
||||
help="If the GPU is not available, exit with an error (useful on shared HPC systems to avoid running out of memory & affecting other users)", action="store_true")
|
||||
|
||||
return subcommand, parser.parse_args(args=sys.argv[2:])
|
||||
|
Loading…
Reference in a new issue