image_to_image.utils.argument_parsing

  1# ---------------------------
  2#        > Imports <
  3# ---------------------------
  4import argparse
  5import textwrap
  6
  7
  8
  9# ---------------------------
 10#     > Argument Parser <
 11# ---------------------------
 12def get_arg_parser():
 13    """
 14    Creates and configures an argument parser for the Image-to-Image framework.
 15
 16    This function defines and groups all command-line arguments required for 
 17    training, validation, and inference of different image-to-image translation 
 18    models (e.g., ResFCN, Pix2Pix, PhysicsFormer, Residual Design Model).
 19    It includes parameters for model configuration, dataset handling, 
 20    optimization, loss weighting, experiment tracking, and hardware setup.
 21
 22    Returns:
 23    - argparse.ArgumentParser: 
 24        A fully configured argument parser with all options and default values.
 25
 26    Argument Groups:
 27        General:
 28            --mode: Operation mode ('train', 'test', 'inference').
 29
 30        Training:
 31            --checkpoint_save_dir: Directory to store model checkpoints.
 32            --save_only_best_model: If set, saves only the best model based on validation loss.
 33            --checkpoint_interval: Interval (in epochs) to save model checkpoints.
 34            --validation_interval: Interval (in epochs) for validation.
 35            --epochs: Number of total training epochs.
 36            --batch_size: Number of samples per training batch.
 37            --lr: Base learning rate.
 38            --loss: Loss function to use ('l1', 'crossentropy', 'weighted_combined').
 39
 40        Weighted Combined Loss Parameters (for primary and secondary models):
 41            --wc_loss_silog_lambda, --wc_loss_weight_silog, --wc_loss_weight_grad, etc.
 42            Fine-tuning coefficients for each subcomponent of the weighted combined loss.
 43
 44        Optimization:
 45            --optimizer, --optimizer_2: Optimizer types (e.g., Adam, AdamW).
 46            --weight_decay: Enable or disable weight decay regularization.
 47            --weight_decay_rate: Coefficient for weight decay.
 48            --gradient_clipping: Enable or disable gradient clipping.
 49            --gradient_clipping_threshold: Maximum gradient norm threshold.
 50            --scheduler, --scheduler_2: Learning rate scheduler types ('step', 'cosine').
 51            --use_warm_up: Enable warm-up phase for the optimizer.
 52            --warm_up_start_lr, --warm_up_step_duration: Warm-up configuration.
 53
 54        Mixed Precision:
 55            --activate_amp: Enables Automatic Mixed Precision (AMP) training.
 56            --amp_scaler: Type of AMP scaler to use ('grad' or None).
 57
 58        Inference:
 59            --model_params_path: Path to saved model checkpoint for inference.
 60            --image_dir_path: Directory containing images for inference.
 61            --output_dir: Output directory for predicted images.
 62
 63        Model Selection:
 64            --model: Choice of model architecture ('resfcn', 'pix2pix', 'residual_design_model', 'physicsformer').
 65
 66        Model-Specific Parameters:
 67            * ResFCN:
 68                --resfcn_in_channels, --resfcn_hidden_channels, --resfcn_out_channels, --resfcn_num_blocks
 69            * Pix2Pix:
 70                --pix2pix_in_channels, --pix2pix_hidden_channels, --pix2pix_out_channels, --pix2pix_second_loss_lambda
 71            * PhysicsFormer:
 72                --physicsformer_in_channels, --physicsformer_out_channels, --physicsformer_img_size,
 73                --physicsformer_patch_size, --physicsformer_embedded_dim, --physicsformer_num_blocks,
 74                --physicsformer_heads, --physicsformer_mlp_dim, --physicsformer_dropout
 75            * Residual Design Model:
 76                --base_model, --complex_model, --combine_mode
 77                (includes separate parameters for the second model branch, e.g., *_2)
 78
 79        Data:
 80            --data_variation: Dataset variant to use (e.g., 'sound_baseline', 'sound_reflection').
 81            --input_type, --output_type: Define input/output representation types.
 82            --fake_rgb_output: Converts single-channel input into fake RGB format.
 83            --make_14_dividable_size: Ensures image dimensions are multiples of 14.
 84
 85        Hardware:
 86            --device: Compute device ('cpu' or 'cuda').
 87
 88        Experiment Tracking:
 89            --experiment_name: Group name for MLflow and TensorBoard logging.
 90            --run_name: Specific run name (timestamp is prepended automatically).
 91            --tensorboard_path: Directory for TensorBoard logs.
 92            --save_path: Path for saving generated images during training/inference.
 93            --cmap: Color map used for visualization of images.
 94    """
 95    parser = argparse.ArgumentParser(description="Image-to-Image Framework - train and inferencing")
 96
 97    # General Parameter
 98    parser.add_argument('--mode', type=str, default='train', choices=['train', 'test', 'inference'],
 99                        help='Modus: train, test or inference')
100
101    # Trainingsparameter
102    parser.add_argument('--checkpoint_save_dir', type=str, default='./checkpoints', help='Path to save the model checkpoints. Is builded: checkpoint_save_dir/experiment_name/run_name')
103    parser.add_argument('--save_only_best_model', action='store_true', help='Should every checkpoint be saved or only the best model?')
104    parser.add_argument('--checkpoint_interval', type=int, default=5, help='Every x epochs checkpoint will be saved (if not `save_only_best_model` is active).')
105    parser.add_argument('--validation_interval', type=int, default=5, help='Every x epochs validation will be calculated.')
106    parser.add_argument('--epochs', type=int, default=50, help='Amount of whole data loops.')
107    parser.add_argument('--batch_size', type=int, default=8, help='Size of a batch, data is processed in batches (smaller packages) and the GPU processes then one batch at a time.')
108    parser.add_argument('--lr', type=float, default=1e-4, help='Learnrate of adjusting the weights towards the gradients.')
109    parser.add_argument('--loss', type=str, default='l1', choices=['l1', 'crossentropy', 'weighted_combined'],
110                        help='Loss Function.')
111    # ---> WeightedCombinedLoss parameters
112    parser.add_argument('--wc_loss_silog_lambda', type=float, default=0.5, help='Lambda parameter for SILog loss.')
113    parser.add_argument('--wc_loss_weight_silog', type=float, default=0.5, help='Weight for SILog loss.')
114    parser.add_argument('--wc_loss_weight_grad', type=float, default=10.0, help='Weight for gradient loss.')
115    parser.add_argument('--wc_loss_weight_ssim', type=float, default=5.0, help='Weight for SSIM loss.')
116    parser.add_argument('--wc_loss_weight_edge_aware', type=float, default=10.0, help='Weight for edge-aware loss.')
117    parser.add_argument('--wc_loss_weight_l1', type=float, default=1.0, help='Weight for L1 loss.')
118    parser.add_argument('--wc_loss_weight_var', type=float, default=1.0, help='Weight for variance loss.')
119    parser.add_argument('--wc_loss_weight_range', type=float, default=1.0, help='Weight for range loss.')
120    parser.add_argument('--wc_loss_weight_blur', type=float, default=1.0, help='Weight for blur loss.')
121
122    parser.add_argument('--optimizer', type=str, default="adam", choices=['adam', 'adamw'],
123                        help='Optimizer, which decides how exactly to calculate the loss and weight gradients.')
124    parser.add_argument('--optimizer_2', type=str, default="adam", choices=['adam', 'adamw'],
125                        help='Optimizer, which decides how exactly to calculate the loss and weight gradients -> for the second model(-part).\nFor example for the discriminator part of pix2pix model or the complex part in the residual design model.')
126    parser.add_argument('--weight_decay', action="store_true", help='Whether or not to use weight decay (keeping weights smaller).')
127    parser.add_argument('--weight_decay_rate', type=float, default=0.0005, help='Coefficient of weight decay -> weighting of the penalty.')
128    parser.add_argument('--gradient_clipping', action="store_true", help='Whether or not to use gradient clipping.')
129    parser.add_argument('--gradient_clipping_threshold', type=float, default=0.5, help='Coefficient of gradient clipping -> threshold for clipping.')
130    parser.add_argument('--scheduler', type=str, default="step", choices=['step', 'cosine'],
131                        help='Decides how to update the learnrate dynamically.')
132    parser.add_argument('--scheduler_2', type=str, default="step", choices=['step', 'cosine'],
133                        help='Decides how to update the learnrate dynamically -> for the second model(-part).\nFor example for the discriminator part of pix2pix model or the complex part in the residual design model.')
134    parser.add_argument('--use_warm_up', action="store_true", help="Whether to use warm up for optimizer/lr.")
135    parser.add_argument('--warm_up_start_lr', type=float, default=0.00005, help='Warm-Up Start learning rate will end at the lr.')
136    parser.add_argument('--warm_up_step_duration', type=int, default=1000, help='Duration of increasing learning rate in steps (one step = one batch process).')
137    parser.add_argument('--activate_amp', action="store_true", help='Activates (Automatically) Mixed Precision and use scaler to loose no details because of the smaller float.')
138    parser.add_argument('--amp_scaler', type=str, default='none', choices=['none', 'grad'],
139                        help='Decides whichscaler should be used-')
140
141    # Inference
142    parser.add_argument('--model_params_path', type=str, required=False, help='Path to the model checkpoints.')
143    parser.add_argument('--image_dir_path', type=str, default=None, required=False, help='Path to folder with images for inference.')
144    parser.add_argument('--output_dir', type=str, default='../../data/eval', help='Path to save the real and predicted Images.')
145
146    # Model Loading
147    parser.add_argument('--model', type=str, default="resfcn", choices=['resfcn', 'pix2pix', 'residual_design_model', 'physicsformer'],
148                        help='Which Model should be choosen')
149
150    # ---> ResFCN
151    parser.add_argument('--resfcn_in_channels', type=int, default=1, help='How much channels as input?')
152    parser.add_argument('--resfcn_hidden_channels', type=int, default=64, help='How much channels in the hidden layers?')
153    parser.add_argument('--resfcn_out_channels', type=int, default=1, help='How much channels as output?')
154    parser.add_argument('--resfcn_num_blocks', type=int, default=16, help='How many Residual Blocks should be stacked.')
155
156    # ---> Pix2Pix
157    parser.add_argument('--pix2pix_in_channels', type=int, default=1, help='How much channels as input?')
158    parser.add_argument('--pix2pix_hidden_channels', type=int, default=64, help='How much channels in the hidden layers?')
159    parser.add_argument('--pix2pix_out_channels', type=int, default=1, help='How much channels as output?')
160    parser.add_argument('--pix2pix_second_loss_lambda', type=float, default=100.0, help='Weighting of second loss.')
161
162    # ---> PhysicsFormer
163    parser.add_argument('--physicsformer_in_channels', type=int, default=1, help='How much channels as input?')
164    parser.add_argument('--physicsformer_out_channels', type=int, default=1, help='How much channels as output?')
165    parser.add_argument('--physicsformer_img_size', type=int, default=256, help='Size of the image (width or height).')
166    parser.add_argument('--physicsformer_patch_size', type=int, default=4, help='Size of patches.')
167    parser.add_argument('--physicsformer_embedded_dim', type=int, default=1024, help='Dimension size of embedding.')
168    parser.add_argument('--physicsformer_num_blocks', type=int, default=8, help='Number of transformer blocks.')
169    parser.add_argument('--physicsformer_heads', type=int, default=16)
170    parser.add_argument('--physicsformer_mlp_dim', type=int, default=2048, help='Dimension of MLP.')
171    parser.add_argument('--physicsformer_dropout', type=float, default=0.1, help='Dropout rate.')
172
173    # ---> Residual Design Model
174    parser.add_argument('--base_model', type=str, default="pix2pix", choices=['resfcn', 'pix2pix'],
175                        help='Model to predict base propagation.')
176    parser.add_argument('--complex_model', type=str, default="pix2pix", choices=['resfcn', 'pix2pix'],
177                        help='Model to predict complex part of propagation (e.g. only reflection).')
178    parser.add_argument('--combine_mode', type=str, default="nn", choices=['math', 'nn'],
179                        help='Using math calculation or CNN for combining sub predictions.')
180    
181    parser.add_argument('--loss_2', type=str, default='l1', choices=['l1', 'crossentropy', 'weighted_combined'],
182                        help='Loss Function.')
183    # ---> WeightedCombinedLoss parameters
184    parser.add_argument('--wc_loss_silog_lambda_2', type=float, default=0.5, help='Lambda parameter for SILog loss.')
185    parser.add_argument('--wc_loss_weight_silog_2', type=float, default=0.5, help='Weight for SILog loss.')
186    parser.add_argument('--wc_loss_weight_grad_2', type=float, default=10.0, help='Weight for gradient loss.')
187    parser.add_argument('--wc_loss_weight_ssim_2', type=float, default=5.0, help='Weight for SSIM loss.')
188    parser.add_argument('--wc_loss_weight_edge_aware_2', type=float, default=10.0, help='Weight for edge-aware loss.')
189    parser.add_argument('--wc_loss_weight_l1_2', type=float, default=1.0, help='Weight for L1 loss.')
190    parser.add_argument('--wc_loss_weight_var_2', type=float, default=1.0, help='Weight for variance loss.')
191    parser.add_argument('--wc_loss_weight_range_2', type=float, default=1.0, help='Weight for range loss.')
192    parser.add_argument('--wc_loss_weight_blur_2', type=float, default=1.0, help='Weight for blur loss.')
193    
194    # ---> ResFCN Model 2
195    parser.add_argument('--resfcn_2_in_channels', type=int, default=1, help='How much channels as input?')
196    parser.add_argument('--resfcn_2_hidden_channels', type=int, default=64, help='How much channels in the hidden layers?')
197    parser.add_argument('--resfcn_2_out_channels', type=int, default=1, help='How much channels as output?')
198    parser.add_argument('--resfcn_2_num_blocks', type=int, default=16, help='How many Residual Blocks should be stacked.')
199
200    # ---> Pix2Pix Model 2
201    parser.add_argument('--pix2pix_2_in_channels', type=int, default=1, help='How much channels as input?')
202    parser.add_argument('--pix2pix_2_hidden_channels', type=int, default=64, help='How much channels in the hidden layers?')
203    parser.add_argument('--pix2pix_2_out_channels', type=int, default=1, help='How much channels as output?')
204    parser.add_argument('--pix2pix_2_second_loss_lambda', type=float, default=100.0, help='Weighting of second loss.')
205
206    # ---> PhysicsFormer Model 2
207    parser.add_argument('--physicsformer_in_channels_2', type=int, default=1, help='How much channels as input?')
208    parser.add_argument('--physicsformer_out_channels_2', type=int, default=1, help='How much channels as output?')
209    parser.add_argument('--physicsformer_img_size_2', type=int, default=256, help='Size of the image (width or height).')
210    parser.add_argument('--physicsformer_patch_size_2', type=int, default=4, help='Size of patches.')
211    parser.add_argument('--physicsformer_embedded_dim_2', type=int, default=1026, help='Dimension size of embedding.')
212    parser.add_argument('--physicsformer_num_blocks_2', type=int, default=8, help='Number of transformer blocks.')
213    parser.add_argument('--physicsformer_heads_2', type=int, default=16)
214    parser.add_argument('--physicsformer_mlp_dim_2', type=int, default=2048, help='Dimension of MLP.')
215    parser.add_argument('--physicsformer_dropout_2', type=float, default=0.1, help='Dropout rate.')
216
217
218    # Data
219    parser.add_argument('--data_variation', type=str, default='sound_baseline', choices=['sound_baseline', 'sound_reflection', 'sound_diffraction', 'sound_combined'],
220                        help='Name of the dataset variation.')
221    parser.add_argument('--input_type', type=str, default='osm', choices=['osm', 'base_simulation'],
222                        help='Input type (can be used to get the base simulation/propagation as input).')
223    parser.add_argument('--output_type', type=str, default='standard', choices=['standard', 'complex_only'],
224                        help='Output Type (can be used to get only reflexion as target).')
225    parser.add_argument('--fake_rgb_output', action='store_true',
226                        help='If setted: Input image is putted with 3 channels.')
227    parser.add_argument('--make_14_dividable_size', action='store_true',
228                        help='Adjusts imagesizes to a multiple of 14 if setted (needed for some networks).')
229    parser.add_argument('--reflexion_channels', action='store_true',
230                        help='If ray-traces should add to the input.')
231    parser.add_argument('--reflexion_steps', type=int, default=36,
232                        help='Defines how many traces should get created.')
233    parser.add_argument('--reflexions_as_channels', action='store_true',
234                        help='If True, every trace gets its own channel, else every trace in one channel.')
235
236    # Hardware
237    parser.add_argument('--device', type=str, default='cuda', choices=['cpu', 'cuda'],
238                        help='Rechen-Device')
239    
240    # Experiment Tracking
241    parser.add_argument('--experiment_name', type=str, default="image-to-image", help='Name of the overall experiment (will stay the same over most runs).')
242    parser.add_argument('--run_name', type=str, default="image-to-image", help='Name of the specific run. Will be used for naming but will add "YEAR-MONTH-DAY_HOUR_MINUTE" in front of your choosen name.')
243    parser.add_argument('--tensorboard_path', type=str, default="../tensorboard", help='Where should the results from tensorboard be saved to?')
244    parser.add_argument('--save_path', type=str, default="../train_inference", help='Where should the results from your model be saved to?')
245    parser.add_argument('--cmap', type=str, default="gray", help='Color Map for saving images.')
246
247
248
249    return parser
250
251
252
253# ---------------------------
254#     > Get Arguments <
255# ---------------------------
256def parse_args():
257    """
258    Parses command-line arguments for the Image-to-Image framework.
259
260    This function calls `get_arg_parser()` to construct the full parser, 
261    then reads and returns all command-line arguments provided by the user.
262    The returned namespace contains all configuration parameters for 
263    training, validation, and inference.
264
265    Returns:
266    - argparse.Namespace: 
267        Parsed arguments containing configuration for model setup, 
268        data handling, training hyperparameters, and logging options.
269    """
270    parser = get_arg_parser()
271    args = parser.parse_args()
272    return args