add intesity scale and suffix parameters to predict scripts

predict-gui.py

@@ -1,4 +1,3 @@
-import os
 import torch
 torch.set_grad_enabled(False)
 import numpy as np
@@ -30,7 +29,7 @@ parser.add_argument(
 )
 args = parser.parse_args()
 
-def process_file(input_file, config_file, weights_file, input_spectrometer_frequency=None,reference_spectrum=None):
+def process_file(input_file, config_file, weights_file, input_spectrometer_frequency=None,reference_spectrum=None, scale=None, suffix=None):
     if input_spectrometer_frequency == 0:
         input_spectrometer_frequency = None
     # Load configuration and initialize predictor
@@ -53,7 +52,9 @@ def process_file(input_file, config_file, weights_file, input_spectrometer_frequ
 
     # Scale and process spectrum
     spectrum_tensor = torch.tensor(spectrum).float()
-    scaling_factor = Defaults.SCALE / spectrum_tensor.max()
+    if scale is None:
+        scale = Defaults.SCALE
+    scaling_factor = scale / spectrum_tensor.max()
     spectrum_tensor *= scaling_factor
     prediction = predictor(spectrum_tensor).numpy()
     prediction /= scaling_factor
@@ -63,7 +64,9 @@ def process_file(input_file, config_file, weights_file, input_spectrometer_frequ
 
     # Prepare output data for download
     output_data = np.column_stack((input_freqs_input_ppm, output_prediction))
-    output_file = f"{Path(input_file).stem}_processed{Path(input_file).suffix}"
+    if suffix is None:
+        suffix = Defaults.SUFFIX
+    output_file = f"{Path(input_file).stem}{suffix}{Path(input_file).suffix}"
     np.savetxt(output_file, output_data)
 
     # Create Plotly figure
@@ -141,6 +144,18 @@ with gr.Blocks() as app:
             )
             config_file = gr.File(label="Custom Config File (.yaml)", visible=False, height=120)
             weights_file = gr.File(label="Custom Weights File (.pt)", visible=False, height=120)
+
+            with gr.Accordion("Advanced", open=False):
+                scale_input = gr.Number(
+                    label="Intensity Scale",
+                    value=Defaults.SCALE,
+                    info=f"Adjust the scaling factor for intensity normalization. Default is {Defaults.SCALE}.",
+                )
+                suffix_input = gr.Textbox(
+                    label="Output File Suffix",
+                    value=Defaults.SUFFIX,
+                    info=f"Suffix to add to processed output filenames. Default is '{Defaults.SUFFIX}'.",
+                )
         
         with gr.Column():
             input_file = gr.File(label="Input File (.txt | .csv)", height=120)
@@ -166,7 +181,7 @@ 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):
+    def process_file_with_model(input_file, model_selection, config_file, weights_file, input_spectrometer_frequency, reference_spectrum_file, scale, suffix):
         if model_selection == "600 MHz":
             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")
@@ -177,11 +192,19 @@ with gr.Blocks() as app:
             config_file = config_file.name
             weights_file = weights_file.name
 
-        return process_file(input_file.name, config_file, weights_file, input_spectrometer_frequency, reference_spectrum_file.name if reference_spectrum_file else None)
+        return process_file(
+            input_file.name,
+            config_file,
+            weights_file,
+            input_spectrometer_frequency,
+            reference_spectrum_file.name if reference_spectrum_file else None,
+            scale,
+            suffix
+        )
 
     process_button.click(
         process_file_with_model,
-        inputs=[input_file, model_selection, config_file, weights_file, input_spectrometer_frequency, reference_spectrum_file],
+        inputs=[input_file, model_selection, config_file, weights_file, input_spectrometer_frequency, reference_spectrum_file, scale_input, suffix_input],
         outputs=[plot_output, download_button]
     )
 

predict.py

@@ -14,6 +14,9 @@ from shimnet.predict_utils import Defaults, resample_input_spectrum, resample_ou
 import warnings
 warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
 
+class Defaults:
+    SCALE = 16.0
+    SUFFIX = "_processed"
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -22,6 +25,8 @@ def parse_args():
     parser.add_argument("--weights", help="model weights")
     parser.add_argument("-o", "--output_dir", default=".", help="Output directory")
     parser.add_argument("--input_spectrometer_frequency", default=None, type=float, help="spectrometer frequency in MHz (input sample collection frequency). Empty if the same as in the training data")
+    parser.add_argument("--suffix", default=Defaults.SUFFIX, help=f"Output file suffix (default: {Defaults.SUFFIX})")
+    parser.add_argument("--intensity_scale", default=Defaults.SCALE, type=float, help=f"Intensity scaling factor (default: {Defaults.SCALE})")
     args = parser.parse_args()
     return args
     
@@ -52,7 +57,7 @@ if __name__ == "__main__":
         
         spectrum = torch.tensor(spectrum).float()
         # scale height of the spectrum
-        scaling_factor = Defaults.SCALE / spectrum.max()
+        scaling_factor = args.intensity_scale / spectrum.max()
         spectrum *= scaling_factor
 
         # correct spectrum
@@ -65,7 +70,7 @@ if __name__ == "__main__":
         output_prediction = resample_output_spectrum(input_freqs_model_ppm, freqs, prediction)
 
         # save result
-        output_file = output_dir / f"{Path(input_file).stem}_processed{Path(input_file).suffix}"
+        output_file = output_dir / f"{Path(input_file).stem}{args.suffix}{Path(input_file).suffix}"
 
         np.savetxt(output_file, np.column_stack((input_freqs_input_ppm, output_prediction)))
         print(f"saved to {output_file}")