Published April 28, 2026 | Version v1.1.0
Software Open

dsgoficial/pytorch_segmentation_models_trainer: Release v1.1.0

  • 1. UFMG
  • 2. @anthropics
  • 3. @ultralytics @viddexa

Description

Bug fixes

  • Fixed tests/test_build_mask.py::Test_BuildMask::test_build_output_dirs_raises_exception: changed the output path to be a sub-directory of the input path in the test, ensuring the path validation logic in build_destination_dirs is correctly triggered and the expected exception is raised.
  • Fixed tests/test_configs/predict.yaml Hydra composition: added the @_global_ package override to the train_config_used_in_predict_test default inclusion. This ensures the configuration fields are merged into the root scope, allowing train_dataset.input_csv_path to be successfully overridden during prediction tests.
  • Fixed NameError: name 'DictConfig' is not defined in dataset_loader/dataset.py (load_augmentation_object): DictConfig was used but not imported from omegaconf, causing all augmentation loading to silently fall back to raw OmegaConf objects and crash with albumentations >= 2.x when A.Compose tried to access .available_keys on them.
  • Fixed Model._unpack_batch (model_loader/model.py) to respect the image_key and mask_key config fields (was hardcoded to "image" / "mask"). Also propagated the fix to _shared_step so training/validation steps honour custom keys end-to-end.
  • Fixed ValueError: prefetch_factor option could only be specified in multiprocessing in Model.train_dataloader, val_dataloader, and test_dataloader: when num_workers=0, prefetch_factor is now set to None as required by PyTorch's DataLoader API.
  • Fixed MCDropoutInferenceProcessor (tools/inference/mc_dropout_inference_processor.py) to save uncertainty rasters with suffix _mc_uncertainty instead of the generic _uncertainty from the parent class, matching the test expectation and making it distinguishable from TTA uncertainty maps.
  • Fixed merge_lora_weights (fine_tuning/lora_utils.py) to use is_peft_model() for the PEFT check rather than a local isinstance guard; this makes the function testable without a real PEFT installation and consistent with is_peft_model.
  • Fixed TimmEncoderWithSMPDecoder (custom_models/timm_models.py) for SMP 0.5.0 compatibility: replaced removed use_batchnorm kwarg with use_norm in UnetDecoder.__init__, and changed the decoder forward call to pass features as a single list instead of splatted positional arguments (both UnetDecoder and FPNDecoder now use forward(features: list) in SMP 0.5.x).
  • Updated test_caches.py::TestClassPresenceCacheAutoSave::test_auto_save_json_structure to account for the _config metadata key now present in the class-presence cache JSON, filtering it out before counting data entries.
  • Fixed Model._shared_step (model_loader/model.py) to apply _prepare_preds_for_metrics before computing train/val metrics: binary models output [B, 1, H, W] but torchmetrics expects [B, H, W], causing a shape mismatch RuntimeError during training. The same fix was applied to the per-class IoU path.
  • Fixed test_frame_field_model.py::_make_ff_batch: class_freq was built as torch.ones(B) (1-D) but compute_seg_loss_weights in base_loss.py sums over dim=1, requiring [B, C]. Also added the missing gt_crossfield_angle: torch.zeros(B, 1, H, W) key, required by compute_gt_field when compute_crossfield: true.
  • Fixed mod_polymapper.py validation metric logging: added numel() > 0 guard to skip empty tensors (produced when no polygons are detected in fast_dev_run), and added .float() before .mean() to handle Long-dtype metric values, preventing ValueError: tensor must have a single element and RuntimeError: mean() dtype.
  • Fixed test_polygonizer.py::test_polygonizer_acm_processor: replaced geopandas.testing.geom_almost_equals (fixed 5×10⁻⁷ m tolerance, too strict for GeoJSON coordinate rounding) with a 1 cm tolerance wrapper; regenerated the acm_polygonizer.geojson baseline because the ACM algorithm produces different vertex coordinates with newer shapely/geopandas.
  • Fixed custom_callbacks/image_callbacks.py: CombinedLoader.loaders was renamed to CombinedLoader.iterables in PyTorch Lightning ≥ 2.x; added a getattr fallback so the callback works on both old and new PL versions. Also fixed the follow-up DataLoader.loader.dataset chain: PL ≥ 2.x exposes plain DataLoader objects in iterables, so .dataset is now accessed directly with a fallback to the old .loader.dataset path.

Installation & Environment

  • Migrated to uv as the primary project and dependency manager.
  • Added uv sync as the recommended installation method in README.md and documentation.
  • Updated project requirements to Python 3.12+ and PyTorch 2.0+ (Lightning 2.4+).
  • Updated website documentation (website/docs/getting-started/installation.md) with comprehensive uv installation guide and updated troubleshooting for CUDA 11.8.

Reproducibility (Training Seed)

  • Added seed: Optional[int] and deterministic_cudnn: bool fields to the TrainConfig dataclass (config_definitions/train_config.py). Both default to None / False so all existing configs are fully backward-compatible.
  • Added set_training_seed(seed, deterministic_cudnn=False) utility function (utils/seed_utils.py). A single call seeds all randomness sources before model or dataset creation: Python random, NumPy np.random, torch.manual_seed, torch.cuda.manual_seed_all, and PYTHONHASHSEED. Optionally sets torch.backends.cudnn.deterministic = True and torch.backends.cudnn.benchmark = False. Returns a torch.Generator seeded with the same value for use in DataLoaders.
  • Modified train() (train.py) to call set_training_seed as the very first operation when cfg.seed is present, ensuring model weight initialisation (SMP, timm, HuggingFace, custom architectures) is also reproducible.
  • Modified _worker_init_fn (dataset_loader/dataset.py) to additionally seed Python's random module alongside NumPy. Since PyTorch sets torch.initial_seed() = global_seed + worker_id per worker automatically, both seeds are deterministic and unique per worker once a global seed is set.
  • Modified Model.train_dataloader, Model.val_dataloader, and Model.test_dataloader (model_loader/model.py) to pass a torch.Generator().manual_seed(cfg.seed) to each DataLoader when cfg.seed is set, making the shuffle sampler sequence reproducible. Added Model._make_dataloader_generator() helper method.
  • Added reference YAML config conf/examples/reproducible_training.yaml with inline comments explaining every field and listing all controlled randomness sources.
  • Added reproducibility block (commented) to conf/examples/smp_mit_b2.yaml so users can enable it in one line.
  • Added user documentation website/docs/user-guide/reproducibility.md covering the seed field, deterministic_cudnn trade-offs, Python API usage, and known limitations.

MC Dropout (test-time uncertainty)

  • Added mc_dropout_utils.py (utils/mc_dropout_utils.py): three pure utility functions with no inference-framework dependency. enable_mc_dropout(model) sets all Dropout / Dropout2d / Dropout3d layers to train mode while leaving the rest of the model in eval mode (BatchNorm keeps running statistics, only dropout randomness is re-enabled). warn_if_no_dropout(model) emits a UserWarning if the model has no dropout layers — in that case all T samples are identical and uncertainty is zero. compute_uncertainty(samples, mode) accepts a [T, B, C, H, W] tensor of softmax probabilities and returns [B, 1, H, W] uncertainty: "entropy" computes predictive entropy of the mean distribution (total uncertainty); "mutual_information" computes BALD — the difference between the entropy of the mean and the mean of the individual entropies (epistemic uncertainty only).
  • Added MCDropoutInferenceProcessor (tools/inference/mc_dropout_inference_processor.py): extends MultiClassInferenceProcessor with MC Dropout inference. Overrides predict_and_merge to run n_samples stochastic forward passes per tile (after calling enable_mc_dropout), average the softmax probabilities for the class prediction, and — only when export_uncertainty_map=True — compute and merge the per-pixel uncertainty map. Uncertainty tensors are only allocated when requested, keeping the no-uncertainty path free of extra memory cost. Overrides process to skip the striped inference path (which uses a single TileMerger and cannot merge uncertainty); a warning is logged for large images. Exports {stem}_mc_uncertainty.tif (float32 single-band GeoTIFF, range [0, log C]) alongside the segmentation output when export_uncertainty_map=True.
  • Modified MultiClassInferenceProcessor (tools/inference/inference_processors.py): added three new parameters export_uncertainty_map (default False), uncertainty_mode (default "entropy"), and output_uncertainty_dir (default None). When tta_mode is set and export_uncertainty_map=True, per-sample softmax probabilities are kept during the TTA loop and used to compute uncertainty via compute_uncertainty(); the result is merged via a separate TileMerger and exported as {stem}_uncertainty.tif. The no-uncertainty path (default) is fully unchanged. Added _save_uncertainty_raster helper for writing the float32 GeoTIFF with source CRS and transform preserved.
  • Added MCDropoutInferenceProcessorConfig (config_definitions/mc_dropout_config.py): Hydra dataclass registered in the ConfigStore under group="inference_processor", name="mc_dropout". Fields: n_samples, uncertainty_mode, export_uncertainty_map, num_classes, model_input_shape, step_shape, output_uncertainty_dir.
  • Added reference YAML config mc_dropout_inference.yaml under conf/examples/ with inline comments explaining the decoder_dropout requirement and the export_uncertainty_map flag.

Evidential Deep Learning (EDL)

  • Added EvidentialWrapper (custom_models/edl_wrapper.py): architecture-agnostic wrapper that replaces Softmax with a Dirichlet parameterisation (evidence = Softplus(logits), alpha = evidence + 1). Works with any model that returns [B, K, H, W] tensors including SMP, HuggingFace, timm, and custom models. Forward pass returns a dict with logits, evidence, alpha, probs, and uncertainty keys. Handles tuple, dict, and plain-tensor model outputs automatically.
  • Added EvidentialMSELoss and EvidentialKLLoss (custom_losses/edl_loss.py): two-component EDL loss designed for use with the existing CompoundLoss / MultiLoss epoch-weight scheduling mechanism. EvidentialMSELoss computes the MSE integrated analytically over the Dirichlet distribution (bias² + variance terms). EvidentialKLLoss computes KL[Dir(α̃) || Dir(1,...,1)] after removing evidence for the correct class, penalising residual wrong-class evidence. KL annealing is expressed as a dynamic weight list in the YAML — no custom scheduler needed.
  • Added edl_utils.py (custom_losses/edl_utils.py): analytic helpers for Dirichlet statistics: one_hot_encode (hard + soft label support, ignore_index handling), dirichlet_strength, epistemic_uncertainty (u = K/S), dirichlet_kl_divergence, kl_divergence_to_uniform, edl_kl_regulariser.
  • Modified Model._shared_step to detect EvidentialWrapper output (dict with "alpha" key): probs is extracted for metrics and logging; mean uncertainty edl/{prefix}_uncertainty is logged automatically at each step. The change is backward-compatible — non-EDL models are unaffected.
  • Added EvidentialWarmupCallback (custom_callbacks/edl_callbacks.py): three-phase encoder freeze schedule for fine-tuning. Phase 1 (epoch < warmup_epochs): encoder frozen. Phase 2 (warmup_epochs ≤ epoch < partial_unfreeze_epoch): last two encoder stages unfrozen. Phase 3 (epoch ≥ partial_unfreeze_epoch): full encoder unfrozen. Setting freeze_encoder=False (training from scratch) disables all freezing while preserving the logging hooks.
  • Added EvidentialUncertaintyVisualizationCallback (custom_callbacks/edl_callbacks.py): logs a 4-column diagnostic grid [input | predicted class | uncertainty map | ground truth] every N validation epochs. Compatible with TensorBoard, WandB, and plain file-system fallback.
  • Added EvidentialInferenceProcessor (tools/inference/edl_inference_processor.py): extends SingleImageInfereceProcessor with separate TileMerger instances for probabilities and uncertainty. Exports a single-band float32 GeoTIFF of the uncertainty map (u = K/S) via save_uncertainty_raster, preserving source CRS and transform exactly. Optional alpha band export via export_alpha=True.
  • Added edl_config.py (config_definitions/edl_config.py): Hydra dataclass configs for EvidentialWrapper, EvidentialMSELoss, EvidentialKLLoss, EvidentialWarmupCallback, EvidentialUncertaintyVisualizationCallback, and EvidentialInferenceProcessor, all registered in the ConfigStore.
  • Added reference YAML configs edl_from_scratch.yaml and edl_finetune.yaml under conf/examples/ with inline comments explaining every field.

RandomCropSegmentationDataset

  • Added RandomCropSegmentationDataset: reads large GeoTIFF images on-the-fly using rasterio windowed reads instead of pre-generating tiles on disk. Eliminates the disk-space overhead of a tile library and allows crop size, augmentation, and sampling strategy to be changed without reprocessing data.
  • Per-worker LRU cache (_RasterioLRUCache) keeps a configurable number of open DatasetReader handles (lru_cache_size, default 64) to avoid repeated rasterio open/close overhead.
  • class_balanced_sampling: weights image selection by inverse class frequency so images containing rare classes are sampled more often. Computed once at dataset initialisation from mask histograms in the CSV.
  • Class-aware CutMix (cutmix_prob, cutmix_alpha): pastes a rectangular region from a second crop chosen to maximise class diversity.
  • ClassMix (classmix_prob): copies a randomly selected class region from a second image and pastes it onto the primary crop; particularly effective for rare classes.
  • soft_labels mode: returns float masks in [0, 1] for label-noise and probabilistic annotation workflows. The _shared_step training loop detects soft labels automatically and uses the appropriate loss path.
  • grid_mode / grid_step: switches from random crop positions to a deterministic sliding-window grid for reproducible validation coverage and pseudo-labelling. configure_optimizers accounts for grid mode when computing steps_per_epoch for OneCycleLR.
  • Added RandomCropSegmentationDatasetConfig dataclass.

Mixture of Experts Models

  • Added UPerNetMoE (custom_models/upernet_moe.py): UPerNet variant that replaces fusion and/or FPN convolutions with MoEConv2dReLU blocks. Supports token_choice (each token picks top-k experts) and expert_choice (each expert picks top-k tokens) routing, configurable noise injection, capacity factor, and an optional shared dense expert. Load-balancing auxiliary loss is automatically detected and added to the training loss in _shared_step.
  • Added UPerNetMEDoE (custom_models/upernet_medoe.py): extends UPerNetMoE with structured expert dropout during training (randomly drops a fraction of experts per forward pass) to improve regularisation and reduce reliance on any single expert. _shared_step automatically logs extra/train_medoe_expert_utilization and extra/train_medoe_expert_entropy when a MEDoE model is detected.

Dual-Head Training

  • Added UPerNetDualHead (custom_models/upernet_dual_head.py): UPerNet with two independent decoders sharing a single encoder. Head A is supervised with hard labels (integer class indices); Head B is supervised with soft labels (float probabilities). A consistency loss couples the two heads during training. inference_head controls which head is active at inference time: "A", "B", or "average" (default).

TTA improvements

  • Added tta_mode compact interface: passing tta_mode: "d4" or tta_mode: "flip" to any inference processor or to test_step automatically selects the corresponding augmentation preset (d4 = all 8 D4 symmetries; flip = 4 flip/rotation combinations), without listing augmentations explicitly.
  • SingleImageInfereceProcessor now accepts tta_mode as an alias for use_tta=True + the corresponding augmentations list. Backward-compatible with the existing use_tta + tta_augmentations interface.
  • MultiClassInferenceProcessor now accepts tta_mode (replaces the previous use_tta parameter on that class).
  • _get_tta_augmentations() in Model checks cfg.tta_mode first; falls back to cfg.use_tta + cfg.tta_augmentations for backward compatibility.

MultiClassInferenceProcessor improvements

  • Added striped inference (make_inference_striped): very large images are automatically split into horizontal stripes when pixel count exceeds striped_threshold_pixels (default 50 MP), processed in parallel via ThreadPoolExecutor, and reassembled in-memory. Stripe height is configurable via stripe_height (default 4096 px).
  • Added confidence map output: confidence_mode ("max_prob" or "entropy") computes per-pixel confidence scores alongside the class prediction. Saved to output_probs_dir when provided.
  • Added tile_weight parameter (passed through to AbstractInferenceProcessor) to control how overlapping tile predictions are merged.
  • process() override: automatically routes each image to the striped or standard inference path based on image size.

Callbacks

  • Added EMACallback: maintains an exponential moving average of model weights (configurable decay). During validation the shadow EMA weights are swapped in so validation metrics reflect the averaged model; the original weights are restored immediately after. Checkpoints saved during a validation epoch contain the EMA weights.
  • Added MixStyleCallback: applies MixStyle feature-level domain augmentation via forward hooks on selected encoder stages. stages controls which encoder stage indices receive the hook; p and alpha control application probability and Beta distribution parameter respectively.

LR Warmup

  • Added warmup_epochs support in hyperparameters: a linear LR warmup is prepended to any scheduler (except OneCycleLR, which has its own warmup via pct_start). The framework automatically subtracts warmup_epochs from T_max (or equivalent period parameters) so the total scheduled duration remains correct.

Inference pipeline

  • predict_from_batch.py updated: when inference_processor is present in the config, batch prediction is routed through instantiate_inference_processor, enabling sliding-window, striped, and TTA inference modes. The legacy trainer.predict() path is preserved for backward compatibility.
  • InferenceProcessorConfig and PredictSingleImageConfig updated with new fields: tta_mode, tile_weight, confidence_mode, striped_threshold_pixels, stripe_height, output_probs_dir.

Domain Adaptation — DANN with Gradient Reversal

  • Added GradientReversalFunction and GradientReversalLayer (domain_adaptation/methods/gradient_reversal.py): a torch.autograd.Function implementing the identity forward pass with gradient negation in backward, wrapped in a parameter-free nn.Module with a set_lambda() method.
  • Added DomainClassifier (domain_adaptation/methods/dann.py): resolution-agnostic MLP domain classifier using AdaptiveAvgPool2d(1) so it works with any encoder spatial size.
  • Added DANNMethod (domain_adaptation/methods/dann.py): full DANN implementation via BaseDomainAdaptationMethod. Features: requires_features=True, configurable feature_layer, dedicated discriminator_lr parameter group, and GRL lambda initialized to 0 to avoid adversarial pressure at epoch 0.
  • Added step_mode parameter to DANNMethod ("epoch" or "batch"): controls the granularity at which the lambda schedule is applied. "batch" mode updates λ every training step via the new on_train_batch_start hook, producing smoother growth closer to the original Ganin et al. implementation; "epoch" (default) updates once per epoch.
  • Added _current_lambda attribute to DANNMethod: caches the most recently computed lambda so DomainAdaptationModel._get_lambda_da() always returns the correct value regardless of update granularity.
  • Added on_train_batch_start lifecycle hook to BaseDomainAdaptationMethod (no-op default) and forwarded it in DomainAdaptationModel.on_train_batch_start.
  • Updated DomainAdaptationModel._get_lambda_da() to prefer method._current_lambda when present, falling back to the epoch-level schedule query.
  • Refactored BaseLambdaScheduler.get_lambda signature from (epoch, total_epochs) to (step, total_steps) — the same formula, now granularity-agnostic. All scheduler subclasses (ConstantScheduler, LinearScheduler, DANNScheduler) updated accordingly.
  • Added full working example config (conf/examples/dann_domain_adaptation.yaml) for U-Net ResNet-34 with DANN.
  • Added test suite for GRL (tests/test_gradient_reversal.py, 17 tests) and DANN (tests/test_dann_method.py, 49 tests including 12 new TestDANNMethodStepMode tests).
  • Added dedicated DANN user guide (website/docs/advanced/dann-method.md) covering mechanism, in_channels lookup table, full config reference, step_mode guidance, monitoring, and limitations.
  • Updated website/docs/advanced/domain-adaptation.md with a "Built-in Methods" section linking to the DANN guide.
  • Updated website/docs/advanced/domain-adaptation-implementing-methods.md Example 2 to use the built-in GradientReversalLayer and add a callout pointing to the dedicated DANN guide.

Domain Adaptation — Initial structure

  • Added initial domain adaptation module (domain_adaptation/) with an extensible base class (BaseDomainAdaptationMethod), feature hook extraction (feature_hooks.py), adaptation schedulers (schedulers.py), and a monitoring callback (callbacks/monitor_callback.py).
  • Added DomainAdaptationModel (inherits Model) that orchestrates source/target dataloaders, adaptation loss weighting, and per-epoch scheduler stepping.
  • Added DomainAdaptationConfig dataclass in config_definitions/domain_adaptation_config.py.
  • Added comprehensive test suite for domain adaptation (test_base_method.py, test_domain_adaptation_config.py, test_domain_adaptation_model.py, test_feature_hooks.py, test_schedulers.py).
  • Added documentation: user guide (advanced/domain-adaptation.md), config reference (advanced/domain-adaptation-config-reference.md), and implementing custom methods guide (advanced/domain-adaptation-implementing-methods.md).

Test Time Augmentation (TTA)

  • Added tools/tta/tta.py implementing TTA with all 8 symmetries of the D4 dihedral group (4 rotations × 2 flips). Each augmentation has an exact inverse so predictions are de-augmented and averaged with no spatial artifacts.
  • Added apply_tta() helper for applying TTA to any segmentation model callable.
  • Exposed use_tta and tta_augmentations fields on all inference processor classes and InferenceProcessorConfig/PredictSingleImageConfig. SingleImageFromFrameFieldProcessor automatically skips the crossfield output during TTA de-augmentation.
  • test_step() in Model now applies TTA when cfg.use_tta=True.
  • Added TTA user guide (website/docs/advanced/tta.md) and inference documentation section.

Transformer & Foundation Model Support

  • Added HuggingFaceSegmentationWrapper (custom_models/huggingface_models.py): loads any AutoModelForSemanticSegmentation from the Hub or local path, bypasses the HF processor, and upsamples logits back to input resolution.
  • Added TimmEncoderWithSMPDecoder (custom_models/timm_models.py): combines a timm features_only backbone with SMP UNet/FPN/PAN decoders.
  • Added TerraTorchSegmentationWrapper (custom_models/terratorch_models.py): bridges TerraTorch foundation model encoders (Prithvi, Clay, SatMAE) with a linear or FPN segmentation head; supports single- and multi-temporal inputs.
  • Added ModelOutputAdapter (custom_models/transformer_adapters.py): normalises any model output (HF dataclass, dict, tuple, or plain tensor) to a (B, C, H, W) tensor with optional bilinear upsampling.
  • Added LoRA / PEFT fine-tuning support (fine_tuning/lora_utils.py): apply_fine_tuning_strategy() supports full, freeze_backbone, linear_probe, and lora strategies; LoraAdapterConfig and FineTuningConfig dataclasses added; merge_lora_weights() for deployment.
  • predict.py now auto-merges LoRA adapter weights before inference; a keep_lora_adapters flag skips the merge for fine-tuning resumption.
  • Hardened training loop in Model: _unpack_batch() replaces fragile batch.values() unpacking with configurable image_key/mask_key; set_encoder_trainable() no longer assumes a .encoder attribute; _prepare_preds_for_metrics() guards metric calls against malformed outputs.
  • Added 6 example YAML configs: smp_mit_b2.yaml, smp_tu_convnext.yaml, segformer_hf.yaml, segformer_lora.yaml, vit_linear_probe.yaml, prithvi_terratorch.yaml.
  • Added [transformers] pip extras group and a dedicated CI job for the transformer test suite.

RasterPatchDataset

  • Added RasterPatchDataset (dataset_loader/raster_patch_dataset.py): scans image/mask directory pairs recursively and exposes every patch_size × patch_size window (with configurable stride) as an independent dataset item. Global index to (image, row, col) mapping runs in O(log N) via bisect over cumulative patch counts; rasterio windowed reads ensure full images never enter RAM.
  • Supports augmentations, selected_bands, image_dtype, mask_extension, n_classes (binary binarisation when n_classes=2), and reset_augmentation_function.
  • Emits UserWarning for orphaned mask files (mask without a corresponding image) to surface dataset misconfiguration.
  • Added RasterPatchDatasetConfig dataclass and example YAML (conf/examples/raster_patch_segmentation.yaml).

Dataset improvements

  • Added test_dataset support with test_step(), test_dataloader(), and test_metrics (prefixed test/) in Model and FrameFieldSegmentationPLModel. trainer.test() is now called automatically after trainer.fit() when test_dataset is present in the config. This completes the three-way dataset split: train_dataset → training loop; val_dataset → per-epoch monitoring during fit; test_dataset → final held-out evaluation after fit.
  • Added test_dataset field to TrainConfig dataclass and test_dataset block to all example YAML configs.
  • Added SegmentationDatasetFromFolder: a new dataset class that discovers image/mask pairs recursively from two root folders, without requiring a CSV file. Matching is done by relative subfolder path and file stem. Supports all parameters of SegmentationDataset. Raises ValueError when no valid pairs are found.
  • SegmentationDataset.__init__ now accepts an optional df parameter (pre-built pd.DataFrame) in addition to input_csv_path, enabling programmatic dataset creation without a CSV file on disk. Fully backwards-compatible.
  • Added configurable image_dtype field to SegmentationDataset, RandomCropSegmentationDataset, and their configs, accepting uint8 (default), uint16, float32, or native. Auto-normalization scales correctly per dtype (/255, /65535, or no division). Fully backwards-compatible.

Inference improvements

  • Added normalize_max_value parameter to all inference processor classes (AbstractInferenceProcessor and all subclasses), exposing Albumentations' max_pixel_value for the normalization step. Default None preserves the previous behaviour (255.0). Use normalize_max_value: 65535.0 for uint16 imagery or 1.0 for pre-normalized float32.

Bug fixes

  • Fixed bug in Model.__init__: gpu_val_transform and gpu_train_transform were accessed via cfg.val_dataset/cfg.train_dataset without checking if those keys exist, causing AttributeError when the corresponding dataset config was omitted.
  • Fixed val_dataloader() to return None gracefully when val_ds is None (i.e. val_dataset absent from config), allowing training-only runs without a validation loop.
  • Removed the orphan set_test_dataset() method from FrameFieldSegmentationPLModel (superseded by the new test_ds attribute set in Model.__init__).

Files

dsgoficial/pytorch_segmentation_models_trainer-v1.1.0.zip

Files (77.7 MB)

Additional details