dietr 0.2.0__py3-none-any.whl

dietr/__init__.py

@@ -0,0 +1 @@
+from .wrapper import DIETR

dietr/data/box_conversion.py

@@ -0,0 +1,101 @@
+"""
+------------------------------------------------------------------------
+DIETR
+Copyright (c) 2026 Koen Botermans
+Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+------------------------------------------------------------------------
+"""
+import numpy as np
+import torch
+
+def box_coco_to_yolo(box: list[int], img_hw: tuple[int]) -> list[float]:
+    """
+    Convert Coco pixel formal boxes to yolo normalized.
+    Coco/Pixel format:      [x_min, y_min, width, height]
+    YOLO format normalized: [x_center, y_center, width, height]
+    xyxy format:            [x_min, y_min, x_max, y_max]
+    This is often called cxcywh format.
+    """
+    x, y, w, h = box[:4]
+    x /= img_hw[1]
+    y /= img_hw[0]
+    w /= img_hw[1]
+    h /= img_hw[0]
+    return [x + w / 2, y + h / 2, w, h]
+
+
+def box_yolo_to_coco(box_np: np.ndarray, img_hw: tuple[int, int]) -> np.ndarray:
+    """
+    Convert yolo format to coco format.
+    Coco/Pixel format:      [x_min, y_min, width, height]
+    YOLO format normalized: [x_center, y_center, width, height]
+    xyxy format:            [x_min, y_min, x_max, y_max]
+    """
+    x_min = (box_np[:, 0] - box_np[:, 2] / 2) * img_hw[1]
+    y_min = (box_np[:, 1] - box_np[:, 3] / 2) * img_hw[0]
+    w = (box_np[:, 2] * img_hw[1])
+    h = (box_np[:, 3] * img_hw[0])
+    return np.floor(np.vstack([x_min, y_min, w, h]).T).astype(int)
+
+
+def box_tensor_yolo_to_xyxy(box: torch.Tensor) -> torch.Tensor:
+    """
+    Coco/Pixel format:      [x_min, y_min, width, height]
+    YOLO format normalized: [x_center, y_center, width, height]
+    xyxy format:            [x_min, y_min, x_max, y_max]
+    Can also be the pascal voc format!
+    """
+    x_c, y_c, w, h = box.unbind(-1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
+    return torch.stack(b, dim=-1)
+
+
+def box_tensor_xyxy_to_yolo(box: torch.Tensor) -> torch.Tensor:
+    """
+
+    Coco/Pixel format:      [x_min, y_min, width, height]
+    YOLO format normalized: [x_center, y_center, width, height]
+    xyxy format:            [x_min, y_min, x_max, y_max]
+    Can also be the pascal voc format!
+    """
+    xmin, ymin, xmax, ymax = box.unbind(-1)
+    box = [(xmax - xmin) / 2, (ymax - ymin) / 2, (xmax - xmin), (ymax - ymin)]
+    return torch.stack(box, dim=-1)
+
+
+def box_xyxy_area(box: torch.Tensor) -> torch.Tensor:
+    """Computes the area of a box in normalized pascal format"""
+    return (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1])
+
+
+def box_xyxy_iou(box1: torch.Tensor, box2: torch.Tensor):
+    area1 = box_xyxy_area(box1)
+    area2 = box_xyxy_area(box2)
+
+    lt = torch.max(box1[:, None, :2], box2[:, :2])  # [N,M,2]
+    rb = torch.min(box1[:, None, 2:], box2[:, 2:])  # [N,M,2]
+
+    wh = (rb - lt).clamp(min=0)  # [N,M,2]
+    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
+
+    union = area1[:, None] + area2 - inter
+
+    iou = inter / union
+    return iou, union
+
+
+def box_yolo_iou(box1: torch.Tensor, box2: torch.Tensor):
+    return box_xyxy_iou(
+        box_tensor_yolo_to_xyxy(box=box1), box_tensor_yolo_to_xyxy(box=box2)
+    )
+
+
+def box_yolo_giou(box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
+    iou, union = box_yolo_iou(box1, box2)
+    box1 = box_tensor_yolo_to_xyxy(box=box1)
+    box2 = box_tensor_yolo_to_xyxy(box=box2)
+    lt = torch.min(box1[:, None, :2], box2[:, :2])  # [N, 2]
+    rb = torch.max(box1[:, None, 2:], box2[:, 2:])  # [N, 2]
+    wh = (rb - lt).clamp(min=0)  # [N, 2]
+    area = wh[:, :, 0] * wh[:, :, 1]
+    return iou - (area - union) / area

dietr/data/coco_definitions.py

@@ -0,0 +1,287 @@
+"""
+------------------------------------------------------------------------
+DIETR
+Copyright (c) 2026 Koen Botermans
+Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+------------------------------------------------------------------------
+"""
+
+coco_classes = (
+    "person",
+    "bicycle",
+    "car",
+    "motorcycle",
+    "airplane",
+    "bus",
+    "train",
+    "truck",
+    "boat",
+    "traffic light",
+    "fire hydrant",
+    "stop sign",
+    "parking meter",
+    "bench",
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+    "backpack",
+    "umbrella",
+    "handbag",
+    "tie",
+    "suitcase",
+    "frisbee",
+    "skis",
+    "snowboard",
+    "sports ball",
+    "kite",
+    "baseball bat",
+    "baseball glove",
+    "skateboard",
+    "surfboard",
+    "tennis racket",
+    "bottle",
+    "wine glass",
+    "cup",
+    "fork",
+    "knife",
+    "spoon",
+    "bowl",
+    "banana",
+    "apple",
+    "sandwich",
+    "orange",
+    "broccoli",
+    "carrot",
+    "hot dog",
+    "pizza",
+    "donut",
+    "cake",
+    "chair",
+    "couch",
+    "potted plant",
+    "bed",
+    "dining table",
+    "toilet",
+    "tv",
+    "laptop",
+    "mouse",
+    "remote",
+    "keyboard",
+    "cell phone",
+    "microwave",
+    "oven",
+    "toaster",
+    "sink",
+    "refrigerator",
+    "book",
+    "clock",
+    "vase",
+    "scissors",
+    "teddy bear",
+    "hair drier",
+    "toothbrush",
+)
+coco_categories = [
+    {"supercategory": "person", "id": 1, "name": "person"},
+    {"supercategory": "vehicle", "id": 2, "name": "bicycle"},
+    {"supercategory": "vehicle", "id": 3, "name": "car"},
+    {"supercategory": "vehicle", "id": 4, "name": "motorcycle"},
+    {"supercategory": "vehicle", "id": 5, "name": "airplane"},
+    {"supercategory": "vehicle", "id": 6, "name": "bus"},
+    {"supercategory": "vehicle", "id": 7, "name": "train"},
+    {"supercategory": "vehicle", "id": 8, "name": "truck"},
+    {"supercategory": "vehicle", "id": 9, "name": "boat"},
+    {"supercategory": "outdoor", "id": 10, "name": "traffic light"},
+    {"supercategory": "outdoor", "id": 11, "name": "fire hydrant"},
+    {"supercategory": "outdoor", "id": 13, "name": "stop sign"},
+    {"supercategory": "outdoor", "id": 14, "name": "parking meter"},
+    {"supercategory": "outdoor", "id": 15, "name": "bench"},
+    {"supercategory": "animal", "id": 16, "name": "bird"},
+    {"supercategory": "animal", "id": 17, "name": "cat"},
+    {"supercategory": "animal", "id": 18, "name": "dog"},
+    {"supercategory": "animal", "id": 19, "name": "horse"},
+    {"supercategory": "animal", "id": 20, "name": "sheep"},
+    {"supercategory": "animal", "id": 21, "name": "cow"},
+    {"supercategory": "animal", "id": 22, "name": "elephant"},
+    {"supercategory": "animal", "id": 23, "name": "bear"},
+    {"supercategory": "animal", "id": 24, "name": "zebra"},
+    {"supercategory": "animal", "id": 25, "name": "giraffe"},
+    {"supercategory": "accessory", "id": 27, "name": "backpack"},
+    {"supercategory": "accessory", "id": 28, "name": "umbrella"},
+    {"supercategory": "accessory", "id": 31, "name": "handbag"},
+    {"supercategory": "accessory", "id": 32, "name": "tie"},
+    {"supercategory": "accessory", "id": 33, "name": "suitcase"},
+    {"supercategory": "sports", "id": 34, "name": "frisbee"},
+    {"supercategory": "sports", "id": 35, "name": "skis"},
+    {"supercategory": "sports", "id": 36, "name": "snowboard"},
+    {"supercategory": "sports", "id": 37, "name": "sports ball"},
+    {"supercategory": "sports", "id": 38, "name": "kite"},
+    {"supercategory": "sports", "id": 39, "name": "baseball bat"},
+    {"supercategory": "sports", "id": 40, "name": "baseball glove"},
+    {"supercategory": "sports", "id": 41, "name": "skateboard"},
+    {"supercategory": "sports", "id": 42, "name": "surfboard"},
+    {"supercategory": "sports", "id": 43, "name": "tennis racket"},
+    {"supercategory": "kitchen", "id": 44, "name": "bottle"},
+    {"supercategory": "kitchen", "id": 46, "name": "wine glass"},
+    {"supercategory": "kitchen", "id": 47, "name": "cup"},
+    {"supercategory": "kitchen", "id": 48, "name": "fork"},
+    {"supercategory": "kitchen", "id": 49, "name": "knife"},
+    {"supercategory": "kitchen", "id": 50, "name": "spoon"},
+    {"supercategory": "kitchen", "id": 51, "name": "bowl"},
+    {"supercategory": "food", "id": 52, "name": "banana"},
+    {"supercategory": "food", "id": 53, "name": "apple"},
+    {"supercategory": "food", "id": 54, "name": "sandwich"},
+    {"supercategory": "food", "id": 55, "name": "orange"},
+    {"supercategory": "food", "id": 56, "name": "broccoli"},
+    {"supercategory": "food", "id": 57, "name": "carrot"},
+    {"supercategory": "food", "id": 58, "name": "hot dog"},
+    {"supercategory": "food", "id": 59, "name": "pizza"},
+    {"supercategory": "food", "id": 60, "name": "donut"},
+    {"supercategory": "food", "id": 61, "name": "cake"},
+    {"supercategory": "furniture", "id": 62, "name": "chair"},
+    {"supercategory": "furniture", "id": 63, "name": "couch"},
+    {"supercategory": "furniture", "id": 64, "name": "potted plant"},
+    {"supercategory": "furniture", "id": 65, "name": "bed"},
+    {"supercategory": "furniture", "id": 67, "name": "dining table"},
+    {"supercategory": "furniture", "id": 70, "name": "toilet"},
+    {"supercategory": "electronic", "id": 72, "name": "tv"},
+    {"supercategory": "electronic", "id": 73, "name": "laptop"},
+    {"supercategory": "electronic", "id": 74, "name": "mouse"},
+    {"supercategory": "electronic", "id": 75, "name": "remote"},
+    {"supercategory": "electronic", "id": 76, "name": "keyboard"},
+    {"supercategory": "electronic", "id": 77, "name": "cell phone"},
+    {"supercategory": "appliance", "id": 78, "name": "microwave"},
+    {"supercategory": "appliance", "id": 79, "name": "oven"},
+    {"supercategory": "appliance", "id": 80, "name": "toaster"},
+    {"supercategory": "appliance", "id": 81, "name": "sink"},
+    {"supercategory": "appliance", "id": 82, "name": "refrigerator"},
+    {"supercategory": "indoor", "id": 84, "name": "book"},
+    {"supercategory": "indoor", "id": 85, "name": "clock"},
+    {"supercategory": "indoor", "id": 86, "name": "vase"},
+    {"supercategory": "indoor", "id": 87, "name": "scissors"},
+    {"supercategory": "indoor", "id": 88, "name": "teddy bear"},
+    {"supercategory": "indoor", "id": 89, "name": "hair drier"},
+    {"supercategory": "indoor", "id": 90, "name": "toothbrush"},
+]
+coco_to_label = {
+    1: 'person',
+    2: 'bicycle',
+    3: 'car',
+    4: 'motorcycle',
+    5: 'airplane',
+    6: 'bus',
+    7: 'train',
+    8: 'truck',
+    9: 'boat',
+    10: 'traffic light',
+    11: 'fire hydrant',
+    13: 'stop sign',
+    14: 'parking meter',
+    15: 'bench',
+    16: 'bird',
+    17: 'cat',
+    18: 'dog',
+    19: 'horse',
+    20: 'sheep',
+    21: 'cow',
+    22: 'elephant',
+    23: 'bear',
+    24: 'zebra',
+    25: 'giraffe',
+    27: 'backpack',
+    28: 'umbrella',
+    31: 'handbag',
+    32: 'tie',
+    33: 'suitcase',
+    34: 'frisbee',
+    35: 'skis',
+    36: 'snowboard',
+    37: 'sports ball',
+    38: 'kite',
+    39: 'baseball bat',
+    40: 'baseball glove',
+    41: 'skateboard',
+    42: 'surfboard',
+    43: 'tennis racket',
+    44: 'bottle',
+    46: 'wine glass',
+    47: 'cup',
+    48: 'fork',
+    49: 'knife',
+    50: 'spoon',
+    51: 'bowl',
+    52: 'banana',
+    53: 'apple',
+    54: 'sandwich',
+    55: 'orange',
+    56: 'broccoli',
+    57: 'carrot',
+    58: 'hot dog',
+    59: 'pizza',
+    60: 'donut',
+    61: 'cake',
+    62: 'chair',
+    63: 'couch',
+    64: 'potted plant',
+    65: 'bed',
+    67: 'dining table',
+    70: 'toilet',
+    72: 'tv',
+    73: 'laptop',
+    74: 'mouse',
+    75: 'remote',
+    76: 'keyboard',
+    77: 'cell phone',
+    78: 'microwave',
+    79: 'oven',
+    80: 'toaster',
+    81: 'sink',
+    82: 'refrigerator',
+    84: 'book',
+    85: 'clock',
+    86: 'vase',
+    87: 'scissors',
+    88: 'teddy bear',
+    89: 'hair drier',
+    90: 'toothbrush'
+}
+
+coco_label_to_normal = {k: i for i, k in enumerate(coco_to_label.keys())}
+normal_label_to_coco = {v: k for k, v in coco_label_to_normal.items()}
+
+coco_result_names = [
+    "AP-0.95",
+    "AP-0.50",
+    "AP-0.75",
+    "AP-small",
+    "AP-medium",
+    "AP-large",
+    "AR-1det",
+    "AR-10det",
+    "AR-0.95",
+    "AR-small",
+    "AR-medium",
+    "AR-large",
+]
+
+mapping = {
+    "Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ]": "AP-0.95",
+    "Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ]": "AP-0.50",
+    "Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ]": "AP-0.75",
+    "Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ]": "AP-small",
+    "Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ]": "AP-medium",
+    "Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ]": "AP-large",
+    "Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ]": "AR-1det",
+    "Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ]": "AR-10det",
+    "Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ]": "AR-0.95",
+    "Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ]": "AR-small",
+    "Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ]": "AR-medium",
+    "Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ]": "AR-large",
+}

dietr/data/conversion.py

@@ -0,0 +1,128 @@
+"""
+------------------------------------------------------------------------
+DIETR
+Copyright (c) 2026 Koen Botermans
+Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+------------------------------------------------------------------------
+"""
+import torch
+import numpy as np
+from dietr.data.coco_definitions import (
+    coco_label_to_normal,
+    normal_label_to_coco,
+    coco_classes,
+)
+from dietr.data.msk_conversion import (
+    msk_rle_to_msk_np,
+    msk_rle_to_compressed_rle,
+    msk_polygon_to_msk_np,
+    msk_np_to_msk_crp_np,
+)
+
+from dietr.data.box_conversion import box_coco_to_yolo
+
+
+def inverse_sigmoid(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
+    x = x.clip(min=0.0, max=1.0)
+    return torch.log(x.clip(min=eps) / (1 - x).clip(min=eps))
+
+
+def coco_seg_to_msk_np(coco_seg: dict, img_hw: tuple[int, int]) -> np.ndarray:
+    if not isinstance(coco_seg, dict):
+        msk_np = msk_polygon_to_msk_np(polygon=coco_seg, img_hw=img_hw)
+    else:
+        if isinstance(coco_seg["counts"], str):
+            msk_np = msk_rle_to_msk_np(rle=coco_seg)
+        else:
+            msk_np = msk_rle_to_msk_np(msk_rle_to_compressed_rle(coco_ann=coco_seg))
+    return msk_np.astype(np.bool_)
+
+
+def coco_ann_to_sample_np(
+    coco_ann: list[any],
+    img_hw: tuple[int, int],
+    min_visibility: float,
+    coco_dataset: bool,
+    msk: bool,
+) -> dict[str, np.ndarray]:
+    """Convert anntations from the coco annotation format to the internal data format, which is just a dictionary of numpy arrays..
+
+    Args:
+        coco_ann (list[coco_annotation]): list of coco annotations.
+        img_hw (tuple[int, int]): the original height and with of the original image, this is need to scale the bounding boxes.
+
+    Returns:
+        sample_np (dict[str, np.ndarray]): combined targets that can be use for training.
+
+        It has the following elements:
+        box_np      (np.ndarray):  Boxes in normalized yolo format. (or cxcywh).
+        cls_np      (np.ndarray):  Labels in normal format, thus 0: person, etc...
+        msk_np      (np.ndarray):  Masks in np.ndarray.
+        crp_np      (np.ndarray):  Cropping of the masks in a numpy array, this cropping is used when computing the loss, to stabilize training the segmentation outputs.
+    """
+    box_np = []
+    cls_np = []
+    if msk:
+        msk_np = []
+        crp_np = []
+    for annotation in coco_ann:
+        if annotation["iscrowd"] == 1:
+            continue
+
+        box_yolo = box_coco_to_yolo(annotation["bbox"], img_hw=img_hw)
+        if (box_yolo[2] <= min_visibility) or (box_yolo[3] <= min_visibility):
+            continue
+        box_np.append(box_yolo)
+
+        if coco_dataset:
+            cls_np.append(label_coco_to_normal(annotation["category_id"]))
+        else:
+            cls_np.append(annotation["category_id"])
+
+        if msk:
+            msk_np.append(
+                coco_seg_to_msk_np(coco_seg=annotation["segmentation"], img_hw=img_hw)
+            )
+            crp_np.append(msk_np_to_msk_crp_np(msk_np=msk_np[-1]))
+
+    if len(box_np) == 0:
+        if msk:
+            return {
+                "box_np": np.array([]).reshape(-1, 4),
+                "msk_np": np.array([]).reshape(-1, 4, 4),
+                "crp_np": np.array([]).reshape(-1, 4, 4),
+                "cls_np": np.array([]),
+            }
+        return {
+            "box_np": np.array([]).reshape(-1, 4),
+            "cls_np": np.array([]),
+        }
+    if msk:
+        return {
+            "box_np": np.stack(box_np),
+            "cls_np": np.stack(cls_np),
+            "msk_np": np.stack(msk_np),
+            "crp_np": np.stack(crp_np),
+        }
+    return {
+        "box_np": np.stack(box_np),
+        "cls_np": np.stack(cls_np),
+    }
+
+
+def label_coco_to_normal(coco_label: int) -> int:
+    """In the coco format the labels are from 1 to 90, this function just maps these labels to label_id,
+    with id=0, is person."""
+    return coco_label_to_normal[coco_label]
+
+
+def label_normal_to_coco(normal_label: int) -> int:
+    return normal_label_to_coco[normal_label]
+
+
+def label_normal_to_name(normal_label: int) -> str:
+    return coco_classes[normal_label - 1]
+
+
+def label_coco_to_name(normal_label: int) -> str:
+    return coco_classes[label_coco_to_normal(normal_label) - 1]

dietr/data/dataloaders.py

@@ -0,0 +1,107 @@
+"""
+------------------------------------------------------------------------
+DIETR
+Copyright (c) 2026 Koen Botermans
+Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+------------------------------------------------------------------------
+"""
+import torch
+import numpy as np
+from torch.utils.data.distributed import DistributedSampler
+from dietr.data.trn_dataset import CocoDataset
+from dietr.data.val_dataset import CocoValDataset
+from dietr.data.trn_batch_collate import BatchImageCollateFunction
+from dietr.data.val_batch_collate import ValBatchImageCollateFunction
+
+from collections.abc import Iterator
+from typing import Any
+
+
+
+def get_val_dataloader(
+    config, multi_gpu: bool = False, tta: bool = False
+) -> torch.utils.data.DataLoader:
+    coco_dataset = CocoValDataset(
+        coco_ann_file=config["val_ann_file"],
+        coco_img_root=config["val_img_root"],
+        img_base_size=config["val_img_base_size"],
+    )
+    batch_collate = ValBatchImageCollateFunction(
+        base_size=config["val_img_base_size"],
+        val_sizes=[config["val_img_base_size"]]
+        if not tta
+        else config["val_tta_img_sizes"],
+    )
+    return torch.utils.data.DataLoader(
+        dataset=coco_dataset,
+        batch_size=config["val_batch_size"],
+        shuffle=False,
+        collate_fn=batch_collate,
+        persistent_workers=True,
+        pin_memory=True,
+        num_workers=config["n_val_workers"],
+        sampler=DistributedSampler(coco_dataset, shuffle=False) if multi_gpu else None,
+    )
+
+def contains_targets(y_trn_batch: dict[str, list[torch.Tensor]]) -> bool:
+    for box in y_trn_batch["box"]:
+        if len(box) == 0:
+            return False
+    return True
+
+
+def cycle(iterable: Any) -> Iterator[Any]:
+    """Create a dataloader-safe cyclical iterator.
+    From: https://github.com/pytorch/pytorch/issues/23900#issuecomment-518858050
+    """
+    iterator = iter(iterable)
+    while True:
+        try:
+            (x_trn_batch, y_trn_batch) =  next(iterator)
+            # while not contains_targets(y_trn_batch=y_trn_batch):
+            #     (x_trn_batch, y_trn_batch) =  next(iterator)
+            yield (x_trn_batch, y_trn_batch)
+        except StopIteration:
+            iterator = iter(iterable)
+
+def get_trn_dataloader(config, multi_gpu: bool, cycle_dataloader: bool) -> Iterator[Any]:
+    coco_dataset = CocoDataset(
+        msk=config["msk"],
+        coco_ann_file=config["trn_ann_file"],
+        coco_img_root=config["trn_img_root"],
+        img_base_size=config["trn_img_base_size"],
+        min_visibility=config["min_visibility"],
+        td_flip_p=config["dataset_dt_flip_p"],
+        lr_flip_p=config["dataset_lr_flip_p"],
+        mask_scaling=config["mask_scaling"],
+        photo_transform_p=config["dataset_photo_transform_p"],
+        filter_empty=config["dataset_filter_empy"],
+        random_zoom_p=config["dataset_random_zoom_p"],
+        random_crop_p=config["dataset_random_crop_p"],
+        mosaic_p=config["dataset_mosaic_p"],
+        coco_dataset=config["coco_dataset"],
+    )
+    batch_collate = BatchImageCollateFunction(
+        base_size=config["trn_img_base_size"],
+        multi_scale_p=config.get("multi_scale_p", 1.0),
+        multi_input_size=config["multi_input_size"],
+        double_multi_input_size=config["double_multi_input_size"],
+        mask_scaling=config["mask_scaling"],
+        dtype=np.float32,
+        both_sides=config.get("both_sides", True),
+    )
+    dataloader = torch.utils.data.DataLoader(
+            dataset=coco_dataset,
+            batch_size=config["trn_batch_size"],
+            shuffle=not multi_gpu,
+            collate_fn=batch_collate,
+            persistent_workers=True,
+            pin_memory=True,
+            num_workers=config["n_trn_workers"],
+            sampler=DistributedSampler(coco_dataset) if multi_gpu else None,
+        )
+    if cycle_dataloader:
+        return cycle(dataloader)
+    return dataloader
+
+

dietr/data/loading.py

@@ -0,0 +1,32 @@
+"""
+------------------------------------------------------------------------
+DIETR
+Copyright (c) 2026 Koen Botermans
+Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+------------------------------------------------------------------------
+"""
+import os
+import cv2
+import pycocotools
+import numpy as np
+
+
+def load_coco_img_np(
+    coco_gt: pycocotools.coco.COCO, img_id: int, coco_img_path: str
+) -> tuple[np.ndarray, tuple[int, int]]:
+    """Given a pycocotools and img_id,"""
+    img_info = coco_gt.loadImgs(ids=[img_id])[0]
+    img_path = img_info["file_name"]
+    combined_img_path = os.path.join(coco_img_path, img_path)
+
+    return (
+        cv2.cvtColor(cv2.imread(combined_img_path), cv2.COLOR_RGB2BGR),
+        (
+            img_info["height"],
+            img_info["width"],
+        )    )
+
+
+def load_coco_ann(coco_gt: pycocotools.coco.COCO, img_id: int) -> list[any]:
+    ann_ids = coco_gt.getAnnIds(imgIds=[img_id])
+    return coco_gt.loadAnns(ann_ids)