make shimnet a package

.gitignore

@@ -11,3 +11,7 @@ data/
 
 # gradio
 .gradio/
+
+# package building
+build/
+*shimnet.egg-info/

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.10 as build
+FROM python:3.10 AS build
 
 RUN useradd -m -u 1000 user
 USER user
@@ -10,16 +10,18 @@ ENV HOME=/home/user \
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app
 
+# copy installation files
+COPY --chown=user shimnet shimnet/
+COPY --chown=user pyproject.toml ./
+# install shimnet (cpu version + GUI)
+RUN pip install --no-cache-dir .[cpu,gui] --extra-index-url https://download.pytorch.org/whl/cpu 
 
-COPY --chown=user requirements-cpu.txt requirements-gui.txt ./
-RUN pip install --no-cache-dir -r requirements-cpu.txt -r requirements-gui.txt --extra-index-url https://download.pytorch.org/whl/cpu 
-
-FROM build as final
+FROM build AS final
 
 COPY --chown=user . .
 
 # download weights
 RUN python download_files.py --overwrite
 
-CMD [ "python", "./predict-gui.py"]
+CMD [ "python", "./predict-gui.py", "--server_name", "0.0.0.0" ]
 

Readme.md

@@ -9,17 +9,18 @@ Web service: [![Open in Hugging Face Spaces](https://huggingface.co/datasets/hug
 
 ## Installation
 
-Python 3.9+ (3.10+ for GUI)
+Python 3.10+ is required
 
-GPU version (for training and inference)
-```
-pip install -r requirements-gpu.txt
-```
+You may install CPU-only version for inference only. If you need both training and inference, GPU version is strongly recommended.
 
-CPU version (for inference, not recommended for training)
-```
-pip install -r requirements-cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu
-```
+In both CPU-only and GPU versions you may also install GUI (graphical user interface for inference)
+
+- CPU-only version
+  - with GUI (recommended): `pip install .[cpu,gui] --extra-index-url https://download.pytorch.org/whl/cpu`
+  - without GUI: `pip install .[cpu] --extra-index-url https://download.pytorch.org/whl/cpu`
+- GPU version (strongly recommended for training)
+  - with GUI:  `pip install .[gpu,gui]`
+  - without GUI: `pip install .[gpu]`
 
 ## Usage
 To correct spectra presented in the paper:
@@ -50,7 +51,7 @@ python predict.py sample_data/CresolRed_after_styrene_600MHz.csv -o output --con
 
 ### input format
 
-The spectrum file for reconstruction should be in the format of two columns separated by a space and without the sign at the end of the line at the end of the file(example below):
+The spectrum file for reconstruction should be in the format of two columns separated by a space and without the sign at the end of the line at the end of the file. The first column is frequency in ppm, the second is the intensity. The frequency values ​​should be in ascending order (example below):
 ```csv
 -1.97134	0.0167137	
 -1.97085	-0.00778748	
@@ -188,14 +189,6 @@ If you want to train the network using the calibration data from our paper, foll
 
 ## GUI
 
-### Installation
-
-To use the ShimNet GUI, ensure you have Python 3.10 installed (not tested with Python 3.11+). After installing the ShimNet requirements (CPU/GPU), install the additional dependencies for the GUI:
-
-```bash
-pip install -r requirements-gui.txt
-```
-
 ### Launching the GUI
 
 The ShimNet GUI is built using Gradio. To start the application, run:
@@ -221,3 +214,17 @@ python predict-gui.py --share
 ```
 
 A public web address will be displayed in the terminal, which you can use to access the GUI remotely or share with others.
+
+### GUI inference with Docker
+
+Create docker image:
+```bash
+docker build -t shimnetgui .
+```
+
+Run the container:
+```bash
+docker run -it -p 7860:7860 shimnetgui
+```
+
+The GUI should be working at `http://127.0.0.1:7860`

predict-gui.py

@@ -1,3 +1,4 @@
+import os
 import torch
 torch.set_grad_enabled(False)
 import numpy as np
@@ -6,8 +7,7 @@ from omegaconf import OmegaConf
 import gradio as gr
 import plotly.graph_objects as go
 
-from src.models import ShimNetWithSCRF, Predictor
-from predict import Defaults, resample_input_spectrum, resample_output_spectrum, initialize_predictor
+from shimnet.predict_utils import Defaults, resample_input_spectrum, resample_output_spectrum, initialize_predictor
 
 # silent deprecation warnings
 import warnings
@@ -168,11 +168,11 @@ with gr.Blocks() as app:
     # Process button click logic
     def process_file_with_model(input_file, model_selection, config_file, weights_file, input_spectrometer_frequency, reference_spectrum_file):
         if model_selection == "600 MHz":
-            config_file = "configs/shimnet_600.yaml"
-            weights_file = "weights/shimnet_600MHz.pt"
+            config_file = os.path.join(os.path.dirname(__file__), "configs/shimnet_600.yaml")
+            weights_file = os.path.join(os.path.dirname(__file__), "weights/shimnet_600MHz.pt")
         elif model_selection == "700 MHz":
-            config_file = "configs/shimnet_700.yaml"
-            weights_file = "weights/shimnet_700MHz.pt"
+            config_file = os.path.join(os.path.dirname(__file__), "configs/shimnet_700.yaml")
+            weights_file = os.path.join(os.path.dirname(__file__), "weights/shimnet_700MHz.pt")
         else:
             config_file = config_file.name
             weights_file = weights_file.name
@@ -186,14 +186,3 @@ with gr.Blocks() as app:
     )
 
 app.launch(share=args.share, server_name=args.server_name)
-
-# '#636efa',
-#  '#EF553B',
-#  '#00cc96',
-#  '#ab63fa',
-#  '#FFA15A',
-#  '#19d3f3',
-#  '#FF6692',
-#  '#B6E880',
-#  '#FF97FF',
-#  '#FECB52'

predict.py

@@ -6,15 +6,14 @@ from pathlib import Path
 import sys, os
 from omegaconf import OmegaConf
 
-from src.models import ShimNetWithSCRF, Predictor
+from shimnet.predict_utils import Defaults, resample_input_spectrum, resample_output_spectrum, initialize_predictor
+
 
 # silent deprecation warnings
 # https://github.com/pytorch/pytorch/issues/97207#issuecomment-1494781560
 import warnings
 warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
 
-class Defaults:
-    SCALE = 16.0
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -26,23 +25,6 @@ def parse_args():
     args = parser.parse_args()
     return args
     
-# functions
-def resample_input_spectrum(input_freqs, input_spectrum, Mhz_per_point):
-    """resample input spectrum to match the model's frequency range"""
-    freqs = np.arange(input_freqs.min(), input_freqs.max(), Mhz_per_point)
-    spectrum = np.interp(freqs, input_freqs, input_spectrum)
-    return freqs, spectrum
-
-def resample_output_spectrum(input_freqs, freqs, prediction):
-    """resample prediction to match the input spectrum's frequency range"""
-    prediction = np.interp(input_freqs, freqs, prediction)
-    return prediction
-
-def initialize_predictor(config, weights_file):
-    model = ShimNetWithSCRF(**config.model.kwargs)
-    predictor = Predictor(model, weights_file)
-    return predictor
-
 # run
 if __name__ == "__main__":
     args = parse_args()

pyproject.toml

@@ -0,0 +1,47 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+packages = ["shimnet"]   # only include this package
+
+[project]
+name = "shimnet"
+version = "0.1.0"
+description = "Package for ShimNet: A Neural Network for Postacquisition Improvement of NMR Spectra Distorted by Magnetic-Field Inhomogeneity https://pubs.acs.org/doi/full/10.1021/acs.jpcb.5c02632"
+authors = [
+    {name = "Marek Bukowicki", email = "m.bukowicki@uw.edu.pl"},
+]
+readme = "Readme.md"
+requires-python = ">=3.10"
+dependencies = [
+    "nmrglue==0.11",
+    "torchdata==0.9.0",
+    "numpy==2.0.2",
+    "matplotlib==3.9.3",
+    "pandas==2.2.3",
+    "tqdm==4.67.1",
+    "hydra-core==1.3.2"
+]
+
+[project.optional-dependencies]
+# GPU installs the default torch build from PyPI (CUDA-enabled by default)
+gpu = [
+    "torch==2.5.1",
+    "torchaudio==2.5.1"
+]
+
+# CPU needs extra index URL: pip install shimnet[cpu] --extra-index-url https://download.pytorch.org/whl/cpu
+cpu = [
+    "torch==2.5.1+cpu",
+    "torchaudio==2.5.1+cpu"
+]
+
+# predictions with GUI
+gui = [
+    "gradio==5.23.2",
+    "plotly==6.0.1",
+    "ipykernel",
+    "psutil",
+    "pexpect"
+]

requirements-cpu.txt

@@ -1,9 +0,0 @@
-torch==2.4.1+cpu
-torchaudio==2.4.1+cpu
-nmrglue==0.11
-torchdata==0.9.0
-numpy==2.0.2
-matplotlib==3.9.3
-pandas==2.2.3
-tqdm==4.67.1
-hydra-core==1.3.2

requirements-gpu.txt

@@ -1,9 +0,0 @@
-torch==2.4.1
-torchaudio==2.4.1
-nmrglue==0.11
-torchdata==0.9.0
-numpy==2.0.2
-matplotlib==3.9.3
-pandas==2.2.3
-tqdm==4.67.1
-hydra-core==1.3.2

requirements-gui.txt

@@ -1,2 +0,0 @@
-gradio==5.23.2
-plotly==6.0.1

shimnet/__init__.py

@@ -0,0 +1 @@
+from . import generators, models, nmr_utils, predict_utils

shimnet/nmr_utils.py

@@ -0,0 +1,18 @@
+import nmrglue as ng
+
+def fid_to_spectrum(varian_fid_path, ph0_correction, ph1_correction, autophase_fn, target_length=None, sin_pod=False):
+    dic, data = ng.varian.read(varian_fid_path)
+    data[0] *= 0.5
+    if sin_pod:
+        data = ng.proc_base.sp(data, end=0.98)
+
+    if target_length is not None:
+        if (pad_length := target_length - len(data)) > 0:
+            data = ng.proc_base.zf(data, pad_length)
+        else:
+            data = data[:target_length]
+
+    spec=ng.proc_base.fft(data)
+    spec = ng.process.proc_autophase.autops(spec, autophase_fn, p0=ph0_correction, p1=ph1_correction, disp=False)
+
+    return spec

shimnet/predict_utils.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+from .models import ShimNetWithSCRF, Predictor
+
+class Defaults:
+    SCALE = 16.0
+
+# functions
+def resample_input_spectrum(input_freqs, input_spectrum, Mhz_per_point):
+    """resample input spectrum to match the model's frequency range"""
+    freqs = np.arange(input_freqs.min(), input_freqs.max(), Mhz_per_point)
+    spectrum = np.interp(freqs, input_freqs, input_spectrum)
+    return freqs, spectrum
+
+def resample_output_spectrum(input_freqs, freqs, prediction):
+    """resample prediction to match the input spectrum's frequency range"""
+    prediction = np.interp(input_freqs, freqs, prediction)
+    return prediction
+
+def initialize_predictor(config, weights_file):
+    model = ShimNetWithSCRF(**config.model.kwargs)
+    predictor = Predictor(model, weights_file)
+    return predictor

train.py

@@ -15,8 +15,8 @@ matplotlib.use('Agg')
 import warnings
 warnings.filterwarnings("ignore", category=UserWarning, module='torchdata')
 
-from src import models
-from src.generators import get_datapipe
+# from shiment import models
+from shiment.generators import get_datapipe
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 if len(sys.argv) < 2:
@@ -32,7 +32,7 @@ else:
     minimum = float("inf")
 
 # initialization        
-model = instantiate({"_target_": f"__main__.models.{config.model.name}", **config.model.kwargs}).to(device)
+model = instantiate({"_target_": f"shimnet.models.{config.model.name}", **config.model.kwargs}).to(device)
 model_weights_file = run_dir / f'model.pt'
 optimizer = torch.optim.Adam(model.parameters())
 optimizer_weights_file = run_dir / f'optimizer.pt'