OSATTA: One-Shot Automatic Test Time Augmentation for Domain Adaptation

Fundamental models (FM) are reshaping the research paradigm by providing ready-to-use solutions to many challenging tasks, such as image classification, registration, or segmentation. Yet, their performance on new dataset cohorts significantly drops, particularly due to domain gaps between the training (source) and testing (target) data. Recently, test time augmentation strategies aim at finding target-to-source- mappings (t2sm), which improve the performance of the FM on the target dataset, by inspecting the the FM weights, thus assuming access to them. While this assumption holds for open research models, it does not for commercial ones (e.g., Chat-GPT). These are provided as black boxes, and the training data and the model weights are not available. In our work, we propose a new generic few-shot method that enables the computation of a target-to-source mapping by only using the black-box model's outputs. We start by defining a parametric family of functions for the t2sm, and with a simple loss function, we optimize the t2sm parameters based on a single labeled image volume. This effectively provides a mapping between the source domain and the target domain. In our experiments, we investigate how to improve the segmentation performance of a given FM (a UNet), and we outperform state-of-the-art accuracy in the 1-shot setting, with further improvement in a few-shot setting. Our approach is invariant to the model architecture as the FM is treated as a black box, which significantly increases our method's practical utility in real-world scenarios. We make the code available for reproducibility purposes.