v1.5.0: Custom augmentations

🚀 Added

• Custom augmentations via Albumentations. You can now control training augmentations through the aug_config parameter in train(). Pass a dictionary of Albumentations transforms, choose a built-in named preset, or disable augmentations entirely. Bounding boxes and segmentation masks are automatically transformed alongside images. (#263, #702)

  from rfdetr import RFDETRSmall
  from rfdetr.datasets.aug_config import AUG_CONSERVATIVE, AUG_AGGRESSIVE, AUG_AERIAL, AUG_INDUSTRIAL
  
  model = RFDETRSmall()
  
  # Use a built-in preset
  model.train(dataset_dir="...", aug_config=AUG_AGGRESSIVE, progress_bar=True)
  
  # Or define transforms explicitly
  model.train(
      dataset_dir="...",
      aug_config={
          "HorizontalFlip": {"p": 0.5},
          "RandomBrightnessContrast": {"brightness_limit": 0.2, "p": 0.4},
          "GaussianBlur": {"blur_limit": 3, "p": 0.2},
      },
      progress_bar=True,
  )
  
  # Disable all augmentations
  model.train(dataset_dir="...", aug_config={})

  Preset	Best for
  AUG_CONSERVATIVE	Small datasets (under 500 images)
  AUG_AGGRESSIVE	Large datasets (2000+ images)
  AUG_AERIAL	Satellite / overhead imagery
  AUG_INDUSTRIAL	Manufacturing / inspection data

• Save augmented training image samples. Enable save_dataset_grids=True in TrainConfig to write 3×3 JPEG grids of augmented training and validation images to your output directory before training begins, making it easy to verify your augmentation pipeline without running a full epoch. (#153)

  from rfdetr import RFDETRSmall
  
  model = RFDETRSmall()
  model.train(dataset_dir="...", save_dataset_grids=True, output_dir="output/")
  # Grids are saved to output/:
  #   train_batch0_grid.jpg, train_batch1_grid.jpg, train_batch2_grid.jpg
  #   val_batch0_grid.jpg,   val_batch1_grid.jpg,   val_batch2_grid.jpg

• ClearML training logger. Set clearml=True in TrainConfig to stream per-epoch metrics directly to your ClearML project. (#520)

  from rfdetr import RFDETRSmall
  
  model = RFDETRSmall()
  model.train(dataset_dir="...", clearml=True)

• MLflow training logger. Set mlflow=True in TrainConfig to log runs and metrics to MLflow, with support for custom tracking URIs and system metrics. (#109)

  from rfdetr import RFDETRSmall
  
  model = RFDETRSmall()
  model.train(dataset_dir="...", mlflow=True)

• Progress bar for training and validation. A live progress bar now shows batch-level progress during training and validation, and on-screen logs are structured for easier reading. (#204)

• device field added to TrainConfig, allowing explicit device selection when configuring training programmatically. (#687)

• ModelConfig now raises an error on unknown parameters, preventing silent misconfiguration from typos or stale config keys. (#196)

• TensorRT export guide. New documentation section covering how to convert an exported ONNX model to a TensorRT engine for maximum inference throughput. (#175)

🌱 Changed

• OPEN_SOURCE_MODELS constant deprecated in favour of the ModelWeights enum for cleaner model weight references. (#696)
• Added MD5 checksum validation for pretrained weight downloads, preventing silent use of corrupted files. (#679)

🔧 Fixed

• Fixed Albumentations bool-mask crash that occurred during segmentation training. (#706)
• Fixed UnboundLocalError when resuming training from a completed checkpoint. (#707)
• Prevented corruption of checkpoint_best_total.pth via atomic checkpoint stripping. (#708)
• Fixed PyTorch 2.9+ compatibility issue with CUDA capability detection. (#686)
• Fixed dtype mismatch error when use_position_supervised_loss=True. (#447)
• Fixed inconsistent return values from build_model. (#519)
• Fixed positional_encoding_size type annotation from bool to int. (#524)
• Fixed ONNX export output_names to include masks when exporting segmentation models. (#402)
• Fixed num_select not being correctly updated during segmentation model fine-tuning. (#399)
• Fixed np.argwhere → np.argmax misuse. (#536)
• Fixed COCO sparse category ID remapping logic for non-contiguous or offset category IDs are correctly handled. (#712)
• Fixed segmentation mask filtering when using aggressive augmentations. (#717)

🏆 Contributors

A special welcome to our new contributors and a big thank you to everyone who helped with this release:

• Panagiotis Moraitis (@panagiotamoraiti) (LinkedIn) – Custom Albumentations augmentation wrapper
• Shubham Rajvanshi (@shubsraj) (LinkedIn) – Progress bar and structured training logs
• Clement (@CorporalCleg) – ClearML logger integration
• Lakshman (@lab176344) – MLflow logger integration
• Mattia Di Giusto (@picjul) (LinkedIn) – Save augmented training image samples
• Juan Cobos (@juan-cobos) – device field in TrainConfig
• Ahmed Samir (@Ahmed-Samir11) – Error on unknown ModelConfig parameters
• Dominik Baran (@Yozer) (LinkedIn) – Fix segmentation mask filtering with aggressive augmentations
• Sungchul Kim (@sungchul2) (LinkedIn) – Fix num_select during segmentation fine-tuning
• Abdul Mukit (@Abdul-Mukit) (LinkedIn) – Fix ONNX export output names for segmentation
• Alarmod (@Alarmod) – PyTorch 2.9+ compatibility fix
• lixiaolei1982 (@lixiaolei1982) – Fix build_model return values & positional_encoding_size type
• kawabe-jiw (@kawabe-jiw) – Fix dtype mismatch with use_position_supervised_loss=True
• Andrei Moraru (@AndreiMoraru123) (LinkedIn) – np.argwhere → np.argmax fix
• Niels Teunissen (@DatSplit) – TensorRT export documentation
• stop1one (@stop1one) (LinkedIn) – Stabilize distributed training & test reliability
• Jirka Borovec (@Borda) (LinkedIn) – Augmentation presets, MD5 weight validation, ModelWeights enum, CI/testing infrastructure, docs