birder 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

birder/common/lib.py

@@ -1,11 +1,7 @@
 import os
-import random
 from typing import Any
 from typing import Optional
 
-import numpy as np
-import torch
-
 from birder.conf import settings
 from birder.data.transforms.classification import RGBType
 from birder.model_registry import registry
@@ -19,11 +15,8 @@ from birder.net.ssl.base import SSLBaseNet
 from birder.version import __version__
 
 
-def set_random_seeds(seed: int) -> None:
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-    random.seed(seed)
+def env_bool(name: str) -> bool:
+    return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"}
 
 
 def get_size_from_signature(signature: SignatureType | DetectionSignatureType) -> tuple[int, int]:

birder/common/training_cli.py

@@ -5,6 +5,7 @@ import typing
 from typing import Optional
 from typing import get_args
 
+from birder.common.cli import FlexibleDictAction
 from birder.common.cli import ValidationError
 from birder.common.training_utils import OptimizerType
 from birder.common.training_utils import SchedulerType
@@ -82,11 +83,23 @@ def add_lr_wd_args(parser: argparse.ArgumentParser, backbone_lr: bool = False, w
         metavar="WD",
         help="weight decay for embedding parameters for vision transformer models",
     )
+    group.add_argument(
+        "--custom-layer-wd",
+        action=FlexibleDictAction,
+        metavar="LAYER=WD",
+        help="custom weight decay for specific layers by name (e.g., offset_conv=0.0)",
+    )
     group.add_argument("--layer-decay", type=float, help="layer-wise learning rate decay (LLRD)")
     group.add_argument("--layer-decay-min-scale", type=float, help="minimum layer scale factor clamp value")
     group.add_argument(
         "--layer-decay-no-opt-scale", type=float, help="layer scale threshold below which parameters are frozen"
     )
+    group.add_argument(
+        "--custom-layer-lr-scale",
+        action=FlexibleDictAction,
+        metavar="LAYER=SCALE",
+        help="custom lr_scale for specific layers by name (e.g., offset_conv=0.01,attention=0.5)",
+    )
 
 
 def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
@@ -185,6 +198,11 @@ def add_detection_input_args(parser: argparse.ArgumentParser) -> None:
         action="store_true",
         help="enable random square resize once per batch (capped by max(--size))",
     )
+    group.add_argument(
+        "--multiscale-min-size",
+        type=int,
+        help="minimum short-edge size for multiscale lists (rounded up to nearest multiple of 32)",
+    )
 
 
 def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs: int = 100) -> None:

birder/common/training_utils.py

@@ -3,8 +3,10 @@ import contextlib
 import logging
 import math
 import os
+import random
 import re
 import subprocess
+import sys
 from collections import deque
 from collections.abc import Callable
 from collections.abc import Generator
@@ -29,12 +31,25 @@ from birder.data.transforms.classification import training_preset
 from birder.optim import Lamb
 from birder.optim import Lars
 from birder.scheduler import CooldownLR
+from birder.version import __version__ as birder_version
 
 logger = logging.getLogger(__name__)
 
 OptimizerType = Literal["sgd", "rmsprop", "adam", "adamw", "nadam", "nadamw", "lamb", "lambw", "lars"]
 SchedulerType = Literal["constant", "step", "multistep", "cosine", "polynomial"]
 
+###############################################################################
+# Core Utilities
+###############################################################################
+
+
+def set_random_seeds(seed: int) -> None:
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
 ###############################################################################
 # Data Sampling
 ###############################################################################
@@ -207,13 +222,16 @@ def count_layers(model: torch.nn.Module) -> int:
 def optimizer_parameter_groups(
     model: torch.nn.Module,
     weight_decay: float,
+    base_lr: float,
     norm_weight_decay: Optional[float] = None,
     custom_keys_weight_decay: Optional[list[tuple[str, float]]] = None,
+    custom_layer_weight_decay: Optional[dict[str, float]] = None,
     layer_decay: Optional[float] = None,
     layer_decay_min_scale: Optional[float] = None,
     layer_decay_no_opt_scale: Optional[float] = None,
     bias_lr: Optional[float] = None,
     backbone_lr: Optional[float] = None,
+    custom_layer_lr_scale: Optional[dict[str, float]] = None,
 ) -> list[dict[str, Any]]:
     """
     Return parameter groups for optimizers with per-parameter group weight decay.
@@ -233,11 +251,16 @@ def optimizer_parameter_groups(
         The PyTorch model whose parameters will be grouped for optimization.
     weight_decay
         Default weight decay (L2 regularization) value applied to parameters.
+    base_lr
+        Base learning rate that will be scaled by lr_scale factors for each parameter group.
     norm_weight_decay
         Weight decay value specifically for normalization layers. If None, uses weight_decay.
     custom_keys_weight_decay
         List of (parameter_name, weight_decay) tuples for applying custom weight decay
         values to specific parameters by name matching.
+    custom_layer_weight_decay
+        Dictionary mapping layer name substrings to custom weight decay values.
+        Applied to parameters whose names contain the specified keys.
     layer_decay
         Layer-wise learning rate decay factor.
     layer_decay_min_scale
@@ -248,6 +271,9 @@ def optimizer_parameter_groups(
         Custom learning rate for bias parameters (parameters ending with '.bias').
     backbone_lr
         Custom learning rate for backbone parameters (parameters starting with 'backbone.').
+    custom_layer_lr_scale
+        Dictionary mapping layer name substrings to custom lr_scale values.
+        Applied to parameters whose names contain the specified keys.
 
     Returns
     -------
@@ -291,14 +317,14 @@ def optimizer_parameter_groups(
     if layer_decay is not None:
         layer_max = num_layers - 1
         layer_scales = [max(layer_decay_min_scale, layer_decay ** (layer_max - i)) for i in range(num_layers)]
-        logger.info(f"Layer scaling in range of {min(layer_scales)} - {max(layer_scales)} on {num_layers} layers")
+        logger.info(f"Layer scaling ranges from {min(layer_scales)} to {max(layer_scales)} across {num_layers} layers")
 
     # Set weight decay and layer decay
     idx = 0
     params = []
     module_stack_with_prefix = [(model, "")]
     visited_modules = []
-    while len(module_stack_with_prefix) > 0:
+    while len(module_stack_with_prefix) > 0:  # pylint: disable=too-many-nested-blocks
         skip_module = False
         (module, prefix) = module_stack_with_prefix.pop()
         if id(module) in visited_modules:
@@ -324,13 +350,35 @@ def optimizer_parameter_groups(
                 for key, custom_wd in custom_keys_weight_decay:
                     target_name_for_custom_key = f"{prefix}.{name}" if prefix != "" and "." in key else name
                     if key == target_name_for_custom_key:
+                        # Calculate lr_scale (from layer_decay or custom_layer_lr_scale)
+                        lr_scale = 1.0 if layer_decay is None else layer_scales[idx]
+                        if custom_layer_lr_scale is not None:
+                            for layer_name_key, custom_scale in custom_layer_lr_scale.items():
+                                if layer_name_key in target_name:
+                                    lr_scale = custom_scale
+                                    break
+
+                        # Apply custom layer weight decay (substring matching)
+                        wd = custom_wd
+                        if custom_layer_weight_decay is not None:
+                            for layer_name_key, custom_wd_value in custom_layer_weight_decay.items():
+                                if layer_name_key in target_name:
+                                    wd = custom_wd_value
+                                    break
+
                         d = {
                             "params": p,
-                            "weight_decay": custom_wd,
-                            "lr_scale": 1.0 if layer_decay is None else layer_scales[idx],
+                            "weight_decay": wd,
+                            "lr_scale": lr_scale,  # Used only for reference/debugging
                         }
-                        if backbone_lr is not None and target_name.startswith("backbone.") is True:
+
+                        # Apply learning rate based on priority: bias_lr > backbone_lr > lr_scale
+                        if bias_lr is not None and target_name.endswith(".bias") is True:
+                            d["lr"] = bias_lr
+                        elif backbone_lr is not None and target_name.startswith("backbone.") is True:
                             d["lr"] = backbone_lr
+                        elif lr_scale != 1.0:
+                            d["lr"] = base_lr * lr_scale
 
                         params.append(d)
                         is_custom_key = True
@@ -342,16 +390,34 @@ def optimizer_parameter_groups(
                 else:
                     wd = weight_decay
 
+                # Apply custom layer weight decay (substring matching)
+                if custom_layer_weight_decay is not None:
+                    for layer_name_key, custom_wd_value in custom_layer_weight_decay.items():
+                        if layer_name_key in target_name:
+                            wd = custom_wd_value
+                            break
+
+                # Calculate lr_scale (from layer_decay or custom_layer_lr_scale)
+                lr_scale = 1.0 if layer_decay is None else layer_scales[idx]
+                if custom_layer_lr_scale is not None:
+                    for layer_name_key, custom_scale in custom_layer_lr_scale.items():
+                        if layer_name_key in target_name:
+                            lr_scale = custom_scale
+                            break
+
                 d = {
                     "params": p,
                     "weight_decay": wd,
-                    "lr_scale": 1.0 if layer_decay is None else layer_scales[idx],
+                    "lr_scale": lr_scale,  # Used only for reference/debugging
                 }
-                if backbone_lr is not None and target_name.startswith("backbone.") is True:
-                    d["lr"] = backbone_lr
 
+                # Apply learning rate based on priority: bias_lr > backbone_lr > lr_scale
                 if bias_lr is not None and target_name.endswith(".bias") is True:
                     d["lr"] = bias_lr
+                elif backbone_lr is not None and target_name.startswith("backbone.") is True:
+                    d["lr"] = backbone_lr
+                elif lr_scale != 1.0:
+                    d["lr"] = base_lr * lr_scale
 
                 params.append(d)
 
@@ -442,6 +508,8 @@ def get_optimizer(parameters: list[dict[str, Any]], l_rate: float, args: argpars
     else:
         raise ValueError("Unknown optimizer")
 
+    logger.debug(f"Created {opt} optimizer with lr={lr}, weight_decay={args.wd}")
+
     return optimizer
 
 
@@ -477,10 +545,10 @@ def get_scheduler(
 
     main_steps = steps - begin_step - remaining_warmup - remaining_cooldown - 1
 
-    logger.debug(f"Using {steps_per_epoch} steps per epoch")
+    logger.debug(f"Scheduler using {steps_per_epoch} steps per epoch")
     logger.debug(
         f"Scheduler {args.lr_scheduler} set for {steps} steps of which {warmup_steps} "
-        f"are warmup and {cooldown_steps} cooldown"
+        f"are warmup and {cooldown_steps} are cooldown"
     )
     logger.debug(
         f"Currently starting from step {begin_step} with {remaining_warmup} remaining warmup steps "
@@ -810,6 +878,51 @@ def is_local_primary(args: argparse.Namespace) -> bool:
     return args.local_rank == 0  # type: ignore[no-any-return]
 
 
+def init_training(
+    args: argparse.Namespace,
+    log: logging.Logger,
+    *,
+    cudnn_dynamic_size: bool = False,
+) -> tuple[torch.device, int, bool]:
+    init_distributed_mode(args)
+
+    log.info(f"Starting training, birder version: {birder_version}, pytorch version: {torch.__version__}")
+
+    log_git_info()
+
+    if args.cpu is True:
+        device = torch.device("cpu")
+        device_id = 0
+    else:
+        device = torch.device("cuda")
+        device_id = torch.cuda.current_device()
+
+    if args.use_deterministic_algorithms is True:
+        torch.backends.cudnn.benchmark = False
+        torch.use_deterministic_algorithms(True)
+    elif cudnn_dynamic_size is True:
+        # Dynamic sizes: avoid per-size algorithm selection overhead.
+        torch.backends.cudnn.enabled = False
+    else:
+        torch.backends.cudnn.enabled = True
+        torch.backends.cudnn.benchmark = True
+
+    if args.seed is not None:
+        set_random_seeds(args.seed)
+
+    if args.non_interactive is True or is_local_primary(args) is False:
+        disable_tqdm = True
+    elif sys.stderr.isatty() is False:
+        disable_tqdm = True
+    else:
+        disable_tqdm = False
+
+    # Enable or disable the autograd anomaly detection.
+    torch.autograd.set_detect_anomaly(args.grad_anomaly_detection)
+
+    return (device, device_id, disable_tqdm)
+
+
 ###############################################################################
 # Utility Functions
 ###############################################################################

birder/data/collators/detection.py

@@ -1,13 +1,14 @@
 import math
 import random
 from typing import Any
+from typing import Optional
 
 import torch
 from torchvision import tv_tensors
 from torchvision.transforms import v2
 from torchvision.transforms.v2 import functional as F
 
-BATCH_MULTISCALE_SIZES = (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
+from birder.data.transforms.detection import build_multiscale_sizes
 
 
 def collate_fn(batch: list[tuple[Any, ...]]) -> tuple[Any, ...]:
@@ -63,13 +64,19 @@ class DetectionCollator:
 
 
 class BatchRandomResizeCollator(DetectionCollator):
-    def __init__(self, input_offset: int, size: tuple[int, int], size_divisible: int = 32) -> None:
+    def __init__(
+        self,
+        input_offset: int,
+        size: tuple[int, int],
+        size_divisible: int = 32,
+        multiscale_min_size: Optional[int] = None,
+    ) -> None:
         super().__init__(input_offset, size_divisible=size_divisible)
         if size is None:
             raise ValueError("size must be provided for batch multiscale")
 
         max_side = max(size)
-        sizes = [side for side in BATCH_MULTISCALE_SIZES if side <= max_side]
+        sizes = [side for side in build_multiscale_sizes(multiscale_min_size) if side <= max_side]
         if len(sizes) == 0:
             sizes = [max_side]
 

birder/data/datasets/coco.py

@@ -98,10 +98,14 @@ class CocoTraining(CocoBase):
 class CocoInference(CocoBase):
     def __getitem__(self, index: int) -> tuple[str, torch.Tensor, Any, list[int]]:
         coco_id = self.dataset.ids[index]
-        path = self.dataset.coco.loadImgs(coco_id)[0]["file_name"]
+        img_info = self.dataset.coco.loadImgs(coco_id)[0]
+        path = img_info["file_name"]
         (sample, labels) = self.dataset[index]
 
-        return (path, sample, labels, F.get_size(sample))
+        # Get original image size (height, width) before transforms
+        orig_size = [img_info["height"], img_info["width"]]
+
+        return (path, sample, labels, orig_size)
 
 
 class CocoMosaicTraining(CocoBase):
@@ -127,9 +131,7 @@ class CocoMosaicTraining(CocoBase):
         self._mosaic_decay_epochs: Optional[int] = None
         self._mosaic_decay_start: Optional[int] = None
 
-    def configure_mosaic_linear_decay(
-        self, base_prob: float, total_epochs: int, decay_fraction: float = 0.1
-    ) -> None:
+    def configure_mosaic_linear_decay(self, base_prob: float, total_epochs: int, decay_fraction: float = 0.1) -> None:
         if total_epochs <= 0:
             raise ValueError("total_epochs must be positive")
         if decay_fraction <= 0.0 or decay_fraction > 1.0:
@@ -141,11 +143,7 @@ class CocoMosaicTraining(CocoBase):
         self._mosaic_decay_start = max(1, total_epochs - decay_epochs + 1)
 
     def update_mosaic_prob(self, epoch: int) -> Optional[float]:
-        if (
-            self._mosaic_base_prob is None
-            or self._mosaic_decay_epochs is None
-            or self._mosaic_decay_start is None
-        ):
+        if self._mosaic_base_prob is None or self._mosaic_decay_epochs is None or self._mosaic_decay_start is None:
             return None
 
         if epoch >= self._mosaic_decay_start:

birder/data/transforms/detection.py

@@ -1,3 +1,4 @@
+import math
 import random
 from collections.abc import Callable
 from typing import Any
@@ -10,6 +11,24 @@ from torchvision.transforms import v2
 
 from birder.data.transforms.classification import RGBType
 
+MULTISCALE_STEP = 32
+DEFAULT_MULTISCALE_MIN_SIZE = 480
+DEFAULT_MULTISCALE_MAX_SIZE = 800
+
+
+def build_multiscale_sizes(
+    min_size: Optional[int] = None, max_size: int = DEFAULT_MULTISCALE_MAX_SIZE
+) -> tuple[int, ...]:
+    if min_size is None:
+        min_size = DEFAULT_MULTISCALE_MIN_SIZE
+
+    start = int(math.ceil(min_size / MULTISCALE_STEP) * MULTISCALE_STEP)
+    end = int(math.floor(max_size / MULTISCALE_STEP) * MULTISCALE_STEP)
+    if end < start:
+        return (start,)
+
+    return tuple(range(start, end + 1, MULTISCALE_STEP))
+
 
 class ResizeWithRandomInterpolation(nn.Module):
     def __init__(
@@ -39,6 +58,7 @@ def get_birder_augment(
     dynamic_size: bool,
     multiscale: bool,
     max_size: Optional[int],
+    multiscale_min_size: Optional[int],
     post_mosaic: bool = False,
 ) -> Callable[..., torch.Tensor]:
     if dynamic_size is True:
@@ -78,9 +98,7 @@ def get_birder_augment(
     # Resize
     if multiscale is True:
         transformations.append(
-            v2.RandomShortestSize(
-                min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-            ),
+            v2.RandomShortestSize(min_size=build_multiscale_sizes(multiscale_min_size), max_size=max_size or 1333),
         )
     else:
         transformations.append(
@@ -132,6 +150,7 @@ def get_birder_augment(
 AugType = Literal["birder", "lsj", "multiscale", "ssd", "ssdlite", "yolo", "detr"]
 
 
+# pylint: disable=too-many-return-statements
 def training_preset(
     size: tuple[int, int],
     aug_type: AugType,
@@ -140,6 +159,7 @@ def training_preset(
     dynamic_size: bool = False,
     multiscale: bool = False,
     max_size: Optional[int] = None,
+    multiscale_min_size: Optional[int] = None,
     post_mosaic: bool = False,
 ) -> Callable[..., torch.Tensor]:
     mean = rgv_values["mean"]
@@ -159,7 +179,9 @@ def training_preset(
         return v2.Compose(  # type:ignore
             [
                 v2.ToImage(),
-                get_birder_augment(size, level, fill_value, dynamic_size, multiscale, max_size, post_mosaic),
+                get_birder_augment(
+                    size, level, fill_value, dynamic_size, multiscale, max_size, multiscale_min_size, post_mosaic
+                ),
                 v2.ToDtype(torch.float32, scale=True),
                 v2.Normalize(mean=mean, std=std),
                 v2.ToPureTensor(),
@@ -190,9 +212,7 @@ def training_preset(
         return v2.Compose(  # type: ignore
             [
                 v2.ToImage(),
-                v2.RandomShortestSize(
-                    min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-                ),
+                v2.RandomShortestSize(min_size=build_multiscale_sizes(multiscale_min_size), max_size=max_size or 1333),
                 v2.RandomHorizontalFlip(0.5),
                 v2.SanitizeBoundingBoxes(),
                 v2.ToDtype(torch.float32, scale=True),
@@ -264,21 +284,18 @@ def training_preset(
         )
 
     if aug_type == "detr":
+        multiscale_sizes = build_multiscale_sizes(multiscale_min_size)
         return v2.Compose(  # type: ignore
             [
                 v2.ToImage(),
                 v2.RandomChoice(
                     [
-                        v2.RandomShortestSize(
-                            (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-                        ),
+                        v2.RandomShortestSize(min_size=multiscale_sizes, max_size=max_size or 1333),
                         v2.Compose(
                             [
                                 v2.RandomShortestSize((400, 500, 600)),
                                 v2.RandomIoUCrop() if post_mosaic is False else v2.Identity(),  # RandomSizeCrop
-                                v2.RandomShortestSize(
-                                    (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-                                ),
+                                v2.RandomShortestSize(min_size=multiscale_sizes, max_size=max_size or 1333),
                             ]
                         ),
                     ]

birder/inference/detection.py

@@ -5,17 +5,99 @@ from typing import Optional
 import torch
 import torch.amp
 from PIL import Image
+from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from birder.conf import settings
+from birder.data.collators.detection import batch_images
 from birder.data.transforms.detection import InferenceTransform
+from birder.inference.wbf import fuse_detections_wbf
+from birder.net.base import make_divisible
+
+
+def _normalize_image_sizes(inputs: torch.Tensor, image_sizes: Optional[list[list[int]]]) -> list[list[int]]:
+    if image_sizes is not None:
+        return image_sizes
+
+    (_, _, height, width) = inputs.shape
+    return [[height, width] for _ in range(inputs.size(0))]
+
+
+def _hflip_inputs(inputs: torch.Tensor, image_sizes: list[list[int]]) -> torch.Tensor:
+    # Detection collator pads on the right/bottom, so flip only the valid region to keep padding aligned.
+    flipped = inputs.clone()
+    for idx, (height, width) in enumerate(image_sizes):
+        flipped[idx, :, :height, :width] = torch.flip(inputs[idx, :, :height, :width], dims=[2])
+
+    return flipped
+
+
+def _resize_batch(
+    inputs: torch.Tensor, image_sizes: list[list[int]], scale: float, size_divisible: int
+) -> tuple[torch.Tensor, torch.Tensor, list[list[int]]]:
+    resized_images: list[torch.Tensor] = []
+    for idx, (height, width) in enumerate(image_sizes):
+        target_h = make_divisible(height * scale, size_divisible)
+        target_w = make_divisible(width * scale, size_divisible)
+        image = inputs[idx, :, :height, :width]
+        resized = F.interpolate(image.unsqueeze(0), size=(target_h, target_w), mode="bilinear", align_corners=False)
+        resized_images.append(resized.squeeze(0))
+
+    return batch_images(resized_images, size_divisible)
+
+
+def _rescale_boxes(boxes: torch.Tensor, from_size: list[int], to_size: list[int]) -> torch.Tensor:
+    scale_w = to_size[1] / from_size[1]
+    scale_h = to_size[0] / from_size[0]
+    scale = boxes.new_tensor([scale_w, scale_h, scale_w, scale_h])
+    return boxes * scale
+
+
+def _rescale_detections(
+    detections: list[dict[str, torch.Tensor]],
+    from_sizes: list[list[int]],
+    to_sizes: list[list[int]],
+) -> list[dict[str, torch.Tensor]]:
+    for idx, (detection, from_size, to_size) in enumerate(zip(detections, from_sizes, to_sizes)):
+        boxes = detection["boxes"]
+        if boxes.numel() == 0:
+            continue
+
+        detections[idx]["boxes"] = _rescale_boxes(boxes, from_size, to_size)
+
+    return detections
+
+
+def _invert_hflip_boxes(boxes: torch.Tensor, image_size: list[int]) -> torch.Tensor:
+    width = boxes.new_tensor(image_size[1])
+    x1 = boxes[:, 0]
+    x2 = boxes[:, 2]
+    flipped = boxes.clone()
+    flipped[:, 0] = width - x2
+    flipped[:, 2] = width - x1
+
+    return flipped
+
+
+def _invert_detections(
+    detections: list[dict[str, torch.Tensor]], image_sizes: list[list[int]]
+) -> list[dict[str, torch.Tensor]]:
+    for idx, (detection, image_size) in enumerate(zip(detections, image_sizes)):
+        boxes = detection["boxes"]
+        if boxes.numel() == 0:
+            continue
+
+        detections[idx]["boxes"] = _invert_hflip_boxes(boxes, image_size)
+
+    return detections
 
 
 def infer_image(
     net: torch.nn.Module | torch.ScriptModule,
     sample: Image.Image | str,
     transform: Callable[..., torch.Tensor],
+    tta: bool = False,
     device: Optional[torch.device] = None,
     score_threshold: Optional[float] = None,
     **kwargs: Any,
@@ -43,7 +125,7 @@ def infer_image(
         device = torch.device("cpu")
 
     input_tensor = transform(image).unsqueeze(dim=0).to(device)
-    detections = infer_batch(net, input_tensor, **kwargs)
+    detections = infer_batch(net, input_tensor, tta=tta, **kwargs)
     if score_threshold is not None:
         for i, detection in enumerate(detections):
             idxs = torch.where(detection["scores"] > score_threshold)
@@ -63,16 +145,36 @@ def infer_batch(
     inputs: torch.Tensor,
     masks: Optional[torch.Tensor] = None,
     image_sizes: Optional[list[list[int]]] = None,
+    tta: bool = False,
     **kwargs: Any,
 ) -> list[dict[str, torch.Tensor]]:
-    (detections, _) = net(inputs, masks=masks, image_sizes=image_sizes, **kwargs)
-    return detections  # type: ignore[no-any-return]
+    if tta is False:
+        (detections, _) = net(inputs, masks=masks, image_sizes=image_sizes, **kwargs)
+        return detections  # type: ignore[no-any-return]
+
+    normalized_sizes = _normalize_image_sizes(inputs, image_sizes)
+    detections_list: list[list[dict[str, torch.Tensor]]] = []
+
+    for scale in (0.8, 1.0, 1.2):
+        (scaled_inputs, scaled_masks, scaled_sizes) = _resize_batch(inputs, normalized_sizes, scale, size_divisible=32)
+        (detections, _) = net(scaled_inputs, masks=scaled_masks, image_sizes=scaled_sizes, **kwargs)
+        detections = _rescale_detections(detections, scaled_sizes, normalized_sizes)
+        detections_list.append(detections)
+
+        flipped_inputs = _hflip_inputs(scaled_inputs, scaled_sizes)
+        (flipped_detections, _) = net(flipped_inputs, masks=scaled_masks, image_sizes=scaled_sizes, **kwargs)
+        flipped_detections = _invert_detections(flipped_detections, scaled_sizes)
+        flipped_detections = _rescale_detections(flipped_detections, scaled_sizes, normalized_sizes)
+        detections_list.append(flipped_detections)
+
+    return fuse_detections_wbf(detections_list, iou_thr=0.55, conf_type="avg")
 
 
 def infer_dataloader(
     device: torch.device,
     net: torch.nn.Module | torch.ScriptModule,
     dataloader: DataLoader,
+    tta: bool = False,
     model_dtype: torch.dtype = torch.float32,
     amp: bool = False,
     amp_dtype: Optional[torch.dtype] = None,
@@ -97,6 +199,8 @@ def infer_dataloader(
         The model to use for inference.
     dataloader
         The DataLoader containing the dataset to perform inference on.
+    tta
+        Run inference with multi-scale and horizontal flip test time augmentation and fuse results with WBF.
     model_dtype
         The base dtype to use.
     amp
@@ -142,7 +246,7 @@ def infer_dataloader(
             masks = masks.to(device, non_blocking=True)
 
             with torch.amp.autocast(device.type, enabled=amp, dtype=amp_dtype):
-                detections = infer_batch(net, inputs, masks, image_sizes)
+                detections = infer_batch(net, inputs, masks=masks, image_sizes=image_sizes, tta=tta)
 
             detections = InferenceTransform.postprocess(detections, image_sizes, orig_sizes)
             if targets[0] != settings.NO_LABEL: