Upload 11 files

README.md

@@ -1,3 +1,45 @@
----
-license: mit
----
+# Deep learning for speed of sound inversion in ultrasound imaging
+
+This repository contains the code and models for the following papers:
+
+
+1. Feigin M, Freedman D, Anthony B. W. A Deep Learning Framework for Single-Sided Sound Speed Inversion in Medical Ultrasound. IEEE Trans Biomed Eng. 2020;67(4):1142-1151. doi:10.1109/TBME.2019.2931195
+2. Feigin M, Zwecker M, Freedman D, Anthony BW. Detecting muscle activation using ultrasound speed of sound inversion with deep learning. In: 2020 42nd Annual International Conference of the IEEE Engineering in Medicine & Biology Society (EMBC). IEEE; 2020:2092-2095. doi:10.1109/EMBC44109.2020.9175237
+3. Feigin M, Freedman D, Anthony BW. Computing Speed-of-Sound from ultrasound: user-agnostic recovery and a new benchmark. IEEE Trans Biomed Eng. 2023; doi:TBF
+
+This repository contain the network code and models for the algorithms and results contained in the paper.
+
+The code was tested under python 3.9. The anaconda environment is defined in environment.yml (setup environment with the `command conda env create -f environment.yml`)
+
+## Data
+
+The dataset used is available on huggingface at https://huggingface.co/datasets/laughingrice/Ultrasound_planewave_sos_inversion
+
+Variables in the files are `[sample, layer, x/channel, y/sample]` order
+
+* `alpha_coeff` -- Alpha coefficient used for simulations, full resolution
+* `c0`  -- Speed-of-sound used for simulations, full resolution
+* `data` -- Channel data (first 2048 samples, 64 active channels, first layer with flat plane wave, to
+  match existing physical hardware were used for the results in the paper)
+* `dx` -- spatial dx value of `c0` and `alpha_coef`
+* `f` -- temporal sampling frequency of channel data (40MHz)
+
+## Models
+
+Model files appearing under the `models` directory for results presented in the paper with teh matching
+execution parameters are as follows:
+
+* `tbme_sos.pt` -- network weights for the network presented in [1]
+  * `python . --test_files data/supplamentary_sample.mat --test_fname tbme_sos.h5 --load_ver models/tbme_sos.pt --net_type tbme`
+* `embc_sos.pt` -- network weights for the network presented in [2] 
+  * `python . --test_files data/supplamentary_sample.mat --test_fname embc_sos.h5 --load_ver models/embc_sos.pt --net_type embc`
+* `tbme2_sos.pt` -- network weights for the network presented in [3]
+  * `python . --test_files data/supplamentary_sample.mat --test_fname tbme2_sos.h5 --load_ver models/tbme2_sos.pt`
+* `tbme2_sos_rand_gain.pt` -- [3] trained to recover the speed-of-sound map with random gain profile and scaling
+  * `python . --test_files data/supplamentary_sample.mat --test_fname tbme2_sos_gain.h5 --load_ver models/tbme2_sos_rand_gain.pt`
+* `tbme2_attn.pt` -- [3] trained to recover the attenuation coefficient
+  * `python . --test_files data/supplamentary_sample.mat --test_fname tbme2_attn.h5 --load_ver models/tbme2_attn.pt --label_vars alpha_coeff`
+* `tbme2_sos_attn.pt` -- [3] trained to recover both the speed-of-sound map and attenuation coefficient
+  * `python . --test_files data/supplamentary_sample.mat --test_fname tbme2_sos_attn.h5 --load_ver models/tbme2_sos_attn.pt --label_vars c0 alpha_coeff`
+* `tbme2_phase_sos.pt` -- [3] trained to recover the speed-of-sound map using the IQ phase component
+  * `python . --test_files data/supplamentary_sample.mat --test_fname tbme2_phase_sos.h5 --load_ver models/tbme2_phase_sos.pt --phase_inv 1`

__main__.py

@@ -0,0 +1,170 @@
+"""
+Deep learning framework for sound speed inversion
+"""
+
+import json
+import git
+import argparse
+import pathlib
+import glob
+import os
+import h5py
+
+import loader
+import run_logger
+import net
+
+import torch
+import torch.utils.data as td
+import pytorch_lightning as pl
+
+
+# ----------------------------
+# Setup command line arguments
+# ----------------------------
+
+parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+parser.add_argument('--test_files', nargs='?', help='Test data (file pattern) to process / data to evaluate ')
+parser.add_argument('--train_files', nargs='?', help='Train data (file pattern) to process, only evaluate test if empty')
+parser.add_argument('--test_fname', default='output.h5', help='Filename into which to write testing output -- will be overwritten')
+
+parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
+
+parser.add_argument('--experiment', default='DeepLearning US', help='experiment name')
+parser.add_argument('--tags', nargs='?', help='Optional run tags, should evaluate to dictionary via json.loads')
+
+parser.add_argument('--load_ver', type=str, help='Network weights to load')
+
+parser.add_argument('--conf', type=str, action='append', help='Config file(s) to import (overridden by command line arguments)')
+parser.add_argument('--conf_export', type=str, help='Filename where to store settings')
+
+
+parser = pl.Trainer.add_argparse_args(parser)
+parser = loader.Loader.add_argparse_args(parser)
+parser = net.Net.add_model_specific_args(parser)
+parser = run_logger.ImgCB.add_argparse_args(parser)
+
+args = parser.parse_args()
+
+if args.conf is not None:
+    for conf_fname in args.conf:
+        with open(conf_fname, 'r') as f:
+            parser.set_defaults(**json.load(f))
+
+    # Reload arguments to override config file values with command line values
+    args = parser.parse_args()
+
+if args.conf_export is not None:
+    with open(args.conf_export, 'w') as f:
+        json.dump(vars(args), f, indent=4, sort_keys=True)
+
+if args.test_files is None and args.train_files is None:
+    raise ValueError('At least one of train files or test files is required')
+
+# ----------------------------
+# Load data
+# ----------------------------
+
+ld = loader.Loader(**vars(args))
+test_input, test_label, train_input, train_label = ld.load_data(test_file_pattern=args.test_files, train_file_pattern=args.train_files)
+
+for name, tensor in (
+        ('test_input', test_input),
+        ('test_label', test_label),
+        ('train_input', train_input),
+        ('train_label', train_label)):
+    print(f'{name}: {tensor.shape if tensor is not None else None} -- {tensor.dtype if tensor is not None else None}')
+
+loaders = []
+
+if args.train_files is not None:
+    if train_input is None or train_label is None or (test_input is not None and test_label is None):
+        raise ValueError('Training requires labeled data')
+
+    train_ds = td.TensorDataset(train_input, train_label)
+    loaders.append(td.DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, pin_memory=True))
+
+if args.test_files is not None:
+    ds = [test_input]
+    if test_label is not None:
+        ds.append(test_label)
+
+    test_ds = td.TensorDataset(*ds)
+    loaders.append(td.DataLoader(test_ds, args.batch_size, shuffle=False, pin_memory=True))
+
+# ----------------------------
+# Run
+# ----------------------------
+
+if args.train_files is not None:
+    if args.tags is None:
+        args.tags = {}
+    elif type(args.tags) == str:
+        args.tags = json.loads(args.tags)
+
+    try:
+        repo = git.Repo(search_parent_directories=True)
+        sha = repo.head.object.hexsha
+        args.tags.update({'commit': sha})
+    except:
+        print('Not a git repo, not logging commit ID')
+
+    mfl = pl.loggers.MLFlowLogger(experiment_name=args.experiment, tags=args.tags)
+    mfl.log_hyperparams(args)
+
+    path = pathlib.Path(__file__).parent.absolute()
+    files = glob.glob(str(path) + os.sep + '*.py')
+    for f in files:
+        mfl.experiment.log_artifact(mfl.run_id, f, 'source')
+
+    chkpnt_cb = pl.callbacks.ModelCheckpoint(
+        monitor='validate_mean',
+        verbose=True,
+        save_top_k=1,
+        save_weights_only=True,
+        mode='min',
+        every_n_train_steps=1,
+        filename='{epoch}-{validate_mean}-{train_mean}',
+    )
+
+    img_cb = run_logger.ImgCB(**vars(args))
+    lr_logger = pl.callbacks.LearningRateMonitor()
+
+    args.__dict__.update({'logger': mfl, 'callbacks': [chkpnt_cb, img_cb, lr_logger]})
+else:
+    if os.path.exists(args.test_fname):
+        os.remove(args.test_fname)
+
+    args.__dict__.update({'callbacks': [run_logger.TestLogger(args.test_fname)]})
+
+
+if test_label is not None:
+    args.n_outputs = test_label.shape[1]
+elif train_label is not None:
+    args.n_outputs = train_label.shape[1]
+
+if test_input is not None:
+    args.n_inputs = test_input.shape[1]
+elif train_input is not None:
+    args.n_inputs = train_input.shape[1]
+
+
+n = net.Net(**vars(args))
+if args.load_ver is not None:
+    t = torch.load(args.load_ver, map_location='cpu')['state_dict']
+    n.load_state_dict(t)
+
+trainer = pl.Trainer.from_argparse_args(args)
+
+if args.train_files is not None:
+    trainer.fit(n, *loaders)
+
+    print(chkpnt_cb.best_model_path)
+elif args.label_vars:
+     trainer.test(n, *loaders)
+else:
+     predictions = trainer.predict(n, *loaders)
+     with h5py.File(args.test_fname, "w") as F:
+         F["predictions"] = torch.cat(predictions).numpy()
+

environment.yml

@@ -0,0 +1,286 @@
+name: DL-US-inversion
+channels:
+  - conda-forge
+  - defaults
+  - pytorch
+  - gimli
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_kmp_llvm
+  - abseil-cpp=20211102.0=h27087fc_1
+  - absl-py=1.1.0=pyhd8ed1ab_0
+  - aiohttp=3.8.1=py310h5764c6d_1
+  - aiosignal=1.2.0=pyhd8ed1ab_0
+  - alembic=1.8.1=pyhd8ed1ab_0
+  - aom=3.4.0=h27087fc_1
+  - appdirs=1.4.4=pyh9f0ad1d_0
+  - arrow-cpp=8.0.0=py310h893e394_4_cpu
+  - asn1crypto=1.5.1=pyhd8ed1ab_0
+  - asttokens=2.0.5=pyhd8ed1ab_0
+  - async-timeout=4.0.2=pyhd8ed1ab_0
+  - attrs=21.4.0=pyhd8ed1ab_0
+  - aws-c-cal=0.5.11=h95a6274_0
+  - aws-c-common=0.6.2=h7f98852_0
+  - aws-c-event-stream=0.2.7=h3541f99_13
+  - aws-c-io=0.10.5=hfb6a706_0
+  - aws-checksums=0.1.11=ha31a3da_7
+  - aws-sdk-cpp=1.8.186=hb4091e7_3
+  - backcall=0.2.0=pyh9f0ad1d_0
+  - backports=1.1=pyhd3eb1b0_0
+  - backports.functools_lru_cache=1.6.4=pyhd8ed1ab_0
+  - blas=2.115=mkl
+  - blas-devel=3.9.0=15_linux64_mkl
+  - blinker=1.4=py_1
+  - boost-cpp=1.79.0=h75c5d50_0
+  - brotli=1.0.9=h166bdaf_7
+  - brotli-bin=1.0.9=h166bdaf_7
+  - brotlipy=0.7.0=py310h5764c6d_1004
+  - bzip2=1.0.8=h7f98852_4
+  - c-ares=1.18.1=h7f98852_0
+  - ca-certificates=2022.6.15=ha878542_0
+  - cached-property=1.5.2=hd8ed1ab_1
+  - cached_property=1.5.2=pyha770c72_1
+  - cachetools=5.0.0=pyhd8ed1ab_0
+  - certifi=2022.6.15=py310hff52083_0
+  - cffi=1.15.1=py310h255011f_0
+  - charset-normalizer=2.1.0=pyhd8ed1ab_0
+  - click=8.1.3=py310hff52083_0
+  - cloudpickle=2.1.0=pyhd8ed1ab_0
+  - colorama=0.4.5=pyhd8ed1ab_0
+  - configparser=5.2.0=pyhd8ed1ab_0
+  - cryptography=37.0.1=py310h9ce1e76_0
+  - cudatoolkit=11.6.0=hecad31d_10
+  - cudnn=8.4.1.50=hed8a83a_0
+  - cycler=0.11.0=pyhd8ed1ab_0
+  - databricks-cli=0.17.0=pyhd8ed1ab_0
+  - decorator=5.1.1=pyhd8ed1ab_0
+  - docker-py=5.0.3=py310hff52083_2
+  - docker-pycreds=0.4.0=py_0
+  - entrypoints=0.4=pyhd8ed1ab_0
+  - executing=0.8.3=pyhd8ed1ab_0
+  - expat=2.4.8=h27087fc_0
+  - ffmpeg=5.0.1=gpl_h512afef_107
+  - flask=2.1.3=pyhd8ed1ab_0
+  - font-ttf-dejavu-sans-mono=2.37=hab24e00_0
+  - font-ttf-inconsolata=3.000=h77eed37_0
+  - font-ttf-source-code-pro=2.038=h77eed37_0
+  - font-ttf-ubuntu=0.83=hab24e00_0
+  - fontconfig=2.14.0=h8e229c2_0
+  - fonts-conda-ecosystem=1=0
+  - fonts-conda-forge=1=0
+  - fonttools=4.34.4=py310h5764c6d_0
+  - freetype=2.11.0=h70c0345_0
+  - frozenlist=1.3.0=py310h5764c6d_1
+  - fsspec=2022.5.0=pyhd8ed1ab_0
+  - future=0.18.2=py310hff52083_5
+  - gettext=0.21.0=hf68c758_0
+  - gflags=2.2.2=he1b5a44_1004
+  - giflib=5.2.1=h36c2ea0_2
+  - gitdb=4.0.9=pyhd8ed1ab_0
+  - gitpython=3.1.27=pyhd8ed1ab_0
+  - glog=0.6.0=h6f12383_0
+  - gmp=6.2.1=h58526e2_0
+  - gnutls=3.7.6=hf3e180e_5
+  - google-auth=2.9.1=pyh6c4a22f_0
+  - google-auth-oauthlib=0.4.1=py_2
+  - greenlet=1.1.2=py310hd8f1fbe_2
+  - grpc-cpp=1.46.3=hbd84cd8_2
+  - grpcio=1.46.3=py310ha0b7d45_2
+  - gunicorn=20.1.0=py310hff52083_2
+  - h5py=3.7.0=nompi_py310h06dffec_100
+  - hdf5=1.12.1=nompi_h2386368_104
+  - htmlmin=0.1.12=py_1
+  - icu=70.1=h27087fc_0
+  - idna=3.3=pyhd8ed1ab_0
+  - imagehash=4.2.1=pyhd8ed1ab_0
+  - importlib-metadata=4.11.4=py310hff52083_0
+  - importlib_resources=5.8.0=pyhd8ed1ab_0
+  - ipython=8.4.0=py310hff52083_0
+  - itsdangerous=2.1.2=pyhd8ed1ab_0
+  - jedi=0.18.1=py310hff52083_1
+  - jinja2=3.1.2=pyhd8ed1ab_1
+  - joblib=1.1.0=pyhd8ed1ab_0
+  - jpeg=9e=h166bdaf_2
+  - keyutils=1.6.1=h166bdaf_0
+  - kiwisolver=1.4.4=py310hbf28c38_0
+  - krb5=1.19.3=h3790be6_0
+  - lame=3.100=h7f98852_1001
+  - lcms2=2.12=hddcbb42_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h9c3ff4c_0
+  - libblas=3.9.0=15_linux64_mkl
+  - libbrotlicommon=1.0.9=h166bdaf_7
+  - libbrotlidec=1.0.9=h166bdaf_7
+  - libbrotlienc=1.0.9=h166bdaf_7
+  - libcblas=3.9.0=15_linux64_mkl
+  - libcrc32c=1.1.2=h9c3ff4c_0
+  - libcurl=7.83.1=h7bff187_0
+  - libdeflate=1.12=h166bdaf_0
+  - libdrm=2.4.112=h166bdaf_0
+  - libedit=3.1.20210910=h7f8727e_0
+  - libev=4.33=h516909a_1
+  - libevent=2.1.10=h9b69904_4
+  - libffi=3.4.2=h7f98852_5
+  - libgcc-ng=12.1.0=h8d9b700_16
+  - libgfortran-ng=12.1.0=h69a702a_16
+  - libgfortran5=12.1.0=hdcd56e2_16
+  - libgomp=12.1.0=h8d9b700_16
+  - libgoogle-cloud=1.40.2=hefc27d0_0
+  - libiconv=1.16=h516909a_0
+  - libidn2=2.3.3=h166bdaf_0
+  - liblapack=3.9.0=15_linux64_mkl
+  - liblapacke=3.9.0=15_linux64_mkl
+  - libllvm11=11.1.0=hf817b99_3
+  - libnghttp2=1.47.0=h727a467_0
+  - libnsl=2.0.0=h7f98852_0
+  - libpciaccess=0.16=h516909a_0
+  - libpng=1.6.37=h753d276_3
+  - libprotobuf=3.20.1=h6239696_0
+  - libssh2=1.10.0=ha56f1ee_2
+  - libstdcxx-ng=12.1.0=ha89aaad_16
+  - libtasn1=4.18.0=h166bdaf_1
+  - libthrift=0.16.0=h519c5ea_1
+  - libtiff=4.4.0=hc85c160_1
+  - libunistring=0.9.10=h7f98852_0
+  - libutf8proc=2.7.0=h7f98852_0
+  - libuuid=2.32.1=h7f98852_1000
+  - libva=2.15.0=h166bdaf_0
+  - libvpx=1.11.0=h9c3ff4c_3
+  - libwebp=1.2.2=h3452ae3_0
+  - libwebp-base=1.2.2=h7f98852_1
+  - libxcb=1.13=h7f98852_1004
+  - libxml2=2.9.14=h22db469_3
+  - libzlib=1.2.12=h166bdaf_2
+  - llvm-openmp=14.0.4=he0ac6c6_0
+  - llvmlite=0.38.1=py310h58363a5_0
+  - lz4-c=1.9.3=h9c3ff4c_1
+  - magma=2.5.4=h6103c52_2
+  - mako=1.2.1=pyhd8ed1ab_0
+  - markdown=3.4.1=pyhd8ed1ab_0
+  - markupsafe=2.1.1=py310h5764c6d_1
+  - matplotlib-base=3.5.2=py310h5701ce4_0
+  - matplotlib-inline=0.1.3=pyhd8ed1ab_0
+  - missingno=0.4.2=py_1
+  - mkl=2022.1.0=h84fe81f_915
+  - mkl-devel=2022.1.0=ha770c72_916
+  - mkl-include=2022.1.0=h84fe81f_915
+  - mlflow=1.27.0=py310ha13cd29_0
+  - multidict=6.0.2=py310h5764c6d_1
+  - multimethod=1.4=py_0
+  - munkres=1.1.4=pyh9f0ad1d_0
+  - nccl=2.12.12.1=h0800d71_0
+  - ncurses=6.3=h27087fc_1
+  - nettle=3.8=hc379101_0
+  - networkx=2.8.4=pyhd8ed1ab_0
+  - ninja=1.11.0=h924138e_0
+  - numba=0.55.0=py310h00e6091_0
+  - numpy=1.23.1=py310h53a5b5f_0
+  - oauthlib=3.2.0=pyhd8ed1ab_0
+  - openh264=2.2.0=h6239696_1
+  - openjpeg=2.4.0=hb52868f_1
+  - openssl=1.1.1q=h166bdaf_0
+  - orc=1.7.5=h6c59b99_0
+  - p11-kit=0.24.1=hc5aa10d_0
+  - packaging=21.3=pyhd8ed1ab_0
+  - pandas=1.4.3=py310h769672d_0
+  - pandas-profiling=3.2.0=pyhd8ed1ab_0
+  - parso=0.8.3=pyhd8ed1ab_0
+  - patsy=0.5.2=pyhd8ed1ab_0
+  - pexpect=4.8.0=pyh9f0ad1d_2
+  - phik=0.12.2=py310h7c64c84_0
+  - pickleshare=0.7.5=py_1003
+  - pillow=9.2.0=py310he619898_0
+  - pip=22.1.2=pyhd8ed1ab_0
+  - prometheus_client=0.14.1=pyhd8ed1ab_0
+  - prometheus_flask_exporter=0.20.2=pyhd8ed1ab_0
+  - prompt-toolkit=3.0.30=pyha770c72_0
+  - protobuf=3.20.1=py310hd8f1fbe_0
+  - pthread-stubs=0.4=h36c2ea0_1001
+  - ptyprocess=0.7.0=pyhd3deb0d_0
+  - pure_eval=0.2.2=pyhd8ed1ab_0
+  - pyarrow=8.0.0=py310h468efa6_0
+  - pyasn1=0.4.8=py_0
+  - pyasn1-modules=0.2.8=py_0
+  - pybind11-abi=4=hd8ed1ab_3
+  - pycparser=2.21=pyhd8ed1ab_0
+  - pydantic=1.9.1=py310h5764c6d_0
+  - pydeprecate=0.3.2=pyhd8ed1ab_0
+  - pygments=2.12.0=pyhd8ed1ab_0
+  - pyjwt=2.4.0=pyhd8ed1ab_0
+  - pyopenssl=22.0.0=pyhd8ed1ab_0
+  - pyparsing=3.0.9=pyhd8ed1ab_0
+  - pysocks=1.7.1=py310hff52083_5
+  - python=3.10.5=h582c2e5_0_cpython
+  - python-dateutil=2.8.2=pyhd8ed1ab_0
+  - python_abi=3.10=2_cp310
+  - pytorch=1.12.0=py3.10_cuda11.6_cudnn8.3.2_0
+  - pytorch-lightning=1.6.5=pyhd8ed1ab_0
+  - pytorch-mutex=1.0=cuda
+  - pytz=2022.1=pyhd8ed1ab_0
+  - pyu2f=0.1.5=pyhd8ed1ab_0
+  - pywavelets=1.3.0=py310hde88566_1
+  - pyyaml=6.0=py310h5764c6d_4
+  - querystring_parser=1.2.4=py_0
+  - re2=2022.06.01=h27087fc_0
+  - readline=8.1.2=h0f457ee_0
+  - requests=2.28.1=pyhd8ed1ab_0
+  - requests-oauthlib=1.3.1=pyhd8ed1ab_0
+  - rsa=4.8=pyhd8ed1ab_0
+  - s2n=1.0.10=h9b69904_0
+  - scikit-learn=1.1.1=py310hffb9edd_0
+  - scipy=1.8.1=py310h7612f91_0
+  - seaborn=0.11.2=hd8ed1ab_0
+  - seaborn-base=0.11.2=pyhd8ed1ab_0
+  - setuptools=59.5.0=py310hff52083_0
+  - shap=0.41.0=py310h769672d_0
+  - six=1.16.0=pyh6c4a22f_0
+  - sleef=3.5.1=h9b69904_2
+  - slicer=0.0.7=pyhd8ed1ab_0
+  - smmap=3.0.5=pyh44b312d_0
+  - snappy=1.1.9=hbd366e4_1
+  - sqlalchemy=1.4.39=py310h5764c6d_0
+  - sqlite=3.39.1=h4ff8645_0
+  - sqlparse=0.4.2=pyhd8ed1ab_0
+  - stack_data=0.3.0=pyhd8ed1ab_0
+  - statsmodels=0.13.2=py310hde88566_0
+  - svt-av1=1.1.0=h27087fc_1
+  - tabulate=0.8.10=pyhd8ed1ab_0
+  - tangled-up-in-unicode=0.2.0=pyhd8ed1ab_0
+  - tbb=2021.5.0=h924138e_1
+  - tensorboard=2.6.0=py_0
+  - tensorboard-plugin-wit=1.8.1=pyhd8ed1ab_0
+  - threadpoolctl=3.1.0=pyh8a188c0_0
+  - tk=8.6.12=h27826a3_0
+  - torchaudio=0.12.0=py310_cu116
+  - torchmetrics=0.9.2=pyhd8ed1ab_0
+  - torchvision=0.13.0=py310_cu116
+  - tqdm=4.64.0=pyhd8ed1ab_0
+  - traitlets=5.3.0=pyhd8ed1ab_0
+  - typing-extensions=4.3.0=hd8ed1ab_0
+  - typing_extensions=4.3.0=pyha770c72_0
+  - tzdata=2022a=h191b570_0
+  - unicodedata2=14.0.0=py310h5764c6d_1
+  - urllib3=1.26.10=pyhd8ed1ab_0
+  - visions=0.7.4=pyhd8ed1ab_0
+  - wcwidth=0.2.5=pyh9f0ad1d_2
+  - websocket-client=1.3.3=pyhd8ed1ab_0
+  - werkzeug=2.1.2=pyhd8ed1ab_1
+  - wheel=0.37.1=pyhd8ed1ab_0
+  - x264=1!161.3030=h7f98852_1
+  - x265=3.5=h924138e_3
+  - xorg-fixesproto=5.0=h7f98852_1002
+  - xorg-kbproto=1.0.7=h7f98852_1002
+  - xorg-libx11=1.7.2=h7f98852_0
+  - xorg-libxau=1.0.9=h7f98852_0
+  - xorg-libxdmcp=1.1.3=h7f98852_0
+  - xorg-libxext=1.3.4=h7f98852_1
+  - xorg-libxfixes=5.0.3=h7f98852_1004
+  - xorg-xextproto=7.3.0=h7f98852_1002
+  - xorg-xproto=7.0.31=h7f98852_1007
+  - xz=5.2.5=h516909a_1
+  - yaml=0.2.5=h7f98852_2
+  - yarl=1.7.2=py310h5764c6d_2
+  - zipp=3.8.0=pyhd8ed1ab_0
+  - zlib=1.2.12=h166bdaf_2
+  - zstd=1.5.2=h8a70e8d_2
+prefix: /home/micha/.conda/envs/DL-US-inversion

loader.py

@@ -0,0 +1,176 @@
+"""
+Defines a Loader class to load data from a file or file wildcard
+"""
+
+import argparse
+
+import h5py
+import torch
+import numpy as np
+import glob
+from typing import Tuple
+
+
+class Loader:
+    """
+    Data loader class
+    """
+
+    def __init__(self, **kwargs):
+        parser = Loader.add_argparse_args()
+        for action in parser._actions:
+            if action.dest in kwargs:
+                action.default = kwargs[action.dest]
+        args = parser.parse_args([])
+        self.__dict__.update(vars(args))
+
+        if type(self.label_vars) is str:
+            self.label_vars = [self.label_vars]
+
+    @staticmethod
+    def add_argparse_args(parent_parser=None):
+        """
+        Add argeparse argument for the data loader
+        """
+        parser = argparse.ArgumentParser(
+            prog='Loader',
+            usage=Loader.__doc__,
+            parents=[parent_parser] if parent_parser is not None else [],
+            add_help=False)
+
+        parser.add_argument('--input_var', default='p_f5.0_o0', help='Variable name for the label data')
+        parser.add_argument('--label_vars', nargs='*', default='c0', help='Variable name(s) for the label data')
+
+        parser.add_argument('--inputs_crop', type=int, default=[0, 1, 32, 96, 42, 2090], nargs='*',
+                            help='Crop input data on load [layer_min layer_max x_min x_max y_min y_max]')
+        parser.add_argument('--labels_crop', type=int, default=[322, 830, 60, 1076], nargs='*', help='Crop label data on load [x_min x_max y_min y_max]')
+        parser.add_argument('--labels_resize', type=float, default=256.0 / 1016.0, help='scaling factor for labels image')
+
+        parser.add_argument('--data_scale', type=float, default=1.0, help='Data scaling factor')
+        parser.add_argument('--data_gain', type=float, default=1.8, help='Data gain factor in dB/20 at farthest point in data.')
+
+        return parser
+
+    def load_data(self, test_file_pattern: str, train_file_pattern: str = None) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        """Loads training/testing data from file(s)
+
+        Arguments:
+            test_file_pattern  {str} -- testing dataset(s) pattern
+            train_file_pattern  {str} -- training dataset(s) pattern
+
+        Returns:
+            (test_inputs, test_labels, train_inputs, train_labels) -- None for values that are not loaded
+        """
+
+        test_inputs, test_labels = self._load_data_files(test_file_pattern)
+        train_inputs, train_labels = self._load_data_files(train_file_pattern)
+
+        if train_file_pattern is not None and train_inputs is None:
+            raise ValueError('Failed to load train set')
+
+        if test_file_pattern is not None and test_inputs is None:
+            raise ValueError('Failed to load train set')
+
+        return test_inputs, test_labels, train_inputs, train_labels
+
+    def _load_data_files(self, file_pattern: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        """ Perform actual data loading
+
+        Args:
+            file_pattern: file name pattern
+
+        Returns:
+            inputs and labels tensors
+        """
+
+        inputs, labels = None, None
+
+        if file_pattern is None:
+            return inputs, labels
+
+        files = glob.glob(file_pattern)
+
+        if len(files) == 0:
+            raise ValueError(f'{file_pattern=} comes up empty')
+
+        # Load first file to get output dimensions
+        with h5py.File(files[0], 'r') as f:
+            if self.input_var not in f:
+                raise ValueError(f'input data key not in file: {self.input_var=}')
+
+            shape = list(f[self.input_var].shape)
+            if self.inputs_crop is not None:
+                for i in range(len(self.inputs_crop) // 2):
+                    shape[-i - 1] = self.inputs_crop[-i * 2 - 1] - self.inputs_crop[-i * 2 - 2]
+
+            shape[0] *= len(files)
+
+            inputs = np.empty(shape, np.single)
+
+            if len(self.label_vars):
+                if not all([v in f for v in self.label_vars]):
+                    raise ValueError(f'labels data key(s) not in file: {self.label_vars=}')
+
+                shape = list(f[self.label_vars[0]].shape)
+                shape[1] *= len(self.label_vars)
+                if self.labels_crop is not None:
+                    for i in range(len(self.labels_crop) // 2):
+                        shape[-i - 1] = self.labels_crop[-i * 2 - 1] - self.labels_crop[-i * 2 - 2]
+
+                shape[-1] = int(shape[-1] * self.labels_resize)
+                shape[-2] = int(shape[-2] * self.labels_resize)
+                shape[0] *= len(files)
+
+                labels = np.empty(shape, np.single)
+
+        # Load data from files
+        pos = 0
+        for file in files:
+            with h5py.File(files[0], 'r') as f:
+                tmp_inputs = np.array(f[self.input_var])
+
+                if self.inputs_crop is not None:
+                    slc = [slice(None)] * 4
+                    for i in range(len(self.inputs_crop) // 2):
+                        slc[-i - 1] = slice(self.inputs_crop[-i * 2 - 2], self.inputs_crop[-i * 2 - 1])
+                    tmp_inputs = tmp_inputs[tuple(slc)]
+
+                inputs[pos:pos + tmp_inputs.shape[0], ...] = tmp_inputs
+
+                if len(self.label_vars):
+                    tmp_labels = []
+                    for v in self.label_vars:
+                        tmp_labels.append(np.array(f[v]))
+                    tmp_labels = np.concatenate(tmp_labels, axis=1)
+
+                    if self.labels_crop is not None and self.labels_crop:
+                        slc = [slice(None)] * 4
+                        for i in range(len(self.labels_crop) // 2):
+                            slc[-i - 1] = slice(self.labels_crop[-i * 2 - 2], self.labels_crop[-i * 2 - 1])
+                        tmp_labels = tmp_labels[tuple(slc)]
+
+                    if self.labels_resize != 1.0:
+                        tmp_labels = torch.nn.Upsample(scale_factor=self.labels_resize, mode='nearest')(torch.from_numpy(tmp_labels)).numpy()
+
+                    labels[pos:pos + tmp_labels.shape[0], ...] = tmp_labels
+
+                pos += tmp_inputs.shape[0]
+
+        inputs = inputs[:pos, ...]
+        if len(self.label_vars):
+            labels = labels[:pos, ...]
+
+        if self.data_scale != 1.0:
+            inputs *= self.data_scale
+
+        if self.data_gain != 0.0:
+            gain = 10.0 ** np.linspace(0, self.data_gain, inputs.shape[-1], np.single).reshape((1, 1, 1, -1))
+            inputs *= gain
+
+        # Required when inputs is non-continuous due to transpose
+        # TODO: Could probably use a check on strides and do a conditional copy.
+        inputs = torch.from_numpy(inputs.copy())
+        if len(self.label_vars):
+            labels = torch.from_numpy(labels)
+
+        return inputs, labels

models/embc_sos.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f16952eb255caf2c845fd8d0cbdc17e5ea8d8a63cc7669699f8003cc388b22e6
+size 17301623

models/tbme2_attn.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6888856204f0137b3c396fd554b1e5a30b97dcc066638f6297267804f0fd44c
+size 9633441

models/tbme2_phase_sos.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5582ded80ec25c459a2307f698fccd0ea7d2c2537914fc65aa0ef0777ee27759
+size 9639457

models/tbme2_sos.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9636c47b19f376e7236149584b715477249190f996f9f0f44a11692b02cf28c4
+size 9633441

models/tbme_sos.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de4663eeb6b655d021fd69381703db2e0e90ba20ceb033586a6097f33d7883d4
+size 6100085

net.py

@@ -0,0 +1,494 @@
+"""
+Network definition file
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchaudio.functional import lfilter
+
+from pytorch_lightning import LightningModule
+
+import numpy as np
+from scipy.signal import butter, gaussian
+from copy import deepcopy
+import argparse
+
+
+class Net(LightningModule):
+    def __init__(self, **kwargs):
+        super().__init__()
+
+        parser = Net.add_model_specific_args()
+        for action in parser._actions:
+            if action.dest in kwargs:
+                action.default = kwargs[action.dest]
+
+        args = parser.parse_args([])
+        self.hparams.update(vars(args))
+
+        if not hasattr(self, f"_init_{self.hparams.net_type}_net"):
+            raise ValueError(f"Unknown net type {self.hparams.net_type}")
+
+        self._net = eval(f"self._init_{self.hparams.net_type}_net(n_inputs={self.hparams.n_inputs}, n_outputs={self.hparams.n_outputs})")
+
+        if self.hparams.bias is not None:
+            if hasattr(self.hparams.bias, "__iter__"):
+                for i in range(len(self.hparams.bias)):
+                    self._net[-1].c.bias[i].data.fill_(self.hparams.bias[i])
+            else:
+                self._net[-1].c.bias.data.fill_(self.hparams.bias)
+
+    @staticmethod
+    def _init_tbme2_net(n_inputs: int = 1, n_outputs: int = 1):
+        return nn.Sequential(
+            # Encoder
+            DownBlock(n_inputs, 32, 32, 3, stride=[1, 2], pool=None,   push=False, layers=3),
+            DownBlock(32,  32,  32,     3, stride=[1, 2], pool=None,   push=False, layers=3),
+            DownBlock(32,  32,  32,     3, stride=[1, 2], pool=None,   push=False, layers=3),
+            DownBlock(32,  32,  32,     3, stride=[1, 2], pool=None,   push=True,  layers=3),
+            DownBlock(32,  32,  64,     3, stride=1,      pool=[2, 2], push=True,  layers=3),
+            DownBlock(64,  64,  128,    3, stride=1,      pool=[2, 2], push=True,  layers=3),
+            DownBlock(128, 128, 512,    3, stride=1,      pool=[2, 2], push=False, layers=3),
+            # Decoder
+            UpBlock(512, 128, 3, scale_factor=2, pop=False, layers=3),
+            UpBlock(256, 64, 3,  scale_factor=2, pop=True,  layers=3),
+            UpBlock(128, 32, 3,  scale_factor=2, pop=True,  layers=3),
+            UpBlock(64, 32, 3,   scale_factor=2, pop=True,  layers=3),
+            UpStep(32, 32, 3,    scale_factor=1),
+            Compress(32, n_outputs))
+
+    @staticmethod
+    def _init_embc_net(n_inputs: int = 1, n_outputs: int = 1):
+        return nn.Sequential(
+            # Encoder
+            DownBlock(n_inputs, 32, 32, 15, [1, 2], None, layers=1),
+            DownBlock(32, 32, 32, 13, [1, 2], None, layers=1),
+            DownBlock(32, 32, 32, 11, [1, 2], None, layers=1),
+            DownBlock(32, 32, 32, 9, [1, 2], None, True, layers=1),
+            DownBlock(32, 32, 64, 7, 1, [2, 2], True, layers=1),
+            DownBlock(64, 64, 128, 5, 1, [2, 2], True, layers=1),
+            DownBlock(128, 128, 512, 3, 1, [2, 2], layers=1),
+            # Decoder
+            UpBlock(512, 128, 5, 2, layers=1),
+            UpBlock(256, 64, 7, 2, True, layers=1),
+            UpBlock(128, 32, 9, 2, True, layers=1),
+            UpBlock(64, 32, 11, 2, True, layers=1),
+            UpStep(32, 32, 3, 1),
+            Compress(32, n_outputs))
+
+    @staticmethod
+    def _init_tbme_net(n_inputs: int = 1, n_outputs: int = 1):
+        return nn.Sequential(
+            # Encoder
+            DownBlock(n_inputs, 32, 32, 3, [1, 2], None, layers=1),
+            DownBlock(32, 32, 32, 3, [1, 2], None, layers=1),
+            DownBlock(32, 32, 32, 3, [1, 2], None, layers=1),
+            DownBlock(32, 32, 32, 3, [1, 2], None, True, layers=1),
+            DownBlock(32, 32, 64, 3, 1, [2, 2], True, layers=1),
+            DownBlock(64, 64, 128, 3, 1, [2, 2], True, layers=1),
+            DownBlock(128, 128, 512, 3, 1, [2, 2], layers=1),
+            # Decoder
+            UpBlock(512, 128, 3, 2, layers=1),
+            UpBlock(256, 64, 3, 2, True, layers=1),
+            UpBlock(128, 32, 3, 2, True, layers=1),
+            UpBlock(64, 32, 3, 2, True, layers=1),
+            UpStep(32, 32, 3, 1),
+            Compress(32, n_outputs))
+
+    @staticmethod
+    def add_model_specific_args(parent_parser=None):
+        parser = argparse.ArgumentParser(
+            prog="Net",
+            usage=Net.__doc__,
+            parents=[parent_parser] if parent_parser is not None else [],
+            add_help=False)
+
+        parser.add_argument("--random_mirror", type=int, nargs="?", default=1, help="Randomly mirror data to increase diversity when using flat plate wave")
+        parser.add_argument("--noise_std", type=float, nargs="*", help="range of std of random noise to add to the input signal [0 val] or [min max]")
+        parser.add_argument("--quantization", type=float, nargs="?", help="Quantization noise")
+        parser.add_argument("--rand_drop", type=int, nargs="*", help="Random drop lines, between 0 and value lines if single value, or between two values")
+        parser.add_argument("--normalize_net", type=float, default=0.0, help="Coefficient for normalizing network weights")
+
+        parser.add_argument("--learning_rate", type=float, default=5e-3, help="Learning rate to use for optimizer")
+        parser.add_argument("--lr_sched_step", type=int, default=15, help="Learning decay, update step size")
+        parser.add_argument("--lr_sched_gamma", type=float, default=0.65, help="Learning decay gamma")
+
+        parser.add_argument("--net_type", default="tbme2", help="The network to use [tbme2/embc/tbme]")
+        parser.add_argument("--bias", type=float, nargs="*", help="Set bias on last layer, set to 1500 when training from scratch on SoS output")
+        parser.add_argument("--decimation", type=int, help="Subsample phase signal")
+        parser.add_argument("--phase_inv", type=int, default=0, help="Use phase for inversion")
+
+        parser.add_argument("--center_freq", type=float, default=5e6, help="Matched filter and IQ demodulation frequency")
+        parser.add_argument("--n_periods", type=float, default=5, help="Matched filter length")
+        parser.add_argument("--matched_filter", type=int, nargs="?", default=0, help="Apply matched filter, set to 1 to run during forward pass, 2 to run during preprocessing phase (before adding noise)")
+
+        parser.add_argument("--rand_output_crop", type=int, help="Subsample phase signal")
+        parser.add_argument("--rand_scale", type=float, nargs="*", help="Random scaling range [min max] -- (10 ** rand_scale)")
+        parser.add_argument("--rand_gain", type=float, nargs="*", help="Random gain coefficient range [min max] -- (10 ** rand_gain)")
+
+        parser.add_argument("--n_inputs", type=int, default=1, help="Number of input layers")
+        parser.add_argument("--n_outputs", type=int, default=1, help="Number of output layers")
+        parser.add_argument("--scale_losses", type=float, nargs="*", help="Scale each layer of the loss function by given value")
+
+        return parser
+
+    def forward(self, x) -> torch.Tensor:
+        # Matched filter
+        if self.hparams.matched_filter == 1:
+            x = self._matched_filter(x)
+
+        # compute IQ phase if in phase_inv mode
+        if self.hparams.phase_inv:
+            x = self._phase(x)
+
+        # Decimation
+        if self.hparams.decimation != 1:
+            x = x[..., ::self.hparams.decimation]
+
+        # Apply network
+        x = self._net((x, []))
+
+        return x
+
+    def _matched_filter(self, x):
+        sampling_freq = 40e6
+
+        samples_per_cycle = sampling_freq / self.hparams.center_freq
+        n_samples = np.ceil(samples_per_cycle * self.hparams.n_periods + 1)
+
+        signal = torch.sin(torch.arange(n_samples, device=x.device) / samples_per_cycle * 2 * np.pi) * torch.from_numpy(gaussian(n_samples, (n_samples - 1) / 6).astype(np.single)).to(x.device)
+
+        return torch.nn.functional.conv1d(x.reshape(x.shape[:2] + (-1,)), signal.reshape(1, 1, -1), padding="same").reshape(x.shape)
+
+    def _phase(self, x):
+        f = self.hparams.center_freq
+        F = 40e6
+        N = x.shape[-1]
+
+        n = int(round(f * N / F))
+
+        X = torch.fft.fft(x, dim=-1)
+        X[..., (2 * n + 1):] = 0
+        X[..., :(2 * n + 1)] *= torch.from_numpy(gaussian(2 * n + 1, 2 * n  / 6).astype(np.single)).to(x.device)
+        X = X.roll(-n, dims=-1)
+        x = torch.fft.ifft(X, dim=-1)
+
+        return x.angle()
+
+    def _preprocess(self, x):
+        # Matched filter
+        if self.hparams.matched_filter == 2:
+            x = self._matched_filter(x)
+
+        #  Gaussian (normal) noise - random scaling, normalized to signal STD
+        if (ns := self.hparams.noise_std) and len(ns):
+            scl = ns[0] if len(ns) == 1 else torch.rand([x.shape[0]] + [1] * 3).to(x.device) * (ns[-1] - ns[-2]) + ns[-2]
+            scl *= x.std()
+            x += torch.empty_like(x).normal_() * scl
+
+        # Random multiplicative scaling
+        if (rs := self.hparams.rand_scale) and len(rs):
+            x *= 10 ** (torch.rand([x.shape[0]] + [1] * 3).to(x.device) * (rs[-1] - rs[-2]) + rs[-2])
+
+        # Random exponential gain
+        if (gs := self.hparams.rand_gain) and len(gs):
+            gain = torch.FloatTensor([10.0]).to(x.device) ** \
+                (torch.rand([x.shape[0]] + [1] * 3).to(x.device) * ((gs[-1] - gs[-2]) + gs[-2]) *
+                 torch.linspace(0, 1, x.shape[-1]).to(x.device).view(1, 1, 1, -1))
+            x *= gain
+
+        # Quantization noise, to emulated ADC
+        if (quantization := self.hparams.quantization) is not None:
+            x = (x * quantization).round() * (1.0 / quantization)
+
+        # Randomly zero out some of the channels
+        if (rand_drop := self.hparams.rand_drop) and len(rand_drop):
+            if len(rand_drop) == 1:
+                rand_drop = [0, ] + rand_drop
+
+            for i in range(x.shape[0]):
+                lines = np.random.randint(0, x.shape[2], np.random.randint(rand_drop[0], rand_drop[1] + 1))
+                x[i, :, lines, :] = 0.
+
+        return x
+
+    def _log_losses(self, outputs: torch.Tensor, labels: torch.Tensor, prefix: str = ""):
+        diff = torch.abs(labels.detach() - outputs.detach())
+
+        s1 = int(diff.shape[-1] * (1.0 / 3.0))
+        s2 = int(diff.shape[-1] * (2.0 / 3.0))
+
+        for i in range(diff.shape[1]):
+            tag = f"{i}_" if diff.shape[1] > 1 else ""
+
+            losses = {
+                f"{prefix + tag}rmse":  torch.sqrt(torch.mean(diff[:, i, ...] * diff[:, i, ...])).item(),
+                f"{prefix + tag}mean":  torch.mean(diff[:, i, ...]).item(),
+                f"{prefix + tag}short": torch.mean(diff[:, i, :, :s1]).item(),
+                f"{prefix + tag}med":   torch.mean(diff[:, i, :, s1:s2]).item(),
+                f"{prefix + tag}long":  torch.mean(diff[:, i, :, s2:]).item()}
+
+            self.log_dict(losses, prog_bar=True)
+
+    def training_step(self, batch, batch_idx):
+        if self.hparams.random_mirror:
+            mirror = np.random.randint(0, 2, batch[0].shape[0])
+
+            for b in batch:
+                for i, m in enumerate(mirror):
+                    if not m:
+                        continue
+
+                    b[i, ...] = b[i, :, range(b.shape[-2] - 1, -1, -1), :]  # Pytorch does not handle negative steps
+
+        loss = self._common_step(batch, batch_idx, "train_")
+
+        if self.hparams.normalize_net:
+            for W in self.parameters():
+                loss += self.hparams.normalize_net * W.norm(2)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        return self._common_step(batch, batch_idx, "validate_")
+
+    def test_step(self, batch, batch_idx):
+        return self._common_step(batch, batch_idx, "test_")
+
+    def predict_step(self, batch, batch_idx):
+        x = batch[0]
+
+        x = self._preprocess(x)
+        z = self(x)
+
+        if isinstance(z, tuple):
+            z = z[0]
+
+        return z
+
+    def _common_step(self, batch, batch_idx, prefix):
+        x, y = batch
+
+        if self.hparams.rand_output_crop:
+            crop = np.random.randint(0, self.hparams.rand_output_crop, batch[0].shape[0])
+
+            for i, c in enumerate(crop):
+                if not c:
+                    continue
+
+                x[i, :, :-c, :] = x[i, :, c:, :].clone()
+                y[i, :, :-c*2, :] = \
+                    y[i, :, c*2-1:-1, :].clone() if np.random.randint(2) else \
+                        y[i, :, c*2:, :].clone()
+
+            x = x[..., :-self.hparams.rand_output_crop, :]
+            y = y[..., :-self.hparams.rand_output_crop*2, :]
+
+        x = self._preprocess(x)
+        z = self(x)
+
+        outputs = z[0] if isinstance(z, tuple) or isinstance(z, list) else z
+        self._log_losses(outputs, y, prefix)
+
+        if (self.hparams.scale_losses) and len(self.hparams.scale_losses):
+            s = torch.FloatTensor(self.hparams.scale_losses).to(y.device).view(1, -1, 1, 1)
+            loss = F.mse_loss(s * z, s * y)
+        else:
+            loss = F.mse_loss(y, outputs)
+
+        self.log(prefix + "loss", np.sqrt(loss.item()))
+
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, self.hparams.lr_sched_step, self.hparams.lr_sched_gamma)
+
+        return [optimizer], [scheduler]
+
+
+class DownStep(nn.Module):
+    """
+    Down scaling step in the encoder decoder network
+    """
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: int = 1, pool: tuple = None) -> None:
+        """Constructor
+
+        Arguments:
+            in_channels {int} -- Number of input channels for 2D convolution
+            out_channels {int} -- Number of output channels for 2D convolution
+            kernel_size {tuple} -- Convolution kernel size
+
+        Keyword Arguments:
+            stride {int} -- Stride of convolution, set to 1 to disable (default: {1})
+            pool {tuple} -- max pulling size, set to None to disable (default: {None})
+        """
+        super(DownStep, self).__init__()
+
+        self.c = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=kernel_size // 2)
+        self.n = nn.BatchNorm2d(out_channels)
+        self.pool = pool
+
+    def forward(self, x: torch.tensor) -> torch.tensor:
+        """Run the forward step
+
+        Arguments:
+            x {torch.tensor} -- input tensor
+
+        Returns:
+            torch.tensor -- output tensor
+        """
+        x = self.c(x)
+        x = F.relu(x)
+        if self.pool is not None:
+            x = F.max_pool2d(x, self.pool)
+        x = self.n(x)
+
+        return x
+
+
+class UpStep(nn.Module):
+    """
+    Up scaling step in the encoder decoder network
+    """
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, scale_factor: int = 2) -> None:
+        """Constructor
+
+        Arguments:
+            in_channels {int} -- Number of input channels for 2D convolution
+            out_channels {int} -- Number of output channels for 2D convolution
+            kernel_size {int} -- Convolution kernel size
+
+        Keyword Arguments:
+            scale_factor {int} -- Upsampling scaling factor (default: {2})
+        """
+        super(UpStep, self).__init__()
+
+        self.c = nn.Conv2d(in_channels, out_channels, kernel_size, padding=kernel_size // 2)
+        self.n = nn.BatchNorm2d(out_channels)
+        self.scale_factor = scale_factor
+
+    def forward(self, x: torch.tensor) -> torch.tensor:
+        """Run the forward step
+
+        Arguments:
+            x {torch.tensor} -- input tensor
+
+        Returns:
+            torch.tensor -- output tensor
+        """
+        if isinstance(x, tuple):
+            x = x[0]
+
+        if self.scale_factor != 1:
+            x = F.interpolate(x, scale_factor=self.scale_factor)
+
+        x = self.c(x)
+        x = F.relu(x)
+        x = self.n(x)
+
+        return x
+
+
+class Compress(nn.Module):
+    """
+    Up scaling step in the encoder decoder network
+    """
+    def __init__(self, in_channels: int, out_channels: int = 1, kernel_size: int = 1, scale_factor: int = 1) -> None:
+        """Constructor
+
+        Arguments:
+            in_channels {int} -- [description]
+
+        Keyword Arguments:
+            out_channels {int} -- [description] (default: {1})
+            kernel_size {int} -- [description] (default: {1})
+        """
+        super(Compress, self).__init__()
+
+        self.scale_factor = scale_factor
+
+        self.c = nn.Conv2d(in_channels, out_channels, kernel_size, padding=kernel_size // 2)
+
+    def forward(self, x: torch.tensor) -> torch.tensor:
+        """Run the forward step
+
+        Arguments:
+            x {torch.tensor} -- input tensor
+
+        Returns:
+            torch.tensor -- output tensor
+        """
+        if isinstance(x, tuple) or isinstance(x, list):
+            x = x[0]
+
+        x = self.c(x)
+
+        if self.scale_factor != 1:
+            x = F.interpolate(x, scale_factor=self.scale_factor)
+
+        return x
+
+
+class DownBlock(nn.Module):
+    def __init__(
+            self,
+            in_chan: int, inter_chan: int, out_chan: int,
+            kernel_size: int = 3, stride: int = 1, pool: tuple = None,
+            push: bool = False,
+            layers: int = 3):
+        super().__init__()
+
+        self.s = []
+        for i in range(layers):
+            self.s.append(deepcopy(DownStep(
+                in_chan if i == 0 else inter_chan,
+                inter_chan if i < layers - 1 else out_chan,
+                kernel_size,
+                1 if i < layers - 1 else stride,
+                None if i < layers - 1 else pool)))
+        self.s = nn.Sequential(*self.s)
+
+        self.push = push
+
+    def forward(self, x: torch.tensor) -> torch.tensor:
+        i, s = x
+
+        i = self.s(i)
+
+        if self.push:
+            s.append(i)
+
+        return i, s
+
+
+class UpBlock(nn.Module):
+    def __init__(
+            self,
+            in_chan: int, out_chan: int,
+            kernel_size: int, scale_factor: int = 2,
+            pop: bool = False,
+            layers: int = 3):
+        super().__init__()
+
+        self.s = []
+        for i in range(layers):
+            self.s.append(deepcopy(UpStep(
+                in_chan if i == 0 else out_chan,
+                out_chan,
+                kernel_size,
+                1 if i < layers - 1 else scale_factor)))
+        self.s = nn.Sequential(*self.s)
+
+        self.pop = pop
+
+    def forward(self, x: torch.tensor) -> torch.tensor:
+        i, s = x
+
+        if self.pop:
+            i = torch.cat((i, s.pop()), dim=1)
+
+        i = self.s(i)
+
+        return i, s

run_logger.py

@@ -0,0 +1,120 @@
+"""
+Support log functions
+
+TODO: log model using mlflow.pytorch in parallel / addition to checkpointing
+"""
+
+import numpy as np
+import h5py
+import os
+import argparse
+
+import torch
+import torchvision.utils as vutils
+import pytorch_lightning as pl
+
+
+class ImgCB(pl.Callback):
+    def __init__(self, **kwargs):
+        parser = ImgCB.add_argparse_args()
+        for action in parser._actions:
+            if action.dest in kwargs:
+                action.default = kwargs[action.dest]
+        args = parser.parse_args([])
+        self.__dict__.update(vars(args))
+
+    @staticmethod
+    def add_argparse_args(parent_parser=None):
+        parser = argparse.ArgumentParser(
+            prog='ImgCB',
+            usage=ImgCB.__doc__,
+            parents=[parent_parser] if parent_parser is not None else [],
+            add_help=False)
+
+        parser.add_argument('--img_ranges', default=[1300, 1800], nargs='*', help='Scaling range on output image, either pair, or set of pairs')
+        parser.add_argument('--err_ranges', default=[0, 50], nargs='*', help='Scaling range on error images, either pair, or set of pairs')
+
+        return parser
+
+    def log_images(self, mfl_logger, y, z, prefix):
+        img_ranges = tuple(self.img_ranges)
+        err_ranges = tuple(self.err_ranges)
+        # 
+        for i in range(y.shape[1]):
+            if y.shape[1] > 1:
+                tag = f'_{i}_'
+
+                if len(self.img_ranges) > 2:
+                    img_ranges = tuple(self.img_ranges[2*i, 2*i + 1])
+                if len(self.err_ranges) > 2:
+                    err_ranges = tuple(self.err_ranges[2*i, 2*i + 1])
+            else:
+                tag = ''
+
+            mfl_logger.experiment.log_image(
+                mfl_logger.run_id,
+                (np.array(vutils.make_grid(
+                    y[:, [i], ...].detach(),
+                    normalize=True, value_range=img_ranges, nrow=6).cpu())[0, ...] * 255.).astype(np.int),
+                prefix + tag + '_labels.png')
+
+            mfl_logger.experiment.log_image(
+                mfl_logger.run_id,
+                (np.array(vutils.make_grid(
+                    z[:, [i], ...].detach(),
+                    normalize=True, value_range=img_ranges, nrow=6).cpu())[0, ...] * 255.).astype(np.int),
+                prefix + tag + '_outputs.png')
+
+            mfl_logger.experiment.log_image(
+                mfl_logger.run_id,
+                (np.array(vutils.make_grid(
+                    torch.abs(y[:, [i], ...].detach() - z[:, [i], ...].detach()),
+                    normalize=True, value_range=err_ranges, nrow=6).cpu())[0, ...] * 255.).astype(np.int),
+                prefix + tag + '_errors.png')
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        if batch_idx == 0:
+            with torch.no_grad():
+                x, y = batch
+
+                if pl_module.hparams.rand_output_crop:
+                    x = x[..., :-pl_module.hparams.rand_output_crop, :]
+                    y = y[..., :-pl_module.hparams.rand_output_crop * 2, :]
+
+                z = pl_module(x.to(pl_module.device))
+
+                if isinstance(z, tuple) or isinstance(z, list):
+                    z = z[0]
+
+                self.log_images(pl_module.logger, y.to(pl_module.device), z, 'train_')
+
+    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
+        if batch_idx == 0:
+            with torch.no_grad():
+                x, y = batch
+
+                if pl_module.hparams.rand_output_crop:
+                    x = x[..., :-pl_module.hparams.rand_output_crop, :]
+                    y = y[..., :-pl_module.hparams.rand_output_crop * 2, :]
+
+                z = pl_module(x.to(pl_module.device))
+
+                if isinstance(z, tuple) or isinstance(z, list):
+                    z = z[0]
+
+                self.log_images(pl_module.logger, y.to(pl_module.device), z, 'validate_')
+
+
+class TestLogger(pl.Callback):
+    """
+    pytorch_lightning Data saving logger for testing output
+    Warning !!! : this function is not multi GPU / multi device safe -- only run on a single gpu / device
+    """
+    def __init__(self, fname: str = 'output.h5'):
+        self.fname = fname
+
+    def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
+        with h5py.File(self.fname, 'a') as f:
+            f[f'batch_{batch_idx:05}'] = outputs.to('cpu').numpy()
+            if len(batch) > 1:
+                f[f'labels_{batch_idx:05}'] = batch[1].to('cpu').numpy()