From 315a360b6f67e0f200e452021cfe663cd29ef01e Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Fri, 4 Apr 2025 11:30:52 +0200
Subject: [PATCH 01/22] Remove some desynced test file pointers

---
 src/flat_bug/tests/assets/pyramid_tps_1.pt      | 1 -
 src/flat_bug/tests/assets/pyramid_tps_2.pt      | 1 -
 src/flat_bug/tests/assets/pyramid_tps_3.pt      | 1 -
 src/flat_bug/tests/assets/pyramid_tps_4.pt      | 1 -
 src/flat_bug/tests/assets/pyramid_tps_5.pt      | 1 -
 src/flat_bug/tests/assets/single_scale_tps_1.pt | 1 -
 6 files changed, 6 deletions(-)
 delete mode 100644 src/flat_bug/tests/assets/pyramid_tps_1.pt
 delete mode 100644 src/flat_bug/tests/assets/pyramid_tps_2.pt
 delete mode 100644 src/flat_bug/tests/assets/pyramid_tps_3.pt
 delete mode 100644 src/flat_bug/tests/assets/pyramid_tps_4.pt
 delete mode 100644 src/flat_bug/tests/assets/pyramid_tps_5.pt
 delete mode 100644 src/flat_bug/tests/assets/single_scale_tps_1.pt

diff --git a/src/flat_bug/tests/assets/pyramid_tps_1.pt b/src/flat_bug/tests/assets/pyramid_tps_1.pt
deleted file mode 100644
index 5943933..0000000
--- a/src/flat_bug/tests/assets/pyramid_tps_1.pt
+++ /dev/null
@@ -1 +0,0 @@
-ERDA Pointer
\ No newline at end of file
diff --git a/src/flat_bug/tests/assets/pyramid_tps_2.pt b/src/flat_bug/tests/assets/pyramid_tps_2.pt
deleted file mode 100644
index 5943933..0000000
--- a/src/flat_bug/tests/assets/pyramid_tps_2.pt
+++ /dev/null
@@ -1 +0,0 @@
-ERDA Pointer
\ No newline at end of file
diff --git a/src/flat_bug/tests/assets/pyramid_tps_3.pt b/src/flat_bug/tests/assets/pyramid_tps_3.pt
deleted file mode 100644
index 3241d37..0000000
--- a/src/flat_bug/tests/assets/pyramid_tps_3.pt
+++ /dev/null
@@ -1 +0,0 @@
-ERDA Pointer
diff --git a/src/flat_bug/tests/assets/pyramid_tps_4.pt b/src/flat_bug/tests/assets/pyramid_tps_4.pt
deleted file mode 100644
index 3241d37..0000000
--- a/src/flat_bug/tests/assets/pyramid_tps_4.pt
+++ /dev/null
@@ -1 +0,0 @@
-ERDA Pointer
diff --git a/src/flat_bug/tests/assets/pyramid_tps_5.pt b/src/flat_bug/tests/assets/pyramid_tps_5.pt
deleted file mode 100644
index 3241d37..0000000
--- a/src/flat_bug/tests/assets/pyramid_tps_5.pt
+++ /dev/null
@@ -1 +0,0 @@
-ERDA Pointer
diff --git a/src/flat_bug/tests/assets/single_scale_tps_1.pt b/src/flat_bug/tests/assets/single_scale_tps_1.pt
deleted file mode 100644
index 5943933..0000000
--- a/src/flat_bug/tests/assets/single_scale_tps_1.pt
+++ /dev/null
@@ -1 +0,0 @@
-ERDA Pointer
\ No newline at end of file

From b0d6974cc1631330d144f59b90a824de98037855 Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Mon, 7 Apr 2025 10:25:22 +0200
Subject: [PATCH 02/22] Fixed error on saving with SVG due to invalid execution
 flow logic in TensorPredictions.plot

---
 src/flat_bug/predictor.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/flat_bug/predictor.py b/src/flat_bug/predictor.py
index 80a9d45..ef6f253 100644
--- a/src/flat_bug/predictor.py
+++ b/src/flat_bug/predictor.py
@@ -586,7 +586,8 @@ def plot(
         params.pop("self", None)
         if outpath not in [None, ""] and outpath.lower().endswith(".svg"):
             retval = self._plot_svg(**params)
-        retval = self._plot_image(**params)
+        else:
+            retval = self._plot_image(**params)
         if retval is None:
             return outpath
         return retval

From 6f9aee66719ba5d800b86eeedb45cee99a15a861 Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Mon, 7 Apr 2025 16:44:42 +0200
Subject: [PATCH 03/22] Fix --long_format inference

---
 src/bin/fb_predict.py     | 2 +-
 src/flat_bug/predictor.py | 5 +----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/bin/fb_predict.py b/src/bin/fb_predict.py
index 8763d99..b0ea160 100644
--- a/src/bin/fb_predict.py
+++ b/src/bin/fb_predict.py
@@ -291,7 +291,7 @@ def predict(
                     identifier = UUID, #str(uuid.uuid4()),
                 )
                 if not result_directory is None:
-                    json_files = [f for f in glob.glob(os.path.join(glob.escape(result_directory), "*.json"))]
+                    json_files = [f for f in glob.glob(os.path.join(glob.escape(metadata if isinstance(metadata, str) else result_directory), f"*{os.path.splitext(os.path.basename(f))[0]}*.json"))]
                     assert len(json_files) == 1
                     all_json_results.append(json_files[0])
                 if isVideo and overviews:
diff --git a/src/flat_bug/predictor.py b/src/flat_bug/predictor.py
index ef6f253..00706dd 100644
--- a/src/flat_bug/predictor.py
+++ b/src/flat_bug/predictor.py
@@ -1077,9 +1077,6 @@ def save(
         if prediction_directory_is_used:
             if not os.path.exists(prediction_directory):
                 os.makedirs(prediction_directory)
-        else:
-            # If the prediction directory is not used set it to None
-            return None
 
         # Save overview
         if overview:
@@ -1129,7 +1126,7 @@ def save(
             # Serialize the data to the metadata path
             self.serialize(outpath=metadata_path, identifier=identifier)
 
-        return prediction_directory
+        return prediction_directory if prediction_directory_is_used else None
 
 def _process_batch(
             image : torch.Tensor, 

From ab5361bd6f91a027c1167ea4623eca8d02f3b3cc Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Fri, 6 Jun 2025 15:42:52 +0200
Subject: [PATCH 04/22] Stricter version requirement for ultralytics dependency
 due to a breaking change in 8.3.125

---
 pyproject.toml        | 2 +-
 src/bin/fb_predict.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7f7d88f..31c16d7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,7 +19,7 @@ classifiers = [
 dependencies = [
     "torch>=2.2.0",
     "torchvision>=0.17.0",
-    "ultralytics>=8.2.16",
+    "ultralytics>=8.2.16,<=8.3.124",
     "shapely>=2.0.2",
     "scikit-optimize>=0.10.1"
 ]
diff --git a/src/bin/fb_predict.py b/src/bin/fb_predict.py
index b0ea160..873a62f 100644
--- a/src/bin/fb_predict.py
+++ b/src/bin/fb_predict.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""
+r"""
 Inference CLI script for ``flatbug``. 
 
 A comprehensive CLI API for ``flatbug`` inference with support for hyperparameter configuration, flexible input parsing, output format specification, and hardware specification.

From f59cec5fcd67d88eca21001e293118b758857221 Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Fri, 6 Jun 2025 15:46:04 +0200
Subject: [PATCH 05/22] Change README install guide from git clone with SSH to
 HTTPS (re: #121)

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index fcddfb2..4256172 100644
--- a/README.md
+++ b/README.md
@@ -33,7 +33,7 @@ micromamba install flat-bug -c conda-forge
 #### Source/development
 Or a development version can be installed from source by cloning this repository:
 ```sh
-git clone git@github.com:darsa-group/flat-bug.git
+git clone https://github.com/darsa-group/flat-bug.git
 cd flat-bug
 pip install -e .
 ```

From 72f407aa78ab2886467368db2d95fa6864936658 Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Wed, 9 Jul 2025 16:15:17 +0200
Subject: [PATCH 06/22] Remove tile and image tensors from intermediate
 `Results` objects to avoid memory-leak and OOM on large images

---
 src/flat_bug/yolo_helpers.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/flat_bug/yolo_helpers.py b/src/flat_bug/yolo_helpers.py
index 5ddaba8..f76203d 100644
--- a/src/flat_bug/yolo_helpers.py
+++ b/src/flat_bug/yolo_helpers.py
@@ -1,3 +1,4 @@
+from argparse import Namespace
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
@@ -132,7 +133,7 @@ def merge_tile_results(
         raise NotImplementedError("'Probs' not implemented yet")
     if not all([r.keypoints is None for r in results]):
         raise NotImplementedError("'Keypoints' not implemented yet")
-    return ResultsWithTiles(tiles=tile_indices, orig_img=orig_img, path=path, names=names, boxes=boxes, masks=masks, polygons=polygons, probs=None, keypoints=None)
+    return ResultsWithTiles(tiles=tile_indices, orig_img=Namespace(shape=orig_img.shape), path=path, names=names, boxes=boxes, masks=masks, polygons=polygons, probs=None, keypoints=None)
 
 def stack_masks(
         masks : List["Masks"], 
@@ -417,5 +418,5 @@ def postprocess(
         boxes = boxes[~too_small]
         masks = masks[~too_small]
         pred[:, :4] = boxes
-        results.append({"orig_img" : imgs[i].clone().permute(1,2,0), "path" : "", "names" : ["insect"], "boxes" : pred[:, :6], "masks" : masks})
+        results.append({"orig_img" : Namespace(shape=imgs[i].permute(1,2,0).shape), "path" : "", "names" : ["insect"], "boxes" : pred[:, :6], "masks" : masks})
     return results
\ No newline at end of file

From 8e23f7f53e8fc3bc3af50fa8aadc611ca5bdd587 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Wed, 8 Oct 2025 19:07:18 +0200
Subject: [PATCH 07/22] work on #134

---
 prototypes/mask_refiner/fb_refine.py | 225 +++++++++++++++++++++++++++
 1 file changed, 225 insertions(+)
 create mode 100644 prototypes/mask_refiner/fb_refine.py

diff --git a/prototypes/mask_refiner/fb_refine.py b/prototypes/mask_refiner/fb_refine.py
new file mode 100644
index 0000000..f3d6dfd
--- /dev/null
+++ b/prototypes/mask_refiner/fb_refine.py
@@ -0,0 +1,225 @@
+import json
+import argparse
+import os
+import cv2 as cv
+import numpy as np
+
+
+# fixme, resume should continue on the same "run folder"
+def main():
+    args_parse = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
+    args_parse.add_argument("-d", "--data-dir", dest="data_dir",
+                            help="The directory containing the prepared data (i.e., the output of  `fb_prepare.py`",
+                            type=str)
+
+    args_parse.add_argument("-c", "--config-file", dest="config_file",
+                            help="A YAML-formatted config file that overrides the default training meta-parameters",
+                            default=None)
+    args_parse.add_argument("-r", "--resume", dest="resume",
+                            help="resume training",
+                            action='store_true')
+
+    args, extra = args_parse.parse_known_args()
+
+# def best_new_contour(cnt):
+
+def best_yolo_contour(pts, roi_padded, conf_thresh=0.25):
+    """
+    Choose the best YOLOv8 polygon (on roi_padded) by IoU against `pts` (N,2).
+    Uses polygon IoU via shapely if available; otherwise falls back to mask IoU.
+    Returns an (M,2) int32 array of the best polygon in ROI coordinates.
+    """
+
+    # --- get global YOLOv8 model ---
+
+    H, W = roi_padded.shape[:2]
+    pts = np.asarray(pts, dtype=np.float32)
+    if pts.ndim != 2 or pts.shape[1] != 2 or len(pts) < 3:
+        return pts.astype(np.int32)
+
+    # --- try shapely for polygon IoU ---
+    try:
+        from shapely.geometry import Polygon
+        from shapely.errors import TopologicalError
+        use_shapely = True
+    except Exception:
+        use_shapely = False
+
+    # helper: polygon IoU with shapely
+    def _poly_iou_shapely(a_xy, b_xy):
+        try:
+            pa = Polygon(a_xy).buffer(0)  # buffer(0) fixes minor self-intersections
+            pb = Polygon(b_xy).buffer(0)
+            if not pa.is_valid or not pb.is_valid:
+                return 0.0
+            inter = pa.intersection(pb).area
+            union = pa.union(pb).area
+            return float(inter / union) if union > 0 else 0.0
+        except TopologicalError:
+            return 0.0
+
+    # fallback: mask IoU if shapely not available
+    def _poly_iou_mask(a_xy, b_xy):
+        am = np.zeros((H, W), dtype=np.uint8)
+        bm = np.zeros((H, W), dtype=np.uint8)
+        a_int = np.round(a_xy).astype(np.int32)
+        b_int = np.round(b_xy).astype(np.int32)
+        a_int[:, 0] = np.clip(a_int[:, 0], 0, W - 1)
+        a_int[:, 1] = np.clip(a_int[:, 1], 0, H - 1)
+        b_int[:, 0] = np.clip(b_int[:, 0], 0, W - 1)
+        b_int[:, 1] = np.clip(b_int[:, 1], 0, H - 1)
+        cv.fillPoly(am, [a_int], 1)
+        cv.fillPoly(bm, [b_int], 1)
+        inter = np.logical_and(am, bm).sum()
+        union = np.logical_or(am, bm).sum()
+        return float(inter / union) if union > 0 else 0.0
+
+    # --- run YOLOv8 ---
+    try:
+        res = yolo(roi_padded, verbose=False)[0]
+    except Exception:
+        return pts.astype(np.int32)
+
+    # no masks predicted
+    if not hasattr(res, "masks") or res.masks is None or getattr(res.masks, "xy", None) is None:
+        return pts.astype(np.int32)
+
+    # confidences (aligned with masks)
+    try:
+        confs = res.boxes.conf.detach().cpu().numpy()
+    except Exception:
+        confs = np.ones(len(res.masks.xy), dtype=np.float32)
+
+    # choose IoU function
+    iou_fn = _poly_iou_shapely if use_shapely else _poly_iou_mask
+
+    best_iou = -1.0
+    best_poly = None
+
+    # ground-truth polygon (float)
+    gt = pts
+
+    for i, poly in enumerate(res.masks.xy):
+        if i < len(confs) and confs[i] < conf_thresh:
+            continue
+        if poly is None or len(poly) < 3:
+            continue
+
+        poly = np.asarray(poly, dtype=np.float32)
+
+        # clip to image bounds to be safe (also helps mask fallback)
+        poly[:, 0] = np.clip(poly[:, 0], 0, W - 1)
+        poly[:, 1] = np.clip(poly[:, 1], 0, H - 1)
+
+        iou = iou_fn(gt, poly)
+        if iou > best_iou:
+            best_iou = iou
+            best_poly = poly
+
+    if best_poly is None or best_iou <= 0.0:
+        return pts.astype(np.int32)
+
+    return np.round(best_poly).astype(np.int32)
+
+
+def refine_instance(im, box, cnt, cls):
+    h, w = im.shape[:2]
+    x1, y1, x2, y2 = box
+
+    # --- 1️⃣ Expand box by 10% ---
+    box_w = x2 - x1
+    box_h = y2 - y1
+    expand_x = int(0.1 * box_w)
+    expand_y = int(0.1 * box_h)
+
+    x1_exp, y1_exp = x1 - expand_x, y1 - expand_y
+    x2_exp, y2_exp = x2 + expand_x, y2 + expand_y
+
+    # --- 2️⃣ Clip to image bounds ---
+    x1_clip, y1_clip = max(x1_exp, 0), max(y1_exp, 0)
+    x2_clip, y2_clip = min(x2_exp, w), min(y2_exp, h)
+
+    # --- 3️⃣ Extract valid region ---
+    roi = im[y1_clip:y2_clip, x1_clip:x2_clip]
+
+    # White background of expanded size
+    roi_h, roi_w = y2_exp - y1_exp, x2_exp - x1_exp
+    white_bg = np.ones((roi_h, roi_w, 3), dtype=np.uint8) * 255
+
+    y_offset = y1_clip - y1_exp
+    x_offset = x1_clip - x1_exp
+    white_bg[y_offset:y_offset + roi.shape[0], x_offset:x_offset + roi.shape[1]] = roi
+
+    # --- 4️⃣ Contour alignment (relative to expanded ROI) ---
+    xs, ys = cnt[0], cnt[1]
+    # assert xs.shape == ys.shape, "xs and ys must have same length"
+    cnt_np = np.stack([xs, ys], axis=1).astype(np.float32)
+
+    cnt_np[:, 0] -= x1_exp
+    cnt_np[:, 1] -= y1_exp
+
+    # --- 5️⃣ Scale so that longest side = 1024 ---
+    long_side = max(roi_w, roi_h)
+    scale = 1024 / long_side
+    new_w = int(roi_w * scale)
+    new_h = int(roi_h * scale)
+
+    roi_scaled = cv.resize(white_bg, (new_w, new_h), interpolation=cv.INTER_LINEAR)
+    cnt_scaled = cnt_np * scale
+
+    # --- 6️⃣ Pad to 1024×1024 with white margins ---
+    pad_x = (1024 - new_w) // 2
+    pad_y = (1024 - new_h) // 2
+
+    roi_padded = np.ones((1024, 1024, 3), dtype=np.uint8) * 255
+    roi_padded[pad_y:pad_y + new_h, pad_x:pad_x + new_w] = roi_scaled
+
+    # Adjust contour for padding (relative to final ROI)
+    cnt_final = cnt_scaled + np.array([pad_x, pad_y], dtype=np.float32)
+
+    # --- 7️⃣ Keep mapping info to reconstruct original coordinates ---
+    transform_info = {
+        "x1_exp": x1_exp,
+        "y1_exp": y1_exp,
+        "scale": scale,
+        "pad_x": pad_x,
+        "pad_y": pad_y
+    }
+
+
+    pts = np.array(cnt_final, dtype=np.int32).reshape(-1, 2)
+
+    refined_candidate = best_yolo_contour(pts, roi_padded)
+
+    cv.polylines(roi_padded, [pts], isClosed=True, color=(0, 0, 255), thickness=2)
+    cv.polylines(roi_padded, [refined_candidate], isClosed=True, color=(255, 0, 0), thickness=2)
+    cv.imshow("test", roi_padded)
+    cv.waitKey(-1)
+
+
+
+def refine_file(path):
+    with open(path, 'r') as f:
+        data = json.load(f)
+
+    dir_name = os.path.dirname(path)
+    parent_image_rel_path = data["image_path"]
+    # parent_image_abs_path = os.path.join(dir_name, parent_image_rel_path)
+    parent_image_abs_path = os.path.join( parent_image_rel_path)
+    assert os.path.isfile(parent_image_abs_path), parent_image_abs_path
+
+    im = cv.imread(parent_image_abs_path)
+
+
+    for box, cnt, cls in zip(data["boxes"], data["contours"], data["classes"]):
+        refine_instance(im, box,  cnt, cls)
+
+
+if __name__ == "__main__":
+    # main()
+    from ultralytics import YOLO
+    from ultralytics.engine.results import Results
+    model_file = "flat_bug_S.pt"
+    result_file = "data/metadata_mask-refiner-test_UUID_ChangeThisTEMPORARY.json"
+    yolo = YOLO(model_file, "segment", verbose=True)
+    refine_file(result_file)

From 22acf67a859faa856f5e72ef89479ab4639beb09 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Sun, 12 Oct 2025 13:09:18 +0200
Subject: [PATCH 08/22] improved

---
 prototypes/mask_refiner/fb_refine.py | 176 +++++++++++++--------------
 1 file changed, 85 insertions(+), 91 deletions(-)

diff --git a/prototypes/mask_refiner/fb_refine.py b/prototypes/mask_refiner/fb_refine.py
index f3d6dfd..53e66c5 100644
--- a/prototypes/mask_refiner/fb_refine.py
+++ b/prototypes/mask_refiner/fb_refine.py
@@ -1,8 +1,11 @@
 import json
 import argparse
 import os
+
+import cv2
 import cv2 as cv
 import numpy as np
+from torchgen.gen_functionalization_type import return_from_mutable_noop_redispatch
 
 
 # fixme, resume should continue on the same "run folder"
@@ -23,103 +26,94 @@ def main():
 
 # def best_new_contour(cnt):
 
-def best_yolo_contour(pts, roi_padded, conf_thresh=0.25):
-    """
-    Choose the best YOLOv8 polygon (on roi_padded) by IoU against `pts` (N,2).
-    Uses polygon IoU via shapely if available; otherwise falls back to mask IoU.
-    Returns an (M,2) int32 array of the best polygon in ROI coordinates.
-    """
-
-    # --- get global YOLOv8 model ---
-
-    H, W = roi_padded.shape[:2]
-    pts = np.asarray(pts, dtype=np.float32)
-    if pts.ndim != 2 or pts.shape[1] != 2 or len(pts) < 3:
-        return pts.astype(np.int32)
-
-    # --- try shapely for polygon IoU ---
-    try:
-        from shapely.geometry import Polygon
-        from shapely.errors import TopologicalError
-        use_shapely = True
-    except Exception:
-        use_shapely = False
-
-    # helper: polygon IoU with shapely
-    def _poly_iou_shapely(a_xy, b_xy):
-        try:
-            pa = Polygon(a_xy).buffer(0)  # buffer(0) fixes minor self-intersections
-            pb = Polygon(b_xy).buffer(0)
-            if not pa.is_valid or not pb.is_valid:
-                return 0.0
-            inter = pa.intersection(pb).area
-            union = pa.union(pb).area
-            return float(inter / union) if union > 0 else 0.0
-        except TopologicalError:
-            return 0.0
-
-    # fallback: mask IoU if shapely not available
-    def _poly_iou_mask(a_xy, b_xy):
-        am = np.zeros((H, W), dtype=np.uint8)
-        bm = np.zeros((H, W), dtype=np.uint8)
-        a_int = np.round(a_xy).astype(np.int32)
-        b_int = np.round(b_xy).astype(np.int32)
-        a_int[:, 0] = np.clip(a_int[:, 0], 0, W - 1)
-        a_int[:, 1] = np.clip(a_int[:, 1], 0, H - 1)
-        b_int[:, 0] = np.clip(b_int[:, 0], 0, W - 1)
-        b_int[:, 1] = np.clip(b_int[:, 1], 0, H - 1)
-        cv.fillPoly(am, [a_int], 1)
-        cv.fillPoly(bm, [b_int], 1)
-        inter = np.logical_and(am, bm).sum()
-        union = np.logical_or(am, bm).sum()
-        return float(inter / union) if union > 0 else 0.0
-
-    # --- run YOLOv8 ---
-    try:
-        res = yolo(roi_padded, verbose=False)[0]
-    except Exception:
-        return pts.astype(np.int32)
-
-    # no masks predicted
-    if not hasattr(res, "masks") or res.masks is None or getattr(res.masks, "xy", None) is None:
-        return pts.astype(np.int32)
-
-    # confidences (aligned with masks)
-    try:
-        confs = res.boxes.conf.detach().cpu().numpy()
-    except Exception:
-        confs = np.ones(len(res.masks.xy), dtype=np.float32)
 
-    # choose IoU function
-    iou_fn = _poly_iou_shapely if use_shapely else _poly_iou_mask
 
-    best_iou = -1.0
-    best_poly = None
+def _iou_from_masks(a, b):
+    inter = np.logical_and(a, b).sum()
+    union = np.logical_or(a, b).sum()
+    return float(inter) / float(union) if union > 0 else 0.0
 
-    # ground-truth polygon (float)
-    gt = pts
+def yolo_ensemble_contour(image_bgr, cnt, conf_threshold=0.25, iou_threshold=0.1, mask_threshold=1):
 
-    for i, poly in enumerate(res.masks.xy):
-        if i < len(confs) and confs[i] < conf_thresh:
-            continue
-        if poly is None or len(poly) < 3:
+
+    # get global YOLOv8 model
+
+    H, W = image_bgr.shape[:2]
+    accum = np.zeros((H, W), dtype=np.float32)
+    original_mask = np.zeros_like(accum).astype(np.uint8)
+    cv.fillPoly(original_mask, [cnt.astype(np.int32)], 1)
+    n=4
+    for i in range(n):
+        print(i)
+        if i ==0:
+            image_in = np.copy(image_bgr)
+        elif i == 1:
+            image_in =  cv.rotate(image_bgr, cv.ROTATE_90_CLOCKWISE)
+        elif i == 2:
+            image_in = cv.flip(image_bgr, 0)
+        elif i == 3:
+            image_in = cv.flip(cv.rotate(image_bgr, cv.ROTATE_90_CLOCKWISE), 0)
+        else:
+            raise ValueError("aug_idx must be in {0,1,2,3}")
+
+        # image_in =cv.medianBlur(image_bgr,i*2+1)
+        res = yolo(image_in, verbose=False)[0]
+
+        # no masks predicted this run
+        if not hasattr(res, "masks") or res.masks is None or getattr(res.masks, "data", None) is None:
             continue
 
-        poly = np.asarray(poly, dtype=np.float32)
+        confs = res.boxes.conf.detach().cpu().numpy()
+        masks_t = res.masks.data  # torch.Tensor [num, h, w]
+        num_masks = masks_t.shape[0]
+
+        valid_masks = []
+        mask_areas = []
 
-        # clip to image bounds to be safe (also helps mask fallback)
-        poly[:, 0] = np.clip(poly[:, 0], 0, W - 1)
-        poly[:, 1] = np.clip(poly[:, 1], 0, H - 1)
+        for j in range(num_masks):
+            if confs is not None and j < len(confs) and confs[j] < conf_threshold:
+                continue
 
-        iou = iou_fn(gt, poly)
-        if iou > best_iou:
-            best_iou = iou
-            best_poly = poly
+            m = masks_t[j].detach().cpu().numpy().astype(np.float32)  # float mask (h, w) in [0,1]
 
-    if best_poly is None or best_iou <= 0.0:
-        return pts.astype(np.int32)
+            iou = _iou_from_masks(m, original_mask)
+            print(j, iou)
+            if iou < iou_threshold :
+                continue
+            valid_masks.append(m)
+            mask_areas.append(np.sum(m))
+
+        if len(mask_areas) == 0:
+            continue
 
-    return np.round(best_poly).astype(np.int32)
+        m = valid_masks[np.argmax(mask_areas)]
+
+        if i == 0:
+            m = m
+        elif i == 1:
+            m = cv.rotate(m, cv.ROTATE_90_COUNTERCLOCKWISE)
+        elif i == 2:
+            m = cv.flip(m, 0)
+        elif i == 3:
+            m = cv.rotate(cv.flip(m,0), cv.ROTATE_90_COUNTERCLOCKWISE)
+        else:
+            raise ValueError("aug_idx must be in {0,1,2,3}")
+        print (i,m.shape)
+        accum += m
+
+    frac = accum / float(n)
+    cv2.imshow("test2", frac)
+    final_mask = (frac >= mask_threshold).astype(np.uint8) * 255
+    if np.count_nonzero(final_mask) == 0:
+        return cnt
+    # find largest contour
+    cnts, _ = cv.findContours(final_mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
+    if not cnts:
+        return None
+
+    best = max(cnts, key=cv.contourArea)  # largest by area
+    best = cv.approxPolyDP(best, 0.0005 * cv.arcLength(best, True), True)
+    return best
 
 
 def refine_instance(im, box, cnt, cls):
@@ -129,8 +123,8 @@ def refine_instance(im, box, cnt, cls):
     # --- 1️⃣ Expand box by 10% ---
     box_w = x2 - x1
     box_h = y2 - y1
-    expand_x = int(0.1 * box_w)
-    expand_y = int(0.1 * box_h)
+    expand_x = int(0.2 * box_w)
+    expand_y = int(0.2 * box_h)
 
     x1_exp, y1_exp = x1 - expand_x, y1 - expand_y
     x2_exp, y2_exp = x2 + expand_x, y2 + expand_y
@@ -189,7 +183,7 @@ def refine_instance(im, box, cnt, cls):
 
     pts = np.array(cnt_final, dtype=np.int32).reshape(-1, 2)
 
-    refined_candidate = best_yolo_contour(pts, roi_padded)
+    refined_candidate = yolo_ensemble_contour( roi_padded, pts)
 
     cv.polylines(roi_padded, [pts], isClosed=True, color=(0, 0, 255), thickness=2)
     cv.polylines(roi_padded, [refined_candidate], isClosed=True, color=(255, 0, 0), thickness=2)
@@ -220,6 +214,6 @@ def refine_file(path):
     from ultralytics import YOLO
     from ultralytics.engine.results import Results
     model_file = "flat_bug_S.pt"
-    result_file = "data/metadata_mask-refiner-test_UUID_ChangeThisTEMPORARY.json"
+    result_file = "data/metadata_mask-refiner-test3_UUID_ChangeThisTEMPORARY.json"
     yolo = YOLO(model_file, "segment", verbose=True)
     refine_file(result_file)

From e5d4da1290b04b9023183e915a6dedb89fdbf1b8 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Wed, 13 May 2026 12:03:03 +0200
Subject: [PATCH 09/22] Add YOLOv26 support for training and inference

- Bump ultralytics upper bound to 8.4.49 (required for YOLOv26)
- Update postprocess() in yolo_helpers.py to handle the end2end output
  format introduced by YOLOv26 (post-NMS xyxy detections vs legacy
  pre-NMS xywh anchors); also handle the 8.4.x tuple output wrapping
  for legacy YOLO11 models
- Fix trainers.py to use load_checkpoint (ultralytics >= 8.4 renamed
  attempt_load_one_weight) and yaml_load (replaced by YAML.load)
- Add scripts/training/fb_config_yolo26n.yaml for training with yolo26n-seg

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 pyproject.toml                          |   2 +-
 scripts/training/fb_config_yolo26n.yaml |  11 +++
 src/flat_bug/trainers.py                |  24 +++++-
 src/flat_bug/yolo_helpers.py            | 108 +++++++++++++++++-------
 4 files changed, 108 insertions(+), 37 deletions(-)
 create mode 100644 scripts/training/fb_config_yolo26n.yaml

diff --git a/pyproject.toml b/pyproject.toml
index bfdeba3..095fcf8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,7 +19,7 @@ classifiers = [
 dependencies = [
     "torch>=2.11",
     "torchvision>=0.17.0",
-    "ultralytics>=8.2.16,<=8.3.124",
+    "ultralytics>=8.2.16,<=8.4.49",
     "shapely>=2.0.2",
     "scikit-optimize>=0.10.1",
     "scipy>=1.14.1"
diff --git a/scripts/training/fb_config_yolo26n.yaml b/scripts/training/fb_config_yolo26n.yaml
new file mode 100644
index 0000000..15d98d9
--- /dev/null
+++ b/scripts/training/fb_config_yolo26n.yaml
@@ -0,0 +1,11 @@
+model: "yolo26n-seg.pt"
+batch: 8
+epochs: 5000
+device: "cuda"
+patience: 9999
+workers: 4
+lr0: 0.001
+lrf: 0.00001
+optimizer: "SGD"
+fb_max_instances: 150
+name: "yolo26n"
diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index b6883f1..8f4f4d1 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -15,8 +15,24 @@
 from ultralytics.data.build import InfiniteDataLoader
 from ultralytics.models import yolo
 from ultralytics.models.yolo.segment import SegmentationTrainer
-from ultralytics.nn.tasks import attempt_load_one_weight
-from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK, IterableSimpleNamespace, yaml_load
+try:
+    from ultralytics.nn.tasks import attempt_load_one_weight as _attempt_load_one_weight
+    def _load_checkpoint(model):
+        weights, ckpt = _attempt_load_one_weight(model)
+        return weights, ckpt
+except ImportError:
+    from ultralytics.nn.tasks import load_checkpoint as _load_checkpoint_raw  # ultralytics >= 8.4
+    def _load_checkpoint(model):
+        weights, ckpt = _load_checkpoint_raw(model)
+        return weights, ckpt
+from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK, IterableSimpleNamespace
+try:
+    from ultralytics.utils import yaml_load  # ultralytics < 8.4
+except ImportError:
+    from ultralytics.utils import YAML as _YAML  # ultralytics >= 8.4
+    def yaml_load(file):
+        """Load a YAML file."""
+        return _YAML.load(file)
 from ultralytics.utils.files import increment_path
 from ultralytics.utils.torch_utils import smart_inference_mode, torch_distributed_zero_first
 
@@ -293,8 +309,8 @@ def setup_model(self) -> dict | None:  # noqa: D102
         model, weights = self.model, None
         ckpt = None
         if str(model).endswith('.pt'):
-            weights, ckpt = attempt_load_one_weight(model)
-            if hasattr(ckpt['model'], 'yaml'):
+            weights, ckpt = _load_checkpoint(model)
+            if ckpt is not None and hasattr(ckpt.get('model', None), 'yaml'):
                 cfg = ckpt['model'].yaml
             else:
                 cfg = weights.yaml
diff --git a/src/flat_bug/yolo_helpers.py b/src/flat_bug/yolo_helpers.py
index e04efe1..185ea18 100644
--- a/src/flat_bug/yolo_helpers.py
+++ b/src/flat_bug/yolo_helpers.py
@@ -388,15 +388,32 @@ def scale_boxes(
     return clip_boxes(boxes, img0_shape)
 
 # Revised from ultralytics
+def _is_end2end_output(preds) -> bool:
+    """Return True if preds is from an end2end model (e.g. YOLOv26).
+
+    Three formats exist:
+    - ultralytics < 8.4 legacy: preds[0] = tensor(b, features, n_anchors)
+    - ultralytics >= 8.4 legacy (e.g. YOLO11): preds[0] = (tensor(b, features, n_anchors), protos)
+    - ultralytics >= 8.4 end2end (e.g. YOLOv26): preds[0] = (tensor(b, n_dets, features), protos)
+
+    The end2end case is identified by preds[0][0].shape[1] > preds[0][0].shape[2]
+    (n_dets > n_features), which is the opposite of the legacy layout.
+    """
+    if not isinstance(preds[0], (list, tuple)):
+        return False  # pre-8.4 legacy tensor format
+    det = preds[0][0]
+    return det.shape[1] > det.shape[2]  # end2end: n_dets > n_features
+
+
 def postprocess(
-        preds, 
-        imgs : Sequence[torch.Tensor] | torch.Tensor, 
-        max_det : int=300, 
-        min_confidence : float=0, 
+        preds,
+        imgs : Sequence[torch.Tensor] | torch.Tensor,
+        max_det : int=300,
+        min_confidence : float=0,
         overlap_threshold : float=0.1,
-        overlap_metric : str="IoU", 
-        nms : int=0, 
-        valid_size_range : tuple[int, int] | list[int] | None=None, 
+        overlap_metric : str="IoU",
+        nms : int=0,
+        valid_size_range : tuple[int, int] | list[int] | None=None,
         edge_margin : int | None=None
     ) -> list[Results]:
     """Postprocesses the predictions of the model.
@@ -410,7 +427,7 @@ def postprocess(
         overlap_metric: Overlap metric to use for NMS. Default is "IoU".
         nms: The type of non-maximum suppression to use. Defaults to 0. 0 is no NMS, 1 is standard NMS, 2 is fancy NMS and 3 is mask NMS.
         valid_size_range: The range of valid sizes for the bounding boxes in pixels. Defaults to None (no valid size range).
-        edge_margin: The minimum gap between the edge of the image and the bounding box in pixels for a prediction to be considered valid. 
+        edge_margin: The minimum gap between the edge of the image and the bounding box in pixels for a prediction to be considered valid.
             Defaults to None (no edge margin).
 
     Returns:
@@ -418,34 +435,61 @@ def postprocess(
 
     """
     tile_size = imgs[0].shape[-1]
-    p = preds[0]
-    assert isinstance(p, torch.Tensor)
-    # Convert from xywh to xyxy
-    p[:, :4, :] = torch.cat((
-            p[:, 0:2, :] - p[:, 2:4, :] / 2,  # x_min, y_min
-            p[:, 0:2, :] + p[:, 2:4, :] / 2   # x_max, y_max
-        ),  
-        dim=1)
+
     if min_confidence < 0 or min_confidence > 1:
         raise ValueError("min_confidence must be between 0 and 1.")
-    if min_confidence > 0:
-        num_above_min_conf = (p[:, 4, :] > min_confidence).sum(dim=1)
-        max_det = min(max_det, int(num_above_min_conf.max().item()))
-    # Filter top-`max_det` predictions
-    if max_det != 0:
-        # Filter out the predictions with the lowest confidence
-        p = p.gather(2, torch.argsort(p[:, 4, :], dim=1, descending=True)[:, :max_det].unsqueeze(1).expand(-1, p.size(1), -1))
-
-    # Change shape from (batch, xyxy + cls + masks, n) to (batch, n, xyxy + cls + masks)
-    p = p.transpose(-2, -1)
-    
+
+    if _is_end2end_output(preds):
+        # YOLOv26 / end2end models: preds[0] = (dets[batch, n_dets, 4+1+1+32], protos[batch, 32, h, w])
+        # Boxes are already in xyxy format and post-NMS; conf at dim 4, class at dim 5.
+        p = preds[0][0].clone()  # (batch, n_dets, 38)
+        protos = preds[0][1]     # (batch, 32, h, w)
+        if len(protos.shape) == 3:
+            protos = protos.unsqueeze(0)
+        # Filter zero-conf padding slots and apply min_confidence
+        if min_confidence > 0:
+            conf_mask = p[:, :, 4] > min_confidence
+        else:
+            conf_mask = p[:, :, 4] > 0
+        # Cap to max_det
+        if max_det > 0 and p.shape[1] > max_det:
+            p = p[:, :max_det, :]
+            conf_mask = conf_mask[:, :max_det]
+    else:
+        # Legacy format (ultralytics < 8.4): preds[0] = tensor(batch, features, n_anchors)
+        # Legacy format (ultralytics >= 8.4): preds[0] = (tensor(batch, features, n_anchors), protos)
+        if isinstance(preds[0], (list, tuple)):
+            raw_dets = preds[0][0]
+            protos = preds[0][1]
+        else:
+            raw_dets = preds[0]
+            protos = preds[1][-1]
+        p = raw_dets.clone()
+        assert isinstance(p, torch.Tensor)
+        # Convert from xywh to xyxy
+        p[:, :4, :] = torch.cat((
+                p[:, 0:2, :] - p[:, 2:4, :] / 2,  # x_min, y_min
+                p[:, 0:2, :] + p[:, 2:4, :] / 2   # x_max, y_max
+            ),
+            dim=1)
+        if min_confidence > 0:
+            num_above_min_conf = (p[:, 4, :] > min_confidence).sum(dim=1)
+            max_det = min(max_det, int(num_above_min_conf.max().item()))
+        # Filter top-`max_det` predictions
+        if max_det != 0:
+            p = p.gather(2, torch.argsort(p[:, 4, :], dim=1, descending=True)[:, :max_det].unsqueeze(1).expand(-1, p.size(1), -1))
+        # Change shape from (batch, xyxy + cls + masks, n) to (batch, n, xyxy + cls + masks)
+        p = p.transpose(-2, -1)
+        conf_mask = None
+
+        if len(protos.shape) == 3:
+            protos = protos.unsqueeze(0)
     results = []
-    protos = preds[1][-1]
-    if len(protos.shape) == 3:
-        protos = protos.unsqueeze(0)
     for i, (pred, _) in enumerate(zip(p, range(len(imgs)))):
-        # Remove predictions with a confidence below min_confidence
-        if min_confidence != 0:
+        # Remove empty slots (end2end) or low-confidence predictions (legacy)
+        if conf_mask is not None:
+            pred = pred[conf_mask[i]]
+        elif min_confidence != 0:
             pred = pred[pred[:, 4] > min_confidence]
         boxes = scale_boxes(
             (tile_size, tile_size), 

From a2e7c6dfe64d2cc2e4db8365a031d07a4a3363fb Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Wed, 13 May 2026 12:07:23 +0200
Subject: [PATCH 10/22] Simplify Predictor model loading, fix ultralytics 8.4.x
 compatibility

Replace the internal-API hack (pred.setup_model(self=pred, ...)) with
direct access to yolo.model. The setup_model call in ultralytics 8.4
now requires args.end2end which the dict2attr stub did not provide.
For .pt models the AutoBackend wrapper is unnecessary since flat-bug
preprocesses images itself and only uses PyTorch tensors.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/predictor.py | 20 +-------------------
 1 file changed, 1 insertion(+), 19 deletions(-)

diff --git a/src/flat_bug/predictor.py b/src/flat_bug/predictor.py
index bc23036..65ab0dc 100644
--- a/src/flat_bug/predictor.py
+++ b/src/flat_bug/predictor.py
@@ -1572,25 +1572,7 @@ def __init__(
                     raise FileNotFoundError(f"No such model or file: '{model}'")
             
             yolo = YOLO(model, "segment", verbose=True)
-            pred = yolo._smart_load("predictor")
-            class dict2attr:
-                def __init__(self, d):
-                    self.__dict__ = d
-            args = dict2attr({
-                "device": self._device,
-                "half": self._dtype == torch.float16,
-                "batch": self.BATCH_SIZE,
-                "model": yolo.model,
-                "fp16" : self._dtype == torch.float16,
-                "dnn" : False,
-                # If we want to support multiclass inference, 
-                # this needs to point to "Path to the additional data.yaml file containing class names. Optional." 
-                # see: https://github.com/ultralytics/ultralytics/blob/bc9fd45cdf10ebe8009037aaf8def2353761c9ed/ultralytics/nn/autobackend.py#L53
-                "data" : None
-            })
-            pred.args = args
-            pred.setup_model(self=pred, model=yolo.model, verbose=True)
-            self._model = pred.model
+            self._model = yolo.model
             self._model.to(self._device, dtype=self._dtype)
             self._model.eval()
         elif isinstance(model, torch.nn.Module):

From 75c1305326f3e07d70e7f6b5f4d0ca4f60de3634 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Thu, 14 May 2026 14:27:52 +0200
Subject: [PATCH 11/22] Name training runs as fb_{size}_{timestamp} and update
 N40S config

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 scripts/training/fb_config_N40S.yaml |  5 ++---
 scripts/training/train.sh            | 17 +++++++++++++----
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/scripts/training/fb_config_N40S.yaml b/scripts/training/fb_config_N40S.yaml
index df8076d..c1b09c6 100644
--- a/scripts/training/fb_config_N40S.yaml
+++ b/scripts/training/fb_config_N40S.yaml
@@ -1,7 +1,7 @@
 batch: 8
-model: "./yolov8n-seg.pt"
+#model: "yolov26n.pt"
 epochs: 100
-device: [0, 1]
+device: [0]
 patience: 9999
 lr0: 0.01
 lrf: 0.0001
@@ -11,6 +11,5 @@ fb_max_instances: 9999
 fb_max_images: -1
 fb_custom_eval: false
 fb_custom_eval_num_images: 100
-fb_exclude_datasets: ["00-prospective-ALUS-mixed", "00-prospective-crall2023", "00-prospective-chavez2024", "00-prospective-InsectCV"] # "01-partial-NHM-beetles-crops", "01-partial-Diopsis", "01-partial-abram2023", "01-partial-AMI-traps", "01-partial-gernat2018"
 plots: true
 
diff --git a/scripts/training/train.sh b/scripts/training/train.sh
index bbb7036..1ede1da 100644
--- a/scripts/training/train.sh
+++ b/scripts/training/train.sh
@@ -7,9 +7,18 @@
 # # SBATCH -t 96:00:00 
 
 #CONFIG=fb_config_L40S_fine-tune.yaml
-CONFIG=fb_config_S40S.yaml
-ROOT=/home/altair/flat-bug
-# source ${ROOT}/.venv/bin/activate
-# fb_prepare_data.py -i ${ROOT}/flat-bug-data/pre-pro/  -o ${ROOT}/flat-bug-data/yolo/ -f
+CONFIG=fb_config_N40S.yaml
+
+# Derive model size letter (N/S/M/L/X) from config filename, e.g. fb_config_N40S.yaml -> N
+SIZE="N"
+NAME="fb_${SIZE}_$(date +%Y-%m-%d_%H-%M-%S)"
+
+ROOT=~/Desktop/flatbug-dir/
+fb_clone_data -s ~/flat-bug/repos/scripts/training/.secrets.yaml -o ${ROOT}/flat-bug-data/pre-pro/
+fb_prepare_data -i ${ROOT}/flat-bug-data/pre-pro/  -o ${ROOT}/flat-bug-data/yolo/ -f
+fb_train -c ${CONFIG} -d ${ROOT}/flat-bug-data/yolo/ --name ${NAME}
+
+
+fb_prepare_data.py -i ${ROOT}/flat-bug-data/pre-pro/  -o ${ROOT}/flat-bug-data/yolo/ -f
 # fb_train.py -c ${ROOT}/scripts/training/${CONFIG} -d ${ROOT}/flat-bug-data/yolo/
 fb_train -c ${ROOT}/scripts/training/${CONFIG} -d dev/fb_yolo

From 75e0d97f3e3c1f0c4292eb1ad7fbb7ac37b4eea7 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Mon, 18 May 2026 14:44:40 +0200
Subject: [PATCH 12/22] Default to YOLOv26 for training, fix model names in
 configs

- fb_train.py: change default model from yolov8m-seg.pt to yolo26m-seg.pt
- fb_config_M40S.yaml: update to yolo26m-seg.pt
- fb_config_M40S_GHPC.yaml: fix yolov26m.pt (wrong name/variant) to yolo26m-seg.pt

Pretrained weights are auto-downloaded by ultralytics if not present locally.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 scripts/training/fb_config_M40S.yaml      | 2 +-
 scripts/training/fb_config_M40S_GHPC.yaml | 2 +-
 src/flat_bug/cli/fb_train.py              | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/scripts/training/fb_config_M40S.yaml b/scripts/training/fb_config_M40S.yaml
index 949c438..70aaf28 100644
--- a/scripts/training/fb_config_M40S.yaml
+++ b/scripts/training/fb_config_M40S.yaml
@@ -1,5 +1,5 @@
 batch: 8
-model: "./yolov8m-seg.pt"
+model: "yolo26m-seg.pt"
 epochs: 100
 device: [0, 1]
 patience: 9999
diff --git a/scripts/training/fb_config_M40S_GHPC.yaml b/scripts/training/fb_config_M40S_GHPC.yaml
index e3b840f..4d481be 100644
--- a/scripts/training/fb_config_M40S_GHPC.yaml
+++ b/scripts/training/fb_config_M40S_GHPC.yaml
@@ -1,5 +1,5 @@
 batch: 8
-model: "yolov8m-seg.pt"
+model: "yolo26m-seg.pt"
 epochs: 500
 device: 0
 patience: 9999
diff --git a/src/flat_bug/cli/fb_train.py b/src/flat_bug/cli/fb_train.py
index 9a08906..9ae05ff 100644
--- a/src/flat_bug/cli/fb_train.py
+++ b/src/flat_bug/cli/fb_train.py
@@ -36,7 +36,7 @@ def main():  # noqa: D103
     DEFAULT_CONF = {
         "batch": 8,
         "imgsz": 1024,
-        "model": "yolov8m-seg.pt",
+        "model": "yolo26m-seg.pt",
         "task": "segment",
         "epochs": 5000,
         "device": "cuda",

From 3a818f6244c0b5d39a2fb80df31c297f8ed90b26 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Mon, 18 May 2026 14:51:13 +0200
Subject: [PATCH 13/22] Download pretrained weights explicitly when not present
 locally

ultralytics < 8.4 does not call attempt_download_asset inside
torch_safe_load, so yolo26m-seg.pt (and any other asset-list model)
was never downloaded on machines running 8.3.x. Add an explicit
attempt_download_asset call in setup_model before _load_checkpoint.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/trainers.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index 8f4f4d1..7ef15cf 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -309,6 +309,10 @@ def setup_model(self) -> dict | None:  # noqa: D102
         model, weights = self.model, None
         ckpt = None
         if str(model).endswith('.pt'):
+            if not os.path.exists(model):
+                # ultralytics < 8.4 doesn't auto-download in torch_safe_load, so do it explicitly
+                from ultralytics.utils.downloads import attempt_download_asset
+                model = attempt_download_asset(model)
             weights, ckpt = _load_checkpoint(model)
             if ckpt is not None and hasattr(ckpt.get('model', None), 'yaml'):
                 cfg = ckpt['model'].yaml

From e551822a0addf150e78346705d4bb7ffb9fac6c0 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Tue, 19 May 2026 16:36:03 +0200
Subject: [PATCH 14/22] Free GPU cache before validation to avoid OOM from
 fragmentation

YOLOv26's one2many loss uses significantly more memory during validation
than YOLOv8 did. Call torch.cuda.empty_cache() before each validation
pass to release fragmented reserved memory.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/trainers.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index 7ef15cf..7d58235 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -405,6 +405,7 @@ def validate(self) -> tuple[dict, float]:
 
         """
         if self.epoch % self.save_period == 0 or self._val_metrics is None:
+            torch.cuda.empty_cache()
             metrics, fitness = super().validate()
             self._val_metrics, self._val_fitness = metrics, fitness
 

From b00fba5e7284608e1d74e5aac332d2d296d6c158 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Tue, 19 May 2026 17:44:54 +0200
Subject: [PATCH 15/22] Cap validation batch size at training batch size to
 avoid OOM

ultralytics doubles batch_size for validation (non-OBB tasks), but
YOLOv26's one2many loss makes this OOM on memory-constrained GPUs.
Clamp validation batch_size back to self.args.batch.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/trainers.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index 7d58235..a0fac16 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -385,6 +385,10 @@ def get_dataloader(
         ) -> InfiniteDataLoader:
         """Construct and return dataloader."""
         assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}."
+        if mode == "val":
+            # ultralytics doubles batch_size for validation; YOLOv26's one2many loss makes this OOM,
+            # so clamp back to the training batch size.
+            batch_size = self.args.batch
         with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
             dataset = self.build_dataset(dataset_path, mode, batch_size)
         shuffle = mode == "train"

From 8a19b4b9d7c0b9b6e96765ffea6f653dde624055 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Wed, 20 May 2026 09:27:02 +0200
Subject: [PATCH 16/22] Use batch size 1 for validation to avoid OOM

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/trainers.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index a0fac16..20296c3 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -386,9 +386,7 @@ def get_dataloader(
         """Construct and return dataloader."""
         assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}."
         if mode == "val":
-            # ultralytics doubles batch_size for validation; YOLOv26's one2many loss makes this OOM,
-            # so clamp back to the training batch size.
-            batch_size = self.args.batch
+            batch_size = 1
         with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
             dataset = self.build_dataset(dataset_path, mode, batch_size)
         shuffle = mode == "train"

From 87e5692333f3baec35d024ce5c50dcd87d330a48 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Wed, 20 May 2026 10:56:26 +0200
Subject: [PATCH 17/22] Use PNG compress_level=1 for crop saving to reduce
 write time

Default compression level 6 is slow for large batches of crops.
Level 1 is significantly faster at the cost of slightly larger files.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/predictor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/flat_bug/predictor.py b/src/flat_bug/predictor.py
index 65ab0dc..cd2bdf0 100644
--- a/src/flat_bug/predictor.py
+++ b/src/flat_bug/predictor.py
@@ -1085,7 +1085,7 @@ def _save_1_crop(
         Image.fromarray(
             obj=chw2hwc_uint8(crop, mask).detach().cpu().numpy(),
             mode="RGB" if mask is None else "RGBA"
-        ).save(path)
+        ).save(path, compress_level=1)
         return path
 
     def save_crops(

From 463d7c6e07ad0e380efd5e8620f8d07ba447c0e8 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Thu, 4 Jun 2026 17:03:07 +0200
Subject: [PATCH 18/22] ibid

---
 scripts/training/fb_config_N40S.yaml | 2 +-
 scripts/training/train.sh            | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/scripts/training/fb_config_N40S.yaml b/scripts/training/fb_config_N40S.yaml
index c1b09c6..46e6fbf 100644
--- a/scripts/training/fb_config_N40S.yaml
+++ b/scripts/training/fb_config_N40S.yaml
@@ -1,5 +1,5 @@
 batch: 8
-#model: "yolov26n.pt"
+# model: "yolov26n.pt"
 epochs: 100
 device: [0]
 patience: 9999
diff --git a/scripts/training/train.sh b/scripts/training/train.sh
index 1ede1da..f422ea3 100644
--- a/scripts/training/train.sh
+++ b/scripts/training/train.sh
@@ -14,11 +14,11 @@ SIZE="N"
 NAME="fb_${SIZE}_$(date +%Y-%m-%d_%H-%M-%S)"
 
 ROOT=~/Desktop/flatbug-dir/
-fb_clone_data -s ~/flat-bug/repos/scripts/training/.secrets.yaml -o ${ROOT}/flat-bug-data/pre-pro/
-fb_prepare_data -i ${ROOT}/flat-bug-data/pre-pro/  -o ${ROOT}/flat-bug-data/yolo/ -f
+#fb_clone_data -s ~/flat-bug/repos/scripts/training/.secrets.yaml -o ${ROOT}/flat-bug-data/pre-pro/
+#fb_prepare_data -i ${ROOT}/flat-bug-data/pre-pro/  -o ${ROOT}/flat-bug-data/yolo/ -f
 fb_train -c ${CONFIG} -d ${ROOT}/flat-bug-data/yolo/ --name ${NAME}
 
 
-fb_prepare_data.py -i ${ROOT}/flat-bug-data/pre-pro/  -o ${ROOT}/flat-bug-data/yolo/ -f
+#fb_prepare_data.py -i ${ROOT}/flat-bug-data/pre-pro/  -o ${ROOT}/flat-bug-data/yolo/ -f
 # fb_train.py -c ${ROOT}/scripts/training/${CONFIG} -d ${ROOT}/flat-bug-data/yolo/
-fb_train -c ${ROOT}/scripts/training/${CONFIG} -d dev/fb_yolo
+#fb_train -c ${ROOT}/scripts/training/${CONFIG} -d dev/fb_yolo

From 95704939753d6c321fe0b8a46e2e1a1b1a89eeb3 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Thu, 25 Jun 2026 16:08:19 +0200
Subject: [PATCH 19/22] Suppress validation loss to fix OOM with YOLOv26
 one2many head
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

YOLOv26's one2many head assigns many positive anchors per GT instance.
On dense flat-bug images (hundreds of insects per 1024×1024 crop) this
causes single_mask_loss/crop_mask to materialise a [n_pos, 256, 256]
float32 tensor of ~11.5 GB — OOM even at batch_size=1.

The validation loss is only logged and has no effect on fitness, mAP,
or early stopping. We therefore replace model.loss (on both the live
model and the EMA model) with a no-op returning zeros for the duration
of each validation pass, then restore the class method via del.

This is orthogonal to the earlier batch_size and empty_cache mitigations,
which did not address the per-image root cause.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/trainers.py | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index 20296c3..6b851a4 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -399,7 +399,7 @@ def get_dataloader(
     @smart_inference_mode()
     def validate(self) -> tuple[dict, float]:
         """Run validation on test set using self.validator.
-        
+
         The returned dict is expected to contain "fitness" key.
 
         Returns:
@@ -408,7 +408,32 @@ def validate(self) -> tuple[dict, float]:
         """
         if self.epoch % self.save_period == 0 or self._val_metrics is None:
             torch.cuda.empty_cache()
-            metrics, fitness = super().validate()
+
+            # YOLOv26's one2many head assigns many positive anchors per GT instance.
+            # On dense flat-bug images this causes single_mask_loss / crop_mask to
+            # allocate 10+ GB for mask tensors — OOM even at batch_size=1.
+            # Validation loss is only logged, not used for fitness or early stopping,
+            # so we replace model.loss with a no-op for the duration of the validation pass.
+            ema_model = getattr(getattr(self, "ema", None), "ema", None)
+            _patch_targets = [m for m in [ema_model, self.model] if m is not None and hasattr(m, "loss")]
+            _saved_losses = [(m, m.loss) for m in _patch_targets]
+            _zero = torch.zeros(len(self.loss_names), device=self.device)
+            for m, _ in _saved_losses:
+                m.loss = lambda *a, **kw: (_zero, _zero)
+            LOGGER.warning(
+                "Validation loss suppressed to prevent OOM from YOLOv26 one2many head "
+                "on dense-annotation images. mAP/fitness metrics are unaffected."
+            )
+            try:
+                metrics, fitness = super().validate()
+            finally:
+                for m, orig in _saved_losses:
+                    try:
+                        del m.loss  # remove instance attribute, restoring the class method
+                    except AttributeError:
+                        m.loss = orig
+                torch.cuda.empty_cache()
+
             self._val_metrics, self._val_fitness = metrics, fitness
 
             # Custom end-to-end validation

From 289f8b3584813aa7cb6c8541dfbd8306a9c411f9 Mon Sep 17 00:00:00 2001
From: Quentin Geissmann <qgeissmann@gmail.com>
Date: Fri, 26 Jun 2026 14:50:59 +0200
Subject: [PATCH 20/22] Fix tensor shape mismatch in validation loss no-op
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

YOLOv26's validator initialises self.loss with 5 components; our
no-op was returning torch.zeros(4) from len(self.loss_names).
Return scalar 0 instead — adding a scalar to any tensor is always valid.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/flat_bug/trainers.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index 6b851a4..6383f65 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -417,9 +417,8 @@ def validate(self) -> tuple[dict, float]:
             ema_model = getattr(getattr(self, "ema", None), "ema", None)
             _patch_targets = [m for m in [ema_model, self.model] if m is not None and hasattr(m, "loss")]
             _saved_losses = [(m, m.loss) for m in _patch_targets]
-            _zero = torch.zeros(len(self.loss_names), device=self.device)
             for m, _ in _saved_losses:
-                m.loss = lambda *a, **kw: (_zero, _zero)
+                m.loss = lambda *a, **kw: (0, 0)
             LOGGER.warning(
                 "Validation loss suppressed to prevent OOM from YOLOv26 one2many head "
                 "on dense-annotation images. mAP/fitness metrics are unaffected."

From ce66fab3d046726e2e4e8d5a488d3dc783b04f47 Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Tue, 7 Jul 2026 18:17:05 +0200
Subject: [PATCH 21/22] Compatibility fixes and small modernization changes

---
 .github/workflows/ci.yml       |   2 +-
 .github/workflows/docs.yml     |   1 -
 pyproject.toml                 |  23 +++--
 src/flat_bug/augmentations.py  |  12 +--
 src/flat_bug/cli/fb_predict.py |   2 +-
 src/flat_bug/predictor.py      |  84 ++++------------
 src/flat_bug/trainers.py       |   2 +
 tests/test_dataset.py          |   7 +-
 tests/test_predictor.py        |   2 +-
 uv.lock                        | 177 ++++++++++++++++++++-------------
 10 files changed, 157 insertions(+), 155 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index b2b6416..13e62af 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -43,7 +43,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       
       - name: Install the project
-        run: uv sync --locked --all-extras --group tests
+        run: uv sync --locked --all-extras
 
       - name: Lint with Ruff
         run: uv run ruff check . --ignore E501
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 08918e7..028773b 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -22,7 +22,6 @@ jobs:
       - name: Install dependencies
         run: |
           uv sync --all-extras --group docs
-          uv pip install sphinx sphinx_rtd_theme myst_parser furo
       - name: Sphinx build
         run: |
           cd ./docs
diff --git a/pyproject.toml b/pyproject.toml
index 095fcf8..846d7d0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,14 +19,18 @@ classifiers = [
 dependencies = [
     "torch>=2.11",
     "torchvision>=0.17.0",
-    "ultralytics>=8.2.16,<=8.4.49",
+    "ultralytics>=8.4.0,<=8.4.90",
     "shapely>=2.0.2",
     "scikit-optimize>=0.10.1",
-    "scipy>=1.14.1"
+    "scipy>=1.14.1",
+    "tqdm>=4.0.0",
 ]
 keywords=['deep learning', 'object detection', 'instance segmentation', 'arthropods']
 
 [project.optional-dependencies]
+eval = [
+    "pandas",
+]
 erda = [
     "pyremotedata>=0.0.16"
 ]
@@ -48,20 +52,19 @@ fb_prepare_data = "flat_bug.cli.fb_prepare_data:main"
 fb_clone_data = "flat_bug.cli.fb_clone_data:main"
 
 [dependency-groups]
-docs = [
-    "furo>=2025.12.19",
-    "myst-parser>=5.0.0",
-    "sphinx>=9.0.4",
-]
-notebook = [
+dev = [
     "ipykernel>=7.2.0",
     "ipywidgets>=8.1.8",
     "jupyter>=1.0.0",
-]
-tests = [
     "pytest>=9.0.3",
     "ruff>=0.15.12",
 ]
+docs = [
+    "furo>=2025.12.19",
+    "myst-parser>=5.0.0",
+    "sphinx>=9.0.4",
+    "sphinx_rtd_theme"
+]
 
 [tool.ruff]
 target-version = "py311"
diff --git a/src/flat_bug/augmentations.py b/src/flat_bug/augmentations.py
index d1d2f5c..503e476 100644
--- a/src/flat_bug/augmentations.py
+++ b/src/flat_bug/augmentations.py
@@ -402,6 +402,7 @@ def scale_labels(  # noqa: D103
 
 class FlatBugRandomPerspective(RandomPerspective):  # noqa: D101
     fill_value = (0, 0, 0)
+    size : tuple[int, int]
 
     def __init__(self, imgsz : int, *args, **kwargs):  # noqa: D107
         super().__init__(*args, **kwargs)
@@ -466,10 +467,9 @@ def __call__(self, labels : dict):
             labels: a dict of `bboxes`, `segments`, `keypoints`.
 
         """
-        if self.pre_transform and "mosaic_border" not in labels:
-            labels = self.pre_transform(labels)
-        # labels.pop("ratio_pad", None)  # do not need ratio pad
-
+        # if self.pre_transform and "mosaic_border" not in labels:
+        #     labels = self.pre_transform(labels)
+        labels.pop("ratio_pad", None)  # do not need ratio pad
         img = labels["img"]
         cls = labels["cls"]
         instances : Instances = labels.pop("instances")
@@ -479,7 +479,7 @@ def __call__(self, labels : dict):
         if instances.normalized:
             instances.denormalize(*img.shape[:2][::-1])
 
-        border = labels.pop("mosaic_border", self.border)
+        border = labels.pop("mosaic_border", (0, 0))
         self.size = img.shape[1] + border[1] * 2, img.shape[0] + border[0] * 2  # w, h
         # M is affine matrix
         # Scale for func:`box_candidates`
@@ -495,7 +495,7 @@ def __call__(self, labels : dict):
             bboxes, segments = apply_segments(segments, M)
 
         if keypoints is not None:
-            keypoints = self.apply_keypoints(keypoints, M)
+            keypoints = self.apply_keypoints(keypoints, M, self.size)
         new_instances = Instances(bboxes, segments, keypoints, bbox_format="xyxy", normalized=False)
 
         # Filter instances
diff --git a/src/flat_bug/cli/fb_predict.py b/src/flat_bug/cli/fb_predict.py
index d342766..e671b11 100644
--- a/src/flat_bug/cli/fb_predict.py
+++ b/src/flat_bug/cli/fb_predict.py
@@ -77,7 +77,7 @@ def cli_args():
         help="The result directory"
     )
     args_parse.add_argument(
-        "-w", "--model-weights", type=str, dest="model_weights", default="flat_bug_M.pt",
+        "-w", "--model-weights", type=str, dest="model_weights", default="flat_bug_M_v2.pt",
         help="The .pt file"
     )
     args_parse.add_argument(
diff --git a/src/flat_bug/predictor.py b/src/flat_bug/predictor.py
index cd2bdf0..43dd2db 100644
--- a/src/flat_bug/predictor.py
+++ b/src/flat_bug/predictor.py
@@ -191,9 +191,14 @@ class TensorPredictions:
     mask_height = None
     device = None
     dtype = None
-    CONSTANTS = ("image", "image_path", "device", "dtype", "time", "mask_height", "mask_width", "CONSTANTS",
-                 "BOX_IS_EQUAL_MARGIN",
-                 "PREFER_POLYGONS")  # Attributes that should not be changed after initialization - should 'contours' be here?
+    CONSTANTS = (
+        "image", "image_path", 
+        "device", "dtype", 
+        "time", 
+        "mask_height", "mask_width", 
+        "BOX_IS_EQUAL_MARGIN",
+        "PREFER_POLYGONS"
+    )
 
     def __init__(
             self, 
@@ -219,13 +224,10 @@ def __init__(
             kwargs: Additional configuration arguments.
 
         """
-        # Set option flags
         self.time = time
         start = end = None
 
-        # Timing could probably be hidden in a decorator...
         if self.time and predictions is not None and len(predictions) > 0:
-            # Initialize timing calculations
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
             start.record()
@@ -255,7 +257,6 @@ def __init__(
         else:
             self.device, self.dtype = torch.device("cpu"), torch.float32
 
-        # Set attributes
         self.image_path = image_path
         if image is None:
             if self.image_path is None:
@@ -304,12 +305,12 @@ def _combine_predictions(
         """
         start = end = end_duplication_removal = end_mask_combination = None
         if self.time:
-            # Initialize timing calculations
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
             end_duplication_removal = torch.cuda.Event(enable_timing=True)
             end_mask_combination = torch.cuda.Event(enable_timing=True)
             start.record()
+        
         self.boxes = torch.cat([torch.as_tensor(p.boxes) for p in predictions])  # Nx4
         self.confs = torch.cat([torch.as_tensor(p.confs) for p in predictions])  # N
         self.scales = [p.scale for p in predictions for _ in range(len(p))]  # N
@@ -1132,20 +1133,12 @@ def save_crops(
     @property
     def json_data(self):
         """JSON-compatible dictionary with instance state data."""
-        ## Clean up the data
-        # 1. Convert the boxes to list
         boxes = self.boxes.cpu().tolist()
-        # 2. Convert the masks to contours as lists 
         contours = [c.T.cpu().tolist() for c in self.contours]
-        # 3. Convert the confidences to floats in a list
         confs = self.confs.float().cpu().tolist()
-        # 4. Convert the classes to integers in a list
         classes = self.classes.cpu().long().tolist()
-        # 5. Get the scales (already floats in a list)
         scales = self.scales
-        # 6. Get the areas (already floats in a list)
         areas = self.areas
-        # 7. Get mask data
         mdata = self.masks.data
         return {
             "boxes": boxes,
@@ -1192,14 +1185,12 @@ def serialize(
         else:
             outpath = f"{outpath}{ext}"
 
-        # Add the basename to the outpath
         pt_path = f'{outpath}.pt'
         json_path = f'{outpath}.json'
 
         if save_pt:
             if os.path.exists(pt_path):
                 logger.warning(f"Pickle ({pt_path}) already exists, overwriting!")
-            ### First serialize as .pt file
             torch.save(self, pt_path)
 
         if save_json:
@@ -1330,7 +1321,7 @@ def save(
                 raise ValueError("Unable to save prediction with unknown source file, when `basename` is not supplied.")
             # Get the base name of the image
             basename = os.path.splitext(os.path.basename(self.image_path))[0]
-        # Construct the prediction directory path
+        
         prediction_directory = os.path.join(output_directory, basename)
         # Create the prediction directory if it does not exist and it is needed 
         # (i.e. if we are saving crops, overview, or metadata to a standard location)
@@ -1344,10 +1335,7 @@ def save(
             overview_directory = overview if isinstance(overview, str) else prediction_directory
             os.makedirs(overview_directory, exist_ok=True)
             assert os.path.isdir(overview_directory), RuntimeError(f"Invalid path for overview: {overview_directory}")
-            # The overview path is then constructed as a .jpg file
-            # in the overview directory with the name overview_{base_name}.jpg
             overview_path = os.path.join(overview_directory, f"overview_{basename}_UUID_{identifier}.jpg")
-            # Save the overview image to the overview path
             scale, linewidth = 1, 2
             if fast:
                 scale = min(1 / 2, 3072 / max(self.image.shape[1:]))
@@ -1360,7 +1348,6 @@ def save(
             crop_directory = crops if isinstance(crops, str) else os.path.join(prediction_directory, "crops") 
             os.makedirs(crop_directory, exist_ok=True)
             assert os.path.isdir(crop_directory), RuntimeError(f"Invalid path for crops: {crop_directory}")
-            # Save the crops to the crops path
             self.save_crops(outdir=crop_directory, basename=basename, mask=mask_crops, identifier=identifier)
 
         # Save json
@@ -1369,11 +1356,7 @@ def save(
             metadata_directory = metadata if isinstance(metadata, str) else prediction_directory
             os.makedirs(metadata_directory, exist_ok=True)
             assert os.path.isdir(metadata_directory), RuntimeError(f"Invalid path for metadata: {metadata_directory}")
-            # The metadata path is then constructed as a .json file 
-            # in the metadata directory with the name metadata_{base_name}_id_{identifier}.<EXT>
             metadata_path = os.path.join(metadata_directory, f'metadata_{basename}_UUID_{identifier}')
-            # Serialize the data to the metadata path 
-            # (we don't do this as a future since it is fast, and then we don't need to copy data)
             self.serialize(outpath=metadata_path, identifier=identifier)
 
         if wait:
@@ -1396,46 +1379,37 @@ def _process_batch(
     start_batch_event = end_fetch_event = end_forward_event = \
         end_batch_event = start_batch_event = current_device_stream = None
     if time:
-        # Initialize batch timing calculations
         start_batch_event = torch.cuda.Event(enable_timing=True)
         end_fetch_event = torch.cuda.Event(enable_timing=True)
         end_forward_event = torch.cuda.Event(enable_timing=True)
         end_batch_event = torch.cuda.Event(enable_timing=True)
         current_device_stream = torch.cuda.current_stream(device=device)
-        # Record batch start
         start_batch_event.record(current_device_stream)
     
     # Get the offsets for the current batch and extract and stack the corresponding tiles
     batch = torch.stack([
-                image[:, o[0]: (o[0] + tile_size), o[1]: (o[1] + tile_size)] 
-                for (m, n), o in offsets[batch_start_idx:min((batch_start_idx + batch_size), len(offsets))]
-            ], dim=0)
+        image[:, o[0]: (o[0] + tile_size), o[1]: (o[1] + tile_size)] 
+        for (m, n), o in offsets[batch_start_idx:min((batch_start_idx + batch_size), len(offsets))]
+    ], dim=0)
     if time:
-        # Record end of fetch
         assert current_device_stream is not None and end_fetch_event is not None
         end_fetch_event.record(current_device_stream)
+
     # Forward pass the model on the batch tiles
-    with torch.no_grad():
+    with torch.inference_mode():
         batch_outputs = getattr(model, callback)(batch)
+    
     if time:
-        # Record end of forward
         assert current_device_stream is not None and end_forward_event is not None
         end_forward_event.record(current_device_stream)
-        # Record batch end
         assert current_device_stream is not None and end_batch_event is not None
         end_batch_event.record(current_device_stream)
-
-        # Calculate timing
         torch.cuda.synchronize(device=device)
         assert current_device_stream is not None and start_batch_event is not None and end_fetch_event is not None
         batch_time = start_batch_event.elapsed_time(end_batch_event) / 1000  # Convert to seconds
         fetch_time = start_batch_event.elapsed_time(end_fetch_event) / 1000  # Convert to seconds
         forward_time = end_fetch_event.elapsed_time(end_forward_event) / 1000  # Convert to seconds
-        # loggger.info(
-        #   f'Batch time: {batch_time:.3f}s,'
-        #   f' fetch time: {fetch_time:.3f}s,'
-        #   f' forward time: {forward_time:.3f}s'
-        # )
+
     # Return the postprocessed batch outputs and optionally the timing
     if time:
         return batch, batch_outputs, (batch_time, fetch_time, forward_time)  # type: ignore
@@ -1536,7 +1510,7 @@ class Predictor:
 
     def __init__(
             self, 
-            model : str | pathlib.Path="flat_bug_M.pt", 
+            model : str | pathlib.Path="flat_bug_M_v2.pt", 
             cfg : dict | str | Path | None=None, 
             device : str | torch.device | int | list[str | torch.device | int]=torch.device("cpu"), 
             dtype : torch.types._dtype | str=torch.float32
@@ -1617,7 +1591,6 @@ def _detect_instances(
             this_EDGE_CASE_MARGIN = 0
 
         if self.TIME:
-            # Initialize timing calculations
             start_detect = torch.cuda.Event(enable_timing=True)
             end_detect = torch.cuda.Event(enable_timing=True)
             main_stream = torch.cuda.current_stream(device=self._device)
@@ -1751,27 +1724,19 @@ def _detect_instances(
         MASK_TO_IMG_RATIO = MASK_SIZE / torch.tensor(
             [TILE_SIZE, TILE_SIZE], dtype=torch.float32, device=self._device
         ).unsqueeze(0)
-        # For the boxes, we can simply add the offsets (and possibly subtract the padding)
+        
         box_offsetters = torch.tensor(
             [[o[1][0] - pad_lrtb[2], o[1][1] - pad_lrtb[0]] for o in offsets], 
             dtype=torch.float32, device=self._device
         )
-        # However for the masks, we need to create a new mask which can contain every tile, 
-        #   and then add the masks from each tile to the correct area
-        #   - this will of course use some memory, but it's probably not too bad
-        # Since the masks do not have the same size as the tiles, we need to scale the offsets
-        mask_offsetters = box_offsetters * MASK_TO_IMG_RATIO
-        # We also need to round the offsets, since they may not line up with the pixel-grid
-        # RE: Now they do since I made sure the offsets are multiples of 4
-        mask_offsetters = torch.round(mask_offsetters).long()
-        # The padding must also be scaled and subtracted from the new mask size
+        mask_offsetters = torch.round(box_offsetters * MASK_TO_IMG_RATIO).long()
         new_mask_size = (
             (mask_offsetters.max(dim=0).values + MASK_SIZE) - 
             torch.tensor(pad_lrtb[1::2][::-1], dtype=torch.long, device=self._device) * MASK_TO_IMG_RATIO[0]
         ).tolist()
-        # Finally, we can merge the results - this function basically just does what I described above
         orig_img = image[:, pad_lrtb[2]:(-pad_lrtb[3] if pad_lrtb[3] != 0 else None),
                    pad_lrtb[0]:(-pad_lrtb[1] if pad_lrtb[1] != 0 else None)] if padded else image
+
         merged_results = merge_tile_results(
             results = postprocessed_results, 
             orig_img = orig_img.permute(1, 2, 0), 
@@ -1847,7 +1812,6 @@ def pyramid_predictions(
 
         """
         if self.TIME:
-            # Initialize timing calculations
             start_pyramid, end_pyramid = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
             start_pyramid.record()
 
@@ -1875,7 +1839,6 @@ def pyramid_predictions(
             resize = transforms.Resize((h, w), antialias=True)
             transform_list.append(resize)
 
-        # Check if the image has an integer data type
         if tensor_image.dtype in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]:
             transform_list.append(transforms.ConvertImageDtype(self._dtype))
 
@@ -1892,9 +1855,6 @@ def pyramid_predictions(
                 fill=0,
                 padding_mode='constant'
             )
-            # padding_for_edge_cases = InpaintPad(
-            #   padding=self.EDGE_CASE_MARGIN * edge_case_margin_padding_multiplier
-            # )
             transform_list.append(padding_for_edge_cases)
         else:
             padding_offset[:] = 0
@@ -1903,11 +1863,9 @@ def pyramid_predictions(
             transforms.Compose(transform_list)(tensor_image) if transform_list else tensor_image
         ).to(device=self._device,  dtype=self._dtype)
 
-        # Check correct dimensions
         assert len(transformed_image.shape) == 3, RuntimeError(
             f"transformed_image.shape {transformed_image.shape} != 3"
         ) 
-        # Check correct number of channels
         assert transformed_image.shape[0] == 3, RuntimeError(
             f"transformed_image.shape[0] {transformed_image.shape[0]} != 3. "
             "The image is probably supplied in WxHxC instead of CxWxH, try image.permute(2, 1, 0) before passing it."
diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index 6383f65..506e6b9 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -15,6 +15,7 @@
 from ultralytics.data.build import InfiniteDataLoader
 from ultralytics.models import yolo
 from ultralytics.models.yolo.segment import SegmentationTrainer
+
 try:
     from ultralytics.nn.tasks import attempt_load_one_weight as _attempt_load_one_weight
     def _load_checkpoint(model):
@@ -26,6 +27,7 @@ def _load_checkpoint(model):
         weights, ckpt = _load_checkpoint_raw(model)
         return weights, ckpt
 from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK, IterableSimpleNamespace
+
 try:
     from ultralytics.utils import yaml_load  # ultralytics < 8.4
 except ImportError:
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 455eff6..b8bc75d 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -103,11 +103,8 @@ def create_validation_dataset(args : IterableSimpleNamespace) -> FlatBugYOLOVali
 def _test_plot_batch(batch, ni):  # noqa: D103
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as f:
         plot_images(
-            batch["img"],
-            batch["batch_idx"],
-            batch["cls"].squeeze(-1),
-            batch["bboxes"],
-            masks=batch["masks"],
+            labels=batch,
+            images=batch["img"],
             paths=batch["im_file"],
             fname=f.name,
             on_plot=os.remove,
diff --git a/tests/test_predictor.py b/tests/test_predictor.py
index e094b62..d44bc82 100644
--- a/tests/test_predictor.py
+++ b/tests/test_predictor.py
@@ -14,7 +14,7 @@
 from flat_bug.predictor import Predictor, TensorPredictions
 from tests.remote_lfs_fallback import check_file_with_remote_fallback
 
-TEST_MODEL_NAME = "flat_bug_M.pt"
+TEST_MODEL_NAME = "flat_bug_M_v2.pt"
 PYRAMID_SCALE_BEFORE = 0.6
 ASSET_DIR = os.path.join(os.path.dirname(__file__), "assets")
 ASSET_NAME = "ALUS_Non-miteArachnids_Unknown_2020_11_03_4545"
diff --git a/uv.lock b/uv.lock
index e406a4e..7d8c3cc 100644
--- a/uv.lock
+++ b/uv.lock
@@ -476,7 +476,7 @@ name = "cuda-bindings"
 version = "13.2.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cuda-pathfinder" },
+    { name = "cuda-pathfinder", marker = "sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/e0/a9/3a8241c6e19483ac1f1dcf5c10238205dcb8a6e9d0d4d4709240dff28ff4/cuda_bindings-13.2.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:721104c603f059780d287969be3d194a18d0cc3b713ed9049065a1107706759d", size = 5730273, upload-time = "2026-03-11T00:12:37.18Z" },
@@ -509,37 +509,37 @@ wheels = [
 
 [package.optional-dependencies]
 cublas = [
-    { name = "nvidia-cublas", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-cublas", marker = "sys_platform == 'linux'" },
 ]
 cudart = [
-    { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" },
 ]
 cufft = [
-    { name = "nvidia-cufft", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-cufft", marker = "sys_platform == 'linux'" },
 ]
 cufile = [
     { name = "nvidia-cufile", marker = "sys_platform == 'linux'" },
 ]
 cupti = [
-    { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" },
 ]
 curand = [
-    { name = "nvidia-curand", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-curand", marker = "sys_platform == 'linux'" },
 ]
 cusolver = [
-    { name = "nvidia-cusolver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" },
 ]
 cusparse = [
-    { name = "nvidia-cusparse", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" },
 ]
 nvjitlink = [
-    { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" },
 ]
 nvrtc = [
-    { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" },
 ]
 nvtx = [
-    { name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" },
 ]
 
 [[package]]
@@ -659,6 +659,7 @@ dependencies = [
     { name = "shapely" },
     { name = "torch" },
     { name = "torchvision" },
+    { name = "tqdm" },
     { name = "ultralytics" },
 ]
 
@@ -670,53 +671,56 @@ cloud-datasets = [
 erda = [
     { name = "pyremotedata" },
 ]
+eval = [
+    { name = "pandas" },
+]
 
 [package.dev-dependencies]
-docs = [
-    { name = "furo" },
-    { name = "myst-parser" },
-    { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
-    { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
-]
-notebook = [
+dev = [
     { name = "ipykernel" },
     { name = "ipywidgets" },
     { name = "jupyter" },
-]
-tests = [
     { name = "pytest" },
     { name = "ruff" },
 ]
+docs = [
+    { name = "furo" },
+    { name = "myst-parser" },
+    { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+    { name = "sphinx-rtd-theme" },
+]
 
 [package.metadata]
 requires-dist = [
     { name = "boto3", marker = "extra == 'cloud-datasets'", specifier = ">=1.40" },
     { name = "cvat-sdk", marker = "extra == 'cloud-datasets'", specifier = ">=2.47" },
+    { name = "pandas", marker = "extra == 'eval'" },
     { name = "pyremotedata", marker = "extra == 'erda'", specifier = ">=0.0.16" },
     { name = "scikit-optimize", specifier = ">=0.10.1" },
     { name = "scipy", specifier = ">=1.14.1" },
     { name = "shapely", specifier = ">=2.0.2" },
     { name = "torch", specifier = ">=2.11" },
     { name = "torchvision", specifier = ">=0.17.0" },
-    { name = "ultralytics", specifier = ">=8.2.16,<=8.3.124" },
+    { name = "tqdm", specifier = ">=4.0.0" },
+    { name = "ultralytics", specifier = ">=8.4.0,<=8.4.90" },
 ]
-provides-extras = ["erda", "cloud-datasets"]
+provides-extras = ["eval", "erda", "cloud-datasets"]
 
 [package.metadata.requires-dev]
-docs = [
-    { name = "furo", specifier = ">=2025.12.19" },
-    { name = "myst-parser", specifier = ">=5.0.0" },
-    { name = "sphinx", specifier = ">=9.0.4" },
-]
-notebook = [
+dev = [
     { name = "ipykernel", specifier = ">=7.2.0" },
     { name = "ipywidgets", specifier = ">=8.1.8" },
     { name = "jupyter", specifier = ">=1.0.0" },
-]
-tests = [
     { name = "pytest", specifier = ">=9.0.3" },
     { name = "ruff", specifier = ">=0.15.12" },
 ]
+docs = [
+    { name = "furo", specifier = ">=2025.12.19" },
+    { name = "myst-parser", specifier = ">=5.0.0" },
+    { name = "sphinx", specifier = ">=9.0.4" },
+    { name = "sphinx-rtd-theme" },
+]
 
 [[package]]
 name = "fonttools"
@@ -1807,7 +1811,7 @@ name = "nvidia-cudnn-cu13"
 version = "9.19.0.56"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas" },
+    { name = "nvidia-cublas", marker = "sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" },
@@ -1819,7 +1823,7 @@ name = "nvidia-cufft"
 version = "12.0.0.61"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink" },
+    { name = "nvidia-nvjitlink", marker = "sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
@@ -1849,9 +1853,9 @@ name = "nvidia-cusolver"
 version = "12.0.4.66"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas" },
-    { name = "nvidia-cusparse" },
-    { name = "nvidia-nvjitlink" },
+    { name = "nvidia-cublas", marker = "sys_platform != 'win32'" },
+    { name = "nvidia-cusparse", marker = "sys_platform != 'win32'" },
+    { name = "nvidia-nvjitlink", marker = "sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
@@ -1863,7 +1867,7 @@ name = "nvidia-cusparse"
 version = "12.6.3.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink" },
+    { name = "nvidia-nvjitlink", marker = "sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
@@ -1879,6 +1883,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:25e30a8a7323935d4ad0340b95a0b69926eee755767e8e0b1cf8dd85b197d3fd", size = 169884119, upload-time = "2025-08-13T19:23:41.967Z" },
 ]
 
+[[package]]
+name = "nvidia-ml-py"
+version = "13.610.43"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f0/b5/a8fbc356f768fa5c9cfd646668fd7d34bf55bdd1c6e20754642a64d930d4/nvidia_ml_py-13.610.43.tar.gz", hash = "sha256:65437eb73d68d0c62c931ca4d45038472faff03bd0b8729abba4b899f70d60f2", size = 52109, upload-time = "2026-06-01T18:54:08.829Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl", hash = "sha256:f13c72698edef492f985cc225f14faafe68ae065a2e407f45bdf6f4b9b43fde8", size = 53163, upload-time = "2026-06-01T18:54:07.704Z" },
+]
+
 [[package]]
 name = "nvidia-nccl-cu13"
 version = "2.28.9"
@@ -2146,6 +2159,34 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
 ]
 
+[[package]]
+name = "polars"
+version = "1.42.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "polars-runtime-32" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/27/99/fe77f10a13a778705ef05b499fc708c9a0b0a3680d9eb6bc6e1b6a6b9914/polars-1.42.1.tar.gz", hash = "sha256:2fe94f3059334650bd850ae19a9c165dcd5d9cb12cd95ea04de2201662e70e8a", size = 741532, upload-time = "2026-06-30T04:57:51.504Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/be/6a/edd939cc6fa04b6415aaa9bf19720fc74ead81234b3d38542e0005816d4d/polars-1.42.1-py3-none-any.whl", hash = "sha256:3c0c65cdfa21a621650c4bdcbbccf93964d052fd766c3e70e84a55d961c259fd", size = 837622, upload-time = "2026-06-30T04:56:34.686Z" },
+]
+
+[[package]]
+name = "polars-runtime-32"
+version = "1.42.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/1f/59/15bcc4dac380c6d63efa5446d8317f22671cbd6c9dadd576bd17a334c45a/polars_runtime_32-1.42.1.tar.gz", hash = "sha256:4d4809e1c1b9a6611f6944f27b24abea902b5159e6b6fa262fd716e947af5afd", size = 3045460, upload-time = "2026-06-30T04:57:52.866Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/62/29/16ff6e4e91d71e530d3581f45e342a9cc35072ac6b31dcbc2fa33de2569e/polars_runtime_32-1.42.1-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:bbdc26d68ee5b23b0ce227fa0599220aa35b77c826b6b0a6b2d8e7f6c1c36974", size = 53117325, upload-time = "2026-06-30T04:56:37.972Z" },
+    { url = "https://files.pythonhosted.org/packages/04/8e/4f8296fcfd1347f1351342fecf13bf2430d7efbae2f1f45964ec7930a99e/polars_runtime_32-1.42.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:f6c0288be940b607dc4a7476c01e67fb6bbee93f5f1dd42c64970274c71008ba", size = 47446251, upload-time = "2026-06-30T04:56:41.459Z" },
+    { url = "https://files.pythonhosted.org/packages/88/2e/0d66a7deadc453b890c3391034ca8ab4b05d0beaebbb92a7d65199fba61b/polars_runtime_32-1.42.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:635d9dbcae2302ae223afb395d5cd220bffa61a53d0ab6871d17c8bc830101cf", size = 51359402, upload-time = "2026-06-30T04:56:44.595Z" },
+    { url = "https://files.pythonhosted.org/packages/5d/10/ffb85fa380bc9c9000dc35f40f44954dde49023018501c54faab94b3a39e/polars_runtime_32-1.42.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d059e8e53cc114ff82f9bd791fd341dc53534a2c745e6f6aa37594c3a93f01fe", size = 57302723, upload-time = "2026-06-30T04:56:47.609Z" },
+    { url = "https://files.pythonhosted.org/packages/c2/63/ca50adc62e44224ca5c622a842ba6f35ee87d1d40ef0df7ea2ed6c6edb08/polars_runtime_32-1.42.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:f91f0b13588324905682809d270e1de5f1990c908721c8527657d77a044c9919", size = 51515673, upload-time = "2026-06-30T04:56:50.88Z" },
+    { url = "https://files.pythonhosted.org/packages/63/c3/08fbbf38deaa17bf34a601d327cb7451074098673c78b7c1a8538dde9794/polars_runtime_32-1.42.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8b8bf972d99d48aaaa2582e2bce966a6f43bc815bd8725d15f5cab9e2fb15d17", size = 55217259, upload-time = "2026-06-30T04:56:53.834Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/0e/51db89361668fe077a835fc579277f824ba526e7daf7b94d23d25439e0d0/polars_runtime_32-1.42.1-cp310-abi3-win_amd64.whl", hash = "sha256:e9364c26da389a8b7339e4d29e20a3d12af730247e6ed3b7804bddce2477f428", size = 52715432, upload-time = "2026-06-30T04:56:57.109Z" },
+    { url = "https://files.pythonhosted.org/packages/76/c5/2fb8592d691bd114de25d9c84300b23541dca7060eac11d7b4bed0327786/polars_runtime_32-1.42.1-cp310-abi3-win_arm64.whl", hash = "sha256:7051226e6b42ffc395a7a9190377cd28649fbfb991b8f85c6271f4e1cfb736fb", size = 46718300, upload-time = "2026-06-30T04:56:59.855Z" },
+]
+
 [[package]]
 name = "prometheus-client"
 version = "0.25.0"
@@ -2213,15 +2254,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" },
 ]
 
-[[package]]
-name = "py-cpuinfo"
-version = "9.0.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/37/a8/d832f7293ebb21690860d2e01d8115e5ff6f2ae8bbdc953f0eb0fa4bd2c7/py-cpuinfo-9.0.0.tar.gz", hash = "sha256:3cdbbf3fac90dc6f118bfd64384f309edeadd902d7c8fb17f02ffa1fc3f49690", size = 104716, upload-time = "2022-10-25T20:38:06.303Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" },
-]
-
 [[package]]
 name = "pyaml"
 version = "26.2.1"
@@ -2798,20 +2830,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/07/39/338d9219c4e87f3e708f18857ecd24d22a0c3094752393319553096b98af/scipy-1.17.1-cp314-cp314t-win_arm64.whl", hash = "sha256:200e1050faffacc162be6a486a984a0497866ec54149a01270adc8a59b7c7d21", size = 25489165, upload-time = "2026-02-23T00:22:29.563Z" },
 ]
 
-[[package]]
-name = "seaborn"
-version = "0.13.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "matplotlib" },
-    { name = "numpy" },
-    { name = "pandas" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/86/59/a451d7420a77ab0b98f7affa3a1d78a313d2f7281a57afb1a34bae8ab412/seaborn-0.13.2.tar.gz", hash = "sha256:93e60a40988f4d65e9f4885df477e2fdaff6b73a9ded434c1ab356dd57eefff7", size = 1457696, upload-time = "2024-01-25T13:21:52.551Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/83/11/00d3c3dfc25ad54e731d91449895a79e4bf2384dc3ac01809010ba88f6d5/seaborn-0.13.2-py3-none-any.whl", hash = "sha256:636f8336facf092165e27924f223d3c62ca560b1f2bb5dff7ab7fad265361987", size = 294914, upload-time = "2024-01-25T13:21:49.598Z" },
-]
-
 [[package]]
 name = "send2trash"
 version = "2.1.0"
@@ -2995,6 +3013,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/3c/dd/018ce05c532a22007ac58d4f45232514cd9d6dd0ee1dc374e309db830983/sphinx_basic_ng-1.0.0b2-py3-none-any.whl", hash = "sha256:eb09aedbabfb650607e9b4b68c9d240b90b1e1be221d6ad71d61c52e29f7932b", size = 22496, upload-time = "2023-07-08T18:40:52.659Z" },
 ]
 
+[[package]]
+name = "sphinx-rtd-theme"
+version = "3.1.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "docutils" },
+    { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+    { name = "sphinxcontrib-jquery" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/84/68/a1bfbf38c0f7bccc9b10bbf76b94606f64acb1552ae394f0b8285bfaea25/sphinx_rtd_theme-3.1.0.tar.gz", hash = "sha256:b44276f2c276e909239a4f6c955aa667aaafeb78597923b1c60babc76db78e4c", size = 7620915, upload-time = "2026-01-12T16:03:31.17Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/87/c7/b5c8015d823bfda1a346adb2c634a2101d50bb75d421eb6dcb31acd25ebc/sphinx_rtd_theme-3.1.0-py2.py3-none-any.whl", hash = "sha256:1785824ae8e6632060490f67cf3a72d404a85d2d9fc26bce3619944de5682b89", size = 7655617, upload-time = "2026-01-12T16:03:28.101Z" },
+]
+
 [[package]]
 name = "sphinxcontrib-applehelp"
 version = "2.0.0"
@@ -3022,6 +3055,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/0a/7b/18a8c0bcec9182c05a0b3ec2a776bba4ead82750a55ff798e8d406dae604/sphinxcontrib_htmlhelp-2.1.0-py3-none-any.whl", hash = "sha256:166759820b47002d22914d64a075ce08f4c46818e17cfc9470a9786b759b19f8", size = 98705, upload-time = "2024-07-29T01:09:36.407Z" },
 ]
 
+[[package]]
+name = "sphinxcontrib-jquery"
+version = "4.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+    { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/de/f3/aa67467e051df70a6330fe7770894b3e4f09436dea6881ae0b4f3d87cad8/sphinxcontrib-jquery-4.1.tar.gz", hash = "sha256:1620739f04e36a2c779f1a131a2dfd49b2fd07351bf1968ced074365933abc7a", size = 122331, upload-time = "2023-03-14T15:01:01.944Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/76/85/749bd22d1a68db7291c89e2ebca53f4306c3f205853cf31e9de279034c3c/sphinxcontrib_jquery-4.1-py2.py3-none-any.whl", hash = "sha256:f936030d7d0147dd026a4f2b5a57343d233f1fc7b363f68b3d4f1cb0993878ae", size = 121104, upload-time = "2023-03-14T15:01:00.356Z" },
+]
+
 [[package]]
 name = "sphinxcontrib-jsmath"
 version = "1.0.1"
@@ -3270,28 +3316,25 @@ wheels = [
 
 [[package]]
 name = "ultralytics"
-version = "8.3.124"
+version = "8.4.90"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "matplotlib" },
     { name = "numpy" },
+    { name = "nvidia-ml-py" },
     { name = "opencv-python" },
-    { name = "pandas" },
     { name = "pillow" },
+    { name = "polars" },
     { name = "psutil" },
-    { name = "py-cpuinfo" },
     { name = "pyyaml" },
     { name = "requests" },
-    { name = "scipy" },
-    { name = "seaborn" },
     { name = "torch" },
     { name = "torchvision" },
-    { name = "tqdm" },
     { name = "ultralytics-thop" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/ec/8f/163d06fd8154e7ddf7a4b9c311d68e38d600b187cd84bbc3909322eea611/ultralytics-8.3.124.tar.gz", hash = "sha256:5f49cee43ed03252dfd2e929027ccc357e164435ab5b847de3f0831a429f261f", size = 864356, upload-time = "2025-05-02T18:56:10.723Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/6a/b9/ba51bac3dbaebf3f927c7318a2a89012f45f2316aada396f6ae6f82aa43a/ultralytics-8.4.90.tar.gz", hash = "sha256:5c6d89c253aa87eccfcde58463a1fe89e9fbd8f9854dcbeaacf1581275d0b781", size = 1136950, upload-time = "2026-07-06T19:33:51.04Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/d7/aa/417553f825f7a5c50d9e535b40aaa7edd7869462ab27093775eb6f417e48/ultralytics-8.3.124-py3-none-any.whl", hash = "sha256:fdae7406fffef54829bac09af7948c6052f13260b6f1c544dd1b5f7f6d5a341d", size = 1006359, upload-time = "2025-05-02T18:56:08.149Z" },
+    { url = "https://files.pythonhosted.org/packages/4f/76/70ae5baf72939c4d7297a49e5803df3f8e020f94b911d8287a5bf9a50323/ultralytics-8.4.90-py3-none-any.whl", hash = "sha256:e32b8d0e46cca10c21d27918822f24114a10176424402a127831d0a04cd23cb6", size = 1348787, upload-time = "2026-07-06T19:33:46.392Z" },
 ]
 
 [[package]]

From 6fe0637424d11a785ddb2795890f83f4575fd22b Mon Sep 17 00:00:00 2001
From: asgersvenning <asgersvenning@gmail.com>
Date: Tue, 7 Jul 2026 19:52:03 +0200
Subject: [PATCH 22/22] Update lint to adhere more closely to `ruff format`

---
 pyproject.toml                      |   2 +-
 src/flat_bug/__init__.py            |  24 +-
 src/flat_bug/augmentations.py       | 236 ++++---
 src/flat_bug/cli/__init__.py        |   2 +-
 src/flat_bug/cli/fb_clone_data.py   |  49 +-
 src/flat_bug/cli/fb_eval.py         |  49 +-
 src/flat_bug/cli/fb_predict.py      | 170 ++---
 src/flat_bug/cli/fb_prepare_data.py |  50 +-
 src/flat_bug/cli/fb_train.py        |  68 +-
 src/flat_bug/coco_utils.py          |  50 +-
 src/flat_bug/config.py              |  98 ++-
 src/flat_bug/datasets.py            | 185 +++---
 src/flat_bug/eval_utils.py          | 669 ++++++++++---------
 src/flat_bug/geometric.py           | 244 ++++---
 src/flat_bug/nms.py                 | 665 +++++++++----------
 src/flat_bug/predictor.py           | 960 ++++++++++++++--------------
 src/flat_bug/trainers.py            | 248 ++++---
 src/flat_bug/yolo_helpers.py        | 275 ++++----
 tests/__init__.py                   |   2 +-
 tests/conftest.py                   |   4 +-
 tests/generate_model_outputs.py     |  11 +-
 tests/remote_lfs_fallback.py        |   7 +-
 tests/restore_assets.py             |   1 -
 tests/test_augmentations.py         |  45 +-
 tests/test_config.py                |  24 +-
 tests/test_dataset.py               |  24 +-
 tests/test_predictor.py             |  77 ++-
 27 files changed, 2163 insertions(+), 2076 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 846d7d0..47c4433 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -68,7 +68,7 @@ docs = [
 
 [tool.ruff]
 target-version = "py311"
-line-length = 140
+line-length = 120
 extend-exclude = [
     "utils",
     "scripts",
diff --git a/src/flat_bug/__init__.py b/src/flat_bug/__init__.py
index 5f20dcd..38da2b3 100644
--- a/src/flat_bug/__init__.py
+++ b/src/flat_bug/__init__.py
@@ -9,14 +9,16 @@
 REMOTE_REPOSITORY = "https://anon.erda.au.dk/share_redirect/Bb0CR1FHG6/"
 # GUI access: https://anon.erda.au.dk/cgi-sid/ls.py?share_id=Bb0CR1FHG6
 
+
 # Thanks to: https://stackoverflow.com/a/53877507/19104786
 class DownloadProgressBar(tqdm):  # noqa: D101
-    def update_to(self, b : int=1, bsize : int=1, tsize : int | None=None):  # noqa: D102
+    def update_to(self, b: int = 1, bsize: int = 1, tsize: int | None = None):  # noqa: D102
         if tsize is not None:
             self.total = tsize
         self.update(b * bsize - self.n)
 
-def download_from_repository(url : str, output_path : str | None=None, strict : bool=True, progress : bool=True):
+
+def download_from_repository(url: str, output_path: str | None = None, strict: bool = True, progress: bool = True):
     """Download a file from the flatbug "repository."""
     if output_path is None:
         output_path = url
@@ -24,11 +26,13 @@ def download_from_repository(url : str, output_path : str | None=None, strict :
     tmp_dl_file = output_path + ".tmp"
     try:
         if progress:
-            with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=f'Downloading {url} to {output_path}') as t:
+            with DownloadProgressBar(
+                unit="B", unit_scale=True, miniters=1, desc=f"Downloading {url} to {output_path}"
+            ) as t:
                 urllib.request.urlretrieve(url, filename=tmp_dl_file, reporthook=t.update_to)
         else:
-                urllib.request.urlretrieve(url,  filename=tmp_dl_file)
-        
+            urllib.request.urlretrieve(url, filename=tmp_dl_file)
+
         os.rename(tmp_dl_file, output_path)
 
     except Exception as e:
@@ -42,16 +46,16 @@ def download_from_repository(url : str, output_path : str | None=None, strict :
 
     return True
 
-# TODO: Improve this perhaps using https://gist.github.com/aldur/f356f245014523330a7070ab12bcfb1f, 
+
+# TODO: Improve this perhaps using https://gist.github.com/aldur/f356f245014523330a7070ab12bcfb1f,
 # as I have done in PyRemoteData https://github.com/asgersvenning/pyremotedata/blob/f0e3506c1abe2bb20106ffa2a1c3fc0f380f3dd8/src/pyremotedata/__init__.py
 logging.basicConfig(
-    level=logging.WARNING, 
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', 
-    datefmt='%Y-%m-%d %H:%M:%S'
+    level=logging.WARNING, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
 )
 
 logger = logging.getLogger(__name__)
 
+
 def set_log_level(level):  # noqa: D103
     logger.setLevel(level)
-    logger.info(f'Log level set to {level}')
+    logger.info(f"Log level set to {level}")
diff --git a/src/flat_bug/augmentations.py b/src/flat_bug/augmentations.py
index 503e476..38d125f 100644
--- a/src/flat_bug/augmentations.py
+++ b/src/flat_bug/augmentations.py
@@ -1,4 +1,5 @@
 """Augmentations used for flatbug."""
+
 import math
 import random
 from typing import cast, overload
@@ -16,17 +17,13 @@
 
 
 ### From Ultralytics repository, remove clipping from `RandomPerspective` and add `apply_segments` function
-def segment2box(
-        segment : torch.Tensor, 
-        width : int=640, 
-        height : int=640
-    ) -> np.ndarray:
+def segment2box(segment: torch.Tensor, width: int = 640, height: int = 640) -> np.ndarray:
     """Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy).
 
     Args:
         segment: the segment label
         width: OBS: Unused. The width of the image. Defaults to 640.
-        height: OBS: Unused. The height of the image. Defaults to 640. 
+        height: OBS: Unused. The height of the image. Defaults to 640.
 
     Returns:
         The minimum and maximum x and y values of the segment (xyxy).
@@ -35,10 +32,8 @@ def segment2box(
     x, y = segment.T  # segment xy
     return np.array([x.min(), y.min(), x.max(), y.max()], dtype=segment.dtype)  # type: ignore
 
-def apply_segments(
-        segments : np.ndarray, 
-        M : np.ndarray
-    ) -> tuple[np.ndarray, np.ndarray]:
+
+def apply_segments(segments: np.ndarray, M: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     """Apply affine to segments and generate new bboxes from segments.
 
     Args:
@@ -66,14 +61,11 @@ def apply_segments(
     segments[..., 1] = segments[..., 1].clip(bboxes[:, 1:2], bboxes[:, 3:4])
     return bboxes, segments
 
-def low_res_inpaint(
-        img : np.ndarray, 
-        mask : np.ndarray, 
-        scale : int=6
-    ) -> np.ndarray:
+
+def low_res_inpaint(img: np.ndarray, mask: np.ndarray, scale: int = 6) -> np.ndarray:
     """Perform low resolution inpainting in-place.
-    
-    In-painting is done on a low-resolution copy of the image, 
+
+    In-painting is done on a low-resolution copy of the image,
     and then copies the upsampled inpainted image back into the original image.
 
     Args:
@@ -90,25 +82,26 @@ def low_res_inpaint(
     lr_mask = cv2.resize(mask, (mask.shape[1] // scale, mask.shape[0] // scale))
 
     # Perform inpainting on the low-res image
-    lr_inpainted = cv2.inpaint(lr_img, lr_mask, inpaintRadius=7, flags=cv2.INPAINT_TELEA)    
+    lr_inpainted = cv2.inpaint(lr_img, lr_mask, inpaintRadius=7, flags=cv2.INPAINT_TELEA)
 
     # Copy the upsampled inpainted image back into the original image
     img[mask == 1] = cv2.resize(lr_inpainted, (img.shape[1], img.shape[0]))[mask == 1]
 
     return img
 
+
 def telea_inpaint_polys(
-        img : np.ndarray, 
-        polys : list[np.ndarray], 
-        exclude_polys : list[np.ndarray] | None=None, 
-        downscale_factor : int | float=6, 
-        **kwargs
-    ) -> np.ndarray:
+    img: np.ndarray,
+    polys: list[np.ndarray],
+    exclude_polys: list[np.ndarray] | None = None,
+    downscale_factor: int | float = 6,
+    **kwargs,
+) -> np.ndarray:
     """Mutably inpaints the polygons in an image using the Fast Marching method by Alexandru Telea.
 
-    The inpainting algorithm is performed on a downsampled version of the image to speed up the process, 
+    The inpainting algorithm is performed on a downsampled version of the image to speed up the process,
     and the inpainted results are then upsampled and pasted back into the original image.
-    
+
     Args:
         img: The image to inpaint.
         polys: A list of polygons to inpaint.
@@ -122,14 +115,14 @@ def telea_inpaint_polys(
     """
     # Type checking and sanitizing
     check_types(img, np.ndarray)
-    if not ((img.ndim == 3 and img.shape[2] < 5) or img.ndim == 2): 
+    if not ((img.ndim == 3 and img.shape[2] < 5) or img.ndim == 2):
         raise ValueError(f"img must be a 2D or 3D numpy array, of shape (H, W) or (H, W, C), got shape {img.shape}")
     check_types(polys, [list, np.ndarray])
     check_types(exclude_polys, ([list, np.ndarray], None))  # type: ignore
     if exclude_polys is None:
         exclude_polys = []
     check_types(downscale_factor, (int, float))
-    
+
     # Early return on no-op
     if len(polys) == 0:
         return img
@@ -152,9 +145,15 @@ def telea_inpaint_polys(
             **kwargs
         )
     
-    # Dilate the inpaint bitmap to ensure that the inpainting doesn't bleed from the edges of the instances under the polygons
-    cv2.dilate(src=inpaint_bitmap, dst=inpaint_bitmap, kernel=np.ones((3, 3), np.uint8), iterations=1)
-    
+    # Dilate the inpaint bitmap to ensure that the inpainting doesn't bleed
+    # from the edges of the instances under the polygons
+    cv2.dilate(
+        src=inpaint_bitmap,
+        dst=inpaint_bitmap,
+        kernel=np.ones((3, 3), np.uint8),
+        iterations=1
+    )
+
     # Inpaint the low-res image using the Fast Marching algorithm
     cv2.inpaint(
         src=lr_img,
@@ -176,27 +175,21 @@ def telea_inpaint_polys(
     # Upsample the inpainted image and bitmap
     inpaint_bitmap = cv2.resize(inpaint_bitmap, orig_shape)
     lr_img = cv2.resize(lr_img, orig_shape)
-    
+
     # Copy the inpainted low-res image back into the original image
     img[inpaint_bitmap == 1] = lr_img[inpaint_bitmap == 1]
 
     # Return the inpainted image (not necessary, as the inpainting is done in-place)
     return img
 
+
 @overload
-def inpaint_pad(
-        array : torch.Tensor, 
-        padding : int | tuple[int, int] | tuple[int, int, int, int]
-    ) -> torch.Tensor: ...
+def inpaint_pad(array: torch.Tensor, padding: int | tuple[int, int] | tuple[int, int, int, int]) -> torch.Tensor: ...
 @overload
-def inpaint_pad(
-        array : np.ndarray, 
-        padding : int | tuple[int, int] | tuple[int, int, int, int]
-    ) -> np.ndarray: ...
+def inpaint_pad(array: np.ndarray, padding: int | tuple[int, int] | tuple[int, int, int, int]) -> np.ndarray: ...
 def inpaint_pad(  # noqa: D103
-        array : torch.Tensor | np.ndarray, 
-        padding : int | tuple[int, int] | tuple[int, int, int, int]
-    ) -> torch.Tensor | np.ndarray:
+    array: torch.Tensor | np.ndarray, padding: int | tuple[int, int] | tuple[int, int, int, int]
+) -> torch.Tensor | np.ndarray:
     # Ensure padding is a tuple (pad_top, pad_bottom, pad_left, pad_right)
     if isinstance(padding, int):
         padding = (padding, padding, padding, padding)
@@ -210,14 +203,14 @@ def inpaint_pad(  # noqa: D103
     pad_t, pad_b, pad_l, pad_r = padding
     if pad_t == 0 and pad_b == 0 and pad_l == 0 and pad_r == 0:
         return array
-    
-    # Convert to integer whc numpy array 
+
+    # Convert to integer whc numpy array
     is_tensor = isinstance(array, torch.Tensor)
     if is_tensor:
         device = array.device
         dtype = array.dtype
         array = array.cpu().numpy()
-    
+
     # If array is not a integer multiply by 255
     array_is_integer = np.issubdtype(array.dtype, np.integer)
     if not array_is_integer:
@@ -243,7 +236,7 @@ def inpaint_pad(  # noqa: D103
     # mask[pad_h:pad_h + original_h, pad_w:pad_w + original_w] = 0
     h_slice = slice(pad_t, pad_t + original_h)
     w_slice = slice(pad_l, pad_l + original_w)
-    padded_image[h_slice, w_slice] = array # <-- HERE
+    padded_image[h_slice, w_slice] = array  # <-- HERE
     mask[h_slice, w_slice] = 0
 
     # Perform inpainting
@@ -256,22 +249,21 @@ def inpaint_pad(  # noqa: D103
         padded_image = padded_image.astype(np.float32) / 255
     if is_tensor:
         padded_image = torch.tensor(padded_image).to(device, dtype)
-    
+
     return padded_image
 
+
 class InpaintPad:  # noqa: D101
-    def __init__(self, padding : int | tuple[int, int] | tuple[int, int, int, int]):  # noqa: D107
+    def __init__(self, padding: int | tuple[int, int] | tuple[int, int, int, int]):  # noqa: D107
         self.padding = padding
 
-    def __call__(self, tensor : torch.Tensor) -> torch.Tensor:  # noqa: D102
+    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:  # noqa: D102
         return inpaint_pad(tensor, self.padding)
-    
+
+
 def remove_instances(  # noqa: D103
-        labels : dict, 
-        area_thr : float | int=1, 
-        max_targets : float | int | None=1000, 
-        min_size : int=0
-    ) -> dict:
+    labels: dict, area_thr: float | int = 1, max_targets: float | int | None = 1000, min_size: int = 0
+) -> dict:
     instances = cast(Instances, labels.pop("instances"))
     assert instances.segments is not None
     imsize = labels["img"].shape[:2][::-1]
@@ -284,8 +276,9 @@ def remove_instances(  # noqa: D103
     bboxes = instances._bboxes.bboxes
 
     if bboxes.shape[0] == 0:
-        labels["instances"] = Instances(np.empty([0, 4], dtype=np.float32), np.empty([0, 2], dtype=np.float32),
-                                        normalized=False)
+        labels["instances"] = Instances(
+            np.empty([0, 4], dtype=np.float32), np.empty([0, 2], dtype=np.float32), normalized=False
+        )
         labels["cls"] = np.empty((0), dtype=np.int32)
         return labels
 
@@ -296,7 +289,7 @@ def remove_instances(  # noqa: D103
     for i, s in enumerate(instances.segments):
         # Initiate overlap using bounding box
         x, y, w, h = bboxes[i]
-        bbox = box(x - w/2, y - h/2, x + w/2, y + h/2)  # type: ignore
+        bbox = box(x - w / 2, y - h / 2, x + w / 2, y + h / 2)  # type: ignore
         bbox_iarea = bbox.intersection(image_bbox).area
         area_ratios[i] = (bbox_iarea + eps) / (bbox.area + eps) if bbox.area > 0 and bbox_iarea > 0 else 0
         if area_ratios[i] < area_thr:
@@ -307,15 +300,15 @@ def remove_instances(  # noqa: D103
 
     valid = np.all(
         [
-            #(b[:, 0] - b[:, 2] / 2) / self._imsize > 0,
-            #(b[:, 1] - b[:, 3] / 2) / self._imsize > 0,
-            #(b[:, 0] + b[:, 2] / 2) / self._imsize < 1,
-            #(b[:, 1] + b[:, 3] / 2) / self._imsize < 1,
+            # (b[:, 0] - b[:, 2] / 2) / self._imsize > 0,
+            # (b[:, 1] - b[:, 3] / 2) / self._imsize > 0,
+            # (b[:, 0] + b[:, 2] / 2) / self._imsize < 1,
+            # (b[:, 1] + b[:, 3] / 2) / self._imsize < 1,
             area_ratios >= area_thr,
             bboxes[:, 2] > min_size,
-            bboxes[:, 3] > min_size
+            bboxes[:, 3] > min_size,
         ],
-       axis=0
+        axis=0,
     )
 
     if max_targets is not None and np.sum(valid) > max_targets:
@@ -327,7 +320,7 @@ def remove_instances(  # noqa: D103
     # 10% outside is flagged as NOT insect!
 
     invalid = np.bitwise_not(valid)
-    invalid_visible = np.bitwise_and(invalid, area_ratios > 0) # We only need to inpaint polygons within the frame
+    invalid_visible = np.bitwise_and(invalid, area_ratios > 0)  # We only need to inpaint polygons within the frame
 
     invalid_i = np.nonzero(invalid_visible)[0]
     invalid_segments = instances.segments[invalid_i]
@@ -345,28 +338,29 @@ def remove_instances(  # noqa: D103
         #     thickness=-1,
         #     lineType=cv2.LINE_4,
         #     offset=(0,0)
-        # ) 
+        # )
         # Up-to-date inpainting method
         telea_inpaint_polys(
-            img=labels["img"], 
-            polys=invalid_segments, 
+            img=labels["img"],
+            polys=invalid_segments,
             exclude_polys=valid_segments,
-            downscale_factor=6, 
+            downscale_factor=6,
             contourIdx=-1,
             thickness=-1,
             lineType=cv2.LINE_4,
-            offset=(0, 0)
+            offset=(0, 0),
         )
 
     # cv2.imwrite(f"/tmp/{os.path.basename(labels['im_file'])}", or_img)
     valid_i = np.nonzero(valid)[0]
 
     if len(valid_i) == 0:
-        labels["instances"] = Instances(np.empty([0, 4], dtype=np.float32), np.empty([0, 2], dtype=np.float32),
-                                        normalized=False)
+        labels["instances"] = Instances(
+            np.empty([0, 4], dtype=np.float32), np.empty([0, 2], dtype=np.float32), normalized=False
+        )
         labels["cls"] = np.empty_like(labels["cls"])
         return labels
-    
+
     # DEBUG: plot boxes on image
     # for bbox in bboxes[valid_i, :]:
     #     x, y, w, h = bbox
@@ -385,10 +379,10 @@ def remove_instances(  # noqa: D103
     # logger.info(labels)
     return labels
 
+
 def scale_labels(  # noqa: D103
-        labels : dict, 
-        scale : float
-    ) -> dict:
+    labels: dict, scale: float
+) -> dict:
     orig_shape = labels["img"].shape[:2]
     # Scale the image
     labels["img"] = cv2.resize(labels["img"], (0, 0), fx=scale, fy=scale)
@@ -400,19 +394,16 @@ def scale_labels(  # noqa: D103
     labels["instances"].denormalize(*new_shape[::-1])
     return labels
 
+
 class FlatBugRandomPerspective(RandomPerspective):  # noqa: D101
     fill_value = (0, 0, 0)
-    size : tuple[int, int]
+    size: tuple[int, int]
 
-    def __init__(self, imgsz : int, *args, **kwargs):  # noqa: D107
+    def __init__(self, imgsz: int, *args, **kwargs):  # noqa: D107
         super().__init__(*args, **kwargs)
         self.imgsz = imgsz
 
-    def affine_transform(
-            self, 
-            img : np.ndarray, 
-            border : tuple[int, int]
-        ) -> tuple[np.ndarray, np.ndarray, float]:
+    def affine_transform(self, img: np.ndarray, border: tuple[int, int]) -> tuple[np.ndarray, np.ndarray, float]:
         """Center."""
         self.scale = self.imgsz / max(img.shape), 1  # fime hardcoded
         C = np.eye(3, dtype=np.float32)
@@ -455,12 +446,12 @@ def affine_transform(
             else:  # affine
                 img = cv2.warpAffine(img, M[:2], dsize=self.size, borderValue=self.fill_value)
                 img_transform_mask = cv2.warpAffine(img_transform_mask, M[:2], dsize=self.size, borderValue=1)
-        
+
         low_res_inpaint(img, img_transform_mask, scale=6)
 
         return img, M, s
-    
-    def __call__(self, labels : dict):
+
+    def __call__(self, labels: dict):
         """Affine images and targets.
 
         Args:
@@ -472,7 +463,7 @@ def __call__(self, labels : dict):
         labels.pop("ratio_pad", None)  # do not need ratio pad
         img = labels["img"]
         cls = labels["cls"]
-        instances : Instances = labels.pop("instances")
+        instances: Instances = labels.pop("instances")
         # Make sure the coord formats are right
         if instances._bboxes.format != "xyxy":
             instances.convert_bbox(format="xyxy")
@@ -518,13 +509,14 @@ def __call__(self, labels : dict):
         # labels["instances"].normalize(*labels["resized_shape"][::-1])
         return labels
 
+
 class Crop:
     """Abstact crop-related augmentation."""
 
     bg_fill = (0, 0, 0)
     min_size = 0  # px
 
-    def __init__(self, imsize : int | tuple[int, int] | list[int] | np.ndarray):  # noqa: D107
+    def __init__(self, imsize: int | tuple[int, int] | list[int] | np.ndarray):  # noqa: D107
         if isinstance(imsize, int):
             self._imsize = (imsize, imsize)
         elif isinstance(imsize, (tuple, list, np.ndarray)):
@@ -534,18 +526,13 @@ def __init__(self, imsize : int | tuple[int, int] | list[int] | np.ndarray):  #
                 raise ValueError("imsize should be a list of length 2")
             self._imsize = imsize
         else:
-            raise TypeError(f'`imsize` should be of type `int`, `tuple`, or `list`, got {type(imsize)}')
+            raise TypeError(f"`imsize` should be of type `int`, `tuple`, or `list`, got {type(imsize)}")
         self._imsize = tuple([int(i) for i in self._imsize])
         self.xsize, self.ysize = self._imsize
 
     def crop_image(  # noqa: D102
-            self, 
-            labels : dict, 
-            start_x : int, 
-            start_y : int, 
-            size_x : int, 
-            size_y : int
-        ) -> dict:
+        self, labels: dict, start_x: int, start_y: int, size_x: int, size_y: int
+    ) -> dict:
         img = labels["img"]
         orig_shape = img.shape
         h, w = img.shape[:2]
@@ -562,13 +549,13 @@ def crop_image(  # noqa: D102
         py = py0 + py1
         n_size_y = size_y - py
 
-        img = img[n_start_y: n_start_y + n_size_y, n_start_x: n_start_x + n_size_x, :]
+        img = img[n_start_y : n_start_y + n_size_y, n_start_x : n_start_x + n_size_x, :]
 
         if px > 0 or py > 0:
-            img = np.pad(img, pad_width=((py0, py1), (px0, px1), (0, 0)), mode="constant", constant_values=0.)
+            img = np.pad(img, pad_width=((py0, py1), (px0, px1), (0, 0)), mode="constant", constant_values=0.0)
             # img = inpaint_pad(img, (py0, py1, px0, px1)) # Fixme: this is very slow for large images
 
-        if img.shape != (size_x, size_y, 3): 
+        if img.shape != (size_x, size_y, 3):
             logger.info("shape:", img.shape)
             logger.info("or-shape", orig_shape)
             logger.info("x, y:", start_x, start_y)
@@ -588,13 +575,13 @@ def crop_image(  # noqa: D102
         if instances.normalized:
             instances.denormalize(*orig_shape[:2][::-1])
 
-        labels['ratio_pad'] = ((1.0, 1.0), (0.0, 0.0))
+        labels["ratio_pad"] = ((1.0, 1.0), (0.0, 0.0))
         x_offset = -n_start_x + px0
         y_offset = -n_start_y + py0
 
         # positions in the cropped image
         instances._bboxes.add([x_offset, y_offset, 0, 0])
-        
+
         assert instances.segments is not None
         for s in instances.segments:
             s[:, 0] += x_offset
@@ -603,16 +590,17 @@ def crop_image(  # noqa: D102
         labels["instances"] = instances
 
         return labels
-    
+
     def __call__(self, x):
         """Abstract function.
-        
+
         Should be implemented in subclasses.
         """
         raise NotImplementedError("This method should be implemented in a subclass")
 
+
 class CenterCrop(Crop):  # noqa: D101
-    def __call__(self, labels : dict) -> dict:  # noqa: D102
+    def __call__(self, labels: dict) -> dict:  # noqa: D102
         h, w = labels["img"].shape[:2]
 
         start_x = (w - self.xsize) // 2
@@ -620,11 +608,12 @@ def __call__(self, labels : dict) -> dict:  # noqa: D102
 
         return self.crop_image(labels, start_x, start_y, self.xsize, self.ysize)
 
+
 class RandomCrop(Crop):  # noqa: D101
     def __init__(self, *args, **kwargs):  # noqa: D107
         super().__init__(*args, **kwargs)
 
-    def __call__(self, labels : dict) -> dict:  # noqa: D102
+    def __call__(self, labels: dict) -> dict:  # noqa: D102
         # Get the initial image to target crop size ratio
         h, w = labels["img"].shape[:2]
         target_source_ratio_h = self.ysize / h
@@ -637,7 +626,7 @@ def __call__(self, labels : dict) -> dict:  # noqa: D102
             scale = np.random.uniform(min_target_source_ratio, 1) ** 2
         # If the image is smaller than the target size we scale between 1, and crop_dim/image_dim
         else:
-            scale = np.random.uniform(1, min_target_source_ratio) ** (1/2)
+            scale = np.random.uniform(1, min_target_source_ratio) ** (1 / 2)
 
         # When we scale up, this is done before cropping
         do_scale_before = scale > 1
@@ -647,10 +636,10 @@ def __call__(self, labels : dict) -> dict:  # noqa: D102
         else:
             target_size = max(int(w * scale), int(h * scale))
             target_xsize, target_ysize = target_size, target_size
-            # Reset the scale such that when the labels/image are 
+            # Reset the scale such that when the labels/image are
             # scaled after cropping the size is self.xsize, self.ysize (assuming these are equal)
             scale = self.xsize / target_xsize
-        
+
         # Calculate possible crop start positions
         h, w = labels["img"].shape[:2]
         if w <= target_xsize:
@@ -669,30 +658,26 @@ def __call__(self, labels : dict) -> dict:  # noqa: D102
 
         return labels
 
+
 class FixInstances:
-    """A callable class that removes instances that are too small or which overlap less than a certain threshold with the image."""
+    """Removes instances that are too small or which overlap less than a certain threshold with the image."""
 
-    def __init__(
-            self, 
-            area_thr : float | int, 
-            max_targets : int | float | None, 
-            min_size : int
-        ):
-        """"Instantiate.
+    def __init__(self, area_thr: float | int, max_targets: int | float | None, min_size: int):
+        """.
 
         Args:
             area_thr: The minimum proportion of the instance that must be within the image in order for it to be kept.
-            max_targets: The maximum number of instances to keep. If there are more instances than this, 
+            max_targets: The maximum number of instances to keep. If there are more instances than this,
                 a random subset of instances will be kept. If `None`, all instances will be kept.
-            min_size: The minimum size of the bounding box of the instance. 
+            min_size: The minimum size of the bounding box of the instance.
                 Instances with a width or height less than this value will be removed.
 
         """
         self.area_thr = area_thr
         self.max_targets = max_targets if max_targets is None or max_targets > 0 else None
         self.min_size = min_size
-    
-    def __call__(self, labels : dict) -> dict:
+
+    def __call__(self, labels: dict) -> dict:
         """Fix instances.
 
         Args:
@@ -704,8 +689,9 @@ def __call__(self, labels : dict) -> dict:
         """
         return remove_instances(labels, area_thr=self.area_thr, max_targets=self.max_targets, min_size=self.min_size)
 
+
 class RandomColorInv:  # noqa: D101
-    def __init__(self, p : float=0.5):
+    def __init__(self, p: float = 0.5):
         """Invert the colors of an image with a probability p.
 
         Args:
@@ -720,9 +706,9 @@ def __init__(self, p : float=0.5):
             p = 1
         self.p = 1 - p
 
-    def __call__(self, labels : dict) -> dict:  # noqa: D102
-        img = labels['img']
+    def __call__(self, labels: dict) -> dict:  # noqa: D102
+        img = labels["img"]
         if random.uniform(0, 1) > self.p:
             assert img.dtype == np.uint8
-            labels['img'] = 255 - img
+            labels["img"] = 255 - img
         return labels
diff --git a/src/flat_bug/cli/__init__.py b/src/flat_bug/cli/__init__.py
index eb24194..c4cac9c 100644
--- a/src/flat_bug/cli/__init__.py
+++ b/src/flat_bug/cli/__init__.py
@@ -1 +1 @@
-"""flatbug public CLI API scripts."""
\ No newline at end of file
+"""flatbug public CLI API scripts."""
diff --git a/src/flat_bug/cli/fb_clone_data.py b/src/flat_bug/cli/fb_clone_data.py
index 20148b6..6a06991 100644
--- a/src/flat_bug/cli/fb_clone_data.py
+++ b/src/flat_bug/cli/fb_clone_data.py
@@ -15,6 +15,7 @@
 
 try:
     from tqdm import tqdm
+
     _HAVE_TQDM = True
 except Exception:
     _HAVE_TQDM = False
@@ -39,17 +40,17 @@
 """
 
 
-
-
 # TARGET_DIR = Path("/home/quentin/Desktop/flat-bug/flat-bug-data/pre-pro")
 FORMAT_NAME = "COCO 1.0"
 # ------------------ Secrets ------------------
 
+
 # ------------------ Load secrets from YAML ------------------
 def load_secrets_yaml(path):
     with open(path) as f:
         return yaml.safe_load(f)
 
+
 # ------------------ Helpers ------------------
 def safe_segment(name: str) -> str:
     """Make a filesystem-safe folder name (keep common chars; replace others with underscore)."""
@@ -58,6 +59,7 @@ def safe_segment(name: str) -> str:
     # avoid empty folder names
     return cleaned or "unnamed_task"
 
+
 def md5_file(path: Path, chunk=1024 * 1024) -> str:
     """Compute MD5 hex digest of a file (for ETag comparison if single-part)."""
     h = hashlib.md5()
@@ -69,9 +71,10 @@ def md5_file(path: Path, chunk=1024 * 1024) -> str:
             h.update(b)
     return h.hexdigest()
 
+
 def task_is_completed(task) -> bool:
     """Check if a task is done.
-    
+
     Consider a task completed if either:
         - task.status == 'completed', OR
         - all of its jobs are in state == 'completed'
@@ -88,6 +91,7 @@ def task_is_completed(task) -> bool:
     except Exception:
         return False
 
+
 def build_s3_client(s3_access_key, s3_secret_key, s3_region, s3_endpoint):
     session = boto3.session.Session(
         aws_access_key_id=s3_access_key,
@@ -103,6 +107,7 @@ def build_s3_client(s3_access_key, s3_secret_key, s3_region, s3_endpoint):
         ),
     )
 
+
 def list_s3_objects_with_prefix(s3, bucket: str, prefix: str):
     """Yield dicts with 'Key', 'Size', 'ETag' (no quotes), and 'LastModified'."""
     continuation = None
@@ -127,13 +132,16 @@ def list_s3_objects_with_prefix(s3, bucket: str, prefix: str):
         else:
             break
 
+
 def ensure_parent(path: Path):
     path.parent.mkdir(parents=True, exist_ok=True)
+
+
 #
 # ------------------ COCO Export ------------------
 def export_coco_annotations_for_task(task, output_json_path: Path, s3_prefix):
     """Export COCO for task.
-    
+
     Export task dataset (COCO 1.0, annotations only) to a temp zip,
     then extract the COCO annotations json to output_json_path.
     """
@@ -144,7 +152,6 @@ def export_coco_annotations_for_task(task, output_json_path: Path, s3_prefix):
         FORMAT_NAME,
         filename=str(tmp_zip),
         include_images=False,
-
     )
 
     # Find the annotations JSON inside the zip (usually 'annotations/instances_default.json')
@@ -157,7 +164,11 @@ def export_coco_annotations_for_task(task, output_json_path: Path, s3_prefix):
         except KeyError:
             for zi in zf.infolist():
                 name = zi.filename.replace("\\", "/")
-                if name.lower().startswith("annotations/") and name.lower().endswith(".json") and "instance" in name.lower():
+                if (
+                    name.lower().startswith("annotations/")
+                    and name.lower().endswith(".json")
+                    and "instance" in name.lower()
+                ):
                     info = zi
                     break
         if info is None:
@@ -176,7 +187,6 @@ def export_coco_annotations_for_task(task, output_json_path: Path, s3_prefix):
         with zf.open(info, "r") as src:
             coco = json.load(src)
 
-
         for im in coco.get("images", []):
             orig = im.get("file_name", "")
             # Make sure we don't accidentally duplicate prefixes
@@ -197,6 +207,7 @@ def _iter_s3_keys(s3, bucket, prefix):
         for obj in page.get("Contents", []):
             yield obj["Key"]
 
+
 def sync_s3_prefix_to_local(
     s3,
     bucket: str,
@@ -218,7 +229,7 @@ def sync_s3_prefix_to_local(
     # 1) Index upstream
     upstream = {}
     for obj in list_s3_objects_with_prefix(s3, bucket, prefix):
-        rel = obj["Key"][len(prefix):].lstrip("/")
+        rel = obj["Key"][len(prefix) :].lstrip("/")
         if not rel:
             continue
         upstream[rel] = obj
@@ -258,7 +269,7 @@ def sync_s3_prefix_to_local(
                         # Single-part ETag is MD5 (no '-')
                         if remote_etag and "-" not in remote_etag:
                             try:
-                                need_dl = (md5_file(dest) != remote_etag)
+                                need_dl = md5_file(dest) != remote_etag
                             except Exception:
                                 need_dl = True
                         else:
@@ -303,14 +314,12 @@ def sync_s3_prefix_to_local(
                 side.unlink(missing_ok=True)
 
     return True
-def _process_one_task(task_id: int,
-                      task_name: str,
-                      cfg_cvat: dict,
-                      s3,
-                      s3_bucket: str,
-                      s3_prefix_root: str,
-                      target_dir: Path):
-    """Runs in a thread. Returns (task_id, ok, msg).""" # noqa: D401
+
+
+def _process_one_task(
+    task_id: int, task_name: str, cfg_cvat: dict, s3, s3_bucket: str, s3_prefix_root: str, target_dir: Path
+):
+    """Runs in a thread. Returns (task_id, ok, msg)."""  # noqa: D401
     CVAT_HOST = cfg_cvat.get("host", "https://app.cvat.ai")
     USERNAME = cfg_cvat["username"]
     PASSWORD = cfg_cvat["password"]
@@ -344,7 +353,7 @@ def _process_one_task(task_id: int,
         with make_client(host=CVAT_HOST, credentials=(USERNAME, PASSWORD)) as client:
             if ORG_SLUG:
                 client.organization_slug = ORG_SLUG
-            t = client.tasks.retrieve(task_id)   # get fresh task handle
+            t = client.tasks.retrieve(task_id)  # get fresh task handle
             t.fetch()
             if not task_is_completed(t):
                 return task_id, False, "Skipping: not completed"
@@ -357,24 +366,22 @@ def _process_one_task(task_id: int,
     except Exception as e:
         return task_id, False, f"error: {e}"
 
+
 # ------------------ Main ------------------
 def main():
     args_parse = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
-
     args_parse.add_argument(
         "-s", "--secrets-file", dest="secrets_file",
         help=(
             "A YAML files containing credentials for s3 and cvat. It has the following structure"
             f"{secrets_structure}"
     ))
-
     args_parse.add_argument(
         "-o", "--output-dir", dest="output_dir",
         help=(
             "The output directory where all subdatasets are stored. "
             "Each subdirectory is a coco dataset, with a JSON file and a list of images"
     ))
-
     args_parse.add_argument(
         "-f", "--force", dest="delete_target_before",
         help="Delete output directory before, this avoids duplicating data etc",
diff --git a/src/flat_bug/cli/fb_eval.py b/src/flat_bug/cli/fb_eval.py
index 792f174..da73804 100644
--- a/src/flat_bug/cli/fb_eval.py
+++ b/src/flat_bug/cli/fb_eval.py
@@ -45,14 +45,16 @@
 # ruff: disable[D103]
 
 
-def load_json(file : str):
+def load_json(file: str):
     with open(file) as f:
         return json.load(f)
 
+
 # Wrapper function to call compare_groups with a single parameter dictionary for multiprocessing
 def process_image(params):
     return compare_groups(**params)
 
+
 def main():
     # # Development defaults
     # predictions = "dev/**/**.json"
@@ -128,7 +130,7 @@ def main():
     pred_coco = filter_coco(pred_coco, confidence=confidence_threshold, area=min_size, verbose=False)
 
     if not os.path.exists(args.ground_truth):
-        raise ValueError(f'Ground truth file not found: {args.ground_truth}')
+        raise ValueError(f"Ground truth file not found: {args.ground_truth}")
     gt_coco = load_json(args.ground_truth)
     gt_coco = filter_coco(gt_coco, area=min_size)
     gt_annotations, pred_annotations = split_annotations(gt_coco), split_annotations(pred_coco)
@@ -141,33 +143,35 @@ def main():
     shared_keys = gt_keys.intersection(pred_keys)
     if len(gt_diff_keys) > 0:
         show = min(2, len(gt_diff_keys))
-        missing_gt_diff_formatted = ', '.join(['"' + str(i) + '"' for i in gt_diff_keys[:show]])
+        missing_gt_diff_formatted = ", ".join(['"' + str(i) + '"' for i in gt_diff_keys[:show]])
         logger.info(
-            f'Ground truth has {len(gt_diff_keys)} images that are not in the predictions:'
-            f'[{missing_gt_diff_formatted}{", ..." if len(gt_diff_keys) > show else ""}] and {len(gt_diff_keys) - show} more'
+            f"Ground truth has {len(gt_diff_keys)} images that are not in the predictions:"
+            f"[{missing_gt_diff_formatted}{', ...' if len(gt_diff_keys) > show else ''}] "
+            f"and {len(gt_diff_keys) - show} more"
         )
     if len(pred_diff_keys) > 0:
         show = min(2, len(pred_diff_keys))
-        missing_pred_diff_formatted = ', '.join(['"' + str(i) + '"' for i in pred_diff_keys[:show]])
+        missing_pred_diff_formatted = ", ".join(['"' + str(i) + '"' for i in pred_diff_keys[:show]])
         logger.info(
-            f'Predictions has {len(pred_diff_keys)} images that are not in the ground truth:'
-            f'[{missing_pred_diff_formatted} {", ..." if len(pred_diff_keys) > show else ""}] and {len(pred_diff_keys) - show} more'
+            f"Predictions has {len(pred_diff_keys)} images that are not in the ground truth:"
+            f"[{missing_pred_diff_formatted} {', ...' if len(pred_diff_keys) > show else ''}] "
+            f"and {len(pred_diff_keys) - show} more"
         )
     if len(shared_keys) == 0:
-        raise ValueError('No images in common between the ground truth and the predictions')
+        raise ValueError("No images in common between the ground truth and the predictions")
 
     shared_keys = sorted(shared_keys)
     if args.n != -1:
-        logger.info(f'Skipping the evaluation of {len(shared_keys) - args.n} images')
-        shared_keys = shared_keys[:args.n]
+        logger.info(f"Skipping the evaluation of {len(shared_keys) - args.n} images")
+        shared_keys = shared_keys[: args.n]
     if len(shared_keys) == 0:
-        raise ValueError('No images to evaluate')
+        raise ValueError("No images to evaluate")
     if len(shared_keys) < args.workers:
         args.workers = min(args.workers, len(shared_keys))
         logger.info(f"Warning: More workers than images, reducing the number of workers to {args.workers}")
-    
+
     result_files = []
-    
+
     if args.workers <= 1:
         for image in tqdm(shared_keys, desc="Evaluating images", dynamic_ncols=True):
             result_files += [process_image(
@@ -184,9 +188,12 @@ def main():
             )]
     else:
         from multiprocessing import Pool
+
         pool = Pool(args.workers)
         all_params = []
-        for image in tqdm(shared_keys, desc="Generating parameters for multiprocessing", dynamic_ncols=True, leave=False):
+        for image in tqdm(
+            shared_keys, desc="Generating parameters for multiprocessing", dynamic_ncols=True, leave=False
+        ):
             this_params = {
                 "group1"            : gt_annotations[image], 
                 "group2"            : pred_annotations[image], 
@@ -201,8 +208,10 @@ def main():
             }
             all_params.append(this_params)
         for matches in tqdm(
-            pool.imap_unordered(process_image, all_params), 
-            total=len(shared_keys), desc="Evaluating images", dynamic_ncols=True
+            pool.imap_unordered(process_image, all_params),
+            total=len(shared_keys),
+            desc="Evaluating images",
+            dynamic_ncols=True,
         ):
             result_files += [matches]
         pool.close()
@@ -216,8 +225,10 @@ def read_and_add_new_column(f):
             df = pd.read_csv(f, sep=";")
             df.insert(0, "image", os.path.splitext(os.path.basename(f))[0])
             return df
+
         combined_result = pd.concat([read_and_add_new_column(f) for f in result_files])
-        combined_result.to_csv(f"{args.output_directory}{os.sep}combined_results.csv", index=False,sep=";")
+        combined_result.to_csv(f"{args.output_directory}{os.sep}combined_results.csv", index=False, sep=";")
+
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/src/flat_bug/cli/fb_predict.py b/src/flat_bug/cli/fb_predict.py
index e671b11..798e560 100644
--- a/src/flat_bug/cli/fb_predict.py
+++ b/src/flat_bug/cli/fb_predict.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 r"""Inference CLI script for ``flatbug``.
 
-A comprehensive CLI API for ``flatbug`` inference with support for hyperparameter configuration, 
+A comprehensive CLI API for ``flatbug`` inference with support for hyperparameter configuration,
 flexible input parsing, output format specification, and hardware specification.
 
 Usage:
@@ -16,7 +16,7 @@
     -w MODEL_WEIGHTS, --model-weights MODEL_WEIGHTS
                         The .pt file
     -p INPUT_PATTERN, --input-pattern INPUT_PATTERN
-                        The pattern to match the images. 
+                        The pattern to match the images.
                         Default is '[^/]*\.([jJ][pP][eE]{0,1}[gG]|[pP][nN][gG])$' i.e. jpg/jpeg/png case-insensitive.
     -n MAX_IMAGES, --max-images MAX_IMAGES
                         Maximum number of images to process. Default is None. Truncates in alphabetical order.
@@ -65,7 +65,7 @@ def cli_args():
         description="""\
             Perform instance detection and segmentation with flatbug on
             one or more images or a video.""",
-        formatter_class=argparse.RawTextHelpFormatter
+        formatter_class=argparse.RawTextHelpFormatter,
     )
 
     args_parse.add_argument(
@@ -81,10 +81,10 @@ def cli_args():
         help="The .pt file"
     )
     args_parse.add_argument(
-        "-p", "--input-pattern", type=str, dest="input_pattern", default=r"[^/]*\.([jJ][pP][eE]{0,1}[gG]|[pP][nN][gG])$",
+        "-p", "--input-pattern", type=str, dest="input_pattern", default=r"[^/]*\.([jJ][pP][eE]?[gG]|[pP][nN][gG])$",
         help=(
             "The pattern to match the images. "
-            r"Default is '[^/]*\.([jJ][pP][eE]{0,1}[gG]|[pP][nN][gG])$' i.e. jpg/jpeg/png case-insensitive."
+            r"Default is '[^/]*\.([jJ][pP][eE]?[gG]|[pP][nN][gG])$' i.e. jpg/jpeg/png case-insensitive."
     ))
     args_parse.add_argument(
         "-n", "--max-images", type=int, dest="max_images", default=None,
@@ -153,41 +153,42 @@ def cli_args():
     args = args_parse.parse_args()
     return vars(args)
 
+
 def predict(
-        input : str,
-        output_dir : str,
-        model_weights : str,
-        input_pattern : str=r"[^/]*\.([jJ][pP][eE]{0,1}[gG]|[pP][nN][gG])$",
-        max_images : int | None=None,
-        recursive : bool=False,
-        scale_before : float=1.0,
-        single_scale : bool=False,
-        nms_metric : str="IoU",
-        device : str="auto",
-        dtype : str=None,
-        fast : bool=False,
-        config : str | None=None,
-        id : str | None=None,
-        no_crops : bool=False,
-        no_overviews : bool=False,
-        no_metadata : bool=False,
-        only_overviews : bool=False,
-        long_format : bool=False,
-        no_save : bool=False,
-        no_compiled_coco : bool=False,
-        verbose : bool=False
-    ):
+    input: str,
+    output_dir: str,
+    model_weights: str,
+    input_pattern: str = r"[^/]*\.([jJ][pP][eE]{0,1}[gG]|[pP][nN][gG])$",
+    max_images: int | None = None,
+    recursive: bool = False,
+    scale_before: float = 1.0,
+    single_scale: bool = False,
+    nms_metric: str = "IoU",
+    device: str = "auto",
+    dtype: str = None,
+    fast: bool = False,
+    config: str | None = None,
+    id: str | None = None,
+    no_crops: bool = False,
+    no_overviews: bool = False,
+    no_metadata: bool = False,
+    only_overviews: bool = False,
+    long_format: bool = False,
+    no_save: bool = False,
+    no_compiled_coco: bool = False,
+    verbose: bool = False,
+):
     if verbose:
         set_log_level("DEBUG")
-    
+
     torch.set_float32_matmul_precision("medium")
-    
+
     logger.debug(f"OPTIONS: {locals()}")
 
     # Sanitize paths
     isVideo = False
     isERDA = input.startswith("erda://")
-    if not isERDA: 
+    if not isERDA:
         input = os.path.normpath(input)
     output_dir = os.path.normpath(output_dir)
     model_weights = os.path.normpath(model_weights)
@@ -196,6 +197,7 @@ def predict(
 
     if isERDA:
         from pyremotedata.implicit_mount import IOHandler, RemotePathIterator
+
         logger.debug("Assuming directory exists on ERDA")
     else:
         _, ext = os.path.splitext(input)
@@ -211,7 +213,7 @@ def predict(
         else:
             device = "cpu"
             logger.info("CUDA not available, using CPU")
-    
+
     if not torch.cuda.is_available() and "cuda" in device:
         raise ValueError(f"Device(s) '{device}' is/are not available.")
     # Detect if multi-gpu, either by comma or semicolon
@@ -237,7 +239,7 @@ def predict(
         else:
             dtype = "float16"
     dtype = dtype
-    
+
     if config is not None:
         config = read_cfg(config)
     else:
@@ -246,10 +248,10 @@ def predict(
         config["TIME"] = device_type == "cuda"
     if nms_metric is not None:
         config["OVERLAP_METRIC"] = nms_metric
-    
+
     if id is None:
         id = str(uuid.uuid4())
-    
+
     crops = not no_crops
     metadata = not no_metadata
     if no_overviews:
@@ -287,11 +289,11 @@ def predict(
         "licenses": [],
         "images": [],
         "annotations": [],
-        "categories": [categories]  # Your category
+        "categories": [categories],  # Your category
     }
     if isERDA:
         input = input.removeprefix("erda://")
-        io = IOHandler(verbose = False, clean = False)
+        io = IOHandler(verbose=False, clean=False)
         io.start()
         io.cd(input)
         # Check for local file index
@@ -302,44 +304,49 @@ def predict(
             io.cache["file_index"] = file_index
 
         file_iter = RemotePathIterator(
-            io_handler = io,
+            io_handler=io,
             # These are basically network-performance parameters
             # How many files to download at once (larger is faster, but more memory intensive)
-            batch_size = 64,
-            # How many files are downloaded in parallel in during each batch (10 seems to be optimal for my connection, 
+            batch_size=64,
+            # How many files are downloaded in parallel in during each batch (10 seems to be optimal for my connection,
             # this is probably dependent on the amount of cores on the server)
-            batch_parallel = 10,
-            # This relates to how much pre-fetching is done, i.e. how many batches are queued before the download is paused.
-            # This can be as large as you want, the larger the less stuttering you will have, but requires more local *disk* (NOT RAM) space
-            max_queued_batches = 3,
-            # This is parameter basically does the same as the one above, 
-            # but it really needs to larger than batch_size * max_queued_batches, 
-            # otherwise files will be deleted before they are used (This *will* result in an error). 
+            batch_parallel=10,
+            # This relates to how much pre-fetching is done,
+            # i.e. how many batches are queued before the download is paused.
+            # This can be as large as you want, the larger the less stuttering you will have,
+            # but requires more local *disk* (NOT RAM) space
+            max_queued_batches=3,
+            # This is parameter basically does the same as the one above,
+            # but it really needs to larger than batch_size * max_queued_batches,
+            # otherwise files will be deleted before they are used (This *will* result in an error).
             # This parameter should probably be removed from the `pyRemoteData` package...
-            n_local_files = 100 * 3 * 2,
+            n_local_files=100 * 3 * 2,
             # Are local files temporary? I.e. should they be deleted after use?
-            # TODO: This should also cause the previous argument to be ignored, and **never** delete files before internally
-            clear_local = False,
+            # TODO: This should also cause the previous argument to be ignored,
+            # and **never** delete files before internally
+            clear_local=False,
             # These parameters are all related to file-indexing and filtering on the remote server
-            # Should the file-index be re-generated? (has to be False if store is False - otherwise an error will be thrown)
-            override = False,
-            # This is important if we do not want to add files to the remote server (i.e. we only want to read them), 
-            # if this is True, then the function will "cache" the file list in 
+            # Should the file-index be re-generated?
+            # (has to be False if store is False - otherwise an error will be thrown)
+            override=False,
+            # This is important if we do not want to add files to the remote server (i.e. we only want to read them),
+            # if this is True, then the function will "cache" the file list in
             # the directory in a file in the remote directory called ".file_index.txt"
-            store = False,
+            store=False,
             # r"^[^\/\.]+(\.jpg$|\.png$|\.jpeg$|\.JPG$|\.PNG$|\.JPEG)$",
-            # # TODO: Currently as a hack, we skip files in subdirectories 
+            # # TODO: Currently as a hack, we skip files in subdirectories
             # i.e. files with a '/' in their name, this is not ideal, as they are still read from the remote server
-            pattern = input_pattern
+            pattern=input_pattern,
         )
     elif isVideo:
         import tempfile
 
         import cv2
+
         tmp_frame_dir = tempfile.TemporaryDirectory()
         video_output_path = os.path.join(output_dir, os.path.splitext(os.path.basename(input))[0] + ".mp4")
         cap = cv2.VideoCapture(input)
-        fps = cap.get(cv2.CAP_PROP_FPS) 
+        fps = cap.get(cv2.CAP_PROP_FPS)
         frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         duration = frame_count / fps
         video_shape = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
@@ -347,7 +354,7 @@ def predict(
         # Make progress bar that shows the progress in time
         pbar = tqdm(total=duration, desc="Reading video frames", dynamic_ncols=True, unit="s")
         while cap.isOpened():
-            pbar.update(1/fps)
+            pbar.update(1 / fps)
             ret, frame = cap.read()
             if not ret:
                 break
@@ -361,14 +368,16 @@ def predict(
         if not no_save and overviews:
             # Create a video writer
             if fast:
-                video_shape = (video_shape[0]//2, video_shape[1]//2)
-            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+                video_shape = (video_shape[0] // 2, video_shape[1] // 2)
+            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
             video_writer = cv2.VideoWriter(video_output_path, fourcc, fps, video_shape)
     else:
         if os.path.isfile(input):
             file_iter = [input]
         else:
-            file_iter = sorted([f for f in glob.glob(os.path.join(input, "**"), recursive=recursive) if re.search(input_pattern, f)])
+            file_iter = sorted(
+                [f for f in glob.glob(os.path.join(input, "**"), recursive=recursive) if re.search(input_pattern, f)]
+            )
     if max_images is not None:
         if isERDA:
             file_iter.subset(list(range(min(max_images, len(file_iter)))))
@@ -376,7 +385,7 @@ def predict(
             file_iter = file_iter[:max_images]
 
     all_json_results = []
-    
+
     pbar = tqdm(enumerate(file_iter), total=len(file_iter), desc="Processing images", dynamic_ncols=True, unit="image")
     for i, f in pbar:
         if isERDA:
@@ -389,28 +398,29 @@ def predict(
         pbar.set_postfix_str(f"Processing {os.path.basename(f)}")
         try:
             # Run the model
-            prediction = pred.pyramid_predictions(f, scale_increment=2/3, scale_before=scale_before, single_scale=single_scale)
+            prediction = pred.pyramid_predictions(
+                f, scale_increment=2 / 3, scale_before=scale_before, single_scale=single_scale
+            )
             # Save the results
             if not no_save:
                 result_directory = prediction.save(
-                    output_directory = output_dir,
-                    fast = fast,
-                    overview = overviews,
-                    metadata = metadata,
-                    crops = crops,
-                    mask_crops = True,
-                    identifier = id,
+                    output_directory=output_dir,
+                    fast=fast,
+                    overview=overviews,
+                    metadata=metadata,
+                    crops=crops,
+                    mask_crops=True,
+                    identifier=id,
                 )
                 if result_directory is not None:
                     basename = os.path.splitext(os.path.basename(f))[0]
                     metadata_directory = metadata if isinstance(metadata, str) else result_directory
                     overview_directory = overviews if isinstance(overviews, str) else result_directory
                     # crop_directory = crops if isinstance(crops, str) else os.path.join(result_directory, crops)
-                    all_json_results.append(os.path.join(metadata_directory, f'metadata_{basename}_UUID_{id}.json'))
+                    all_json_results.append(os.path.join(metadata_directory, f"metadata_{basename}_UUID_{id}.json"))
                     if isVideo and overviews:
                         frames.append(os.path.join(overview_directory, f"overview_{basename}_UUID_{id}.jpg"))
         except Exception:
-            #fixme, what is going on with /home/quentin/todo/toup/20221008_16-01-04-226084_raw_jpg.rf.0b8d397da3c47408694eeaab2cde06e5.jpg?
             logger.exception(f"Issue whilst processing {f}")
             raise
     if verbose:
@@ -426,11 +436,11 @@ def predict(
 
             compiled_coco = os.path.join(output_dir, "coco_instances.json")
             pred_coco = {}
-            
+
             flat_bug_predictions = [json.load(open(p)) for p in all_json_results]
             for d in flat_bug_predictions:
                 fb_to_coco(d, pred_coco)
-            with open(compiled_coco,"w") as f:
+            with open(compiled_coco, "w") as f:
                 json.dump(pred_coco, f)
     if isVideo and frames and not no_save and overviews:
         for frame in tqdm(frames, desc=f"Writing video ({video_output_path})", unit="frame"):
@@ -447,18 +457,22 @@ def predict(
     if verbose:
         logger.info("All steps done, process cleaning up.")
 
+
 def main():
     kwargs = cli_args()
 
-    if kwargs.get('gpu', None) is not None:
+    if kwargs.get("gpu", None) is not None:
         logger.warning("'gpu' argument is deprecated!")
         if kwargs.get("device", None) not in [None, "auto"]:
-            raise RuntimeError("Supplying both 'gpu' and 'device' is ambigous. Please use only one, preferably 'device'.")
+            raise RuntimeError(
+                "Supplying both 'gpu' and 'device' is ambigous. Please use only one, preferably 'device'."
+            )
         kwargs["device"] = kwargs.pop("gpu")
     else:
         kwargs.pop("gpu", None)
 
     predict(**kwargs)
 
+
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/src/flat_bug/cli/fb_prepare_data.py b/src/flat_bug/cli/fb_prepare_data.py
index 2f138b4..c8ed84a 100644
--- a/src/flat_bug/cli/fb_prepare_data.py
+++ b/src/flat_bug/cli/fb_prepare_data.py
@@ -25,14 +25,12 @@ def collapse_in_parent_dir(child):
     shutil.rmtree(child)
 
 
-
 OUT_COCO_CONVERTER = "labels/default/"
 OUT_COCO_CONVERTER_IMAGES = "images/default/"
 JSON_FILE_BASENAME = "instances_default.json"
 DATASET_NAME = "insects"
 
 
-
 # A help sting
 out_structure = """
 ├── data.yaml
@@ -54,7 +52,7 @@ def merge_cocos(files, out_file, delete=False):
     for c in files:
         with open(c) as f:
             coco = json.load(f)
-        id_map = {} ## old: new
+        id_map = {}  ## old: new
         new_images = []
         for i in coco["images"]:
             id_map[i["id"]] = im_id
@@ -79,6 +77,7 @@ def merge_cocos(files, out_file, delete=False):
         for c in files:
             os.remove(c)
 
+
 def prepare_coco_file(source_file, image_list, out):
     with open(source_file) as f:
         coco = json.load(f)
@@ -94,7 +93,6 @@ def prepare_coco_file(source_file, image_list, out):
             new_image.append(i)
     assert len(images_to_keep) > 0
 
-
     new_annots = []
     for a in coco["annotations"]:
         if a["image_id"] in image_ids_to_keep:
@@ -105,6 +103,7 @@ def prepare_coco_file(source_file, image_list, out):
     with open(out, "w") as f:
         json.dump(coco, f)
 
+
 def main():
     args_parse = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
 
@@ -137,8 +136,6 @@ def main():
     args = args_parse.parse_args()
     option_dict = vars(args)
 
-
-
     data_yaml = {"path": DATASET_NAME,
                  "train": "images/train",
                  "val": "images/val",
@@ -159,7 +156,7 @@ def main():
 
     with open(os.path.join(PREPARED_DATA_TARGET, "data.yaml"), "w") as f:
         yaml.dump(data_yaml, f)
-    datasets = [] # [d for d in os.listdir(COCO_DATA_ROOT) if os.path.isdir(os.path.join(COCO_DATA_ROOT, d))]
+    datasets = []  # [d for d in os.listdir(COCO_DATA_ROOT) if os.path.isdir(os.path.join(COCO_DATA_ROOT, d))]
     for d in os.listdir(COCO_DATA_ROOT):
         source_dir = os.path.join(COCO_DATA_ROOT, d)
         if os.path.isdir(source_dir):
@@ -174,12 +171,10 @@ def main():
         shutil.rmtree(tmp_dir)
 
         try:
-
-            coco_files = [f for f in sorted(glob.glob(os.path.join( source_dir, "*.json")))]
-            #,"Multiple label files, only supporting one"
+            coco_files = [f for f in sorted(glob.glob(os.path.join(source_dir, "*.json")))]
+            # ,"Multiple label files, only supporting one"
             assert len(coco_files) == 1, os.path.join(source_dir, "*.json")
 
-
             convert_coco(labels_dir=source_dir, save_dir=tmp_dir, use_segments=True)
             os.makedirs(os.path.join(tmp_dir, OUT_COCO_CONVERTER, "train"), exist_ok=True)
             os.makedirs(os.path.join(tmp_dir, OUT_COCO_CONVERTER, "val"), exist_ok=True)
@@ -197,7 +192,7 @@ def main():
 
             validation_files = {}
             training_files = {}
-            for f in sorted(glob.glob(os.path.join(tmp_dir,OUT_COCO_CONVERTER, "*.txt"))):
+            for f in sorted(glob.glob(os.path.join(tmp_dir, OUT_COCO_CONVERTER, "*.txt"))):
                 basename_sans_ext = os.path.splitext(os.path.basename(f))[0]
 
                 image_matches = [
@@ -219,10 +214,9 @@ def main():
                 im_path = os.path.join(source_dir, im_basename)
                 assert os.path.isfile(im_path)
 
-                with open(im_path, 'rb') as file_obj:
+                with open(im_path, "rb") as file_obj:
                     file_hash = hashlib.md5(file_obj.read()).hexdigest()
 
-
                 new_bn_se = f"{d}_{basename_sans_ext}"
                 p = int(file_hash[0:4], 16) / int("ffff", 16)
                 if p < option_dict["validation_proportion"]:
@@ -233,14 +227,8 @@ def main():
                     training_files[im_basename] = new_bn_se + ".jpg"
                 logging.info(f"{im_basename} -> {subset}")
 
-                shutil.move(
-                    f,
-                    os.path.join(tmp_dir, OUT_COCO_CONVERTER, os.path.join(subset, new_bn_se + ".txt"))
-                )
-                shutil.copy(
-                    im_path,
-                    os.path.join(tmp_dir, OUT_COCO_CONVERTER_IMAGES, subset, new_bn_se + ".jpg")
-                )
+                shutil.move(f, os.path.join(tmp_dir, OUT_COCO_CONVERTER, os.path.join(subset, new_bn_se + ".txt")))
+                shutil.copy(im_path, os.path.join(tmp_dir, OUT_COCO_CONVERTER_IMAGES, subset, new_bn_se + ".jpg"))
 
             if len(validation_files) == 0:
                 logging.warning(f"No validation files for {d}")
@@ -248,7 +236,7 @@ def main():
                 prepare_coco_file(
                     coco_files[0],
                     validation_files,
-                    os.path.join(tmp_dir, OUT_COCO_CONVERTER, "val", f"{d}"+JSON_FILE_BASENAME)
+                    os.path.join(tmp_dir, OUT_COCO_CONVERTER, "val", f"{d}" + JSON_FILE_BASENAME),
                 )
 
             if len(validation_files) == 0:
@@ -256,22 +244,24 @@ def main():
             else:
                 prepare_coco_file(
                     coco_files[0],
-                    training_files, 
-                    os.path.join(tmp_dir, OUT_COCO_CONVERTER, "train", f"{d}"+JSON_FILE_BASENAME)
+                    training_files,
+                    os.path.join(tmp_dir, OUT_COCO_CONVERTER, "train", f"{d}" + JSON_FILE_BASENAME),
                 )
 
             collapse_in_parent_dir(os.path.join(tmp_dir, OUT_COCO_CONVERTER))
-            collapse_in_parent_dir(os.path.join(tmp_dir,  OUT_COCO_CONVERTER_IMAGES))
-            #fixme here should add a subdir like "insects/" same as the name in data.yaml
+            collapse_in_parent_dir(os.path.join(tmp_dir, OUT_COCO_CONVERTER_IMAGES))
+            # fixme here should add a subdir like "insects/" same as the name in data.yaml
             shutil.copytree(tmp_dir, PREPARED_DATA_TARGET_SUBDIR, dirs_exist_ok=True)
         finally:
             if os.path.isdir(tmp_dir):
                 shutil.rmtree(tmp_dir)
 
-
     for subset in {"val", "train"}:
         all_json = [f for f in sorted(glob.glob(os.path.join(PREPARED_DATA_TARGET_SUBDIR, "labels", subset, "*.json")))]
-        merge_cocos(all_json, os.path.join(PREPARED_DATA_TARGET_SUBDIR, "labels", subset,JSON_FILE_BASENAME), delete=True)
+        merge_cocos(
+            all_json, os.path.join(PREPARED_DATA_TARGET_SUBDIR, "labels", subset, JSON_FILE_BASENAME), delete=True
+        )
+
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/src/flat_bug/cli/fb_train.py b/src/flat_bug/cli/fb_train.py
index 9ae05ff..2853185 100644
--- a/src/flat_bug/cli/fb_train.py
+++ b/src/flat_bug/cli/fb_train.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """``flatbug`` training script.
 
-The ``flatbug`` training script uses a lightly modified YOLO training interface 
+The ``flatbug`` training script uses a lightly modified YOLO training interface
 (https://docs.ultralytics.com/modes/train/), with a few additional parameters.
 
 See `scripts/experiments/best_train/default.yaml` for an example training config.
@@ -41,7 +41,7 @@ def main():  # noqa: D103
         "epochs": 5000,
         "device": "cuda",
         "patience": 500,
-        "optimizer": 'auto',
+        "optimizer": "auto",
         "save_period": 5,
         # "optimizer": 'SGD',
         # "lr0": 0.01,
@@ -52,22 +52,27 @@ def main():  # noqa: D103
         "fb_max_images": -1,
         "fb_custom_eval": False,
         "fb_custom_eval_num_images": -1,
-        "fb_exclude_datasets" : [],
-        "cache": False
+        "fb_exclude_datasets": [],
+        "cache": False,
     }
     args_parse = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
-    args_parse.add_argument("-d", "--data-dir", dest="data_dir",
-                            help="The directory containing the prepared data (i.e., the output of  `fb_prepare.py`",
-                            type=str)
-
-    args_parse.add_argument("-c", "--config-file", dest="config_file",
-                            help="A YAML-formatted config file that overrides the default training meta-parameters",
-                            default=None)
-
-    args_parse.add_argument("-r", "--resume", dest="resume",
-                            help="resume training",
-                            action='store_true')
-
+    args_parse.add_argument(
+        "-d",
+        "--data-dir",
+        dest="data_dir",
+        help="The directory containing the prepared data (i.e., the output of  `fb_prepare.py`",
+        type=str,
+    )
+
+    args_parse.add_argument(
+        "-c",
+        "--config-file",
+        dest="config_file",
+        help="A YAML-formatted config file that overrides the default training meta-parameters",
+        default=None,
+    )
+
+    args_parse.add_argument("-r", "--resume", dest="resume", help="resume training", action="store_true")
 
     args, extra = args_parse.parse_known_args()
     cli_overrides = {}
@@ -78,7 +83,9 @@ def main():  # noqa: D103
         if key not in DEFAULT_CONF:
             raise ValueError(f"Unknown argument: {key}\n" + args_parse.format_help())
         if key.startswith("fb_"):
-            raise ValueError("Options starting with 'fb_' should be specified in the config file, not as command line arguments")
+            raise ValueError(
+                "Options starting with 'fb_' should be specified in the config file, not as command line arguments"
+            )
         # fixme: probably unsafe...
         try:
             value = eval(value)
@@ -91,9 +98,9 @@ def main():  # noqa: D103
     option_dict = vars(args)
 
     option_dict["data_dir"] = os.path.abspath(os.path.normpath(option_dict["data_dir"]))
-    assert os.path.isdir(option_dict["data_dir"]), f'Directory {option_dict["data_dir"]} not found.'
+    assert os.path.isdir(option_dict["data_dir"]), f"Directory {option_dict['data_dir']} not found."
 
-    # I think this should be fixed by resolving the path before passing 
+    # I think this should be fixed by resolving the path before passing
     # it to the trainer and setting DATASETS_DIR in the scope of ultralytics.data.utils
     # (see https://github.com/ultralytics/ultralytics/blob/588bbbe4aed122e3d24353856484148bc5ef05ad/ultralytics/data/utils.py#L301)
     # #fixme issue when providing new dataset path, sill using old one?! see when i used pollen data
@@ -110,13 +117,14 @@ def main():  # noqa: D103
         with open(option_dict["config_file"]) as f:
             yaml_config = yaml.safe_load(f)
             overrides.update(yaml_config)
-    
+
     # Update with cli overrides
     overrides.update(cli_overrides)
 
     # Update data directory and resume flag from the command line
     overrides["data"] = os.path.join(option_dict["data_dir"], "data.yaml")
-    # OBS: This is a *very* cursed hack around the fact that ultralytics have decided that you cannot change the settings at runtime. 
+    # OBS: This is a *very* cursed hack around the fact that ultralytics
+    # have decided that you cannot change the settings at runtime.
     # We technically only need to change it here, but I'll change it both places for consistency
     ultralytics_data_utils.DATASETS_DIR = Path(option_dict["data_dir"])
     ultralytics_utils.DATASETS_DIR = Path(option_dict["data_dir"])
@@ -128,24 +136,25 @@ def main():  # noqa: D103
         overrides["resume"] = overrides["model"]
         if (old_optim := overrides.pop("optimizer", None)) is not None:
             logger.warning(
-                f"Ignored optimizer '{old_optim}' - "
-                "YOLO does not support changing the optimizer while training."
+                f"Ignored optimizer '{old_optim}' - YOLO does not support changing the optimizer while training."
             )
     else:
         overrides["resume"] = False
 
     # ruff: disable[F841] - TODO: fixme, we don't actually support multiple DDP
     # This is just a hack to fix this: https://github.com/pytorch/pytorch/issues/37377 - only relevant for DDP
-    if isinstance(overrides["device"], (tuple, list)) :
+    if isinstance(overrides["device"], (tuple, list)):
         num_devices = len(overrides["device"])
     elif isinstance(overrides["device"], str):
         num_devices = len(overrides["device"].split(","))
     else:
-        num_devices = 1 # Fixme: Is this a real case, or just a type error?
+        num_devices = 1  # Fixme: Is this a real case, or just a type error?
     # ruff: enable[F841]
-    if isinstance(overrides["device"], (tuple, list)) or (isinstance(overrides["device"], str) and len(overrides["device"].split(",")) > 1):
-        os.environ['MKL_THREADING_LAYER'] = 'GNU'
-        os.environ['OMP_NUM_THREADS'] = str(overrides["workers"])
+    if isinstance(overrides["device"], (tuple, list)) or (
+        isinstance(overrides["device"], str) and len(overrides["device"].split(",")) > 1
+    ):
+        os.environ["MKL_THREADING_LAYER"] = "GNU"
+        os.environ["OMP_NUM_THREADS"] = str(overrides["workers"])
 
     # Ensure that `~` is not interpreted literally in arguments
     for k in overrides:
@@ -164,5 +173,6 @@ def main():  # noqa: D103
         trainer.start_epoch = 0
     trainer.train()
 
+
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/src/flat_bug/coco_utils.py b/src/flat_bug/coco_utils.py
index 37f0b91..ffefc2a 100644
--- a/src/flat_bug/coco_utils.py
+++ b/src/flat_bug/coco_utils.py
@@ -76,12 +76,7 @@
 #     }
 
 
-
-
-def fb_to_coco(
-        d: dict, 
-        coco: dict
-    ) -> dict:
+def fb_to_coco(d: dict, coco: dict) -> dict:
     """Convert a FlatBug dataset to a COCO dataset.
 
     Args:
@@ -135,7 +130,10 @@ def fb_to_coco(
     # identifier = d["identifier"]
     image_path = d["image_path"]
     image_width, image_height, mask_width, mask_height = (
-        d["image_width"], d["image_height"], d["mask_width"], d["mask_height"]
+        d["image_width"],
+        d["image_height"],
+        d["mask_width"],
+        d["mask_height"],
     )
 
     # Image
@@ -147,7 +145,7 @@ def fb_to_coco(
         "license": 0,
         "flickr_url": "",
         "coco_url": "",
-        "date_captured": 0
+        "date_captured": 0,
     }
     coco["images"].append(image)
 
@@ -156,8 +154,8 @@ def fb_to_coco(
         box, contour, conf = boxes[i], contours[i], confs[i]
         # class_, scale = classes[i], scales[i]
         x1, y1, x2, y2 = box
-        x,y,w,h = x1, y1, x2 - x1, y2 - y1
-        box=[x,y,w,h]
+        x, y, w, h = x1, y1, x2 - x1, y2 - y1
+        box = [x, y, w, h]
 
         # Scale and restructure the contour
         m2i = [(mask_width - 1) / (image_width - 1), (mask_height - 1) / (image_height - 1)]  # Mask to image ratio
@@ -173,14 +171,14 @@ def fb_to_coco(
             "area": 0.0,
             "bbox": box,
             "iscrowd": 0,
-            "conf": conf
+            "conf": conf,
         }
         coco["annotations"].append(annotation)
 
     return coco
 
 
-def format_contour(c : list) -> np.ndarray:
+def format_contour(c: list) -> np.ndarray:
     """Format a contour to the OpenCV format.
 
     Args:
@@ -201,24 +199,22 @@ def contour_bbox(c: np.ndarray) -> np.ndarray:
         c: Contour.
 
     Returns:
-        oyut: Bounding box.
+        Bounding box.
 
     """
     return np.array([c[:, 0].min(), c[:, 1].min(), c[:, 0].max(), c[:, 1].max()])
 
 
-def split_annotations(
-        coco: dict, 
-        strip_directories: bool = True
-    ) -> dict[str, list]:
+def split_annotations(coco: dict, strip_directories: bool = True) -> dict[str, list]:
     """Split COCO annotations by image ID.
 
     Args:
         coco: COCO dataset.
-        strip_directories: Flag to indicate whether only the basename of the images should be included in the result. Defaults to True.
+        strip_directories: Flag indicating whether only the basename of the images should be included in the result.
+            Defaults to `True`.
 
     Returns:
-        Dict of COCO datasets, split by image ID and keyed by image name.
+        `dict` of COCO datasets, split by image ID and keyed by image name.
 
     """
     img_id = np.array([i["image_id"] for i in coco["annotations"]])
@@ -230,13 +226,13 @@ def split_annotations(
             coco["images"][i]["file_name"] = os.path.basename(coco["images"][i]["file_name"])
 
     result = {coco["images"][id - 1]["file_name"]: [coco["annotations"][i] for i in g] for id, g in zip(ids, groups)}
-    
+
     # Ensure that all images are included in the result, even if they have no annotations/predictions
     for i in range(len(coco["images"])):
         image_name = coco["images"][i]["file_name"]
         if image_name not in result:
             result[image_name] = []
-    
+
     return result
 
 
@@ -286,12 +282,8 @@ def annotations_to_numpy(annotations: list[dict[str, list[int]]]) -> tuple[np.nd
     bboxes = np.array([contour_bbox(c) for c in contours])
     return bboxes, contours
 
-def filter_coco(
-        coco : dict, 
-        confidence : float | None=None, 
-        area : int | None=None, 
-        verbose : bool=False
-    ) -> dict:
+
+def filter_coco(coco: dict, confidence: float | None = None, area: int | None = None, verbose: bool = False) -> dict:
     """Filter COCO annotations by confidence.
 
     Args:
@@ -322,5 +314,5 @@ def filter_coco(
         "licenses": coco["licenses"],
         "images": coco["images"],
         "annotations": filtered_annotations,
-        "categories": coco["categories"]
-    }
\ No newline at end of file
+        "categories": coco["categories"],
+    }
diff --git a/src/flat_bug/config.py b/src/flat_bug/config.py
index 851ae0f..8fcf6a6 100644
--- a/src/flat_bug/config.py
+++ b/src/flat_bug/config.py
@@ -1,4 +1,5 @@
 """Configuration submodule for flatbug."""
+
 import os
 from collections import OrderedDict
 from collections.abc import Iterable
@@ -25,7 +26,7 @@
     "OVERLAP_METRIC",
     "TIME",
     "TILE_SIZE",
-    "BATCH_SIZE"
+    "BATCH_SIZE",
 ]
 
 CFG_DESCRIPTION = {
@@ -40,7 +41,7 @@
     "OVERLAP_METRIC": "Metric to use for NMS. One of 'IOU' or 'IOS', more might be added in the future.",
     "TIME": "Enable to print time taken for each step. Can incur a performance penalty.",
     "TILE_SIZE": "Fixed by the model architecture - do not change unless you know what you are doing.",
-    "BATCH_SIZE": "Used for model initialization and batched tile processing."
+    "BATCH_SIZE": "Used for model initialization and batched tile processing.",
 }
 
 DEFAULT_CFG = {
@@ -52,10 +53,10 @@
     "MAX_MASK_SIZE": 1024,
     "PREFER_POLYGONS": True,
     "EXPERIMENTAL_NMS_OPTIMIZATION": True,
-    "OVERLAP_METRIC" : "IoU",
+    "OVERLAP_METRIC": "IoU",
     "TIME": False,
     "TILE_SIZE": 1024,
-    "BATCH_SIZE": 16
+    "BATCH_SIZE": 16,
 }
 
 LEGACY_CFG = {
@@ -67,27 +68,25 @@
     "MAX_MASK_SIZE": 1024,
     "PREFER_POLYGONS": True,
     "EXPERIMENTAL_NMS_OPTIMIZATION": True,
-    "OVERLAP_METRIC" : "IoU",
+    "OVERLAP_METRIC": "IoU",
     "TIME": False,
     "TILE_SIZE": 1024,
-    "BATCH_SIZE": 16
+    "BATCH_SIZE": 16,
 }
 # ruff: enable[E501]
 
-def get_type_def(
-        obj : Any, 
-        tuple_list_interchangeable : bool=False
-    ) -> Any | list[Any]:
-    """Generate a dynamic type definition for an object.
+
+def get_type_def(obj: Any, tuple_list_interchangeable: bool = False) -> Any | list[Any]:
+    r"""Generate a dynamic type definition for an object.
 
     The type definition schema is defined like this:
-    - If the object is a tuple or a list, the first element is the type of the object, 
+    - If the object is a tuple or a list, the first element is the type of the object,
         and the second element is a list of type definitions for the elements of the object.
     - If the object is not a tuple or a list, the type definition is the type of the object.
 
     For example;
-        - the type definition for the object `(1, "A", True)` would be `[tuple, [int, str, bool]]`.
-        - the type definition for the object `[[2, "B"], [3, "C"]]` would be `[list, [[list, [int, str]], [list, [int, str]]]]`.
+        - `(1, "A", True)`          \:   `[tuple, [int, str, bool]]`
+        - `[[2, "B"], [3, "C"]]`    \:   `[list, [[list, [int, str]], [list, [int, str]]]]`
 
     Args:
         obj: The object to generate a type definition for.
@@ -104,18 +103,17 @@ def get_type_def(
         return [otype, [get_type_def(i, tuple_list_interchangeable) for i in obj]]
     return type(obj)
 
-CFG_TYPES = {k : get_type_def(DEFAULT_CFG[k], tuple_list_interchangeable=True) for k in DEFAULT_CFG}
+
+CFG_TYPES = {k: get_type_def(DEFAULT_CFG[k], tuple_list_interchangeable=True) for k in DEFAULT_CFG}
+
 
 def check_types(
-        value : Any, 
-        expected_type : list[Any] | Iterable[type] | type, 
-        key : str="<Not specified>", 
-        strict : bool=True
-    ) -> bool:
+    value: Any, expected_type: list[Any] | Iterable[type] | type, key: str = "<Not specified>", strict: bool = True
+) -> bool:
     """Recursively check if the type of a value matches the expected type.
 
-    If the expected type is a list, the first element is the type of the value, 
-    and the second element is a list of types that the elements of the value match, 
+    If the expected type is a list, the first element is the type of the value,
+    and the second element is a list of types that the elements of the value match,
     a single type that all elements should match or a tuple/type of types that all elements should match any of.
 
     Args:
@@ -140,12 +138,11 @@ def check_types(
             # Check that an expected type has been supplied for both the value and its elements
             if len(expected_type) != 2:
                 raise ValueError(
-                    f"Expected type list must have exactly 2 elements, "
-                    f"got {len(expected_type)} for key: {key}."
+                    f"Expected type list must have exactly 2 elements, got {len(expected_type)} for key: {key}."
                 )
             # Check that the value matches the expected type
             check_types(value, expected_type[0], key, strict)
-            # If the expected type of the elements is a list, 
+            # If the expected type of the elements is a list,
             # each element of the value should match the corresponding element of the expected type list
             if isinstance(expected_type[1], list):
                 # Check that the number of types in the list matches the number of items in the value
@@ -154,21 +151,21 @@ def check_types(
                         f"Expected number of types ({len(expected_type[1])}) "
                         f"does not match number of items in value ({len(value)}) for key: {key}."
                     )
-                # Check that each item in the value matches 
+                # Check that each item in the value matches
                 # the corresponding type in the expected type list
                 for item, et in zip(value, expected_type[1]):
                     check_types(item, et, key, strict)
-            # If the expected type of the elements is a single type, 
+            # If the expected type of the elements is a single type,
             # each element of the value should match the expected type
             elif isinstance(expected_type[1], type):
                 for item in value:
                     check_types(item, expected_type[1], key, strict)
-            # If the expected type of the elements is a tuple, 
+            # If the expected type of the elements is a tuple,
             # each element of the value should match any of the types in the tuple
             elif isinstance(expected_type[1], tuple):
                 check_types(value, expected_type[1], key, strict)
-            # If the expected type of the elements is an iterable, 
-            # the value should be an iterable and each element of the value 
+            # If the expected type of the elements is an iterable,
+            # the value should be an iterable and each element of the value
             # should match the corresponding type in the expected type iterable
             elif hasattr(expected_type[1], "__iter__") and hasattr(expected_type[1], "__len__"):
                 assert len(expected_type[1]) == len(value), (
@@ -185,8 +182,8 @@ def check_types(
                     raise TypeError("\n  - ".join(errors))
             else:
                 raise TypeError(
-                    "Invalid expected type. Expected 'list', 'type', 'tuple' or an iterable "
-                    f"got {type(expected_type[1])} for key: {key}."
+                    "Invalid expected type. "
+                    f"Expected 'list', 'type', 'tuple' or an iterable got {type(expected_type[1])} for key: {key}."
                 )
         # If the expected type is an iterable, check if the value is an instance of any of the types in the iterable
         elif hasattr(expected_type, "__iter__") and not isinstance(expected_type, type):
@@ -206,8 +203,8 @@ def check_types(
         # If the expected type is not a list, a iterable or a 'type' object raise an error
         else:
             raise TypeError(
-                "Invalid expected type. Expected 'list', an iterable, 'type' or 'typing.Any' "
-                f"got {type(expected_type)} for key: {key}."
+                "Invalid expected type. "
+                f"Expected 'list', an iterable, 'type' or 'typing.Any' got {type(expected_type)} for key: {key}."
             )
         # If no errors are raised, return True
         return True
@@ -218,16 +215,14 @@ def check_types(
         else:
             return False
 
-def check_cfg_types(
-        cfg : dict, 
-        strict : bool = False
-    ) -> bool:
+
+def check_cfg_types(cfg: dict, strict: bool = False) -> bool:
     """Check if the config is a dictionary and that the types of the values in the config dictionary are correct.
 
     Args:
         cfg: The config dictionary to check.
         strict: If True, raise an error if a key is not recognized. Defaults to False.
-    
+
     Returns:
         True if all checks pass, raises an error otherwise.
 
@@ -247,10 +242,8 @@ def check_cfg_types(
     # If no errors are raised, return True
     return True
 
-def read_cfg(
-        path : str | Path, 
-        strict : bool=False
-    ) -> dict:
+
+def read_cfg(path: str | Path, strict: bool = False) -> dict:
     """Load and validate the config file.
 
     Missing keys are replaced with default values.
@@ -288,11 +281,8 @@ def read_cfg(
     # Return config
     return cfg
 
-def write_cfg(
-        cfg : dict, 
-        path : str | os.PathLike, 
-        overwrite : bool=False
-    ) -> str | os.PathLike:
+
+def write_cfg(cfg: dict, path: str | os.PathLike, overwrite: bool = False) -> str | os.PathLike:
     """Save the config dictionary to a YAML file.
 
     Args:
@@ -334,12 +324,13 @@ def write_cfg(
             sorted_cfg[key] = cfg[key]
     # Save config file
     with open(path, "w") as f:
-        # OBS: will fail if not using yaml.SafeDumper (default with yaml.safe_dump). 
+        # OBS: will fail if not using yaml.SafeDumper (default with yaml.safe_dump).
         # If another dumper is to be used, the representer for OrderedDict must be added manually.
         yaml.safe_dump(sorted_cfg, f, sort_keys=False, default_flow_style=None)
     # Return the path to the saved config YAML file
     return path
 
+
 if __name__ == "__main__":
     # Print a helpful message:
     logger.info(
@@ -348,13 +339,14 @@ def write_cfg(
         "####################################################################"
         "\nConfigurable parameters:"
     )
-    logger.info(
-        "\n".join([f"\t- {key} ({CFG_TYPES[key]}): {CFG_DESCRIPTION[key]}" for key in CFG_PARAMS])
-    )
+    logger.info("\n".join(
+        [f"\t- {key} ({CFG_TYPES[key]}): {CFG_DESCRIPTION[key]}"
+         for key in CFG_PARAMS]
+    ))
     logger.info(
         "\nParameters can either be specified with a YAML file or manually:"
         "\t* `fb_predict --config <YML_PATH>`"
         "\t* `flat_bug.predictor.Predictor.__init__(..., cfg=<YML_PATH>, ...)`"
         "\t* `flat_bug.predictor.Predictor.set_hyperparameters(<PARAM_i>=<VALUE_i>, <PARAM_j>=<VALUE_j>, ...)`"
         "\nAny parameters not specified will be set to default values."
-    )
\ No newline at end of file
+    )
diff --git a/src/flat_bug/datasets.py b/src/flat_bug/datasets.py
index d92cf23..a575e41 100644
--- a/src/flat_bug/datasets.py
+++ b/src/flat_bug/datasets.py
@@ -1,4 +1,5 @@
 """Modified YOLO dataset used for training flatbug."""
+
 import os
 import re
 import stat
@@ -16,21 +17,23 @@
 
 from flat_bug.augmentations import CenterCrop, FixInstances, FlatBugRandomPerspective, RandomColorInv, RandomCrop
 
-HELP_URL = 'See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
-IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp', 'pfm'  # include image suffixes
+HELP_URL = "See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data"
+IMG_FORMATS = "bmp", "dng", "jpeg", "jpg", "mpo", "png", "tif", "tiff", "webp", "pfm"  # include image suffixes
+
 
 def get_area(image_path):  # noqa: D103
     with Image.open(image_path) as image:
         return image.size[0] * image.size[1]
 
-def calculate_image_weights(image_paths : list[str]) -> list[float]:
+
+def calculate_image_weights(image_paths: list[str]) -> list[float]:
     """Calculate normalized weights for each image based on the file sizes.
 
     Normalized by the minimum file size, so that the values are between 1 and infinity.
-    
+
     Args:
-        image_paths: List of image file paths.
-    
+        image_paths: `list` of image file paths.
+
     Returns:
         normalized weights for each image.
 
@@ -39,16 +42,14 @@ def calculate_image_weights(image_paths : list[str]) -> list[float]:
     min_size = min(file_sizes)
     return [(size / min_size) for size in file_sizes]
 
-def reweight(
-        weights : list[float], 
-        target_sum : float | int
-    ) -> list[float]:
+
+def reweight(weights: list[float], target_sum: float | int) -> list[float]:
     """Reweights the provided list of weights so that their sum equals the target sum.
-    
+
     Args:
-        weights: List of weights to reweight.
+        weights: `list` of weights to reweight.
         target_sum: Desired sum of the weights.
-    
+
     Returns:
         Reweighted weights.
 
@@ -56,18 +57,17 @@ def reweight(
     sum_weights = sum(weights)
     return [max(round(w * target_sum / sum_weights), 1) for w in weights]
 
-def generate_indices(
-        weights : list[float], 
-        target_size : int | None=None
-    ) -> list[int]:
+
+def generate_indices(weights: list[float], target_size: int | None = None) -> list[int]:
     """Deterministically generates a list of indices based on the provided weights to oversample the items.
-    
+
     Args:
-        weights: List of weights for each item.
-        target_size: Desired size of the output list. If None, the size of the output is approximately the sum of the weights.
+        weights: `list` of weights for each item.
+        target_size: Desired length of the output `list`.
+            If `None`, the size of the output is approximately the sum of the weights.
 
     Returns:
-        List of indices to oversample the items.
+        `list` of indices to oversample the items.
 
     """
     # n = len(weights)
@@ -76,7 +76,7 @@ def generate_indices(
 
     if target_size is not None:
         for _ in range(10):
-            if abs(sum(weights) - target_size)/target_size < 0.01:
+            if abs(sum(weights) - target_size) / target_size < 0.01:
                 break
             weights = reweight(weights, target_size)
 
@@ -85,19 +85,17 @@ def generate_indices(
 
     return indices
 
-def get_datasets(files : list[str]) -> dict[str, list[str]]:  # noqa: D103
+
+def get_datasets(files: list[str]) -> dict[str, list[str]]:  # noqa: D103
     file_dataset = [mtch.group(0) for f in files if (mtch := re.match(r"[^_]+", os.path.basename(f)))]
     datasets = list(set(file_dataset))
-    datasets = {d : [] for d in datasets}
+    datasets = {d: [] for d in datasets}
     for file, fd in zip(files, file_dataset):
         datasets[fd].append(file)
     return datasets
 
-def subset(
-        self : "FlatBugYOLODataset", 
-        n : int | None=None, 
-        pattern : str | None=None
-    ):
+
+def subset(self: "FlatBugYOLODataset", n: int | None = None, pattern: str | None = None):
     """Subsets the dataset to the first 'n' elements that match the pattern.
 
     Args:
@@ -110,12 +108,16 @@ def subset(
         return self
     if pattern is not None:
         cp = re.compile(pattern)
+
         def _match_pattern(x):
             return bool(cp.search(os.path.basename(x)))
+
         match_fn = _match_pattern
     else:
+
         def _match_all(_):
             return True
+
         match_fn = _match_all
     # Get the indices of the elements that match the pattern
     indices = [i for i, f in enumerate(self.im_files) if match_fn(f)]
@@ -125,37 +127,41 @@ def _match_all(_):
     # Subset the images
     self.im_files = [f for i, f in enumerate(self.im_files) if i in indices]
 
+
 def hook_get_labels_with_subset(  # noqa: D103
-        obj : "FlatBugYOLODataset", 
-        args : dict
-    ):
+    obj: "FlatBugYOLODataset", args: dict
+):
     if not isinstance(args, dict):
         raise ValueError("args must be a dictionary")
     if not isinstance(obj, FlatBugYOLODataset):
         raise ValueError("obj must be an instance of FlatBugYOLODataset")
+
     def subset_then_get():
         subset(obj, **args)
         obj.get_labels = getattr(super(type(obj), obj), "get_labels")
         return obj.get_labels()
+
     obj.get_labels = subset_then_get
 
+
 class PrintNumInstances:  # noqa: D101
-    def __init__(self, title : str):  # noqa: D107
-        self.fmt = f'({"{num:>5}"}) ({"{imsize:^10}"}) | {title}'
+    def __init__(self, title: str):  # noqa: D107
+        self.fmt = f"({'{num:>5}'}) ({'{imsize:^10}'}) | {title}"
 
-    def __call__(self, labels : dict):  # noqa: D102
+    def __call__(self, labels: dict):  # noqa: D102
         n = len(labels["instances"]) if "instances" in labels else labels["masks"].max().item()
         print(self.fmt.format(num=n, imsize="x".join([str(d) for d in labels["img"].shape])))
         return labels
 
+
 def train_augmentation_pipeline(  # noqa: D103
-        hyperparameters : IterableSimpleNamespace, 
-        image_size : int, 
-        max_instances : int | float | None, 
-        min_size : int, 
-        use_segments : bool, 
-        use_keypoints : bool
-    ) -> Compose:
+    hyperparameters: IterableSimpleNamespace,
+    image_size: int,
+    max_instances: int | float | None,
+    min_size: int,
+    use_segments: bool,
+    use_keypoints: bool,
+) -> Compose:
     return Compose([
         # Crop to slightly larger than needed for training
         RandomCrop(imsize=int(image_size * 1.5)),
@@ -167,7 +173,7 @@ def train_augmentation_pipeline(  # noqa: D103
         RandomColorInv(p=0.25),
         RandomFlip(direction="vertical", p=hyperparameters.flipud),
         RandomFlip(direction="horizontal", p=hyperparameters.fliplr),
-         # Remove instances outside crop
+        # Remove instances outside crop
         FixInstances(area_thr=0.975, max_targets=max_instances, min_size=min_size),
         # YOLO-native preprocessing
         Format(
@@ -177,16 +183,14 @@ def train_augmentation_pipeline(  # noqa: D103
             return_keypoint=use_keypoints,
             batch_idx=True,
             mask_ratio=hyperparameters.mask_ratio,
-            mask_overlap=hyperparameters.overlap_mask
+            mask_overlap=hyperparameters.overlap_mask,
         ),
     ])
 
+
 def validation_augmentation_pipeline(  # noqa: D103
-        image_size : int, 
-        min_size : int, 
-        use_segments : bool, 
-        use_keypoints : bool
-    ) -> Compose:
+    image_size: int, min_size: int, use_segments: bool, use_keypoints: bool
+) -> Compose:
     return Compose([
         RandomCrop(imsize=int(image_size * 1.5)),
         CenterCrop(image_size),
@@ -198,28 +202,23 @@ def validation_augmentation_pipeline(  # noqa: D103
             return_keypoint=use_keypoints,
             batch_idx=True,
             mask_ratio=1,
-            mask_overlap=True
-        )
+            mask_overlap=True,
+        ),
     ])
 
+
 class FlatBugYOLODataset(YOLODataset):  # noqa: D101
-    
     # What is the minimum size of an instance to be considered (width or height in pixels after augmentations)
-    _min_size : int=32
+    _min_size: int = 32
 
-    # How much do we allow the dataset to grow when oversampling - this is done to ensure larger images are not underrepresented
-    _oversample_factor : int=2
+    # How much do we allow the dataset to grow when oversampling, used to ensure larger images are not underrepresented
+    _oversample_factor: int = 2
 
     def __init__(  # noqa: D107
-            self, 
-            max_instances : int | float | None, 
-            classes : None=None, 
-            subset_args : dict | None=None, 
-            *args, 
-            **kwargs
-        ):
+        self, max_instances: int | float | None, classes: None = None, subset_args: dict | None = None, *args, **kwargs
+    ):
         self._max_instances = max_instances
-        self._include_classes = classes # Only used so the class list is visible in the subset method
+        self._include_classes = classes  # Only used so the class list is visible in the subset method
         if subset_args is not None:
             hook_get_labels_with_subset(self, subset_args)
         if "data" in kwargs:
@@ -243,15 +242,12 @@ def _debug_write_loaded_images(self, out, index):
             n = cv2.rectangle(n, (x - w // 2, y - w // 2), (x + w // 2, y + h // 2), 255, 3)
         cv2.imwrite(f"/tmp/test/{index}-img.jpg", n + m / 3)
 
-    def load_image(
-            self, 
-            i : int | slice
-        ) -> tuple[np.ndarray, tuple[int, int], tuple[int, int]]:
+    def load_image(self, i: int | slice) -> tuple[np.ndarray, tuple[int, int], tuple[int, int]]:
         """Load an image.
 
         Args:
             i: Image index.
-        
+
         Returns:
             im, hw_original, hw_resized
 
@@ -273,43 +269,39 @@ def load_image(
             h0, w0 = im.shape[:2]  # orig hw
 
             return im, (h0, w0), im.shape[:2]  # type: ignore
-        return self.ims[i], self.im_hw0[i], self.im_hw[i] # type: ignore
+        return self.ims[i], self.im_hw0[i], self.im_hw[i]  # type: ignore
 
     def build_transforms(  # noqa: D102
-            self, 
-            hyp : IterableSimpleNamespace
-        ) -> Compose:
+        self, hyp: IterableSimpleNamespace
+    ) -> Compose:
         return train_augmentation_pipeline(
-            hyperparameters=hyp, 
-            image_size=self.imgsz, 
-            max_instances=self._max_instances, 
-            min_size=self._min_size, 
-            use_segments=self.use_segments, 
-            use_keypoints=self.use_keypoints
+            hyperparameters=hyp,
+            image_size=self.imgsz,
+            max_instances=self._max_instances,
+            min_size=self._min_size,
+            use_segments=self.use_segments,
+            use_keypoints=self.use_keypoints,
         )
 
-    def cache_labels(
-            self, 
-            path : Path=Path("./labels.cache")
-        ):
+    def cache_labels(self, path: Path = Path("./labels.cache")):
         """OBS: DO NOT USE THIS FUNCTION MANUALLY."""
         LOGGER.warning("!! OBS !! ==>>== Flat-bug doesn't use the .cache-file! ==<<== !! OBS !!")
 
-        # To bypass the creation of .cache files we use a temporary dummy file, which is set to read-only, 
+        # To bypass the creation of .cache files we use a temporary dummy file, which is set to read-only,
         # causing a check in ultralytics to bail on creating the file
         tmp_file = tempfile.NamedTemporaryFile(delete=False)
         # The path passed to the superclass `cache_labels` method must be a pathlib.Path object
         unwriteable_tmp_path = Path(tmp_file.name)
-        
+
         # Change the file to read-only
         os.chmod(str(unwriteable_tmp_path), stat.S_IREAD)
 
-        # Before calling the superclass `cache_labels` method, 
+        # Before calling the superclass `cache_labels` method,
         # we need to create a dummy `<unwriteable_tmp_path>.cache.npy` file
         temporary_dummy_numpy_cache_file = unwriteable_tmp_path.with_suffix(".cache.npy")
         np.save(temporary_dummy_numpy_cache_file, np.array([0]))
-        
-        # Call the superclass `cache_labels` method with the temporary read-only pathlib.Path object 
+
+        # Call the superclass `cache_labels` method with the temporary read-only pathlib.Path object
         return_val = super().cache_labels(path=unwriteable_tmp_path)
 
         # Remove the temporary file if it still exists
@@ -318,28 +310,27 @@ def cache_labels(
         # Remove the temporary numpy cache file if it still exists
         if os.path.exists(temporary_dummy_numpy_cache_file):
             os.remove(temporary_dummy_numpy_cache_file)
-        
+
         return return_val
 
-    
     def __len__(self):
         return len(self.__indices)
 
     def __getitem__(self, index):
         return self.transforms(self.get_image_and_label(self.__indices[index]))
 
+
 class FlatBugYOLOValidationDataset(FlatBugYOLODataset):  # noqa: D101
-    _resample_n : int= 5
+    _resample_n: int = 5
 
     def build_transforms(  # noqa: D102
-            self, 
-            hyp : IterableSimpleNamespace
-        ) -> Compose:
+        self, hyp: IterableSimpleNamespace
+    ) -> Compose:
         return validation_augmentation_pipeline(
-            image_size=self.imgsz, 
-            min_size=self._min_size, 
-            use_segments=self.use_segments, 
-            use_keypoints=self.use_keypoints
+            image_size=self.imgsz,
+            min_size=self._min_size,
+            use_segments=self.use_segments,
+            use_keypoints=self.use_keypoints,
         )
 
     def __len__(self):
diff --git a/src/flat_bug/eval_utils.py b/src/flat_bug/eval_utils.py
index ed8a76a..551be17 100644
--- a/src/flat_bug/eval_utils.py
+++ b/src/flat_bug/eval_utils.py
@@ -1,4 +1,5 @@
 """Utilities for flatbug evaluation."""
+
 import csv
 import os
 import time
@@ -12,24 +13,23 @@
 from flat_bug.coco_utils import annotations_to_numpy, contour_area, contour_bbox
 
 
-def isfloat(num : str) -> bool:  # noqa: D103
+def isfloat(num: str) -> bool:  # noqa: D103
     try:
         fnum = float(num)
         return not fnum.is_integer()
     except Exception:
         return False
 
-def ispath(path : str) -> bool:  # noqa: D103
+
+def ispath(path: str) -> bool:  # noqa: D103
     return "/" in path or "\\" in path
-    
-def format_cell(
-        cell : str, 
-        digits : int = 3, 
-        max_length : int = 30
-    ) -> str:
+
+
+def format_cell(cell: str, digits: int = 3, max_length: int = 30) -> str:
     """Autoformat a cell for a table.
 
-    Standardizes the number of decimals if the cell is coercible to a float, and truncates the cell if it exceeds the maximum length.
+    Standardizes the number of decimals if the cell is coercible to a float,
+    and truncates the cell if it exceeds the maximum length.
 
     Args:
         cell: The cell to format.
@@ -48,17 +48,15 @@ def format_cell(
         return f"{cell[:left_size]}...{cell[-right_size:]}"
     return cell
 
-def format_row(
-        cells : list[Any], 
-        widths : list[int], 
-        align : str = "center"
-    ) -> str:
+
+def format_row(cells: list[Any], widths: list[int], align: str = "center") -> str:
     """Format a row of a table.
 
     Args:
         cells: The cells of the row. Elements should be compatible with f-strings (`__format__`).
         widths: The widths of each column.
-        align: Alignment of the cell content within each column. Valid options are "center"/"left"/"right". Defaults to "center".
+        align: Alignment of the cell content within each column. Valid options are "center"/"left"/"right".
+            Defaults to "center".
 
     Returns:
         The formatted row.
@@ -75,10 +73,8 @@ def format_row(
                 row += f" {cell:>{width}} |"
     return row
 
-def pretty_print_csv(
-        csv_file : str, 
-        delimiter : str = ","
-    ):
+
+def pretty_print_csv(csv_file: str, delimiter: str = ","):
     """Pretty print the CSV file.
 
     Args:
@@ -95,9 +91,10 @@ def pretty_print_csv(
             print("pretty_print_csv: Empty CSV file.")
             return
         max_lengths = [len(h) for h in headers]
-        rows = []
-        for row in csv_reader:
-            rows.append([format_cell(cell, digits=3, max_length=min(30, header_width * 4)) for cell, header_width in zip(row, max_lengths)])
+        rows = [[
+            format_cell(cell, digits=3, max_length=min(30, header_width * 4))
+            for cell, header_width in zip(row, max_lengths)
+        ] for row in csv_reader]
     # Get the maximum length of each column
     for row in rows:
         for i, cell in enumerate(row):
@@ -108,13 +105,11 @@ def pretty_print_csv(
     print(header)
     print(horizontal_line)
     for row in rows:
-        print(format_row(row, max_lengths, align="right"))   
-    print(horizontal_line)   
+        print(format_row(row, max_lengths, align="right"))
+    print(horizontal_line)
 
-def bbox_intersect(
-        b1 : np.ndarray, 
-        b2s : np.ndarray
-    ) -> np.ndarray:
+
+def bbox_intersect(b1: np.ndarray, b2s: np.ndarray) -> np.ndarray:
     """Calculate the intersecting rectangle between two rectangles. The rectangles must be aligned with the axes.
 
     Args:
@@ -137,11 +132,10 @@ def bbox_intersect(
     return ix
 
 
-def bbox_intersect_area(
-        b1 : np.ndarray, 
-        b2s : np.ndarray
-    ) -> np.ndarray:
-    """Calculate the area of the intersecting rectangle between two rectangles. The rectangles must be aligned with the axes.
+def bbox_intersect_area(b1: np.ndarray, b2s: np.ndarray) -> np.ndarray:
+    """Calculate the area of the intersecting rectangle between two rectangles.
+
+    The rectangles must be aligned with the axes.
 
     Args:
         b1: Bounding box 1.
@@ -158,12 +152,7 @@ def bbox_intersect_area(
     return np.prod((ix_min - ix_max).clip(0), axis=1)
 
 
-def contour_intersection(
-        contour1: np.ndarray, 
-        contour2: np.ndarray, 
-        box1: np.ndarray, 
-        box2: np.ndarray
-    ):
+def contour_intersection(contour1: np.ndarray, contour2: np.ndarray, box1: np.ndarray, box2: np.ndarray):
     """Calculate the intersection of two contours.
 
     Contours should be providedd as [x1, y1, x2, y2, ..., xn, yn]
@@ -212,18 +201,19 @@ def contour_intersection(
 
 
 def pairwise_contour_intersection(
-        contours1: list[np.ndarray], 
-        contours2: list[np.ndarray] | None = None,
-        bboxes1: np.ndarray | None = None, 
-        bboxes2: np.ndarray | None = None,
-        areas1: np.ndarray | None = None,
-        areas2: np.ndarray | None = None
-    ) -> np.ndarray:
+    contours1: list[np.ndarray],
+    contours2: list[np.ndarray] | None = None,
+    bboxes1: np.ndarray | None = None,
+    bboxes2: np.ndarray | None = None,
+    areas1: np.ndarray | None = None,
+    areas2: np.ndarray | None = None,
+) -> np.ndarray:
     """Calculate the pairwise intersection of two groups of contours.
 
     Args:
         contours1: Contours in group 1.
-        contours2: Contours in group 2. If None provided, symmetric intersection is calculated for contours1 instead. Defaults to None.
+        contours2: Contours in group 2.
+            If None provided, symmetric intersection is calculated for contours1 instead. Defaults to None.
         areas1: Areas of contours in group 1. Computed if not None. Defaults to None
         areas2: Areas of contours in group 2. Computed if not None. Defaults to None
         bboxes1: Bounding boxes of contours in group 1. Computed if not None. Defaults to None
@@ -264,18 +254,20 @@ def pairwise_contour_intersection(
 
 
 def match_geoms(
-        contours1: list[np.ndarray], 
-        contours2: list[np.ndarray], 
-        threshold: float = 1 / 4,
-        iou_mat: np.ndarray | None = None, 
-        areas1: np.ndarray | None = None,
-        areas2: np.ndarray | None = None
-    ) -> tuple[np.ndarray, int]:
+    contours1: list[np.ndarray],
+    contours2: list[np.ndarray],
+    threshold: float = 1 / 4,
+    iou_mat: np.ndarray | None = None,
+    areas1: np.ndarray | None = None,
+    areas2: np.ndarray | None = None,
+) -> tuple[np.ndarray, int]:
     """Match geometries (polygons) in group 1 to geometries in group 2.
 
     Args:
-        contours1: Geometries (polygons) in group 1. List of length N, where each element is a Xx2 array of contour coordinates.
-        contours2: Geometries (polygons) in group 2. List of length M, where each element is a Xx2 array of contour coordinates.
+        contours1: Geometries (polygons) in group 1.
+            List of length N, where each element is a Xx2 array of contour coordinates.
+        contours2: Geometries (polygons) in group 2.
+            List of length M, where each element is a Xx2 array of contour coordinates.
         threshold: IoU threshold. Defaults to 1/4.
         iou_mat: IoU matrix of size NxM. Computed if None. Defaults to None.
         areas1: Areas of polygons in group 1. Computed if None. Defaults to None.
@@ -300,7 +292,7 @@ def match_geoms(
         iou = iou_mat.copy()
     # Check the shape of the IoU matrix
     if iou.shape != (n, m):
-        raise ValueError(f'Expected IoU matrix of shape {(n, m)}, got {iou.shape}')
+        raise ValueError(f"Expected IoU matrix of shape {(n, m)}, got {iou.shape}")
     # Initialize the match array
     matches = np.zeros((n, 2), dtype=np.int32)
     matches[:, 0] = np.arange(n, dtype=np.int32)
@@ -310,14 +302,14 @@ def match_geoms(
     # Match the geometries in group 1 to the geometries in group 2
     if n > 0 and m > 0:
         for focus in np.argsort(iou.max(axis=1)):
-            best_match = np.argsort(iou[focus])[::-1][:np.sum(iou[focus] > threshold)]
+            best_match = np.argsort(iou[focus])[::-1][: np.sum(iou[focus] > threshold)]
             if len(best_match) == 0:
                 continue
             # Check if the potential matches have a better focus
             best_match = best_match[iou[:, best_match].argmax(axis=0) == focus]
             if len(best_match) == 0:
                 continue
-            best_match = best_match[0] 
+            best_match = best_match[0]
             if iou[focus, best_match] > threshold:
                 matches[focus, 1] = best_match.astype(np.int32)
                 # Set the intersection to 0 so it doesn't get matched again
@@ -333,18 +325,18 @@ def match_geoms(
     matches = np.concatenate([matches, matches_2], axis=0)
     # Check the shape of the matches array
     if matches.shape != (n + len(unmatched), 2):
-        raise ValueError(f'Expected matches of shape {(n + len(unmatched), 2)}, got {matches.shape}')
+        raise ValueError(f"Expected matches of shape {(n + len(unmatched), 2)}, got {matches.shape}")
     return matches, len(unmatched)
 
 
 def plot_heatmap(
-        mat: np.ndarray, 
-        axis_labels: Sequence[str] | None = None, 
-        breaks: int = 25,
-        dimensions: tuple[int, int] | None = None, 
-        output_path: str | None = None, 
-        scale: float = 1
-    ):
+    mat: np.ndarray,
+    axis_labels: Sequence[str] | None = None,
+    breaks: int = 25,
+    dimensions: tuple[int, int] | None = None,
+    output_path: str | None = None,
+    scale: float = 1,
+):
     """Plot a heatmap of a matrix using OpenCV.
 
     Args:
@@ -363,21 +355,21 @@ def plot_heatmap(
             scale_dims = 1000 / min_dim
             dimensions = tuple([int(d * scale_dims) for d in dimensions])  # type: ignore
     if mat.shape[0] == 0 or mat.shape[1] == 0:
-        logger.warning('Empty matrix. Cannot plot heatmap.')
+        logger.warning("Empty matrix. Cannot plot heatmap.")
         return
     assert dimensions is not None
 
     # Create a colormap for viridis
     colormap = cv2.applyColorMap(
-        src = (mat / (mat.max() or 1) * 255).astype(np.uint8), 
-        colormap = cv2.COLORMAP_VIRIDIS
+        src=(mat / (mat.max() or 1) * 255).astype(np.uint8), 
+        colormap=cv2.COLORMAP_VIRIDIS
     )
     # Expand colormap to 1000x1000
     colormap = cv2.resize(
-        src = colormap, 
-        dsize = dimensions, 
-        dst = colormap,
-        interpolation = cv2.INTER_NEAREST_EXACT
+        src=colormap, 
+        dsize=dimensions, 
+        dst=colormap,
+        interpolation=cv2.INTER_NEAREST_EXACT
     )
     # Get the height of the colormap
     cheight, cwidth = colormap.shape[:2]
@@ -386,25 +378,28 @@ def plot_heatmap(
     # Add numbers to the colorbar
     cmin, cmax = mat.min(), mat.max()
 
-    cmin, cmax = 10 ** np.floor(np.log10(cmin)) if cmin > 0 else cmin, 10 ** np.ceil(np.log10(cmax)) if cmax > 0 else cmax
+    cmin, cmax = (
+        10 ** np.floor(np.log10(cmin)) if cmin > 0 else cmin,
+        10 ** np.ceil(np.log10(cmax)) if cmax > 0 else cmax,
+    )
     cmin, cmax = int(cmin), int(cmax)
     if cmin == cmax:
         nice_breaks = np.array([cmin])
     else:
-        # Add semi-equally spaced numbers to the colorbar at "nice" values, 
-        # "nice" values are defined as integer multiples of powers of 10 
+        # Add semi-equally spaced numbers to the colorbar at "nice" values,
+        # "nice" values are defined as integer multiples of powers of 10
         # to the power of the maximum value - the integer rounded 10 logarithm of the number of breaks
         raw_breaks = np.linspace(cmin, cmax, breaks)
         nice_multiple = 10 ** (np.log10(cmax) - np.ceil(np.log10(breaks)))
         nice_breaks = (raw_breaks / nice_multiple).round() * nice_multiple
         nice_breaks = nice_breaks[nice_breaks <= cmax]
-        # Ensure that the minimum and maximum values are included, 
+        # Ensure that the minimum and maximum values are included,
         # and remove the breaks if they are within 1 "nice_multiple" of any other break
         nice_breaks = nice_breaks[np.abs(nice_breaks - cmin) >= (nice_multiple * 0.9)]
         nice_breaks = nice_breaks[np.abs(nice_breaks - cmax) >= (nice_multiple * 0.9)]
         nice_breaks = np.concatenate([[cmin], nice_breaks, [cmax]])
     # Create the labels
-    labels = [f'{(i * 100):.3g}%' for i in nice_breaks]
+    labels = [f"{(i * 100):.3g}%" for i in nice_breaks]
 
     # Define a target font height
     font_height_target = int(min(min(cheight, cwidth) / 50, max(1, ((cheight * 0.5) / breaks))))
@@ -419,15 +414,15 @@ def plot_heatmap(
 
     # Create a colorbar
     colorbar = cv2.applyColorMap(
-        src = np.arange(256, dtype=np.uint8).reshape(256, 1).repeat(colorbar_width, 1),
-        colormap = cv2.COLORMAP_VIRIDIS
+        src=np.arange(256, dtype=np.uint8).reshape(256, 1).repeat(colorbar_width, 1),
+        colormap=cv2.COLORMAP_VIRIDIS
     )
     # Stretch the colorbar to the height of the colormap
     colorbar = cv2.resize(
-        src = colorbar, 
-        dsize = (colorbar.shape[1], cheight), 
-        dst = colorbar,
-        interpolation = cv2.INTER_LINEAR_EXACT
+        src=colorbar, 
+        dsize=(colorbar.shape[1], cheight), 
+        dst=colorbar,
+        interpolation=cv2.INTER_LINEAR_EXACT
     )
     cbheight = colorbar.shape[0]
     # Add the breaks to the colorbar
@@ -435,30 +430,24 @@ def plot_heatmap(
         label = labels[i]
         text_width, font_height = label_sizes[i]
         cv2.putText(
-            img = colorbar, 
-            text = label, 
-            org = (
-                (colorbar_width - text_width) // 2, 
-                int((i + 0.5) * cbheight / len(nice_breaks) + font_height / 2)
-            ),
-            fontFace = cv2.FONT_HERSHEY_SIMPLEX, 
-            fontScale = font_size, 
-            color = (0, 0, 0), 
-            thickness = (3 * font_height_target) // 15, 
-            lineType = cv2.LINE_AA
+            img=colorbar,
+            text=label,
+            org=((colorbar_width - text_width) // 2, int((i + 0.5) * cbheight / len(nice_breaks) + font_height / 2)),
+            fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+            fontScale=font_size,
+            color=(0, 0, 0),
+            thickness=(3 * font_height_target) // 15,
+            lineType=cv2.LINE_AA,
         )
         cv2.putText(
-            img = colorbar, 
-            text = label, 
-            org = (
-                (colorbar_width - text_width) // 2, 
-                int((i + 0.5) * cbheight / len(nice_breaks) + font_height / 2)
-            ),
-            fontFace = cv2.FONT_HERSHEY_SIMPLEX, 
-            fontScale = font_size, 
-            color = (255, 255, 255), 
-            thickness = font_height_target // 15, 
-            lineType = cv2.LINE_AA
+            img=colorbar,
+            text=label,
+            org=((colorbar_width - text_width) // 2, int((i + 0.5) * cbheight / len(nice_breaks) + font_height / 2)),
+            fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+            fontScale=font_size,
+            color=(255, 255, 255),
+            thickness=font_height_target // 15,
+            lineType=cv2.LINE_AA,
         )
 
     # Concatenate the colormap and the colorbar
@@ -472,9 +461,9 @@ def plot_heatmap(
         # Calculate the font size
         axis_label_font_size = cv2.getFontScaleFromHeight(cv2.FONT_HERSHEY_COMPLEX, axis_box_size // 2, 3)
 
-        # First create the x-axis box
+        # Create the x-axis box
         x_axis_box = np.zeros((axis_box_size, dimensions[0] + colorbar_width, 3), dtype=np.uint8) + 255
-        # Then create the y-axis box, remembering to take into account the extra vertical space taken up by the x-axis label. 
+        # Create the y-axis box, remembering to take into account the extra vertical space taken up by the x-axis label.
         # It is instantiated in the flipped orientation.
         y_axis_box = np.zeros((axis_box_size, dimensions[1] + axis_box_size, 3), dtype=np.uint8) + 255
         # Calculate the midpoint on each box with respect to the heatmap
@@ -485,25 +474,25 @@ def plot_heatmap(
         center_offset = axis_box_size // 2 + axis_box_size // 4
         # Add the x-axis label
         cv2.putText(
-            img = x_axis_box, 
-            text = x_label, 
-            org = (x_midpoint, center_offset), 
-            fontFace = cv2.FONT_HERSHEY_COMPLEX, 
-            fontScale = axis_label_font_size,
-            color = (0, 0, 0), 
-            thickness = axis_box_size // 15, 
-            lineType = cv2.LINE_AA
+            img=x_axis_box,
+            text=x_label,
+            org=(x_midpoint, center_offset),
+            fontFace=cv2.FONT_HERSHEY_COMPLEX,
+            fontScale=axis_label_font_size,
+            color=(0, 0, 0),
+            thickness=axis_box_size // 15,
+            lineType=cv2.LINE_AA,
         )
         # Add the y-axis label
         cv2.putText(
-            img = y_axis_box, 
-            text = y_label, 
-            org = (y_midpoint, center_offset), 
-            fontFace = cv2.FONT_HERSHEY_COMPLEX, 
-            fontScale = axis_label_font_size,
-            color = (0, 0, 0), 
-            thickness = axis_box_size // 15, 
-            lineType = cv2.LINE_AA
+            img=y_axis_box,
+            text=y_label,
+            org=(y_midpoint, center_offset),
+            fontFace=cv2.FONT_HERSHEY_COMPLEX,
+            fontScale=axis_label_font_size,
+            color=(0, 0, 0),
+            thickness=axis_box_size // 15,
+            lineType=cv2.LINE_AA,
         )
         # Flip the y-axis box
         y_axis_box = cv2.rotate(y_axis_box, cv2.ROTATE_90_CLOCKWISE)
@@ -514,31 +503,28 @@ def plot_heatmap(
     if scale != 1:
         # Rescale the colormap to 2x lower resolution
         colormap = cv2.resize(
-            src = colormap, 
-            dsize = (int(colormap.shape[1] * scale), int(colormap.shape[0] * scale)),
-            dst = colormap,
-            interpolation = cv2.INTER_LINEAR
+            src=colormap,
+            dsize=(int(colormap.shape[1] * scale), int(colormap.shape[0] * scale)),
+            dst=colormap,
+            interpolation=cv2.INTER_LINEAR,
         )
 
     if output_path is not None:
         # Save the image
         cv2.imwrite(
-            filename = output_path, 
-            img = colormap, 
-            params = [int(cv2.IMWRITE_JPEG_QUALITY), 95]
+            filename=output_path, 
+            img=colormap, 
+            params=[int(cv2.IMWRITE_JPEG_QUALITY), 95]
         )
     else:
         compatible_display(colormap)
 
-def equal_spaced_cuts(
-        k : int, 
-        start : float | int, 
-        end : float | int
-    ) -> np.ndarray:
+
+def equal_spaced_cuts(k: int, start: float | int, end: float | int) -> np.ndarray:
     """Generate k equal spaced cuts between start and end.
-    
-    The edges are not included, and the distance between the left-most and 
-    right-most cut to the edges is half the distance between the cuts. 
+
+    The edges are not included, and the distance between the left-most and
+    right-most cut to the edges is half the distance between the cuts.
 
     Args:
         k: Number of cuts.
@@ -553,57 +539,57 @@ def equal_spaced_cuts(
 
 
 def plot_matches(
-        matches: np.ndarray, 
-        contours1: list[np.ndarray], 
-        contours2: list[np.ndarray],
-        group_labels: Sequence[str] | None = None, 
-        image_path: str | None = None,
-        output_path: str | None = None, 
-        scale: float = 1, 
-        boxes: bool = True
-    ):
+    matches: np.ndarray,
+    contours1: list[np.ndarray],
+    contours2: list[np.ndarray],
+    group_labels: Sequence[str] | None = None,
+    image_path: str | None = None,
+    output_path: str | None = None,
+    scale: float = 1,
+    boxes: bool = True,
+):
     """Plot the matches between two groups of contours using OpenCV.
 
     Args:
         matches: Matches between group 1 and 2.
         contours1: Contours of group 1.
         contours2: Contours of group 2.
-        group_labels: Labels of the two groups (should have a length of 2). If None the groups are labelled as "1" and "2". 
-            Defaults to None.
+        group_labels: Labels of the two groups (should have a length of 2).
+            If None the groups are labelled as "1" and "2". Defaults to None.
         image_path: Path to the image. If None the matches are plotted on a black background. Defaults to None.
-        output_path: Output path of plot. 
-            If None the rasterized result is displayed or returned as an array, depending in the context. 
+        output_path: Output path of plot.
+            If None the rasterized result is displayed or returned as an array, depending in the context.
             Defaults to None.
         scale: Scale of the output image. Defaults to 1.
-        boxes: Flag to indicate whether to plot bounding boxes. Defaults to True.
+        boxes: Flag indicating whether to plot bounding boxes. Defaults to True.
 
     """
     GROUP_COLORS = [(184, 126, 55), (28, 26, 228)]
     # Type check the input
     if not isinstance(matches, np.ndarray):
-        raise ValueError(f'Expected matches to be a NumPy array, got {type(matches)}')
+        raise ValueError(f"Expected matches to be a NumPy array, got {type(matches)}")
     for i, c1 in enumerate(contours1):
         if not isinstance(c1, np.ndarray):
-            raise ValueError(f'Expected contours1[{i}] to be a NumPy array, got {type(c1).__name__}')
+            raise ValueError(f"Expected contours1[{i}] to be a NumPy array, got {type(c1).__name__}")
     for i, c2 in enumerate(contours2):
         if not isinstance(c2, np.ndarray):
-            raise ValueError(f'Expected contours2[{i}] to be a NumPy array, got {type(c2).__name__}')
+            raise ValueError(f"Expected contours2[{i}] to be a NumPy array, got {type(c2).__name__}")
     if not isinstance(group_labels, list) and group_labels is not None:
-        raise ValueError(f'Expected group_labels to be a list or None, got {type(group_labels)}')
+        raise ValueError(f"Expected group_labels to be a list or None, got {type(group_labels)}")
     elif isinstance(group_labels, list):
         for i, lab in enumerate(group_labels):
             if not isinstance(lab, str):
-                raise ValueError(f'Expected group_labels[{i}] to be a string, got {type(lab).__name__}')
+                raise ValueError(f"Expected group_labels[{i}] to be a string, got {type(lab).__name__}")
     elif group_labels is None:
         group_labels = ["1", "2"]
     if not isinstance(image_path, str) and image_path is not None:
-        raise ValueError(f'Expected image_path to be a string or None, got {type(image_path).__name__}')
+        raise ValueError(f"Expected image_path to be a string or None, got {type(image_path).__name__}")
     elif isinstance(image_path, str) and not os.path.exists(image_path):
-        raise ValueError(f'Expected image_path to be a valid file, got {image_path}')
+        raise ValueError(f"Expected image_path to be a valid file, got {image_path}")
     if not isinstance(output_path, str) and output_path is not None:
-        raise ValueError(f'Expected output_path to be a string or None, got {type(output_path).__name__}')
+        raise ValueError(f"Expected output_path to be a string or None, got {type(output_path).__name__}")
     elif isinstance(output_path, str) and not os.path.exists(os.path.dirname(output_path)):
-        raise ValueError(f'Output directory does not exist: {os.path.dirname(output_path)}')
+        raise ValueError(f"Output directory does not exist: {os.path.dirname(output_path)}")
 
     # If a output path is provided the image is saved, otherwise it is displayed with IPython
     save_plot = isinstance(output_path, str)
@@ -611,12 +597,12 @@ def plot_matches(
         # Check the output path extension
         _, out_ext = os.path.splitext(output_path)
         if out_ext not in [".jpg", ".jpeg", ".JPG", ".JPEG"]:
-            raise ValueError(f'Expected output path to have a .JPG/.jpg/.JPEG/.jpeg extension, got {out_ext}')
+            raise ValueError(f"Expected output path to have a .JPG/.jpg/.JPEG/.jpeg extension, got {out_ext}")
 
     # If the is image path is provided
     if isinstance(image_path, str):
         # Load the image
-        image = cv2.imread(filename = image_path)
+        image = cv2.imread(filename=image_path)
         assert image is not None
     else:
         # Otherwise, create a blank image. The dimensions are dynamically calculated to fit the contours
@@ -646,71 +632,49 @@ def plot_matches(
     for idx, (i, j) in enumerate(matches):
         if i != -1:
             # Draw the first contour mask on the first copy
-            cv2.fillPoly(
-                img = cimg1, 
-                pts = [contours1[i]], 
-                color = GROUP_COLORS[0]
-            )
+            cv2.fillPoly(img=cimg1, pts=[contours1[i]], color=GROUP_COLORS[0])
             cv2.drawContours(
-                image = image, 
-                contours = [contours1[i]], 
-                contourIdx = -1,
-                color = GROUP_COLORS[0],
-                thickness = 4,
-                lineType=cv2.LINE_AA
+                image=image,
+                contours=[contours1[i]],
+                contourIdx=-1,
+                color=GROUP_COLORS[0],
+                thickness=4,
+                lineType=cv2.LINE_AA,
             )
             if boxes:
                 # Draw boxes around the contours on the original image
                 cv2.rectangle(
-                    img = image, 
-                    pt1 = (bboxes1[i][0], bboxes1[i][1]), 
-                    pt2 = (bboxes1[i][2], bboxes1[i][3]), 
-                    color = GROUP_COLORS[0], 
-                    thickness = 8
+                    img=image,
+                    pt1=(bboxes1[i][0], bboxes1[i][1]),
+                    pt2=(bboxes1[i][2], bboxes1[i][3]),
+                    color=GROUP_COLORS[0],
+                    thickness=8,
                 )
         if j != -1:
             # Draw the second contour mask on the second copy
-            cv2.fillPoly(
-                img = cimg2, 
-                pts = [contours2[j]], 
-                color = GROUP_COLORS[1]
-            )
+            cv2.fillPoly(img=cimg2, pts=[contours2[j]], color=GROUP_COLORS[1])
             cv2.drawContours(
-                image = image, 
-                contours = [contours2[j]],
-                contourIdx=-1, 
-                color = GROUP_COLORS[1],
+                image=image,
+                contours=[contours2[j]],
+                contourIdx=-1,
+                color=GROUP_COLORS[1],
                 thickness=4,
-                lineType=cv2.LINE_AA
+                lineType=cv2.LINE_AA,
             )
             if boxes:
                 # Draw boxes around the contours on the original image
                 cv2.rectangle(
-                    img = image, 
-                    pt1 = (bboxes2[j][0], bboxes2[j][1]), 
-                    pt2 = (bboxes2[j][2], bboxes2[j][3]), 
-                    color = GROUP_COLORS[1], 
-                    thickness = 8
+                    img=image,
+                    pt1=(bboxes2[j][0], bboxes2[j][1]),
+                    pt2=(bboxes2[j][2], bboxes2[j][3]),
+                    color=GROUP_COLORS[1],
+                    thickness=8,
                 )
 
     # Blend the image copies with the contour masks together (makes the contours semi-transparent - alpha=0.5)
-    cv2.addWeighted(
-        src1 = cimg1, 
-        alpha = 0.5, 
-        src2 = cimg2, 
-        beta = 0.5, 
-        gamma = 0, 
-        dst = cimg1
-    )
+    cv2.addWeighted(src1=cimg1, alpha=0.5, src2=cimg2, beta=0.5, gamma=0, dst=cimg1)
     # Blend with the blended copies with original image
-    cv2.addWeighted(
-        src1 = image, 
-        alpha = 0.5, 
-        src2 = cimg1, 
-        beta = 0.5, 
-        gamma = 0, 
-        dst = image
-    )
+    cv2.addWeighted(src1=image, alpha=0.5, src2=cimg1, beta=0.5, gamma=0, dst=image)
 
     # Downscale the image
     image = cv2.resize(image, (image.shape[1] // 2, image.shape[0] // 2))
@@ -728,8 +692,9 @@ def plot_matches(
         label_font_color = []
         if i != -1:
             # Draw a text label next to the box
-            label_width = \
-                cv2.getTextSize(f"{idx}", cv2.FONT_HERSHEY_SIMPLEX, label_font_scale, label_font_thickness)[0][0]
+            label_width = cv2.getTextSize(
+                f"{idx}", cv2.FONT_HERSHEY_SIMPLEX, label_font_scale, label_font_thickness
+            )[0][0]
             label_font_color.append(GROUP_COLORS[0])
             label_coord.append((bboxes1[i][0] - label_width, bboxes1[i][1] + label_font_height))
         if j != -1:
@@ -738,41 +703,41 @@ def plot_matches(
             label_coord.append((bboxes2[j][0], bboxes2[j][1] - label_font_height // 4))
         for coord, color in zip(label_coord, label_font_color):
             cv2.putText(
-                img = image, 
-                text = str(idx), 
-                org = coord, 
-                fontFace = cv2.FONT_HERSHEY_SIMPLEX, 
-                fontScale = label_font_scale, 
-                color = (2, 210, 238) if no_match else (0, 0, 0),
-                thickness = label_font_thickness * 3, 
-                lineType = cv2.LINE_8
+                img=image,
+                text=str(idx),
+                org=coord,
+                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                fontScale=label_font_scale,
+                color=(2, 210, 238) if no_match else (0, 0, 0),
+                thickness=label_font_thickness * 3,
+                lineType=cv2.LINE_8,
             )
             cv2.putText(
-                img = image, 
-                text = str(idx), 
-                org = coord, 
-                fontFace = cv2.FONT_HERSHEY_SIMPLEX, 
-                fontScale = label_font_scale, 
-                color = (147, 20, 255) if no_match else color,
-                thickness = label_font_thickness, 
-                lineType = cv2.LINE_AA
+                img=image,
+                text=str(idx),
+                org=coord,
+                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                fontScale=label_font_scale,
+                color=(147, 20, 255) if no_match else color,
+                thickness=label_font_thickness,
+                lineType=cv2.LINE_AA,
             )
-            
 
     # Create a legend
-    LEGEND_TEXT_Y_JUST = 1/2
+    LEGEND_TEXT_Y_JUST = 1 / 2
     legend_font_height = max(min(6, image.shape[0] // 40), 24)
     legend_margin = max(1, int(image.shape[0] * 0.01))
     legend_font_size = cv2.getFontScaleFromHeight(cv2.FONT_HERSHEY_COMPLEX, legend_font_height, 3)
-    legend_label_widths = \
-        [cv2.getTextSize(glabel, cv2.FONT_HERSHEY_COMPLEX, legend_font_size, 3)[0][0] for glabel in group_labels]
+    legend_label_widths = [
+        cv2.getTextSize(glabel, cv2.FONT_HERSHEY_COMPLEX, legend_font_size, 3)[0][0] for glabel in group_labels
+    ]
     legend_font_width = max(legend_label_widths)
     legend_box_height = int((legend_font_height * 1.6) * len(group_labels))
     legend_box_width = int(legend_font_width * 1.25)
     # Extract the legend box
     legend_box = image[
         legend_margin:(legend_box_height + legend_margin),
-        -(legend_box_width + legend_margin):-legend_margin, 
+        -(legend_box_width + legend_margin):-legend_margin,
         :
     ]
     # Whiten legend box
@@ -781,25 +746,25 @@ def plot_matches(
     legend_box += whiten_amount
     # Add a black border to the legend box
     cv2.rectangle(
-        img = legend_box, 
-        pt1 = (1, 1), 
-        pt2 = (legend_box_width - legend_font_height // 15 - 1, legend_box_height - legend_font_height // 15 - 1), 
-        color = (0, 0, 0), 
-        thickness = legend_font_height // 15,
-        lineType=cv2.LINE_AA
+        img=legend_box,
+        pt1=(1, 1),
+        pt2=(legend_box_width - legend_font_height // 15 - 1, legend_box_height - legend_font_height // 15 - 1),
+        color=(0, 0, 0),
+        thickness=legend_font_height // 15,
+        lineType=cv2.LINE_AA,
     )
     # Add the legend labels and items
-    legend_attributes = []
+    legend_attributes: list[tuple[str, int, tuple[int, int, int], int, int]] = []
     item_cut_ys = equal_spaced_cuts(len(group_labels), 0, legend_box_height)
     for i, (glabel, label_width, item_cut_y) in enumerate(zip(group_labels, legend_label_widths, item_cut_ys)):
         item_label_y = int(item_cut_y + legend_font_height * LEGEND_TEXT_Y_JUST)
         # Labels
         label_x = int(legend_box_width * 0.975) - label_width
-        
+
         # Items - positioned to the left of the labels with a margin of 'legend_margin'
         item_x = label_x // 2
         item_color = GROUP_COLORS[i]
-        
+
         # Add the attributes to the legend_attributes list
         legend_attributes.append((glabel, item_x, item_color, label_x, item_label_y))
 
@@ -809,50 +774,52 @@ def plot_matches(
     for glabel, _, item_color, label_x, item_label_y in legend_attributes:
         # Draw the legend item label
         cv2.putText(
-            img = legend_box, 
-            text = glabel, 
-            org = (label_x, item_label_y), 
-            fontFace = cv2.FONT_HERSHEY_COMPLEX, 
-            fontScale = legend_font_size, 
-            color = (0, 0, 0),
-            thickness = legend_font_height // 15, 
-            lineType = cv2.LINE_AA
+            img=legend_box,
+            text=glabel,
+            org=(label_x, item_label_y),
+            fontFace=cv2.FONT_HERSHEY_COMPLEX,
+            fontScale=legend_font_size,
+            color=(0, 0, 0),
+            thickness=legend_font_height // 15,
+            lineType=cv2.LINE_AA,
         )
         # Fill the item circle with the color of the group
         cv2.circle(
-            img = legend_box, 
-            center = (min_item_x, item_label_y - legend_font_height // 2), 
-            radius = legend_font_height // 2, 
-            color = item_color,
-            thickness=cv2.FILLED
+            img=legend_box,
+            center=(min_item_x, item_label_y - legend_font_height // 2),
+            radius=legend_font_height // 2,
+            color=item_color,
+            thickness=cv2.FILLED,
         )
         # Add a black border to the item circle
         cv2.circle(
-            img = legend_box,
-            center = (min_item_x, item_label_y - legend_font_height // 2), 
-            radius = legend_font_height // 2, 
-            color = (0, 0, 0),
-            thickness = legend_font_height // 30,
-            lineType=cv2.LINE_AA
+            img=legend_box,
+            center=(min_item_x, item_label_y - legend_font_height // 2),
+            radius=legend_font_height // 2,
+            color=(0, 0, 0),
+            thickness=legend_font_height // 30,
+            lineType=cv2.LINE_AA,
         )
 
     # Add the legend to the image
-    image[legend_margin:(legend_box_height + legend_margin), -(legend_box_width + legend_margin):-legend_margin, :] = legend_box
+    image[
+        legend_margin : (legend_box_height + legend_margin), -(legend_box_width + legend_margin) : -legend_margin, :
+    ] = legend_box
 
     if scale != 1:
         # Scale the image
         cv2.resize(
-            src = image, 
-            dsize = (int(image.shape[1] * scale), int(image.shape[0] * scale)), 
-            dst = image
+            src=image, 
+            dsize=(int(image.shape[1] * scale), int(image.shape[0] * scale)), 
+            dst=image
         )
 
     if save_plot:
         # Save the image
         cv2.imwrite(
-            filename = output_path, 
-            img = image, 
-            params = [int(cv2.IMWRITE_JPEG_QUALITY), 95]
+            filename=output_path, 
+            img=image, 
+            params=[int(cv2.IMWRITE_JPEG_QUALITY), 95]
         )
     else:
         compatible_display(image)
@@ -861,8 +828,7 @@ def plot_matches(
 def compatible_display(image: np.ndarray):  # noqa: D103
     TIMEOUT = 5  # seconds
     # Check if the image is displayed in a Jupyter notebook
-    if 'get_ipython' in globals():
-        # Only import the necessary modules if the image is displayed in a Jupyter notebook, ensures they are optional dependencies
+    if "get_ipython" in globals():
         import ipywidgets as widgets  # type: ignore
         from IPython.display import clear_output, display  # type: ignore
 
@@ -883,7 +849,7 @@ def on_button_clicked(b):
         button.on_click(on_button_clicked)
 
         # Display the image
-        display(widgets.Image(value=cv2.imencode('.jpg', image_rgb)[1].tobytes(), format='jpg'))
+        display(widgets.Image(value=cv2.imencode(".jpg", image_rgb)[1].tobytes(), format="jpg"))
         # Display the button
         display(button)
 
@@ -891,30 +857,30 @@ def on_button_clicked(b):
         start_time = time.time()
         while not button_clicked and (time.time() - start_time) < TIMEOUT:
             time.sleep(0.01)
-        logger.info('Image display closed')
+        logger.info("Image display closed")
     else:
         # Check if a display is available
-        if os.environ.get('DISPLAY', '') == '':
-            logger.info('No display found, unable to display the image')
+        if os.environ.get("DISPLAY", "") == "":
+            logger.info("No display found, unable to display the image")
         else:
             # Display the image
-            cv2.imshow('Matches', image)
+            cv2.imshow("Matches", image)
             cv2.waitKey(0)
             cv2.destroyAllWindows()
 
 
 def compare_groups(
-        group1: list, 
-        group2: list, 
-        threshold: float=0.1,
-        group_labels: Sequence[str] | None=("Ground Truth", "Predictions"), 
-        plot: bool=False, 
-        plot_scale: float=1, 
-        plot_boxes: bool=True,
-        image_path: str | None=None, 
-        output_identifier: str | None=None, 
-        output_directory: str | None=None
-    ) -> str | dict:
+    group1: list,
+    group2: list,
+    threshold: float = 0.1,
+    group_labels: Sequence[str] | None = ("Ground Truth", "Predictions"),
+    plot: bool = False,
+    plot_scale: float = 1,
+    plot_boxes: bool = True,
+    image_path: str | None = None,
+    output_identifier: str | None = None,
+    output_directory: str | None = None,
+) -> str | dict:
     """Compare group 1 to group 2.
 
     Output is saved to a CSV file with the following columns:
@@ -924,10 +890,10 @@ def compare_groups(
         - contourArea_1 (`int`) is the area of the contour in group 1
         - contourArea_2 (`int`) is the area of the contour in group 2, or 0 if there is no match
         - bbox_1 (`list[int, int, int, int : xmin, ymin, xmax, ymax]`) is the bounding box of the geometry in group 1
-        - bbox_2 (`list[int, int, int, int : xmin, ymin, xmax, ymax]`) is the bounding box of the 
+        - bbox_2 (`list[int, int, int, int : xmin, ymin, xmax, ymax]`) is the bounding box of the
             matched geometry in group 2, or an empty list if there is no match
         - contour_1 (`list[list[int, int : x_i, y_i]]`) is the contour of the geometry in group 1
-        - contour_2 (`list[list[int, int : x_i, y_i]]`) is the contour of the matched geometry in group 2, 
+        - contour_2 (`list[list[int, int : x_i, y_i]]`) is the contour of the matched geometry in group 2,
             or an empty list if there is no match
 
     Args:
@@ -935,7 +901,7 @@ def compare_groups(
         group2: Group 2.
         threshold: IoU threshold for matching elements between groups. Defaults to 0.1.
         group_labels: Group labels. Defaults to `["Ground Truth", "Predictions"]`.
-        plot: Whether to plot the matches and the IoU matrix, 
+        plot: Whether to plot the matches and the IoU matrix,
             usually this is much slower than simply comparing the groups. Defaults to False.
         plot_scale: Scale of the plot. Defaults to 1. Lower values will make the plot smaller, but may be faster.
         plot_boxes: Whether to plot the bounding boxes. Defaults to True.
@@ -944,28 +910,31 @@ def compare_groups(
         output_directory: Output directory. Defaults to None.
 
     Returns:
-        Path to the CSV file or the data that would have been saved to the CSV file as a dictionary, 
+        Path to the CSV file or the data that would have been saved to the CSV file as a dictionary,
         where the keys are the column names and the values are the column values.
 
     """
     # Type check the input
     if not isinstance(group1, list) or not isinstance(group2, list):
-        raise ValueError(f'Expected group1 and group2 to be lists, got {type(group1)} and {type(group2)}')
+        raise ValueError(f"Expected group1 and group2 to be lists, got {type(group1)} and {type(group2)}")
     if not all([isinstance(i, dict) for i in group1]) or not all([isinstance(i, dict) for i in group2]):
         raise ValueError(
-            f'Expected group1 and group2 to be lists of dictionaries, got {type(group1[0])} and {type(group2[0])}')
+            f"Expected group1 and group2 to be lists of dictionaries, got {type(group1[0])} and {type(group2[0])}"
+        )
     if not isinstance(threshold, float):
-        raise ValueError(f'Expected threshold to be a float, got {type(threshold)}')
+        raise ValueError(f"Expected threshold to be a float, got {type(threshold)}")
     if not isinstance(plot, bool):
-        raise ValueError(f'Expected plot to be a bool, got {type(plot)}')
+        raise ValueError(f"Expected plot to be a bool, got {type(plot)}")
     if not (isinstance(image_path, str) or image_path is None):
-        raise ValueError(f'Expected image_path to be a string or None, got {type(image_path)}')
+        raise ValueError(f"Expected image_path to be a string or None, got {type(image_path)}")
     if not isinstance(output_directory, str) and output_directory is not None:
-        raise ValueError(f'Expected output_directory to be a string or None, got {type(output_directory)}')
+        raise ValueError(f"Expected output_directory to be a string or None, got {type(output_directory)}")
     elif isinstance(output_directory, str) and not os.path.isdir(output_directory):
-        raise ValueError(f'Expected output_directory to be a valid directory, got {output_directory}')
+        raise ValueError(f"Expected output_directory to be a valid directory, got {output_directory}")
     if not isinstance(output_identifier, str) and (plot or output_directory is not None):
-        raise ValueError(f'Expected output_identifier to be a string when saving to file or plotting, got {type(output_identifier)}')
+        raise ValueError(
+            f"Expected output_identifier to be a string when saving to file or plotting, got {type(output_identifier)}"
+        )
 
     # Convert the annotations to NumPy arrays, and calculate bounding boxes and areas
     b1, c1 = annotations_to_numpy(group1)
@@ -983,25 +952,29 @@ def compare_groups(
     # Plot the matches and the IoU matrix
     if plot:
         if not isinstance(image_path, str):
-            raise ValueError(f'Expected path to be a string, got {type(image_path)}')
+            raise ValueError(f"Expected path to be a string, got {type(image_path)}")
         elif not os.path.isfile(image_path):
-            raise ValueError(f'Expected path to be a valid file, got {image_path}')
+            raise ValueError(f"Expected path to be a valid file, got {image_path}")
         plot_matches(
-            matches = matches, 
-            contours1 = c1, 
-            contours2 = c2, 
-            group_labels = group_labels, 
-            image_path = image_path,
-            output_path = os.path.join(output_directory, f'{output_identifier}_matches.jpg') if output_directory is not None else None,
-            scale = plot_scale, 
-            boxes = plot_boxes
+            matches=matches,
+            contours1=c1,
+            contours2=c2,
+            group_labels=group_labels,
+            image_path=image_path,
+            output_path=os.path.join(output_directory, f"{output_identifier}_matches.jpg")
+            if output_directory is not None
+            else None,
+            scale=plot_scale,
+            boxes=plot_boxes,
         )
         if not any([dim == 0 for dim in iou.shape]):
             plot_heatmap(
-                mat = iou, 
-                axis_labels = group_labels[::-1] if group_labels is not None else None,
-                output_path = os.path.join(output_directory, f'{output_identifier}_heatmap.jpg') if output_directory is not None else None,
-                scale = plot_scale
+                mat=iou,
+                axis_labels=group_labels[::-1] if group_labels is not None else None,
+                output_path=os.path.join(output_directory, f"{output_identifier}_heatmap.jpg")
+                if output_directory is not None
+                else None,
+                scale=plot_scale,
             )
 
     ## Gather the data for the output
@@ -1063,14 +1036,29 @@ def compare_groups(
     len_boxes1, len_boxes2 = len(boxes1), len(boxes2)
     len_contours1, len_contours2 = len(contours1), len(contours2)
     # Check that the lengths are all the same
-    data_length = set(
-        [len_idx1, len_idx2, len_matched_iou, len_careas1, len_careas2, len_boxes1, len_boxes2, len_contours1,
-         len_contours2])
+    data_length = set([
+        len_idx1,
+        len_idx2,
+        len_matched_iou,
+        len_careas1,
+        len_careas2,
+        len_boxes1,
+        len_boxes2,
+        len_contours1,
+        len_contours2,
+    ])
     if len(data_length) != 1:
         raise ValueError(
             "Lengths of the data are not all the same: {}, {}, {}, {}, {}, {}, {}, {}, {}".format(  # noqa: UP032
-                len_idx1, len_idx2, len_matched_iou, len_careas1, 
-                len_careas2, len_boxes1, len_boxes2, len_contours1, len_contours2
+                len_idx1,
+                len_idx2,
+                len_matched_iou,
+                len_careas1,
+                len_careas2,
+                len_boxes1,
+                len_boxes2,
+                len_contours1,
+                len_contours2,
             )
         )
 
@@ -1086,7 +1074,7 @@ def compare_groups(
         "bbox_1": boxes1,
         "bbox_2": boxes2,
         "contour_1": contours1,
-        "contour_2": contours2
+        "contour_2": contours2,
     }
 
     if output_directory is not None:
@@ -1106,6 +1094,7 @@ def compare_groups(
     else:
         return output
 
+
 def generate_block(min: int, max: int, size: int) -> np.ndarray:
     """Generate a block of integers centered around a random start value within a given range.
 
@@ -1120,12 +1109,12 @@ def generate_block(min: int, max: int, size: int) -> np.ndarray:
     """
     if size <= 0 or min >= max:
         raise ValueError("Size must be positive and min must be less than max.")
-    
+
     start = np.random.randint(min, max)
     left = np.random.randint(0, size)
     right = size - left
     block = np.arange(start - left, start + right)
-    
+
     return block[np.logical_and(block >= min, block < max)]
 
 
@@ -1145,12 +1134,16 @@ def generate_bootstraps(s: int, n: int, block: bool = False) -> list[np.ndarray]
         raise ValueError("The size 's' and the number 'n' of bootstraps must be positive.")
 
     if block:
-        blocks = int(max(1, (s ** 0.5) // 2))
-        return [np.concatenate([generate_block(min=0, max=s, size=s // blocks) for _ in range(blocks)]) for _ in range(n)]
+        blocks = int(max(1, (s**0.5) // 2))
+        return [
+            np.concatenate([generate_block(min=0, max=s, size=s // blocks) for _ in range(blocks)])
+            for _ in range(n)
+        ]
     else:
         return [np.random.choice(s, s, replace=True) for _ in range(n)]
 
-def f1_score(GT : np.ndarray, MP : np.ndarray) -> float:
+
+def f1_score(GT: np.ndarray, MP: np.ndarray) -> float:
     """Calculate the F1 score for a binary classification problem.
 
     Args:
@@ -1163,7 +1156,7 @@ def f1_score(GT : np.ndarray, MP : np.ndarray) -> float:
     """
     if len(GT) != len(MP):
         raise ValueError("Lengths of GT and MP must match.")
-    
+
     TP = np.sum(MP & GT)
     FP = np.sum(MP & ~GT)
     FN = np.sum(~MP & GT)
@@ -1173,17 +1166,13 @@ def f1_score(GT : np.ndarray, MP : np.ndarray) -> float:
 
     return float(2 * precision * recall / (precision + recall) if precision + recall > 0 else 0)
 
-def optimal_threshold_f1(
-    y: np.ndarray,
-    iou: np.ndarray,
-    confidence: np.ndarray,
-    num_thresholds: int = 100
-) -> float:
+
+def optimal_threshold_f1(y: np.ndarray, iou: np.ndarray, confidence: np.ndarray, num_thresholds: int = 100) -> float:
     """Find the optimal threshold for F1 score by iterating over possible thresholds.
 
     Args:
         y: Ground truth binary labels.
-        iou: IoU values, 
+        iou: IoU values,
         confidence: Confidence scores for predictions.
         num_thresholds: Number of thresholds to test. Defaults to 100.
 
@@ -1193,14 +1182,14 @@ def optimal_threshold_f1(
     """
     if len(y) != len(iou) or len(y) != len(confidence):
         raise ValueError("Lengths of y, iou, and confidence must match.")
-    
+
     y = np.asarray(y, dtype=bool)
     best_threshold, best_f1 = 0, 0
 
     for i in range(num_thresholds + 1):
         threshold = i / num_thresholds
         MP = confidence >= threshold
-        
+
         f1 = f1_score(y, MP)
 
         if f1 > best_f1:
@@ -1211,10 +1200,7 @@ def optimal_threshold_f1(
 
 
 def best_confidence_threshold(
-    y: list[int] | np.ndarray,
-    iou: list[float] | np.ndarray,
-    confidence: list[float] | np.ndarray,
-    n: int = 100
+    y: list[int] | np.ndarray, iou: list[float] | np.ndarray, confidence: list[float] | np.ndarray, n: int = 100
 ) -> float:
     """Find the best confidence threshold using bootstrapping and F1 score optimization.
 
@@ -1235,4 +1221,7 @@ def best_confidence_threshold(
     iou = np.asarray([i if isinstance(i, float) else 0 for i in iou], dtype=float)
     confidence = np.asarray([c if isinstance(c, float) else 0 for c in confidence], dtype=float)
 
-    return float(np.mean([optimal_threshold_f1(y[boot], iou[boot], confidence[boot]) for boot in generate_bootstraps(len(iou), n, True)]))
+    return float(np.mean([
+        optimal_threshold_f1(y[boot], iou[boot], confidence[boot])
+        for boot in generate_bootstraps(len(iou), n, True)
+    ]))
diff --git a/src/flat_bug/geometric.py b/src/flat_bug/geometric.py
index 5576fd0..f52d983 100644
--- a/src/flat_bug/geometric.py
+++ b/src/flat_bug/geometric.py
@@ -1,4 +1,5 @@
 """Geometric helper functions for flatbug."""
+
 import math
 from collections.abc import Sequence
 from itertools import accumulate
@@ -16,22 +17,24 @@
 
 
 def equal_allocate_overlaps(total: int, segments: int, size: int) -> list[int]:
-    """Generate cumulative positions for placing segments of a given size within a total length, with controlled overlaps.
+    """Generate cumulative positions for placing fixed-length segments within a total length with controlled overlaps.
 
-    This function divides the specified `total` length into `segments` positions, ensuring each segment (of given `size`) fits
-    evenly by introducing a small overlap between adjacent segments. The overlap is distributed uniformly, with the first few gaps 
-    adjusted slightly to ensure the segments collectively sum to `total`.
+    This function divides the specified `total` length into `segments` positions,
+    ensuring each segment (of given `size`) fits evenly by introducing a small overlap between adjacent segments.
+    The overlap is distributed uniformly, with the first few gaps adjusted slightly
+    to ensure the segments collectively sum to `total`.
 
     Args:
-        total: The total length to be covered by the segments. This is the target cumulative length the segments should fit into.
+        total: The total length to be covered by the segments.
+            This is the target cumulative length the segments should fit into.
         segments: The number of segments to place within the total length.
             Must be greater than or equal to 2.
         size: The desired size of each segment, used to determine the ideal spacing between segments.
-        
+
     Returns:
-        A listt of cumulative positions (starting from 0) where each segment should be placed.
+        A list of cumulative positions (starting from 0) where each segment should be placed.
         These positions are spaced with controlled overlaps to ensure they collectively cover the `total` length.
-            
+
     Example:
         >>> equal_allocate_overlaps(1000, 5, 250)
         [0, 187, 374, 562, 750]
@@ -39,31 +42,28 @@ def equal_allocate_overlaps(total: int, segments: int, size: int) -> list[int]:
     """
     if segments < 2:
         return [0] * segments
-    
+
     overlap = segments * size - total
     partial_overlap, remainder = divmod(overlap, segments - 1)
     distance = size - partial_overlap
 
     return list(accumulate([distance - (1 if i < remainder else 0) for i in range(segments - 1)], initial=0))
 
+
 def calculate_tile_offsets(  # noqa: D103
-        image_size : tuple[int, int],
-        tile_size : int,
-        minimum_overlap : int
-    ) -> list[tuple[tuple[int, int], tuple[int, int]]]:
+    image_size: tuple[int, int], tile_size: int, minimum_overlap: int
+) -> list[tuple[tuple[int, int], tuple[int, int]]]:
     w, h = image_size
     x_n_tiles = math.ceil((w - minimum_overlap) / (tile_size - minimum_overlap)) if w != tile_size else 1
     y_n_tiles = math.ceil((h - minimum_overlap) / (tile_size - minimum_overlap)) if h != tile_size else 1
-    
+
     x_range = equal_allocate_overlaps(w, x_n_tiles, tile_size)
     y_range = equal_allocate_overlaps(h, y_n_tiles, tile_size)
 
     return [((m, n), (j, i)) for n, j in enumerate(y_range) for m, i in enumerate(x_range)]
 
-def create_contour_mask(
-        mask: torch.Tensor, 
-        width: int=1
-    ) -> torch.Tensor:
+
+def create_contour_mask(mask: torch.Tensor, width: int = 1) -> torch.Tensor:
     """Convert a binary mask for a filled polygon to a binary mask for the non-filled polygon.
 
     ```
@@ -78,13 +78,13 @@ def create_contour_mask(
         # (here dashes "-" represent 0s and hashes "#" represent 1s)
     ```
 
-    We call the result ("After") the "contour mask". 
-    
+    We call the result ("After") the "contour mask".
+
     Optionally, the "linewidth" of the contour mask can be increased.
-    
+
     Args:
         mask: a NxM binary tensor with 1s inside the "polygon".
-        width: Width of the contour in the result. 
+        width: Width of the contour in the result.
             Reasonable values are >= 1; Setting to 0 will result in all 0s in the output. Defaults to 1.
 
     Returns:
@@ -105,28 +105,30 @@ def create_contour_mask(
     elif width == 1:
         return contour_mask
     elif width > 1:
-        # Expand the contour mask to include the neighbors (with a distance of less than or equal to width in either direction)
+        # Expand the contour mask to include the neighbors
+        # (with a distance of less than or equal to width in either direction)
         expansion_kernel = torch.ones((1, 1, 1 + 2 * width, 1 + 2 * width), dtype=torch.float, device=device)
-        expanded_contour_mask = F.conv2d(contour_mask.float().unsqueeze(0).unsqueeze(0), expansion_kernel, padding=width).squeeze() > 0.5
+        expanded_contour_mask = (
+            F.conv2d(contour_mask.float().unsqueeze(0).unsqueeze(0), expansion_kernel, padding=width).squeeze() > 0.5
+        )
         return expanded_contour_mask
     else:
         raise ValueError(f"Invalid width: {width}")
 
 
 @overload
-def find_contours(mask : V, largest_only : Literal[True]=True, simplify : bool=True) -> V: ...
+def find_contours(mask: V, largest_only: Literal[True] = True, simplify: bool = True) -> V: ...
 @overload
-def find_contours(mask : V, largest_only : Literal[False]=False, simplify : bool=True) -> list[V]: ...
-def find_contours(
-        mask : V, 
-        largest_only : bool=True, 
-        simplify : bool=True
-    ) -> V | list[V]:
+def find_contours(mask: V, largest_only: Literal[False] = False, simplify: bool = True) -> list[V]: ...
+def find_contours(mask: V, largest_only: bool = True, simplify: bool = True) -> V | list[V]:
     """Extract polygons from a boolean mask."""
-    contour = list(cv2.findContours(
-        mask.to(torch.uint8).cpu().numpy() if isinstance(mask, torch.Tensor) else mask.astype(np.uint8), 
-        cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
-    ))[0]
+    contour = list(
+        cv2.findContours(
+            mask.to(torch.uint8).cpu().numpy() if isinstance(mask, torch.Tensor) else mask.astype(np.uint8),
+            cv2.RETR_EXTERNAL,
+            cv2.CHAIN_APPROX_NONE,
+        )
+    )[0]
     if len(contour) == 0:
         logger.info("No contours found; mask shape:", mask.shape, "mask sum:", mask.sum())
         if isinstance(mask, torch.Tensor):
@@ -140,7 +142,7 @@ def find_contours(
         contour = contour[np.argmax(areas).item()]
     if simplify:
         contour = simplify_contour(contour, tolerance=1 if isinstance(simplify, bool) else simplify)
-    
+
     # Convert to tensor
     if isinstance(contour, list):
         contour = [np.asarray(c).squeeze(axis=1) for c in contour]
@@ -153,19 +155,17 @@ def find_contours(
     else:
         return torch.tensor(contour, dtype=torch.long, device=mask.device)
 
+
 @overload
-def simplify_contour(contour : V, tolerance : float=1.0) -> V: ...
+def simplify_contour(contour: V, tolerance: float = 1.0) -> V: ...
 @overload
-def simplify_contour(contour : Sequence[V], tolerance : float=1.0) -> list[V]: ...
-def simplify_contour(
-        contour : V | Sequence[V], 
-        tolerance : float=1.0
-    ) -> V | list[V]:
-    """Simplify one or more polygons via cv2.approxPolyDP.
-    
-    Wrapper for cv2.approxPolyDP that simplifies a contour by reducing the number of points while keeping the shape of the contour.
-    Only works for simple closed contours without holes.
-    
+def simplify_contour(contour: Sequence[V], tolerance: float = 1.0) -> list[V]: ...
+def simplify_contour(contour: V | Sequence[V], tolerance: float = 1.0) -> V | list[V]:
+    """Simplify one or more polygons via `cv2.approxPolyDP`.
+
+    Wrapper for `cv2.approxPolyDP` that simplifies a contour by reducing the number of points
+    while keeping the shape of the contour. Only works for simple closed contours without holes.
+
     Args:
         contour: The contour to simplify, represented as a Nx2 tensor or a Nx1x2 tensor.
         tolerance: The maximum distance between the original contour and the simplified contour. Defaults to 1.0.
@@ -186,29 +186,20 @@ def simplify_contour(
         elif isinstance(contour, np.ndarray):
             return np.asarray(cv2.approxPolyDP(contour, tolerance, True))
     raise TypeError(
-        f'Unable to simplify contour of type {type(contour).__name__}, '
-        'expected a torch.Tensor or np.ndarray or an iterable of such.'
+        f"Unable to simplify contour of type {type(contour).__name__}, "
+        "expected a torch.Tensor or np.ndarray or an iterable of such."
     )
 
+
 @overload
-def contours_to_masks(
-        contours : Sequence[V], 
-        height : int | torch.Tensor, 
-        width : int | torch.Tensor
-    ) -> V: ...
+def contours_to_masks(contours: Sequence[V], height: int | torch.Tensor, width: int | torch.Tensor) -> V: ...
 @overload
 def contours_to_masks(
-        contours : Sequence[Never],
-        height : int | torch.Tensor, 
-        width : int | torch.Tensor
-    ) -> torch.Tensor: ...
-def contours_to_masks(
-        contours : Sequence[V], 
-        height : int | torch.Tensor, 
-        width : int | torch.Tensor
-    ) -> V | torch.Tensor:
+    contours: Sequence[Never], height: int | torch.Tensor, width: int | torch.Tensor
+) -> torch.Tensor: ...
+def contours_to_masks(contours: Sequence[V], height: int | torch.Tensor, width: int | torch.Tensor) -> V | torch.Tensor:
     """Rasterize a list of countors to a NxHxW boolean tensor/array stack.
-    
+
     Contours should be represented as (i, j) index-coordinates in a Xx2 tensor.
 
     Args:
@@ -221,28 +212,28 @@ def contours_to_masks(
 
     """
     N = len(contours)
-    
+
     if isinstance(height, torch.Tensor):
         assert height.numel() == 1, f"Height must be a scalar tensor not {height.shape}"
         int_height = int(height.item())
     else:
         int_height = int(height)
-        
+
     if isinstance(width, torch.Tensor):
         assert width.numel() == 1, f"Width must be a scalar tensor not {width.shape}"
         int_width = int(width.item())
     else:
         int_width = int(width)
-        
+
     assert int_height > 0 and int_width > 0, f"Height and width must be positive not {int_height} and {int_width}"
 
     # Initialize the masks as UMATs
     masks = np.zeros((N, int_height, int_width), dtype=np.uint8)
-    
+
     # If there are no contours, return the empty masks gracefully
     if N == 0:
         return torch.as_tensor(masks, dtype=torch.bool)
-        
+
     is_tensor = isinstance(contours[0], torch.Tensor)
 
     # Type checking
@@ -263,24 +254,25 @@ def contours_to_masks(
         return torch.as_tensor(masks, dtype=torch.bool, device=contours[0].device)
     return masks.astype(bool)
 
+
 @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")
-def _poly_area_tensor(poly : torch.Tensor):
+def _poly_area_tensor(poly: torch.Tensor):
     if len(poly) < 10e4:
         poly = poly.cpu()
     poly_r = poly.roll(1, 0)
     return (poly[:, 0] @ poly_r[:, 1] - poly[:, 1] @ poly_r[:, 0]) / 2.0
 
 
-def poly_area(poly : torch.Tensor | np.ndarray) -> float:
-    """Calculate the area of a 2D simple polygon represented by a positively oriented (counter clock wise) sequence of points.
+def poly_area(poly: torch.Tensor | np.ndarray) -> float:
+    """Calculate the area of a simple 2D polygon.
 
+    The polygon must be represented by a positively oriented (counter clock-wise) sequence of points.
     See https://en.wikipedia.org/wiki/Shoelace_formula#Shoelace_formula for details.
 
     Args:
-        poly: A tensor or array of shape (n, 2), 
-            where n is the number of vertices 
+        poly: A tensor or array of shape (n, 2), where n is the number of vertices
             and the 2 columns are the x and y coordinates of the vertices.
-    
+
     Returns:
         The area of the polygon
 
@@ -291,12 +283,14 @@ def poly_area(poly : torch.Tensor | np.ndarray) -> float:
         poly_r = np.roll(poly, 1, axis=0)
         return float((poly[:, 0] @ poly_r[:, 1] - poly[:, 1] @ poly_r[:, 0]) / 2.0)
 
-def _poly_normals_arr(polygon : np.ndarray) -> np.ndarray:
+
+def _poly_normals_arr(polygon: np.ndarray) -> np.ndarray:
     v = np.roll(polygon, -1, axis=0) - polygon
     n = np.column_stack([v[:, 1], -v[:, 0]])
     n = (n + np.roll(n, 1, axis=0)) / 2
     return n
 
+
 @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")
 def _poly_normals_tensor(polygon: torch.Tensor) -> torch.Tensor:
     v = torch.roll(polygon, shifts=-1, dims=0) - polygon
@@ -304,14 +298,16 @@ def _poly_normals_tensor(polygon: torch.Tensor) -> torch.Tensor:
     n = (n + torch.roll(n, shifts=1, dims=0)) / 2
     return n
 
-def poly_normals(polygon : V) -> V:
+
+def poly_normals(polygon: V) -> V:
     """Calculate the normals of a polygon.
 
     Args:
-        polygon: A tensor of shape (n, 2), where n is the number of vertices and the 2 columns are the x and y coordinates of the vertices.
+        polygon: A tensor or array of shape `(n, 2)`,
+            where `n` is the number of vertices and the 2 columns are the x and y coordinates of the vertices.
 
     Returns:
-        A tensor of shape (n, 2), where n is the number of vertices and the 2 columns are the x and y coordinates of the normals.
+        A tensor or array equivalent to the input `polygon`.
 
     """
     if isinstance(polygon, torch.Tensor):
@@ -330,10 +326,11 @@ def _linear_interpolate_arr(poly: np.ndarray, scale: int) -> np.ndarray:
 
     new_poly = np.zeros((poly.shape[0] * scale, 2), dtype=np.float32)
     for i in range(poly.shape[0] - 1):
-        new_poly[i*scale:(i+1)*scale] = np.linspace(poly[i], poly[i+1], scale, endpoint=False)
+        new_poly[i * scale : (i + 1) * scale] = np.linspace(poly[i], poly[i + 1], scale, endpoint=False)
     new_poly[-scale:] = np.linspace(poly[-1], poly[0], scale, endpoint=False)
     return new_poly[~(new_poly == np.roll(new_poly, -1, axis=0)).all(axis=1)]
 
+
 @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")
 def _linear_interpolate_tensor(poly: torch.Tensor, scale: int) -> torch.Tensor:
     if scale < 1:
@@ -347,16 +344,14 @@ def _linear_interpolate_tensor(poly: torch.Tensor, scale: int) -> torch.Tensor:
     # Using vector math to precisely replicate np.linspace(..., endpoint=False)
     weights = torch.arange(scale, dtype=torch.float32, device=poly.device).unsqueeze(1) / scale
     for i in range(poly.shape[0] - 1):
-        new_poly[i*scale:(i+1)*scale] = poly[i] + (poly[i+1] - poly[i]) * weights
+        new_poly[i * scale : (i + 1) * scale] = poly[i] + (poly[i + 1] - poly[i]) * weights
     new_poly[-scale:] = poly[-1] + (poly[0] - poly[-1]) * weights
-    
+
     mask = ~(new_poly == torch.roll(new_poly, shifts=-1, dims=0)).all(dim=1)
     return new_poly[mask]
 
-def linear_interpolate(
-        poly: V, 
-        scale: int | np.ndarray | torch.Tensor
-    ) -> V:
+
+def linear_interpolate(poly: V, scale: int | np.ndarray | torch.Tensor) -> V:
     """Linearly interpolates a N x 2 polygon to have N x scale vertices."""
     if not isinstance(scale, int):
         scale = int(scale.item())
@@ -367,16 +362,16 @@ def linear_interpolate(
 
 
 def _scale_contour_arr(
-        contour: np.ndarray, 
-        scale: Sequence[float | int] | np.ndarray | torch.Tensor | float | int, 
-        expand_by_one: bool=False
-    ) -> np.ndarray:
+    contour: np.ndarray,
+    scale: Sequence[float | int] | np.ndarray | torch.Tensor | float | int,
+    expand_by_one: bool = False,
+) -> np.ndarray:
     if len(contour.shape) != 2 or contour.shape[1] != 2:
         if contour.shape[0] == 2:
             contour = contour.reshape(1, 2)
         else:
             raise ValueError(f"Contour must be a Nx2 array, not {contour.shape}")
-            
+
     if isinstance(scale, (int, float)):
         scale = [scale, scale]
     scale = np.asarray(scale, dtype=np.float32)
@@ -390,46 +385,47 @@ def _scale_contour_arr(
         return np.round(contour * scale).astype(np.int32)
     if np.all(scale == 1):
         return contour
-        
+
     contour = contour * scale
     centroid = contour.mean(axis=0)
     n_interp = max(1, int(np.ceil(scale.max())) * 2)
-    
+
     contour = _linear_interpolate_arr(contour, n_interp)
     contour_normals = _poly_normals_arr(contour)
-    
+
     if expand_by_one:
         expand_one = np.sign(contour_normals) * (np.abs(contour_normals) > 0)
         contour -= expand_one
-    
+
     if scale[0] < 1:
         contour[:, 0] += contour_normals[:, 0] / scale[0] / 2
     if scale[1] < 1:
         contour[:, 1] += contour_normals[:, 1] / scale[1] / 2
-    
+
     contour[contour_normals > 0] = np.floor(contour[contour_normals > 0])
     contour[contour_normals < 0] = np.ceil(contour[contour_normals < 0])
     contour = contour.round()
     drift = centroid - contour.mean(axis=0)
-    
-    return (contour + drift).round().astype(np.int32)[(n_interp // 2)::n_interp].copy()
+
+    return (contour + drift).round().astype(np.int32)[(n_interp // 2) :: n_interp].copy()
+
 
 def _scale_contour_tensor(
-        contour: torch.Tensor, 
-        scale: Sequence[float | int] | np.ndarray | torch.Tensor | float | int, 
-        expand_by_one: bool=False
-    ) -> torch.Tensor:
+    contour: torch.Tensor,
+    scale: Sequence[float | int] | np.ndarray | torch.Tensor | float | int,
+    expand_by_one: bool = False,
+) -> torch.Tensor:
     if len(contour.shape) != 2 or contour.shape[1] != 2:
         if contour.shape[0] == 2:
             contour = contour.reshape(1, 2)
         else:
             raise ValueError(f"Contour must be a Nx2 tensor, not {contour.shape}")
-            
+
     if isinstance(scale, (int, float)):
         scale = [scale, scale]
 
     scale = torch.as_tensor(scale, dtype=torch.float32, device=contour.device)
-        
+
     if len(scale) != 2:
         raise ValueError(f"Scale must be a scalar or a list of 2 scalars, not {scale}")
 
@@ -439,46 +435,43 @@ def _scale_contour_tensor(
         return torch.round(contour * scale).to(torch.int32)
     if torch.all(scale == 1):
         return contour
-        
+
     contour = contour * scale
     centroid = contour.mean(dim=0)
     n_interp = max(1, int(torch.ceil(scale.max()).item()) * 2)
-    
+
     contour = _linear_interpolate_tensor(contour, n_interp)
     contour_normals = _poly_normals_tensor(contour)
-    
+
     if expand_by_one:
         expand_one = torch.sign(contour_normals) * (torch.abs(contour_normals) > 0)
         contour -= expand_one
-    
+
     if scale[0] < 1:
         contour[:, 0] += contour_normals[:, 0] / scale[0] / 2
     if scale[1] < 1:
         contour[:, 1] += contour_normals[:, 1] / scale[1] / 2
-    
+
     contour[contour_normals > 0] = torch.floor(contour[contour_normals > 0])
     contour[contour_normals < 0] = torch.ceil(contour[contour_normals < 0])
     contour = contour.round()
     drift = centroid - contour.mean(dim=0)
-    
-    return torch.round(contour + drift).to(torch.int32)[(n_interp // 2)::n_interp].clone()
+
+    return torch.round(contour + drift).to(torch.int32)[(n_interp // 2) :: n_interp].clone()
+
 
 def scale_contour(  # noqa: D103
-        contour: V, 
-        scale: Sequence[float | int] | np.ndarray | torch.Tensor | float | int, 
-        expand_by_one: bool=False
-    ) -> V:
+    contour: V, scale: Sequence[float | int] | np.ndarray | torch.Tensor | float | int, expand_by_one: bool = False
+) -> V:
     if isinstance(contour, torch.Tensor):
         return _scale_contour_tensor(contour, scale, expand_by_one)
     else:
         return _scale_contour_arr(contour, scale, expand_by_one)
 
-def resize_masks(
-        masks : torch.Tensor, 
-        new_shape : tuple[int, int] | list[int] | int
-    ) -> torch.Tensor:
+
+def resize_masks(masks: torch.Tensor, new_shape: tuple[int, int] | list[int] | int) -> torch.Tensor:
     """Resize a mask (or a batch of masks) by scaling the contour coordinates and snapping to the integer grid.
-    
+
     Ensures that snapping is always done towards the outside of the mask.
 
     Args:
@@ -499,14 +492,13 @@ def resize_masks(
     if not isinstance(new_shape, int) and (new_shape[0] <= 1 or new_shape[1] <= 1):
         raise ValueError(f"Target shape must be at least 2x2, not {new_shape}")
     # Resize the mask
-    return F.interpolate(masks.float()[None], new_shape, mode='nearest-exact', antialias=False)[0] > 0.5
+    return F.interpolate(masks.float()[None], new_shape, mode="nearest-exact", antialias=False)[0] > 0.5
+
 
 _to_uint8 = torchvision.transforms.ConvertImageDtype(torch.uint8)
 
-def chw2hwc_uint8(
-        crop : torch.Tensor, 
-        mask : torch.Tensor | None
-    ) -> torch.Tensor:
+
+def chw2hwc_uint8(crop: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
     """Convert a crop from CHW to HWC format, and adds the mask as an alpha channel if it exists.
 
     Args:
@@ -521,4 +513,4 @@ def chw2hwc_uint8(
     if mask is not None:
         mask = mask.bool().to(torch.uint8) * 255
         crop = torch.cat([crop, mask], dim=0)
-    return crop.permute(1, 2, 0)
\ No newline at end of file
+    return crop.permute(1, 2, 0)
diff --git a/src/flat_bug/nms.py b/src/flat_bug/nms.py
index f86e2d0..807e23b 100644
--- a/src/flat_bug/nms.py
+++ b/src/flat_bug/nms.py
@@ -1,4 +1,5 @@
 """Implementations of non-maximum suppression for boxes, polygons and masks used in flatbug inference."""
+
 from collections.abc import Callable
 from functools import partial
 from typing import Any, Literal, cast, overload
@@ -10,21 +11,18 @@
 import torchvision
 
 
-def iou_boxes(
-        rectangles : torch.Tensor,
-        other_rectangles : torch.Tensor | None=None
-    ) -> torch.Tensor:
+def iou_boxes(rectangles: torch.Tensor, other_rectangles: torch.Tensor | None = None) -> torch.Tensor:
     """Calculate the intersection over union (IoU) of a set of rectangles.
 
     Args:
-        rectangles: A tensor of shape (n, 4), where n is the number of rectangles
+        rectangles: A tensor of shape `(n, 4)`, where `n` is the number of rectangles
             and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the rectangles.
-        other_rectangles: A tensor of shape (m, 4), where m is the number of rectangles
-            and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the rectangles. 
-            Defaults to None, in which case the symmetric IoU of the rectangles with themselves is calculated.
+        other_rectangles: Same as `rectangles` or `None`,
+            in which case the symmetric IoU of the rectangles with themselves is calculated.
+            Defaults to None.
 
     Returns:
-        A tensor of shape (n, n), where n is the number of rectangles, 
+        A tensor of shape `(n, m)`, where `n`/`m` is the number of rectangles,
         containing the IoU of each rectangle with each other rectangle.
 
     """
@@ -38,31 +36,30 @@ def iou_boxes(
         raise ValueError(f"Other rectangles must be a tensor, not {type(other_rectangles)}")
     elif not len(other_rectangles.shape) == 2 or other_rectangles.shape[1] != 4:
         raise ValueError(f"Other rectangles must be of shape (n, 4), not {other_rectangles.shape}")
-    
+
     return torchvision.ops.box_iou(rectangles, rectangles if other_rectangles is None else other_rectangles)
 
+
 # Check if 'fmt' is an argument in the current version of torchvision
 try:
-    torchvision.ops.boxes._box_inter_union(torch.empty((0,4)), torch.empty((0,4)), fmt="xyxy")
+    torchvision.ops.boxes._box_inter_union(torch.empty((0, 4)), torch.empty((0, 4)), fmt="xyxy")
     _box_inter_union = partial(torchvision.ops.boxes._box_inter_union, fmt="xyxy")
 except TypeError:
     _box_inter_union = torchvision.ops.boxes._box_inter_union
 
-def ios_boxes(
-        rectangles : torch.Tensor,
-        other_rectangles : torch.Tensor | None=None
-    ) -> torch.Tensor:
+
+def ios_boxes(rectangles: torch.Tensor, other_rectangles: torch.Tensor | None = None) -> torch.Tensor:
     """Calculate the intersection over smaller (IoS) of a set of rectangles.
 
     Args:
-        rectangles: A tensor of shape (n, 4), where n is the number of rectangles
+        rectangles: A tensor of shape `(n, 4)`, where `n` is the number of rectangles
             and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the rectangles.
-        other_rectangles: A tensor of shape (m, 4), where m is the number of rectangles
-            and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the rectangles. 
-            Defaults to None, in which case the symmetric IoS of the rectangles with themselves is calculated.
+        other_rectangles: Same as `rectangles` or `None`,
+            in which case the symmetric IoS of the rectangles with themselves is calculated.
+            Defaults to None.
 
     Returns:
-        A tensor of shape (n, n), where n is the number of rectangles, 
+        A tensor of shape `(n, m)`, where `n`/`m` is the number of rectangles,
         containing the IoS of each rectangle with each other rectangle.
 
     """
@@ -87,18 +84,19 @@ def ios_boxes(
     ios = intersections / (sareas + 1e-6)
     return ios
 
+
 @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")
 def iou_masks(
-        m1s : torch.Tensor, 
-        m2s : torch.Tensor, 
-        a1s : torch.Tensor | None=None, 
-        a2s : torch.Tensor | None=None, 
-        dtype : torch.dtype=torch.float32
-    ) -> torch.Tensor:
+    m1s: torch.Tensor,
+    m2s: torch.Tensor,
+    a1s: torch.Tensor | None = None,
+    a2s: torch.Tensor | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
     """Compute IoU between all pairs between two sets of masks.
 
-    The IoU is calculated using the formula: 
-    
+    The IoU is calculated using the formula:
+
     `IoU[i,j] = intersection[i, j] / (m1s[i].sum() + m2s[j].sum() - intersection[i, j])`
 
     `intersection[i, j] = (m1s[i] * m2s[j]).sum()`
@@ -109,14 +107,16 @@ def iou_masks(
     OBS: Results will only be valid for boolean or masks containing only 0s and 1s.
 
     Args:
-        m1s: A tensor of shape (n, h, w), where n is the number of masks and h and w are the height and width of the masks.
-        m2s: A tensor of shape (m, h, w), where m is the number of masks and h and w are the height and width of the masks.
-        a1s: A tensor of shape (n, ) containing the areas of the masks in m1s. Defaults to None, in which case the areas are calculated.
-        a2s: A tensor of shape (m, ) containing the areas of the masks in m2s. Defaults to None, in which case the areas are calculated.
-        dtype: The data type of the output tensor. Defaults to torch.float32.
-        
+        m1s: A tensor of shape `(n, h, w)`,
+            where `n` is the number of masks and `h` and `w` are the height and width of the masks.
+        m2s: Same as `m1s`.
+        a1s: A tensor of shape `(n, )` containing the areas of the masks in m1s.
+            Defaults to `None`, in which case the areas are calculated.
+        a2s: Same as `a1s`.
+        dtype: The data type of the output tensor. Defaults to `torch.float32`.
+
     Returns:
-        A tensor of shape (n, m) containing the IoU of each pair of masks.
+        A tensor of shape `(n, m)` containing the IoU of each pair of masks.
 
     """
     # 1. Standardize Inputs: Ensure batch dim and flatten spatial dims (N, H, W) -> (N, P)
@@ -124,7 +124,7 @@ def iou_masks(
         m1s = m1s.unsqueeze(0)
     if m2s.dim() == 2:
         m2s = m2s.unsqueeze(0)
-        
+
     m1s_flat = m1s.flatten(1)
     m2s_flat = m2s.flatten(1)
 
@@ -133,25 +133,26 @@ def iou_masks(
         a1s = m1s_flat.sum(dim=1).to(dtype)
     else:
         a1s = a1s.to(dtype)
-        
+
     if a2s is None:
         a2s = m2s_flat.sum(dim=1).to(dtype)
     else:
         a2s = a2s.to(dtype)
-    
+
     intersections = torch.mm(m1s_flat.to(dtype), m2s_flat.t().to(dtype))
     unions = a1s.unsqueeze(1) + a2s.unsqueeze(0) - intersections
-    
+
     return intersections / (unions + 1e-6)
 
+
 @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")
 def ios_masks(  # noqa: D103
-        m1s : torch.Tensor, 
-        m2s : torch.Tensor, 
-        a1s : torch.Tensor | None=None, 
-        a2s : torch.Tensor | None=None, 
-        dtype : torch.dtype=torch.float32
-    ) -> torch.Tensor:
+    m1s: torch.Tensor,
+    m2s: torch.Tensor,
+    a1s: torch.Tensor | None = None,
+    a2s: torch.Tensor | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
     """Compute IoS (Intersection over Smaller area) between all pairs between two sets of masks.
 
     The IoS is calculated using the formula:
@@ -160,20 +161,22 @@ def ios_masks(  # noqa: D103
 
     `intersection[i, j] = (m1s[i] * m2s[j]).sum()`
 
-    The reason the intersection is calculated this way is that it can be vectorized 
+    The reason the intersection is calculated this way is that it can be vectorized
     and calculated in a single matrix multiplication for all pairs of masks.
 
     OBS: Results will only be valid for boolean or masks containing only 0s and 1s.
 
     Args:
-        m1s: A tensor of shape (n, h, w), where n is the number of masks and h and w are the height and width of the masks.
-        m2s: A tensor of shape (m, h, w), where m is the number of masks and h and w are the height and width of the masks.
-        a1s: A tensor of shape (n, ) containing the areas of the masks in m1s. Defaults to None, in which case the areas are calculated.
-        a2s: A tensor of shape (m, ) containing the areas of the masks in m2s. Defaults to None, in which case the areas are calculated.
-        dtype: The data type of the output tensor. Defaults to torch.float32.
+        m1s: A tensor of shape `(n, h, w)`,
+            where `n` is the number of masks and `h` and `w` are the height and width of the masks.
+        m2s: Same as `m1s`.
+        a1s: A tensor of shape `(n, )` containing the areas of the masks in `m1s`.
+            Defaults to `None`, in which case the areas are calculated.
+        a2s: Same as `a1s`.
+        dtype: The data type of the output tensor. Defaults to `torch.float32`.
 
     Returns:
-        A tensor of shape (n, m) containing the IoS of each pair of masks.
+        A tensor of shape `(n, m)` containing the IoS of each pair of masks.
 
     """
     # 1. Standardize Inputs: Ensure batch dim and flatten spatial dims (N, H, W) -> (N, P)
@@ -181,7 +184,7 @@ def ios_masks(  # noqa: D103
         m1s = m1s.unsqueeze(0)
     if m2s.dim() == 2:
         m2s = m2s.unsqueeze(0)
-        
+
     m1s_flat = m1s.flatten(1)
     m2s_flat = m2s.flatten(1)
 
@@ -190,28 +193,30 @@ def ios_masks(  # noqa: D103
         a1s = m1s_flat.sum(dim=1).to(dtype)
     else:
         a1s = a1s.to(dtype)
-        
+
     if a2s is None:
         a2s = m2s_flat.sum(dim=1).to(dtype)
     else:
         a2s = a2s.to(dtype)
-    
+
     intersections = torch.mm(m1s_flat.to(dtype), m2s_flat.t().to(dtype))
     amin = torch.minimum(a1s.unsqueeze(1), a2s.unsqueeze(0))
-    
+
     return intersections / (amin + 1e-6)
 
+
 def iou_polygons(  # noqa: D103
-        polygons1: list[torch.Tensor] | np.ndarray, 
-        polygons2: list[torch.Tensor] | np.ndarray | None = None,
-        *args, **kwargs
-    ) -> np.ndarray:
-    
+    polygons1: list[torch.Tensor] | np.ndarray,
+    polygons2: list[torch.Tensor] | np.ndarray | None = None,
+    *args,
+    **kwargs,
+) -> np.ndarray:
+
     if len(polygons1) == 0:
         return np.empty((0, 0 if polygons2 is None else len(polygons2)), dtype=np.float32)
 
     is_symmetric = polygons2 is None
-    
+
     def ensure_geoms(objs: Any) -> np.ndarray:
         # If it's already an object-dtype numpy array, assume it's shapely geoms
         if isinstance(objs, np.ndarray) and objs.dtype == object:
@@ -224,7 +229,7 @@ def ensure_geoms(objs: Any) -> np.ndarray:
 
     areas1 = shapely.area(geoms1)
     areas2 = areas1 if is_symmetric else shapely.area(geoms2)
-    
+
     intersections = shapely.area(shapely.intersection(geoms1[:, np.newaxis], geoms2[np.newaxis, :]))
 
     unions = areas1[:, np.newaxis] + areas2[np.newaxis, :] - intersections
@@ -237,16 +242,17 @@ def ensure_geoms(objs: Any) -> np.ndarray:
 
 
 def ios_polygons(  # noqa: D103
-        polygons1: list[torch.Tensor] | np.ndarray, 
-        polygons2: list[torch.Tensor] | np.ndarray | None = None,
-        *args, **kwargs
-    ) -> np.ndarray:
-    
+    polygons1: list[torch.Tensor] | np.ndarray,
+    polygons2: list[torch.Tensor] | np.ndarray | None = None,
+    *args,
+    **kwargs,
+) -> np.ndarray:
+
     if len(polygons1) == 0:
         return np.empty((0, 0 if polygons2 is None else len(polygons2)), dtype=np.float32)
 
     is_symmetric = polygons2 is None
-    
+
     def ensure_geoms(objs: Any) -> np.ndarray:
         # If it's already an object-dtype numpy array, assume it's shapely geoms
         if isinstance(objs, np.ndarray) and objs.dtype == object:
@@ -259,7 +265,7 @@ def ensure_geoms(objs: Any) -> np.ndarray:
 
     areas1 = shapely.area(geoms1)
     areas2 = areas1 if is_symmetric else shapely.area(geoms2)
-    
+
     intersections = shapely.area(shapely.intersection(geoms1[:, np.newaxis], geoms2[np.newaxis, :]))
 
     areas_min = np.minimum(areas1[:, np.newaxis], areas2[np.newaxis, :])
@@ -270,59 +276,60 @@ def ensure_geoms(objs: Any) -> np.ndarray:
 
     return ios_mat
 
+
 @overload
 def base_nms_(
-        objects : Any, 
-        overlap_fn : Callable, 
-        scores : torch.Tensor, 
-        collate_fn : Callable | None=None, 
-        overlap_threshold : float=0.5, 
-        strict : bool=True, 
-        return_indices : Literal[False]=False, 
-        **kwargs
-    ) -> tuple[Any, torch.Tensor]: ...
+    objects: Any,
+    overlap_fn: Callable,
+    scores: torch.Tensor,
+    collate_fn: Callable | None = None,
+    overlap_threshold: float = 0.5,
+    strict: bool = True,
+    return_indices: Literal[False] = False,
+    **kwargs,
+) -> tuple[Any, torch.Tensor]: ...
 @overload
 def base_nms_(
-        objects : Any, 
-        overlap_fn : Callable, 
-        scores : torch.Tensor, 
-        collate_fn : Callable | None=None, 
-        overlap_threshold : float=0.5, 
-        strict : bool=True, 
-        return_indices : Literal[True]=True, 
-        **kwargs
-    ) -> torch.Tensor: ...
+    objects: Any,
+    overlap_fn: Callable,
+    scores: torch.Tensor,
+    collate_fn: Callable | None = None,
+    overlap_threshold: float = 0.5,
+    strict: bool = True,
+    return_indices: Literal[True] = True,
+    **kwargs,
+) -> torch.Tensor: ...
 def base_nms_(
-        objects : Any, 
-        overlap_fn : Callable, 
-        scores : torch.Tensor, 
-        collate_fn : Callable | None=None, 
-        overlap_threshold : float=0.5, 
-        strict : bool=True, 
-        return_indices : bool=False, 
-        **kwargs
-    ) -> torch.Tensor | tuple[Any, torch.Tensor]:
+    objects: Any,
+    overlap_fn: Callable,
+    scores: torch.Tensor,
+    collate_fn: Callable | None = None,
+    overlap_threshold: float = 0.5,
+    strict: bool = True,
+    return_indices: bool = False,
+    **kwargs,
+) -> torch.Tensor | tuple[Any, torch.Tensor]:
     """Perform the standard non-maximum suppression algorithm.
 
     Args:
         objects: An object which can be indexed by a tensor of indices.
-        overlap_fn: A function which takes an anchor object and a comparison set (not in the Python sense) of (different) objects 
-            and returns the IoU of the anchor object with each object in the comparison set as a tensor of shape (1, n). 
-            The reason it is not just (n, ) is to allow for implementations of `overlap_fn` functions between two sets, 
-            where the IoU is calculated between each pair of objects from distinct sets.
-        scores: A tensor of shape (n, ) containing the "scores" of the objects, this can merely be though of as a priority score, 
-            where the higher the score, the higher the priority of the object - it does not have to be a probability/confidence.
-        collate_fn: A function which takes a list of objects and returns a single combined object. 
-            Defaults to `torch.cat` if `objects` is a tensor and `lambda x : x` if `objects` is a list, otherwise it has to be specified.
-        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to 0.5.
-        strict: A flag to indicate whether to perform strict checks on the algorithm. Defaults to True.
-        return_indices: A flag to indicate whether to return the indices of the picked objects or the objects themselves. 
-            Defaults to False.  If True, both the picked objects and scores are returned.
-        **kwargs: Additional keyword arguments to be passed to the overlap_fn function.
-    
+        overlap_fn: A function which takes an anchor object and a comparison set (not in the Python sense)
+            of (different) objects and returns the IoU of the anchor object with each object in the comparison set as
+            a tensor of shape `(1, n)`.
+        scores: A tensor of shape `(n, )` containing the "scores" of the objects,
+            this can merely be though of as a priority score, it does not have to be a probability/confidence.
+        collate_fn: A function which takes a list of objects and returns a single combined object.
+            Defaults to `torch.cat` if `objects` is a tensor and `lambda x : x` if `objects` is a `list`,
+            otherwise it has to be specified.
+        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to `0.5`.
+        strict: Flag indicating whether to perform strict checks on the algorithm. Defaults to `True`.
+        return_indices: A flag indicating whether to return the indices of the picked objects or the objects themselves.
+            Defaults to `False`.  If `True`, both the picked objects and scores are returned.
+        **kwargs: Additional keyword arguments to be passed to the `overlap_fn` function.
+
     Returns:
-        Either a tensor of shape `(m,)` containing the indices of the picked objects, 
-        or a tuple (`tuple[Any, torch.Tensor]`) where the first element contains 
+        Either a tensor of shape `(m,)` containing the indices of the picked objects,
+        or a tuple (`tuple[Any, torch.Tensor]`) where the first element contains
         the picked objects and the second element is a tensor of their scores.
 
     """
@@ -330,10 +337,10 @@ def base_nms_(
         if isinstance(objects, torch.Tensor):
             collate_fn = torch.cat
         elif isinstance(objects, list):
-            collate_fn = lambda x : x  # noqa: E731
+            collate_fn = lambda x: x  # noqa: E731
         else:
             raise ValueError(f"collate_fn must be specified for objects of type {type(objects)}")
-    
+
     device = scores.device
     if len(scores.shape) != 1:
         raise ValueError(f"Scores must be of shape (n,), not {scores.shape}")
@@ -344,7 +351,7 @@ def base_nms_(
             return torch.arange(len(objects))
         else:
             return collate_fn([objects[i] for i in range(len(objects))]), scores
-    
+
     # Sort the boxes by score (implicitly)
     indices = torch.argsort(scores, descending=True)
 
@@ -365,7 +372,8 @@ def base_nms_(
         # Remove the picked box from the possible boxes
         possible[possible_idx[0]] = False
         # Calculate the overlaps (e.g. IoU) between the picked box and the remaining possible boxes
-        overlaps = overlap_fn(objects[indices[possible_idx[0]]], objects[indices[possible_idx[1:]]], **kwargs).squeeze(0)
+        overlaps = overlap_fn(objects[indices[possible_idx[0]]], objects[indices[possible_idx[1:]]], **kwargs)
+        overlaps = overlaps.squeeze(0)
         # Get the indices of the boxes with an overlap greater than the threshold
         winner_mask = overlaps <= overlap_threshold
         # Remove the boxes with an overlap greater than the threshold from the possible boxes
@@ -375,55 +383,59 @@ def base_nms_(
             # In/Decrement the counters
             increment = (~winner_mask).sum().item() + 1
             left -= increment
-            assert left == (possible_idx.numel() - 1), f"left ({left}) != possible_idx.numel() - 1 ({possible_idx.numel() - 1})"
+            assert left == (possible_idx.numel() - 1), (
+                f"left ({left}) != possible_idx.numel() - 1 ({possible_idx.numel() - 1})"
+            )
             assert (i + 1) == len(winners), f"n ({i + 1}) != winners.sum() ({len(winners)})"
 
-
-    # Map the indices back to the original indices and sort them (returns boxes, scores & indices in the original order of the input)
+    # Map the indices back to the original indices and sort them
+    # (returns boxes, scores & indices in the original order of the input)
     winners = torch.tensor(winners, dtype=torch.long, device=device)
-    winners = indices[winners].sort().values 
-    
+    winners = indices[winners].sort().values
+
     # Return the boxes and scores that were picked
     if return_indices:
         return winners
     else:
         return collate_fn([objects[ni] for ni in winners]), scores[winners]
 
+
 def fancy_nms(
-        objects : Any, 
-        overlap_fn : Callable, 
-        scores : torch.Tensor, 
-        overlap_threshold : float | int=0.5, 
-        return_indices : bool=False
-    ) -> torch.Tensor | tuple[Any, torch.Tensor]:
+    objects: Any,
+    overlap_fn: Callable,
+    scores: torch.Tensor,
+    overlap_threshold: float | int = 0.5,
+    return_indices: bool = False,
+) -> torch.Tensor | tuple[Any, torch.Tensor]:
     """Perform a 'fancy' implementation of non-maximum suppression (NMS).
-    
-    It is not as fast as the non-maximum suppression algorithm, 
+
+    It is not as fast as the non-maximum suppression algorithm,
     nor does it follow the exact same algorithm, but it is more readable and easier to debug.
 
     The algorithm works as follows:
         1. Sort the objects by score (implicitly)
         2. Calculate the overlap (e.g. IoU) matrix
-        3. Create a boolean matrix where overlap > overlap_threshold 
+        3. Create a boolean matrix where overlap > overlap_threshold
         4. Fold the boolean matrix sequentially (i.e. row_i = row_i + row_i-1 + ... + row_0)
-           (The values on the diagonal of the matrix now correspond to the number 
+           (The values on the diagonal of the matrix now correspond to the number
            of higher-priority objects that suppress the current object, including itself)
         5. objects which are suppressed only by themselves are returned.
 
-    
+
     Args:
-        objects: Any object collection that can be indexed by a tensor, where the first dimension corresponds to the objects.
-        overlap_fn: A function that calculates the symmetric overlap (e.g. IoU) matrix 
-            of a set of objects returned as a `torch.Tensor` of shape (n, n), 
-            where n is the number of objects. The device should match the device of the scores.
-        scores: A tensor of shape (n, ) containing the scores of the objects.
-        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to 0.5.
-        return_indices: A flag to indicate whether to return the indices of the picked objects or the objects themselves. 
-            Defaults to False. If True, both the picked objects and scores are returned.
+        objects: Any object collection that can be indexed by a tensor,
+            where the first dimension corresponds to the objects.
+        overlap_fn: A function that calculates the symmetric overlap (e.g. IoU) matrix
+            of a set of objects returned as a `torch.Tensor` of shape `(n, n)`,
+            where `n` is the number of objects. The device should match the device of the scores.
+        scores: A tensor of shape `(n, )` containing the scores of the objects.
+        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to `0.5`.
+        return_indices: A flag indicating whether to return the indices of the picked objects or the objects themselves.
+            Defaults to `False`. If `True`, both the picked objects and scores are returned.
 
     Returns:
-        Either a tensor containing the indices of the picked objects, 
-        or a tuple (`tuple[Any, torch.Tensor`) where the first element contains 
+        Either a tensor containing the indices of the picked objects,
+        or a tuple (`tuple[Any, torch.Tensor`) where the first element contains
         the picked objects and the second element is a tensor of their scores.
 
     """
@@ -432,14 +444,16 @@ def fancy_nms(
     if not len(scores.shape) == 1:
         raise ValueError(f"Scores must be of shape (n,), not {scores.shape}")
     if not objects.shape[0] == scores.shape[0]:
-        raise ValueError(f"Boxes and scores must have the same number of boxes, not {objects.shape[0]} and {scores.shape[0]}")
+        raise ValueError(
+            f"Boxes and scores must have the same number of boxes, not {objects.shape[0]} and {scores.shape[0]}"
+        )
 
     if len(objects) == 0 or len(objects) == 1:
         if return_indices:
             return torch.arange(len(objects))
         else:
             return objects, scores
-    
+
     # Sort the boxes by score (implicitly)
     indices = torch.argsort(scores, descending=True)
 
@@ -449,7 +463,7 @@ def fancy_nms(
     # Fold the overlap matrix sequentially (i.e. row_i = row_i + row_i-1 + ... + row_0)
     overlaps = (overlaps > overlap_threshold).cumsum(dim=1) <= 1
 
-    # The boxes with an overlap greater than the threshold are the elements on 
+    # The boxes with an overlap greater than the threshold are the elements on
     # the diagonal of the folded overlap matrix which are one (suppressed only by itself)
     indices = indices[torch.where(overlaps.diagonal())[0]]
 
@@ -458,18 +472,22 @@ def fancy_nms(
     else:
         return objects[indices], scores[indices]
 
+
 def nms_masks_(
-        masks : torch.Tensor, 
-        scores : torch.Tensor, 
-        overlap_threshold : float=0.5,
-        overlap_fn : Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor]=iou_masks
-    ) -> torch.Tensor:
+    masks: torch.Tensor,
+    scores: torch.Tensor,
+    overlap_threshold: float = 0.5,
+    overlap_fn: Callable[
+        [torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor
+    ] = iou_masks,
+) -> torch.Tensor:
     """Perform non-maximum suppression (NMS) on a set of masks.
-    
+
     Args:
-        masks: A tensor of shape (n, h, w), where n is the number of masks and h and w are the height and width of the masks.
-        scores: A tensor of shape (n, ) containing the scores of the masks.
-        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to 0.5.
+        masks: A tensor of shape `(n, h, w)`,
+            where `n` is the number of masks and `h` and `w` are the height and width of the masks.
+        scores: A tensor of shape `(n, )` containing the scores of the masks.
+        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to `0.5`.
         overlap_fn: A function to compute overlaps between masks.
 
     Returns:
@@ -493,37 +511,38 @@ def nms_masks_(
     for _ in range(N):
         possible_idx = possible.nonzero().squeeze(1)
         n_possible = possible_idx.numel()
-        
+
         if n_possible < 2:
             if n_possible == 1:
                 possible[possible_idx] = False
                 winners[i] = possible_idx
                 i += 1
             break
-            
+
         winners[i] = possible_idx[0]
         possible[possible_idx[0]] = False
-        
+
         overlaps = overlap_fn(
-            masks[possible_idx[0:1]].unsqueeze(1), 
-            masks[possible_idx[1:]].unsqueeze(1), 
-            areas[possible_idx[0:1]], 
-            areas[possible_idx[1:]], 
-            torch.float32
+            masks[possible_idx[0:1]].unsqueeze(1),
+            masks[possible_idx[1:]].unsqueeze(1),
+            areas[possible_idx[0:1]],
+            areas[possible_idx[1:]],
+            torch.float32,
         ).squeeze(0)
-        
+
         winner_mask = overlaps <= overlap_threshold
         possible[possible_idx[1:]] = winner_mask
         i += 1
 
-    return indices[winners[:i]].sort().values 
+    return indices[winners[:i]].sort().values
+
 
 def nms_polygons_(  # noqa: D103
-        polys : list[torch.Tensor], 
-        scores : torch.Tensor, 
-        overlap_threshold : float=0.5,
-        overlap_fn : Callable[[np.ndarray, np.ndarray], np.ndarray]=iou_polygons
-    ) -> torch.Tensor:
+    polys: list[torch.Tensor],
+    scores: torch.Tensor,
+    overlap_threshold: float = 0.5,
+    overlap_fn: Callable[[np.ndarray, np.ndarray], np.ndarray] = iou_polygons,
+) -> torch.Tensor:
     N, device = len(scores), scores.device
     if N <= 1:
         return torch.arange(N, device=device)
@@ -531,7 +550,7 @@ def nms_polygons_(  # noqa: D103
     scores_np = scores.cpu().numpy()
     geoms = np.array([cast(shapely.Polygon, shapely.polygons(p.cpu().numpy())).buffer(0) for p in polys])
 
-    indices = np.argsort(scores_np)[::-1] # Ascending sort -> reverse for descending
+    indices = np.argsort(scores_np)[::-1]  # Ascending sort -> reverse for descending
     geoms = geoms[indices]
 
     # int64 to ensure compatibility when converting back to torch.long later
@@ -542,7 +561,7 @@ def nms_polygons_(  # noqa: D103
     for _ in range(N):
         possible_idx = np.flatnonzero(possible)
         n_possible = possible_idx.size
-        
+
         if n_possible < 2:
             if n_possible == 1:
                 possible[possible_idx] = False
@@ -556,7 +575,7 @@ def nms_polygons_(  # noqa: D103
         possible[curr_idx] = False
 
         overlaps = overlap_fn(
-            geoms[curr_idx:curr_idx+1], 
+            geoms[curr_idx:curr_idx+1],
             geoms[possible_idx[1:]]
         ).squeeze(0)
         
@@ -569,31 +588,31 @@ def nms_polygons_(  # noqa: D103
 
 
 def cluster_overlap_boxes(
-        boxes: torch.Tensor, 
-        overlap_threshold: float = 0.5,
-        overlap_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = iou_boxes,
-        time: bool = False
-    ) -> tuple[list[torch.Tensor], torch.Tensor]:
+    boxes: torch.Tensor,
+    overlap_threshold: float = 0.5,
+    overlap_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = iou_boxes,
+    time: bool = False,
+) -> tuple[list[torch.Tensor], torch.Tensor]:
     """Cluster boxes via connected components.
-    
+
     Note: Implementation relies on ``overlap_fn`` being symmetric (e.g. IoU/IoS).
     """
     N, device = len(boxes), boxes.device
     if N <= 1:
         return [torch.arange(N, dtype=torch.long, device=device)], torch.zeros(N, dtype=torch.long, device=device)
-    
+
     # Chuncked adjacency matrix (memory-to-compute tradeoff)
-    CHUNK_SIZE = 2500 
+    CHUNK_SIZE = 2500
     rows_list = []
     cols_list = []
-    
+
     for i in range(0, N, CHUNK_SIZE):
         chunk_i = boxes[i : i + CHUNK_SIZE]
         for j in range(i, N, CHUNK_SIZE):
             chunk_j = boxes[j : j + CHUNK_SIZE]
             adj_chunk = overlap_fn(chunk_i, chunk_j) >= overlap_threshold
             local_edges = adj_chunk.nonzero().cpu().numpy()
-            
+
             if local_edges.size > 0:
                 # Offset the local indices to global indices
                 rows_list.append(local_edges[:, 0] + i)
@@ -606,13 +625,13 @@ def cluster_overlap_boxes(
     else:
         row, col = np.concatenate(rows_list), np.concatenate(cols_list)
         data = np.ones(len(row), dtype=bool)
-        
+
         sparse_graph = scipy.sparse.coo_matrix((data, (row, col)), shape=(N, N))
-        
+
         # Find connected components (scipy handles the symmetry implicitly with directed=False)
         _, labels = scipy.sparse.csgraph.connected_components(
-            sparse_graph, 
-            directed=False, 
+            sparse_graph,
+            directed=False,
             return_labels=True
         )
 
@@ -620,106 +639,100 @@ def cluster_overlap_boxes(
     cluster_vec = torch.from_numpy(labels).to(device=device, dtype=torch.long)
     sorted_idx = torch.argsort(cluster_vec)
     sorted_labels = cluster_vec[sorted_idx]
-    
+
     _, counts = torch.unique(sorted_labels, return_counts=True)
     groups = torch.split(sorted_idx, counts.tolist())
 
     return list(groups), cluster_vec
 
 
-OVERLAP_FNS : dict[str, dict[str, Callable]] = {
-    "polygon" : {
-        "iou" : iou_polygons,
-        "ios" : ios_polygons
-    },
-    "mask" : {
-        "iou" : iou_masks,
-        "ios" : ios_masks
-    },
-    "box" : {
-        "iou" : iou_boxes,
-        "ios" : ios_boxes
-    }
+OVERLAP_FNS: dict[str, dict[str, Callable]] = {
+    "polygon": {"iou": iou_polygons, "ios": ios_polygons},
+    "mask": {"iou": iou_masks, "ios": ios_masks},
+    "box": {"iou": iou_boxes, "ios": ios_boxes},
 }
 
-def get_overlap_fn(geometry : str, metric : str):  # noqa: D103
+
+def get_overlap_fn(geometry: str, metric: str):  # noqa: D103
     geometry, metric = geometry.lower().strip(), metric.lower().strip()
     if geometry not in OVERLAP_FNS:
         raise NotImplementedError(
-            f'No overlap metrics implemented for geometry type: "{geometry}", ' +
-            'valid options are [{}]'.format(
-                ", ".join(OVERLAP_FNS.keys())
-            )
+            f'No overlap metrics implemented for geometry type: "{geometry}", '
+            + "valid options are [{}]".format(", ".join(OVERLAP_FNS.keys()))
         )
     options = OVERLAP_FNS[geometry]
     if metric not in options:
         raise NotImplementedError(
-            f'Overlap metric: "{metric}" not implemented for geometry type: "{geometry}", ' +
-            'valid options are [{}]'.format(
-                ", ".join(options.keys())
-            )
+            f'Overlap metric: "{metric}" not implemented for geometry type: "{geometry}", '
+            + "valid options are [{}]".format(", ".join(options.keys()))
         )
     return options[metric]
 
 
 @overload
 def nms_masks(
-        masks : torch.Tensor, 
-        scores : torch.Tensor, 
-        overlap_threshold : float=0.5, 
-        return_indices : Literal[False]=False, 
-        group_first : bool=True, 
-        boxes : torch.Tensor | None=None,
-        overlap_fn : (
-            Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor] | str
-        )=iou_masks,
-        overlap_fn_boxes : Callable[..., torch.Tensor] | str | None=None
-    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    masks: torch.Tensor,
+    scores: torch.Tensor,
+    overlap_threshold: float = 0.5,
+    return_indices: Literal[False] = False,
+    group_first: bool = True,
+    boxes: torch.Tensor | None = None,
+    overlap_fn: (
+        Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor]
+        | str
+    ) = iou_masks,
+    overlap_fn_boxes: Callable[..., torch.Tensor] | str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]: ...
 @overload
 def nms_masks(
-        masks : torch.Tensor, 
-        scores : torch.Tensor, 
-        overlap_threshold : float=0.5, 
-        return_indices : Literal[True]=True, 
-        group_first : bool=True, 
-        boxes : torch.Tensor | None=None,
-        overlap_fn : (
-            Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor] | str
-        )=iou_masks,
-        overlap_fn_boxes : Callable[..., torch.Tensor] | str | None=None
-    ) -> torch.Tensor: ...
+    masks: torch.Tensor,
+    scores: torch.Tensor,
+    overlap_threshold: float = 0.5,
+    return_indices: Literal[True] = True,
+    group_first: bool = True,
+    boxes: torch.Tensor | None = None,
+    overlap_fn: (
+        Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor]
+        | str
+    ) = iou_masks,
+    overlap_fn_boxes: Callable[..., torch.Tensor] | str | None = None,
+) -> torch.Tensor: ...
 def nms_masks(
-        masks : torch.Tensor, 
-        scores : torch.Tensor, 
-        overlap_threshold : float=0.5, 
-        return_indices : bool=False, 
-        group_first : bool=True, 
-        boxes : torch.Tensor | None=None,
-        overlap_fn : (
-            Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor] | str
-        )=iou_masks,
-        overlap_fn_boxes : Callable[..., torch.Tensor] | str | None=None
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    masks: torch.Tensor,
+    scores: torch.Tensor,
+    overlap_threshold: float = 0.5,
+    return_indices: bool = False,
+    group_first: bool = True,
+    boxes: torch.Tensor | None = None,
+    overlap_fn: (
+        Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.dtype], torch.Tensor]
+        | str
+    ) = iou_masks,
+    overlap_fn_boxes: Callable[..., torch.Tensor] | str | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """Efficiently perform non-maximum suppression on a set of boolean masks.
 
-    Defaults to a modified two-stage NMS algorithm, that aims to minimize the number of mask intersection calculations needed.
+    Defaults to a modified two-stage NMS algorithm,
+    that aims to minimize the number of mask intersection calculations needed.
 
     Args:
-        masks: A tensor of shape (n, h, w), where n is the number of masks and h and w are the height and width of the masks.
-        scores: A tensor of shape (n, ) containing the "scores" of the masks, this can merely be though of as a priority score, 
-            where the higher the score, the higher the priority of the object - it does not have to be a probability/confidence.
-        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to 0.5.
-        return_indices: A flag to indicate whether to return the indices of the picked objects or the objects themselves. Defaults to False.
-            If True, both the picked objects and scores are returned.
-        group_first: A flag to indicate whether two use the two-stage NMS method. Defaults to True.
-        boxes: Bounding boxes for the masks. A tensor of shape (n, 4), where n is the number of masks and 
+        masks: A tensor of shape `(n, h, w)`,
+            where `n` is the number of masks and `h` and `w` are the height and width of the masks.
+        scores: A tensor of shape `(n, )` containing the "scores" of the masks,
+            this can merely be though of as a priority score, where the higher the score,
+            the higher the priority of the object - it does not have to be a probability/confidence.
+        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to `0.5`.
+        return_indices: A flag indicating whether to return the indices of the picked objects or the objects themselves.
+            Defaults to `False`. If `True`, both the picked objects and scores are returned.
+        group_first: A flag indicating whether two use the two-stage NMS method. Defaults to `True`.
+        boxes: Bounding boxes for the masks. A tensor of shape `(n, 4)`, where `n` is the number of masks and
             the 4 columns are the x_min, y_min, x_max and y_max coordinates of the bounding boxes.
         overlap_fn: A function to compute overlaps between masks.
         overlap_fn_boxes: A function to compute overlaps between boxes.
-    
+
     Returns:
-        Either a tensor of shape `(m,)` containing the indices of the picked objects, 
-        or a tuple (`tuple[torch.Tensor, torch.Tensor]`) where the first element contains 
+        Either a tensor of shape `(m,)` containing the indices of the picked objects,
+        or a tuple (`tuple[torch.Tensor, torch.Tensor]`) where the first element contains
         the picked masks and the second element is a tensor of their scores.
 
     """
@@ -735,8 +748,10 @@ def nms_masks(
         if boxes is None:
             raise ValueError("'boxes' must be specified for nms_masks when 'group_first' is True")
         if overlap_fn_boxes is None:
-            raise RuntimeError("If an overlap function is manually provided for masks, one must also be provided for boxes.")
-        # We decrease the overlap_threshold for the clustering, 
+            raise RuntimeError(
+                "If an overlap function is manually provided for masks, one must also be provided for boxes."
+            )
+        # We decrease the overlap_threshold for the clustering,
         # since there is no straight-forward relationship between the IoU of the boxes and the IoU of the masks
         groups, _ = cluster_overlap_boxes(
             boxes=boxes,
@@ -750,12 +765,17 @@ def nms_masks(
                 _nms_ind[i] = group
             else:
                 group_boxes = boxes[group].round().long()
-                xmin, ymin, xmax, ymax = group_boxes[:, 0].min(), group_boxes[:, 1].min(), group_boxes[:, 2].max(), group_boxes[:, 3].max()
+                xmin, ymin, xmax, ymax = (
+                    group_boxes[:, 0].min(),
+                    group_boxes[:, 1].min(),
+                    group_boxes[:, 2].max(),
+                    group_boxes[:, 3].max(),
+                )
                 _nms_ind[i] = group[nms_masks_(
-                        masks=masks[group, ymin:(ymax+1), xmin:(xmax+1)],
-                        scores=scores[group],
-                        overlap_threshold=overlap_threshold,
-                        overlap_fn=overlap_fn
+                    masks=masks[group, ymin : (ymax + 1), xmin : (xmax + 1)],
+                    scores=scores[group],
+                    overlap_threshold=overlap_threshold,
+                    overlap_fn=overlap_fn,
                 )]
         if len(_nms_ind) > 0:
             nms_ind = torch.cat(_nms_ind)
@@ -769,62 +789,61 @@ def nms_masks(
 
 @overload
 def nms_polygons(
-        polygons : list[torch.Tensor], 
-        scores : torch.Tensor, 
-        overlap_threshold : float | int=0.5, 
-        return_indices : Literal[False]=False, 
-        group_first : bool=True, 
-        boxes : torch.Tensor | None=None,
-        overlap_fn : Callable | str="IoU",
-        overlap_fn_boxes : Callable[..., torch.Tensor] | str | None=None,
-    ) -> tuple[list[torch.Tensor], torch.Tensor]: ...
+    polygons: list[torch.Tensor],
+    scores: torch.Tensor,
+    overlap_threshold: float | int = 0.5,
+    return_indices: Literal[False] = False,
+    group_first: bool = True,
+    boxes: torch.Tensor | None = None,
+    overlap_fn: Callable | str = "IoU",
+    overlap_fn_boxes: Callable[..., torch.Tensor] | str | None = None,
+) -> tuple[list[torch.Tensor], torch.Tensor]: ...
 @overload
 def nms_polygons(
-        polygons : list[torch.Tensor], 
-        scores : torch.Tensor, 
-        overlap_threshold : float | int=0.5, 
-        return_indices : Literal[True]=True, 
-        group_first : bool=True, 
-        boxes : torch.Tensor | None=None,
-        overlap_fn : Callable | str="IoU",
-        overlap_fn_boxes : Callable[..., torch.Tensor] | str | None=None,
-    ) -> torch.Tensor: ...
+    polygons: list[torch.Tensor],
+    scores: torch.Tensor,
+    overlap_threshold: float | int = 0.5,
+    return_indices: Literal[True] = True,
+    group_first: bool = True,
+    boxes: torch.Tensor | None = None,
+    overlap_fn: Callable | str = "IoU",
+    overlap_fn_boxes: Callable[..., torch.Tensor] | str | None = None,
+) -> torch.Tensor: ...
 def nms_polygons(
-        polygons : list[torch.Tensor], 
-        scores : torch.Tensor, 
-        overlap_threshold : float | int=0.5, 
-        return_indices : bool=False, 
-        group_first : bool=True, 
-        boxes : torch.Tensor | None=None,
-        overlap_fn : Callable | str="IoU",
-        overlap_fn_boxes : Callable[..., torch.Tensor] | str | None=None,
-    ) -> torch.Tensor | tuple[list[torch.Tensor], torch.Tensor]:
+    polygons: list[torch.Tensor],
+    scores: torch.Tensor,
+    overlap_threshold: float | int = 0.5,
+    return_indices: bool = False,
+    group_first: bool = True,
+    boxes: torch.Tensor | None = None,
+    overlap_fn: Callable | str = "IoU",
+    overlap_fn_boxes: Callable[..., torch.Tensor] | str | None = None,
+) -> torch.Tensor | tuple[list[torch.Tensor], torch.Tensor]:
     """Efficiently perform non-maximum suppression on a set of polygons.
 
-    Defaults to a modified two-stage NMS algorithm, 
+    Defaults to a modified two-stage NMS algorithm,
     that aims to minimize the number of polygon intersection calculations needed (very expensive).
 
     Args:
-        polygons: A list of tensors of shape (n, 2), 
-            where n is the number of vertices in the polygon and the 2 columns are the x and y coordinates of the vertices.
-        scores: A tensor of shape (n, ) containing the "scores" of the polygons, this can merely be though of as a priority score,
-            where the higher the score, the higher the priority of the object - it does not have to be a probability/confidence.
-        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression.
-            Defaults to 0.5.
-        return_indices: A flag to indicate whether to return the indices of the picked objects or the objects themselves.
+        polygons: A list of tensors of shape `(n, 2)`, where `n` is the number of vertices in the polygon
+            and the 2 columns are the x and y coordinates of the vertices.
+        scores: A tensor of shape `(n, )` containing the "scores" of the polygons,
+            this can merely be though of as a priority score, where the higher the score,
+            the higher the priority of the object - it does not have to be a probability/confidence.
+        overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to 0.5.
+        return_indices: A flag indicating whether to return the indices of the picked objects or the objects themselves.
             Defaults to False. If True, both the picked objects and scores are returned.
-        group_first: A flag to indicate whether two use the two-stage NMS method. Defaults to True (recommended).
-        boxes: Bounding boxes for the polygons. 
-            A tensor of shape (n, 4), where n is the number of polygons and 
+        group_first: Flag indicating whether two use the two-stage NMS method. Defaults to True (recommended).
+        boxes: Bounding boxes for the polygons. A tensor of shape `(n, 4)`, where `n` is the number of polygons and
             the 4 columns are the x_min, y_min, x_max and y_max coordinates of the bounding boxes.
-        overlap_fn: A callable to compute overlap between polygons. Must accept either one or two lists of tensors and return a tensor. 
-            Can also be a string (e.g., "IoU").
+        overlap_fn: A callable to compute overlap between polygons.
+            Must accept either one or two lists of tensors and return a tensor. Can also be a string (e.g., "IoU").
         overlap_fn_boxes: A callable to compute overlap between a set of bounding boxes.
             Can also be a string (e.g., "IoU").
-    
+
     Returns:
-        Either a tensor of shape `(m,)` containing the indices of the picked polygons, 
-        or a tuple (`tuple[list[torch.Tensor], torch.Tensor]`) where the first element 
+        Either a tensor of shape `(m,)` containing the indices of the picked polygons,
+        or a tuple (`tuple[list[torch.Tensor], torch.Tensor]`) where the first element
         contains the picked polygons and the second element is a tensor of their scores.
 
     """
@@ -836,7 +855,9 @@ def nms_polygons(
         overlap_fn = get_overlap_fn("polygon", overlap_fn)
     else:
         if overlap_fn_boxes is None:
-            raise RuntimeError("If an overlap function is manually provided for polygons, one must also be provided for boxes.")
+            raise RuntimeError(
+                "If an overlap function is manually provided for polygons, one must also be provided for boxes."
+            )
     device = polygons[0].device
     if not group_first or len(polygons) < 10:
         nms_ind = nms_polygons_(
@@ -848,8 +869,8 @@ def nms_polygons(
     else:
         if boxes is None:
             raise ValueError("'boxes' must be specified for nms_masks when 'group_first' is True")
-        # We decrease the overlap_threshold for the clustering, 
-        # since there is no straight-forward relationship between the overlap of the boxes and the overlap of the polygons
+        # We decrease the overlap_threshold for the clustering,
+        # since there isn't a straight-forward relationship between the box- and polygon-overlap
         groups, _ = cluster_overlap_boxes(
             boxes=boxes, 
             overlap_threshold=min(1, overlap_threshold / 4), 
@@ -861,13 +882,12 @@ def nms_polygons(
             if len(group) == 1:
                 nms_ind.append(group)
             else:
-                nms_ind.append(
-                    group[nms_polygons_(
-                        polys=[polygons[gi] for gi in group], 
-                        scores=scores[group], 
-                        overlap_threshold=overlap_threshold, overlap_fn=overlap_fn
-                    )]
-                )
+                nms_ind.append(group[nms_polygons_(
+                    polys=[polygons[gi] for gi in group],
+                    scores=scores[group],
+                    overlap_threshold=overlap_threshold,
+                    overlap_fn=overlap_fn,
+                )])
         if len(nms_ind) > 0:
             nms_ind = torch.cat(nms_ind)
         else:
@@ -877,14 +897,15 @@ def nms_polygons(
     else:
         return [polygons[ni] for ni in nms_ind], scores[nms_ind]
 
+
 def nms_boxes(
-        boxes : torch.Tensor, 
-        scores : torch.Tensor, 
-        overlap_threshold : float | int=0.5,
-        overlap_fn : Callable[[torch.Tensor], torch.Tensor] | str | None=None,
-    ) -> torch.Tensor:
+    boxes: torch.Tensor,
+    scores: torch.Tensor,
+    overlap_threshold: float | int = 0.5,
+    overlap_fn: Callable[[torch.Tensor], torch.Tensor] | str | None = None,
+) -> torch.Tensor:
     """Perform NMS on boxes and return the NMS indexes.
-    
+
     Wraps `torchvision.ops.nms`; the standard non-maximum suppression algorithm.
     """
     if overlap_fn is None or isinstance(overlap_fn, str) and (overlap_fn := overlap_fn.strip().lower()) == "iou":
@@ -895,4 +916,4 @@ def nms_boxes(
             return torchvision.ops.nms(boxes, scores, overlap_threshold).sort().values
     if isinstance(overlap_fn, str):
         overlap_fn = get_overlap_fn("box", overlap_fn)
-    return base_nms_(boxes, overlap_fn=overlap_fn, scores=scores, overlap_threshold=overlap_threshold, return_indices=True)
\ No newline at end of file
+    return base_nms_(boxes, overlap_fn, scores, overlap_threshold=overlap_threshold, return_indices=True)
diff --git a/src/flat_bug/predictor.py b/src/flat_bug/predictor.py
index 43dd2db..98ee644 100644
--- a/src/flat_bug/predictor.py
+++ b/src/flat_bug/predictor.py
@@ -1,4 +1,5 @@
 """Implementation of the public flatbug `Predictor` and `TensorPredictions`."""
+
 import atexit
 import base64
 import json
@@ -24,7 +25,6 @@
 from ultralytics import YOLO
 from ultralytics.engine.results import Results
 
-# from flat_bug.yolo_helpers import *
 from flat_bug import download_from_repository, logger
 from flat_bug.config import CFG_PARAMS, DEFAULT_CFG, read_cfg
 from flat_bug.geometric import (
@@ -72,22 +72,22 @@ def _work(self):
             fn, future = self._queue.get()
             try:
                 if future.set_running_or_notify_cancel():
-                    try: 
+                    try:
                         future.set_result(fn())
-                    except Exception as e: 
+                    except Exception as e:
                         future.set_exception(e)
                         logger.error(f"Async task failed: {e}", exc_info=True)
             finally:
-                with self._lock: 
+                with self._lock:
                     self._active.discard(future)
                 self._queue.task_done()
 
     def submit(self, fn, *args, **kwargs):
         """Submit a call to be executed asynchronously."""
-        if not self._threads: 
+        if not self._threads:
             self._init_pool()
         future = Future()
-        with self._lock: 
+        with self._lock:
             self._active.add(future)
         # Blocks here if backlog is full
         self._queue.put((lambda: fn(*args, **kwargs), future))
@@ -95,25 +95,26 @@ def submit(self, fn, *args, **kwargs):
 
     def flush(self, progress=False):
         """Wait for all pending futures to finish."""
-        with self._lock: 
+        with self._lock:
             pending = list(self._active)
-        if not pending: 
+        if not pending:
             return
 
         if progress and tqdm:
-            for _ in tqdm(as_completed(pending), total=len(pending), desc="Finishing pending executions."): 
+            for _ in tqdm(as_completed(pending), total=len(pending), desc="Finishing pending executions."):
                 pass
         else:
             wait(pending)
 
+
 _executor = AsyncExecutor()
 
 
 class Prepared_Results:
     """Class for containing the results from a single `Predictor._detect_instances` call.
-    
-    This should probably not be its own class, but just a TensorPredictions object with a single element instead, 
-    but this would require altering the `TensorPredictions._combine_predictions` function 
+
+    This should probably not be its own class, but just a TensorPredictions object with a single element instead,
+    but this would require altering the `TensorPredictions._combine_predictions` function
     to handle a single element differently or pass a flag or something.
     """
 
@@ -123,7 +124,7 @@ def __init__(self, predictions: ResultsWithTiles, scale: tuple[float, float], de
         assert self._predictions.boxes is not None and isinstance(self._predictions.boxes.data, torch.Tensor)
         self._predictions.boxes.data[:, :4] /= self.wh_scale.repeat(1, 2)
         self._predictions.polygons = self._predictions.polygons._apply(
-            lambda poly : (poly + torch.roll(poly, 1, dims=0)) / (2 * self.wh_scale)
+            lambda poly: (poly + torch.roll(poly, 1, dims=0)) / (2 * self.wh_scale)
         )
         self.scale = sum(scale) / 2
         self.device = device
@@ -142,7 +143,7 @@ def __getitem__(self, i):
     def contours(self):  # noqa: D102
         assert self._predictions.masks is not None
         return [
-            torch.as_tensor(c) if c is not None else torch.tensor([], dtype=torch.long, device=self.device) 
+            torch.as_tensor(c) if c is not None else torch.tensor([], dtype=torch.long, device=self.device)
             for c in self._predictions.masks.xy
         ]
 
@@ -164,23 +165,24 @@ def confs(self) -> torch.Tensor | np.ndarray:  # noqa: D102
     @property
     def classes(self) -> torch.Tensor:
         """Not implemented properly."""
-        ### OBS: This is not really implemented, but exists just so that the the rest of the code already handles the multiclass case, 
-        # but this function will need to be changed for it to work properly ### 
-        # Currently this function is pretty redundant, since the localizer only has a single class. 
-        # If there were more classes, the function should do some kind of argmax on self._predictions.boxes.cls 
+        ### OBS: This is not really implemented, but exists just so that the the rest of the code already handles
+        # the multiclass case, but this function will need to be changed for it to work properly ###
+        # Currently this function is pretty redundant, since the localizer only has a single class.
+        # If there were more classes, the function should do some kind of argmax on self._predictions.boxes.cls
         # (I assume these are class probabilities).
         assert self._predictions.boxes is not None
         return torch.ones_like(torch.as_tensor(self._predictions.boxes.cls))
 
+
 # Class for containing the results from multiple _detect_instances calls
 class TensorPredictions:
-    """Result handling class for combining the results from multiple YOLOv8 detections at different scales into a single object.
+    """Result class for combining the results from multiple YOLOv8 detections at different scales into a single object.
 
-    `TensorPredictions` handles a rather complex merging procedure, 
-    resizing to remove image padding and scaling effects on the masks and boxes, 
+    `TensorPredictions` handles a rather complex merging procedure,
+    resizing to remove image padding and scaling effects on the masks and boxes,
     and non-maximum suppression using mask-IoU or mask-IoS.
 
-    `TensorPredictions` also allows for easy conversion from mask to contours and back, plotting of the results, 
+    `TensorPredictions` also allows for easy conversion from mask to contours and back, plotting of the results,
     and (de-)serialization to save and load the results to/from disk.
     """
 
@@ -192,22 +194,25 @@ class TensorPredictions:
     device = None
     dtype = None
     CONSTANTS = (
-        "image", "image_path", 
-        "device", "dtype", 
-        "time", 
-        "mask_height", "mask_width", 
+        "image",
+        "image_path",
+        "device",
+        "dtype",
+        "time",
+        "mask_height",
+        "mask_width",
         "BOX_IS_EQUAL_MARGIN",
-        "PREFER_POLYGONS"
+        "PREFER_POLYGONS",
     )
 
     def __init__(
-            self, 
-            predictions : list[Prepared_Results] | None=None,
-            image : torch.Tensor | None=None,
-            image_path : str | None = None, 
-            time : bool=False, 
-            **kwargs
-        ):
+        self,
+        predictions: list[Prepared_Results] | None = None,
+        image: torch.Tensor | None = None,
+        image_path: str | None = None,
+        time: bool = False,
+        **kwargs,
+    ):
         """Create a `TensorPredictions` instance from scratch.
 
         You probably don't want to use this method manually. If you want to load saved results use:
@@ -215,7 +220,7 @@ def __init__(
         ```
         prediction = TensorPredictions.load(...)
         ```
-        
+
         Args:
             predictions: Predictions from multiple `Predictor._detect_instances` calls.
             image: The image where the predictions originate.
@@ -239,7 +244,7 @@ def __init__(
             else:
                 logger.warning(f"WARNING: Unknown keyword argument {k}={v} for TensorPredictions is ignored!")
 
-        # Device and dtype are None by default, but they may be set by the user or 
+        # Device and dtype are None by default, but they may be set by the user or
         # passed by **kwargs, so we check if they are None and if so set them to the default values
         # Then we check that they are the same for all predictions and the image (if they are not None)
         if predictions is not None and len(predictions) > 0:
@@ -250,7 +255,9 @@ def __init__(
             if self.dtype is None:
                 self.dtype = elem.dtype
             for pi, p in enumerate(predictions):
-                assert p.device == self.device, RuntimeError(f"predictions[{pi}].device {p.device} != device {self.device}")
+                assert p.device == self.device, RuntimeError(
+                    f"predictions[{pi}].device {p.device} != device {self.device}"
+                )
                 assert p.dtype == self.dtype, RuntimeError(f"predictions[{pi}].dtype {p.dtype} != dtype {self.dtype}")
             if image is not None:
                 assert image.device == self.device, RuntimeError(f"image.device {image.device} != device {self.device}")
@@ -276,24 +283,21 @@ def __init__(
             # If there are no predictions, set other attributes to empty tensors or lists.
             # Ensures correct type and device for the attributes when there are no predictions
             self.masks, self.polygons, self.boxes, self.confs, self.classes, self.scales = (
-                torch.empty((0, 0), device=self.device, dtype=self.dtype), 
-                [], 
-                torch.empty((0, 4), device=self.device, dtype=self.dtype), 
-                torch.empty((0,), device=self.device, dtype=self.dtype), 
-                torch.empty((0,), device=self.device, dtype=self.dtype), 
-                []
+                torch.empty((0, 0), device=self.device, dtype=self.dtype),
+                [],
+                torch.empty((0, 4), device=self.device, dtype=self.dtype),
+                torch.empty((0,), device=self.device, dtype=self.dtype),
+                torch.empty((0,), device=self.device, dtype=self.dtype),
+                [],
             )
 
         if self.time and predictions is not None and len(predictions) > 0:
             assert end is not None and start is not None
             end.record()
             torch.cuda.synchronize()
-            logger.info(f'Initializing TensorPredictions took {start.elapsed_time(end) / 1000:.3f} s')
+            logger.info(f"Initializing TensorPredictions took {start.elapsed_time(end) / 1000:.3f} s")
 
-    def _combine_predictions(
-            self, 
-            predictions: list[Prepared_Results]
-        ):
+    def _combine_predictions(self, predictions: list[Prepared_Results]):
         """Combine a list of Prepared_Results from multiple `Predictor._detect_instances` calls.
 
         This function is used in-place during initialization of a `TensorPrediction` instance.
@@ -310,7 +314,7 @@ def _combine_predictions(
             end_duplication_removal = torch.cuda.Event(enable_timing=True)
             end_mask_combination = torch.cuda.Event(enable_timing=True)
             start.record()
-        
+
         self.boxes = torch.cat([torch.as_tensor(p.boxes) for p in predictions])  # Nx4
         self.confs = torch.cat([torch.as_tensor(p.confs) for p in predictions])  # N
         self.scales = [p.scale for p in predictions for _ in range(len(p))]  # N
@@ -354,13 +358,14 @@ def _combine_predictions(
         self.masks.orig_shape = self.image.shape[1:]
 
         poly_lists = [p._predictions.polygons.to_list() for p in predictions]
-        self.polygons : list[torch.Tensor] = [
-            p[int(nd_i.item()) if isinstance(nd, torch.Tensor) else int(nd)] 
-            for p, nd in zip(poly_lists, valid_chunked) for nd_i in nd
+        self.polygons: list[torch.Tensor] = [
+            p[int(nd_i.item()) if isinstance(nd, torch.Tensor) else int(nd)]
+            for p, nd in zip(poly_lists, valid_chunked)
+            for nd_i in nd
         ]
         self.classes = torch.cat([p.classes[nd] for p, nd in zip(predictions, valid_chunked)])  # N
         self.scales = [predictions[i].scale for i, p in enumerate(valid_chunked) for _ in range(len(p))]  # N
-        
+
         # Sort the polygons, masks, boxes, classes, scales and confidences by confidence
         sorted_indices = self.confs.argsort(descending=True)
         self.masks = self.masks[sorted_indices]
@@ -371,37 +376,38 @@ def _combine_predictions(
         self.confs = self.confs[sorted_indices]
 
         # # Check that everything is the correct size
-        assert len(self) == len(self.boxes), RuntimeError(f"len(self) {len(self)} != len(self.boxes) {len(self.boxes)}")
-        assert len(self) == len(self.confs), RuntimeError(f"len(self) {len(self)} != len(self.confs) {len(self.confs)}")
-        assert len(self) == len(self.classes), RuntimeError(f"len(self) {len(self)} != len(self.classes) {len(self.classes)}")
-        assert len(self) == len(self.scales), RuntimeError(f"len(self) {len(self)} != len(self.scales) {len(self.scales)}")
+        assert len(self) == len(self.boxes), RuntimeError(f"{len(self)=} != {len(self.boxes)=}")
+        assert len(self) == len(self.confs), RuntimeError(f"{len(self)=} != {len(self.confs)=}")
+        assert len(self) == len(self.classes), RuntimeError(f"{len(self)=} != {len(self.classes)=}")
+        assert len(self) == len(self.scales), RuntimeError(f"{len(self)=} != {len(self.scales)=}")
         if self.time:
-            assert start is not None and end is not None and end_duplication_removal is not None and end_mask_combination is not None
+            assert (
+                start is not None
+                and end is not None
+                and end_duplication_removal is not None
+                and end_mask_combination is not None
+            )
             end.record()
             torch.cuda.synchronize()
             total = start.elapsed_time(end) / 1000
             duplication_removal = start.elapsed_time(end_duplication_removal) / 1000
             mask_combination = end_duplication_removal.elapsed_time(end_mask_combination) / 1000
             logger.info(
-                f'Combining {len(predictions)} predictions into a single TensorPredictions object took {total:.3f} s |'
-                f' Duplication removal: {duplication_removal:.3f} s | Mask combination: {mask_combination:.3f} s'
+                f"Combining {len(predictions)} predictions into a single TensorPredictions object took {total:.3f} s |"
+                f" Duplication removal: {duplication_removal:.3f} s | Mask combination: {mask_combination:.3f} s"
             )
 
-    def offset_scale_pad(
-            self, 
-            offset: torch.Tensor, 
-            scale: float, 
-            pad: int = 0
-        ):
+    def offset_scale_pad(self, offset: torch.Tensor, scale: float, pad: int = 0):
         """Scale and offset the detections to real image coordinates in-place.
-        
-        Since the image may be padded, the masks and boxes should be offset by the padding-width and scaled 
+
+        Since the image may be padded, the masks and boxes should be offset by the padding-width and scaled
         by the `scale_before` factor to match the original image size. Also pads the boxes by pad pixels to be safe.
 
         Args:
-            offset: A vector of length 2 containing the x and y offset of the image. Useful for removing image-padding effects.
+            offset: A vector of length 2 containing the x and y offset of the image.
+                Useful for removing image-padding effects.
             scale: The scale factor of the image.
-            pad: The number of pixels to pad the boxes by. Defaults to 0. (Not to be confused with image-padding, 
+            pad: The number of pixels to pad the boxes by. Defaults to 0. (Not to be confused with image-padding,
                 this is about expanding the boxes a bit to ensure they cover the entire mask)
 
         Returns:
@@ -432,7 +438,7 @@ def offset_scale_pad(
             self.polygons = [(poly + offset.unsqueeze(0)) * scale for poly in self.polygons]
 
             # However masks are more complicated since they don't have the same size as the image
-            image_shape = torch.tensor( # Get the shape of the original image
+            image_shape = torch.tensor(  # Get the shape of the original image
                 [self.image.shape[1], self.image.shape[2]], 
                 device=self.device,
                 dtype=self.dtype
@@ -442,21 +448,24 @@ def offset_scale_pad(
             # here the scaled and padded image size is calculated from the original image shape
             # (but it would probably be easier just to pass it...)
             offset_norm = -offset / (image_shape / scale - 2 * offset)
-            orig_mask_shape = torch.tensor([self.masks.shape[1], self.masks.shape[2]], device=self.device, dtype=self.dtype) - 1
+            orig_mask_shape = torch.tensor(
+                [self.masks.shape[1], self.masks.shape[2]],
+                device=self.device, dtype=self.dtype
+            ) - 1
             # Convert the normalized offset to the coordinates of the masks
             offset_mask_coords = offset_norm * orig_mask_shape
             # Round the coordinates to the nearest integer and convert to long (needed for indexing)
             offset_mask_coords = torch.round(offset_mask_coords).long()
             self.masks.data = torch.as_tensor(self.masks.data)[
-                :, 
-                offset_mask_coords[0]:(-(offset_mask_coords[0] + 1) if offset_mask_coords[0] != 0 else None), 
-                offset_mask_coords[1]:(-(offset_mask_coords[1] + 1) if offset_mask_coords[1] != 0 else None)
+                :,
+                offset_mask_coords[0] : (-(offset_mask_coords[0] + 1) if offset_mask_coords[0] != 0 else None),
+                offset_mask_coords[1] : (-(offset_mask_coords[1] + 1) if offset_mask_coords[1] != 0 else None),
             ]  # Slice out the padded parts of the masks
 
         if self.time:
             end.record()
             torch.cuda.synchronize()
-            logger.info(f'Offsetting, scaling and padding took {start.elapsed_time(end) / 1000:.3f} s')
+            logger.info(f"Offsetting, scaling and padding took {start.elapsed_time(end) / 1000:.3f} s")
 
         return self
 
@@ -480,12 +489,15 @@ def fix_boxes(self):
             if len(this_mask_nz) == 0:
                 self.boxes[i] = torch.tensor([0, 0, 0, 0], device=self.device, dtype=self.dtype)
             else:
-                self.boxes[i] = torch.tensor([
-                        this_mask_nz[:, 1].min(), 
-                        this_mask_nz[:, 0].min(), 
+                self.boxes[i] = torch.tensor(
+                    [
+                        this_mask_nz[:, 1].min(),
+                        this_mask_nz[:, 0].min(),
                         this_mask_nz[:, 1].max(),
-                        this_mask_nz[:, 0].max()
-                    ], device=self.device, dtype=self.dtype
+                        this_mask_nz[:, 0].max(),
+                    ],
+                    device=self.device,
+                    dtype=self.dtype,
                 ) * mask_to_image_scale.repeat(2)
         self.boxes[:, :2] = self.boxes[:, :2].floor()
         self.boxes[:, 2:] = self.boxes[:, 2:].ceil()
@@ -493,14 +505,9 @@ def fix_boxes(self):
         self.boxes[:, 1:4:2] = self.boxes[:, 1:4:2].clamp(0, self.image.shape[1])
         return self
 
-    def non_max_suppression(
-            self, 
-            overlap_threshold : float, 
-            metric : str,
-            **kwargs
-        ):
+    def non_max_suppression(self, overlap_threshold: float, metric: str, **kwargs):
         """Perform non-max suppression (NMS) in-place.
-        
+
         Either uses polygons (most likely) or masks.
         """
         if self.time:
@@ -519,56 +526,59 @@ def non_max_suppression(
                 nms_ind = nms_polygons(
                     polygons=self.polygons,
                     scores=self.confs,
-                    overlap_threshold=overlap_threshold, 
+                    overlap_threshold=overlap_threshold,
                     overlap_fn=metric,
-                    return_indices=True, 
-                    boxes=self.boxes, 
-                    **kwargs
+                    return_indices=True,
+                    boxes=self.boxes,
+                    **kwargs,
                 )
             else:
                 image_to_mask_scale = torch.tensor(
                     [self.image.shape[1] / self.masks.data.shape[1], self.image.shape[2] / self.masks.data.shape[2]],
-                    device=self.device, dtype=self.dtype
+                    device=self.device,
+                    dtype=self.dtype,
                 )
-                nms_ind : torch.Tensor = nms_masks(
+                nms_ind: torch.Tensor = nms_masks(
                     masks=torch.as_tensor(self.masks.data),
                     scores=self.confs,
-                    overlap_threshold=overlap_threshold, 
+                    overlap_threshold=overlap_threshold,
                     overlap_fn=metric,
                     return_indices=True,
-                    boxes=self.boxes / image_to_mask_scale.repeat(2).unsqueeze(0), 
-                    **kwargs
+                    boxes=self.boxes / image_to_mask_scale.repeat(2).unsqueeze(0),
+                    **kwargs,
                 )
             # Remove the instances that were not selected
             self = self[nms_ind.sort().values]
         else:
             nms_ind = torch.empty((0,))
-        
+
         if self.time:
             end.record()
             torch.cuda.synchronize()
             logger.info(
-                f'Non-maximum suppression took {start.elapsed_time(end) / 1000:.3f}s '
-                f'for removing {len_before - len(nms_ind)} elements of {len_before} elements'
+                f"Non-maximum suppression took {start.elapsed_time(end) / 1000:.3f}s "
+                f"for removing {len_before - len(nms_ind)} elements of {len_before} elements"
             )
         return self
 
     @property
     def contours(self) -> list[torch.Tensor]:
-        """Wraps the openCV.findContours function, and uses openCV.contourArea to select the largest contour for each mask."""
+        """Wraps the `openCV.findContours` function.
+
+        `openCV.contourArea` is used to select the largest contour for each mask.
+        """
         if self.PREFER_POLYGONS:
             return self.polygons
         else:
             return [
-                self.contour_to_image_coordinates(find_contours(create_contour_mask(mask), largest_only=True, simplify=False)) 
+                self.contour_to_image_coordinates(
+                    find_contours(create_contour_mask(mask), largest_only=True, simplify=False)
+                )
                 for mask in self.masks.data
             ]
 
     @contours.setter
-    def contours(
-            self, 
-            value : list[torch.Tensor | np.ndarray]
-        ):
+    def contours(self, value: list[torch.Tensor | np.ndarray]):
         assert self.mask_height is not None and self.mask_width is not None
         if self.PREFER_POLYGONS:
             if not isinstance(value, list):
@@ -585,7 +595,7 @@ def contours(
                         raise RuntimeError(f"Unknown shape `{value[i].shape}` for `contours[{i}]` - should be (N, 2)")
                 value[i] = torch.from_numpy(
                     scale_contour(
-                        contour=np.asarray(value[i]), 
+                        contour=np.asarray(value[i]),
                         scale=contour_scaling,
                         expand_by_one=True
                     )
@@ -596,7 +606,10 @@ def contours(
                 for _ in range(len(value))
             ])  # Initialize empty masks
         else:
-            self.masks = contours_to_masks(list(map(torch.as_tensor, value)), self.mask_height, self.mask_width).to(self.device)
+            self.masks = contours_to_masks(
+                list(map(torch.as_tensor, value)),
+                self.mask_height, self.mask_width
+            ).to(self.device)
 
     @property
     def areas(self):
@@ -606,11 +619,7 @@ def areas(self):
         else:
             return self.masks.sum(1).sum(1).tolist()
 
-    def contour_to_image_coordinates(
-            self, 
-            contour: torch.Tensor, 
-            scale: float = 1
-        ) -> torch.Tensor:
+    def contour_to_image_coordinates(self, contour: torch.Tensor, scale: float = 1) -> torch.Tensor:
         """Convert a contour from mask coordinates to image coordinates.
 
         Args:
@@ -631,10 +640,7 @@ def contour_to_image_coordinates(
 
         return scaled_contour
 
-    def flip(
-            self, 
-            direction : str="vertical"
-        ):
+    def flip(self, direction: str = "vertical"):
         """Flips the masks, polygons and boxes along the specified axis in-place.
 
         Args:
@@ -672,7 +678,7 @@ def flip(
         if self.time:
             end.record()
             torch.cuda.synchronize()
-            logger.info(f'Flipping masks, polygons and boxes {direction} took {start.elapsed_time(end) / 1000:.3f} s')
+            logger.info(f"Flipping masks, polygons and boxes {direction} took {start.elapsed_time(end) / 1000:.3f} s")
 
         return self
 
@@ -684,7 +690,7 @@ def new(self):  # noqa: D102
 
     def __getitem__(self, i):
         """Flexible indexing for TensorPredictions.
-        
+
         Can be used to get a single element, a slice, or an iterable of indices (e.g. a list, tuple, tensor).
         """
         new_tp = self.new()
@@ -698,7 +704,7 @@ def __getitem__(self, i):
                     if isinstance(i, torch.Tensor):
                         # Just to be super safe we cast to float, then round, then cast to long, then to list
                         i = i.float().round().long().tolist()
-                    if not all([isinstance(j, int) for j in i]) or all([isinstance(j, float) and (j % 1) == 0 for j in i]): # type: ignore
+                    if not all(isinstance(j, int) for j in i) or all(isinstance(j, float) and (j % 1) == 0 for j in i): # type: ignore
                         raise RuntimeError(f"Unknown type or non-integer float for {i}: {type(i).__name__}")
                     i = [int(j) for j in i]  # type: ignore
                     # If v is a tensor, we can just index it with the list
@@ -712,7 +718,9 @@ def __getitem__(self, i):
                         try:
                             new_value = v[i]
                         except Exception as e:
-                            raise RuntimeError(f"Unknown type for {k}: {type(v)} does not support flexible indexing") from e
+                            raise RuntimeError(
+                                f"Unknown type for {k}: {type(v)} does not support flexible indexing"
+                            ) from e
                 else:
                     # Otherwise, assume it's an index
                     if isinstance(i, torch.Tensor) and len(i) == 1:
@@ -731,37 +739,36 @@ def __getitem__(self, i):
                         raise RuntimeError(f"Unknown type for {k}: {type(v)}")
                 setattr(new_tp, k, new_value)
         return new_tp
-    
+
     def plot(
-            self, 
-            linewidth : int=2, 
-            masks : bool=True, 
-            boxes : bool=True, 
-            confidence : bool=True, 
-            outpath : str | None=None, 
-            scale : float=1,
-            contour_color : tuple[int, int, int]=(255, 0, 0),
-            box_color : tuple[int, int, int]=(0, 0, 0),
-            alpha : float=0.3,
-            wait : bool=False
-        ):
+        self,
+        linewidth: int = 2,
+        masks: bool = True,
+        boxes: bool = True,
+        confidence: bool = True,
+        outpath: str | None = None,
+        scale: float = 1,
+        contour_color: tuple[int, int, int] = (255, 0, 0),
+        box_color: tuple[int, int, int] = (0, 0, 0),
+        alpha: float = 0.3,
+        wait: bool = False,
+    ):
         """Visualizes `flatbug` predictions from a `TensorPredictions` object.
-        
+
         Args:
             linewidth: Linewidth of the segmentation countours and bounding boxes.
                 Default to 2.
-            masks: Flag to indicate whether segmentation contours should be included.
+            masks: Flag indicating whether segmentation contours should be included.
                 Default to True.
-            boxes: Flag to indicate whether bounding boxes should be included, if False confidences are also omitted.
-                Defaults to True.
-            confidence: Flag to indicate whether detection confidences should be included, if boxes is False, this argument is ignored.
+            boxes: Flag indicating whether bounding boxes should be included, if False confidences are also omitted.
                 Defaults to True.
-            outpath: Where should the visualization be saved? 
-                If outpath is None, then the rasterized visualization is returned as a `cv2.UMat`/`np.ndarray` (shape: HWC, colors: BGR).
-                Defaults to None.
-            scale: Render the visualization at a scale relative to the image size (from which the predictions originate). 
-                **OBS**: Large images and/or scales can be very slow to render.
-                    Defaults to 1. 
+            confidence: Flag indicating whether detection confidences should be included,
+                if boxes is False, this argument is ignored. Defaults to True.
+            outpath: Where should the visualization be saved? If outpath is None, then the rasterized visualization is
+                returned as a `cv2.UMat`/`np.ndarray` (shape: HWC, colors: BGR). Defaults to None.
+            scale: Render the visualization at a scale relative to the image size
+                (from which the predictions originate).
+                **OBS**: Large images and/or scales can be very slow to render. Defaults to 1.
             contour_color: RGB color ([0, 255]) to use for contour border and fill.
                 Defaults to `(255, 0, 0)` (red).
             box_color: RGB color ([0, 255]) to use for bounding box and confidence text color.
@@ -771,8 +778,8 @@ def plot(
             wait: If `False` (default) returns a future immediately, otherwise block and return the actual result.
 
         Returns:
-            If outpath is supplied, it is returned.
-            Otherwise the rasterized visualization is returned as as a `cv2.UMat`/`np.ndarray` (shape: HWC, colors: BGR).
+            If outpath is supplied, it is returned. Otherwise the rasterized visualization is returned
+            as as a `cv2.UMat`/`np.ndarray` (shape: HWC, colors: BGR).
             **OBS**: If `wait=True` then a future is returned instead.
 
         """
@@ -780,10 +787,10 @@ def plot(
         params.pop("self", None)
         params.pop("wait", None)
         data = {
-            "image" : self.image_path or self.image.detach().cpu().clone(),
-            "bboxes" : self.boxes.detach().cpu().clone(),
-            "contours" : [poly.detach().cpu().clone() for poly in self.polygons],
-            "confs" : self.confs.detach().cpu().clone(),
+            "image": self.image_path or self.image.detach().cpu().clone(),
+            "bboxes": self.boxes.detach().cpu().clone(),
+            "contours": [poly.detach().cpu().clone() for poly in self.polygons],
+            "confs": self.confs.detach().cpu().clone(),
         }
         if outpath and outpath.lower().endswith(".svg"):
             retval = _executor.submit(TensorPredictions._plot_svg, **data, **params)
@@ -795,16 +802,16 @@ def plot(
 
     @staticmethod
     def _box_to_svg_element(
-            box : torch.Tensor, 
-            scale : float=1.0, 
-            color : tuple[int, int, int]=(0, 0, 0), 
-            linewidth : float | int=2, 
-            label : str | None=None,
-            label_fontsize : float | int=12,
-            background_image : Any | None=None  # expected to be a NumPy array in BGR
-        ) -> str:
+        box: torch.Tensor,
+        scale: float = 1.0,
+        color: tuple[int, int, int] = (0, 0, 0),
+        linewidth: float | int = 2,
+        label: str | None = None,
+        label_fontsize: float | int = 12,
+        background_image: Any | None = None,  # expected to be a NumPy array in BGR
+    ) -> str:
         # Convert box color (RGB tuple) to hex.
-        hex_color = f'#{"".join(hs if len(hs)==2 else hs+"0" for v in color if len(hs:=hex(v)[2:]))}'
+        hex_color = f"#{''.join(hs if len(hs) == 2 else hs + '0' for v in color if len(hs := hex(v)[2:]))}"
 
         if scale != 1:
             box = (box.float() * scale).round().long()
@@ -813,13 +820,10 @@ def _box_to_svg_element(
         height = ymax - ymin
 
         # Build the rectangle SVG element.
-        rect_svg = (
-            f'<rect x="{xmin}" y="{ymin}" width="{width}" height="{height}" '
-            f'style="stroke:{hex_color};stroke-width:{linewidth};fill:none"/>'
-        )
+        rect_svg = f'<rect x="{xmin}" y="{ymin}" width="{width}" height="{height}" style="stroke:{hex_color};stroke-width:{linewidth};fill:none"/>'  # noqa: E501
 
         if label is not None:
-            avg_char_width = label_fontsize * 0.6 # 0.6: Arbitrarly chosen value for ~avg. character aspect ratio 
+            avg_char_width = label_fontsize * 0.6  # 0.6: Arbitrarly chosen value for ~avg. character aspect ratio
             text_width = int(len(label) * avg_char_width)
             text_height = label_fontsize
             offset = (linewidth * 3) // 2  # offset in pixels above the box
@@ -842,13 +846,13 @@ def _box_to_svg_element(
                     label_color = (0, 0, 0) if avg_brightness > 150 else (255, 255, 255)
 
             # Convert the label color to hex.
-            label_hex = f'#{"".join(hs if len(hs)==2 else hs+"0" for v in label_color if len(hs:=hex(v)[2:]))}'
+            label_hex = f"#{''.join(hs if len(hs) == 2 else hs + '0' for v in label_color if len(hs := hex(v)[2:]))}"
             # Create a <text> element. Note that we use a fixed font size (12px) and family.
             text_svg = (
                 f'<text x="{xmin}" y="{text_y}" fill="{label_hex}" '
                 f'style="font-size:{label_fontsize}px;font-family:Arial;font-weight:bold;">{label}</text>'
             )
-            out = f'<g>{rect_svg}{text_svg}</g>'
+            out = f"<g>{rect_svg}{text_svg}</g>"
         else:
             out = rect_svg
 
@@ -856,24 +860,24 @@ def _box_to_svg_element(
 
     @staticmethod
     def _contour_to_svg_element(
-            contour : torch.Tensor | Any, 
-            scale : float=1.0,
-            color : tuple[int, int, int]=(255, 0, 0),
-            linewidth : int | float=2,
-            alpha : int | float=0.33
-        ):
+        contour: torch.Tensor | Any,
+        scale: float = 1.0,
+        color: tuple[int, int, int] = (255, 0, 0),
+        linewidth: int | float = 2,
+        alpha: int | float = 0.33,
+    ):
         d_list = []
-        hex_color = f'#{"".join(hs if len(hs) == 2 else hs + "0" for v in color if len(hs := hex(v)[2:]))}'
+        hex_color = f"#{''.join(hs if len(hs) == 2 else hs + '0' for v in color if len(hs := hex(v)[2:]))}"
 
         stroke_colour = hex_color
-        fill_colour   = hex_color
+        fill_colour = hex_color
         if alpha > 1:
             alpha = alpha / 255
         for i in range(len(contour)):
             name = i
             x, y = (contour[i] * scale).round().long().tolist()
             d_list.append(f"{x},{y}")
-        d_str = ' '.join(d_list)
+        d_str = " ".join(d_list)
         return (
             f'<path name="{name}" '
             f'style="stroke:{stroke_colour};stroke-width:{linewidth};stroke-opacity:1;'
@@ -882,28 +886,24 @@ def _contour_to_svg_element(
 
     @staticmethod
     def _plot_svg(
-            image : torch.Tensor | str,
-            bboxes : torch.Tensor,
-            contours : torch.Tensor,
-            confs : torch.Tensor,
-            linewidth : int=2, 
-            masks : bool=True, 
-            boxes : bool=True, 
-            confidence : bool=True, 
-            outpath : str | None=None, 
-            scale : float=1,
-            contour_color : tuple[int, int, int] = (255, 0, 0),
-            box_color : tuple[int, int, int] = (0, 0, 0),
-            alpha : float = 0.3
-        ):
+        image: torch.Tensor | str,
+        bboxes: torch.Tensor,
+        contours: torch.Tensor,
+        confs: torch.Tensor,
+        linewidth: int = 2,
+        masks: bool = True,
+        boxes: bool = True,
+        confidence: bool = True,
+        outpath: str | None = None,
+        scale: float = 1,
+        contour_color: tuple[int, int, int] = (255, 0, 0),
+        box_color: tuple[int, int, int] = (0, 0, 0),
+        alpha: float = 0.3,
+    ):
         embed_jpeg = True
 
         if isinstance(image, str):
-            tensor_image = decode_image(
-                input=image, 
-                mode=ImageReadMode.RGB, 
-                apply_exif_orientation=True
-            )
+            tensor_image = decode_image(input=image, mode=ImageReadMode.RGB, apply_exif_orientation=True)
         else:
             tensor_image = image
         np_image = torchvision.transforms.ConvertImageDtype(torch.uint8)(tensor_image).permute(1, 2, 0).cpu().numpy()
@@ -915,15 +915,13 @@ def _plot_svg(
             height, width = np_image.shape[0:2]
             ((_, text_height), _) = cv2.getTextSize("0", cv2.FONT_HERSHEY_SIMPLEX, 1 * scale, 2)
 
-            encoded_string = base64.b64encode(cv2.imencode('.jpg', np_image)[1])
-            desc = ''
+            encoded_string = base64.b64encode(cv2.imencode(".jpg", np_image)[1])
+            desc = ""
             content = []
             content.append(
-                f'<svg width="{width}" height="{height}" '
-                'xmlns:xlink="http://www.w3.org/1999/xlink" '
-                'xmlns="http://www.w3.org/2000/svg">'
+                f'<svg width="{width}" height="{height}" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns="http://www.w3.org/2000/svg">'
             )
-            
+
             # Embed the background image.
             if embed_jpeg:
                 content.append(
@@ -934,56 +932,60 @@ def _plot_svg(
                 for cont in contours:
                     content.append(TensorPredictions._contour_to_svg_element(
                         cont,
-                        scale=scale, color=contour_color,
-                        linewidth=linewidth, alpha=alpha
+                        scale=scale,
+                        color=contour_color,
+                        linewidth=linewidth,
+                        alpha=alpha
                     ))
             if boxes:
                 for box, conf in zip(bboxes, confs):
-                    lbl = f'{conf.item():.1%}' if confidence else None
+                    lbl = f"{conf.item():.1%}" if confidence else None
                     # Pass the background image so the function can sample the area behind the label.
                     content.append(TensorPredictions._box_to_svg_element(
-                        box, 
-                        scale=scale, color=box_color, 
-                        linewidth=linewidth, label=lbl, 
-                        background_image=image, label_fontsize=text_height
+                        box,
+                        scale=scale,
+                        color=box_color,
+                        linewidth=linewidth,
+                        label=lbl,
+                        background_image=image,
+                        label_fontsize=text_height,
                     ))
-            content.append('</svg>')
-            
+            content.append("</svg>")
+
             if outpath:
-                with open(outpath, 'w+') as f:
+                with open(outpath, "w+") as f:
                     f.writelines(content)
                 return None
 
         except Exception as e:
             raise e
-        
+
         return content
 
     @staticmethod
     def _plot_image(
-            image : torch.Tensor | str,
-            bboxes : torch.Tensor,
-            contours : torch.Tensor,
-            confs : torch.Tensor,
-            linewidth : int=2, 
-            masks : bool=True, 
-            boxes : bool=True, 
-            confidence : bool=True, 
-            outpath : str | None=None, 
-            scale : float=1,
-            contour_color : tuple[int, int, int] = (255, 0, 0),
-            box_color : tuple[int, int, int] = (0, 0, 0),
-            alpha : float = 0.3
-        ) -> cv2.UMat | None:
+        image: torch.Tensor | str,
+        bboxes: torch.Tensor,
+        contours: torch.Tensor,
+        confs: torch.Tensor,
+        linewidth: int = 2,
+        masks: bool = True,
+        boxes: bool = True,
+        confidence: bool = True,
+        outpath: str | None = None,
+        scale: float = 1,
+        contour_color: tuple[int, int, int] = (255, 0, 0),
+        box_color: tuple[int, int, int] = (0, 0, 0),
+        alpha: float = 0.3,
+    ) -> cv2.UMat | None:
         if isinstance(image, str):
-            tensor_image = decode_image(
-                input=image, 
-                mode=ImageReadMode.RGB, 
-                apply_exif_orientation=True
-            )
+            tensor_image = decode_image(input=image, mode=ImageReadMode.RGB, apply_exif_orientation=True)
         else:
             tensor_image = image
-        np_image = cast(np.ndarray, torchvision.transforms.ConvertImageDtype(torch.uint8)(tensor_image).permute(1, 2, 0).cpu().numpy())
+        np_image = cast(
+            np.ndarray,
+            torchvision.transforms.ConvertImageDtype(torch.uint8)(tensor_image).permute(1, 2, 0).cpu().numpy(),
+        )
         if scale != 1:
             np_image = cv2.resize(np_image, (0, 0), fx=scale, fy=scale)
         np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
@@ -995,7 +997,10 @@ def _plot_image(
         if len(contours) > 0:
             # Draw masks
             if masks:
-                smpl_contours = [simplify_contour((c * scale).round().to(torch.int32).cpu().numpy(), scale / 2) for c in contours]
+                smpl_contours = [
+                    simplify_contour((c * scale).round().to(torch.int32).cpu().numpy(), scale / 2)
+                    for c in contours
+                ]
                 ih, iw = np_image.shape[:2]
                 _alpha = int(255 * alpha)
 
@@ -1005,13 +1010,14 @@ def _plot_image(
                     cv2.drawContours(this_poly_alpha, [c], -1, 1, -1)
                     poly_alpha += this_poly_alpha * _alpha
                 poly_alpha = poly_alpha.clip(0, 255) / 255
-                
+
                 # Create a red fill for the polygons
                 poly_fill = np.zeros_like(np_image)
                 for i, channel_color in enumerate(contour_color):
                     poly_fill[:, :, i] = channel_color
                 # Add the polygons to the image by blending the fill and the image using the alpha mask
-                np_image = (np_image.astype(np.float32) * (1 - poly_alpha) + poly_fill * poly_alpha).round().astype(np.uint8)
+                np_image = (np_image.astype(np.float32) * (1 - poly_alpha) + poly_fill * poly_alpha)
+                np_image = np_image.round().astype(np.uint8)
                 # Draw the contours
                 for i, c in enumerate(smpl_contours):
                     cv2.drawContours(np_image, [c], -1, contour_color, linewidth)
@@ -1019,7 +1025,7 @@ def _plot_image(
             # Draw boxes and confidences
             if boxes:
                 for box, conf in zip(bboxes, confs):
-                    box = (box * scale)
+                    box = box * scale
                     box[:2] = box[:2].floor()
                     box[2:] = box[2:].ceil()
                     box = box.long()
@@ -1032,23 +1038,23 @@ def _plot_image(
                             f"{conf * 100:.3g}%",
                             cv2.FONT_HERSHEY_SIMPLEX,
                             fontScale=1 * scale,
-                            thickness=max(1, round(2 * scale))
+                            thickness=max(1, round(2 * scale)),
                         )
                         # Calculate the text position
                         xp, yp = start_point[0], start_point[1] - linewidth * 2
                         if yp < text_height:
                             yp = end_point[1] + text_height + linewidth * 2
                         # Get the average color intensity behind the text
-                        avg_color = np.mean(np_image[yp:yp + text_height, xp:xp + text_width])
+                        avg_color = np.mean(np_image[yp : yp + text_height, xp : xp + text_width])
                         # Draw the text
                         cv2.putText(
                             img=np_image,
                             text=f"{conf * 100:.3g}%",
-                            org=(xp, yp), 
+                            org=(xp, yp),
                             fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                             fontScale=1 * scale,
                             color=(0, 0, 0) if avg_color > 150 else (255, 255, 255),
-                            thickness=max(1, round(2 * scale))
+                            thickness=max(1, round(2 * scale)),
                         )
 
         # Save or show the image
@@ -1056,8 +1062,8 @@ def _plot_image(
             cv2.imwrite(outpath, np_image)
             return None
         else:
-            return cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB) # type: ignore
-    
+            return cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)  # type: ignore
+
     @property
     def crops(self) -> list[torch.Tensor]:
         """Detection crops."""
@@ -1073,30 +1079,29 @@ def crop_masks(self) -> list[torch.Tensor]:
             ]
         else:
             return [
-                resize_masks(mask, self.image.shape[1:])[box[1]:box[3], box[0]:box[2]] # type: ignore - TODO: fixme
-                for mask, box in zip(self.masks, self.boxes.long()) # type: ignore - TODO: fixme
+                resize_masks(mask, self.image.shape[1:])[box[1] : box[3], box[0] : box[2]]  # type: ignore - TODO: fixme
+                for mask, box in zip(self.masks, self.boxes.long())  # type: ignore - TODO: fixme
             ]
 
-    @staticmethod    
+    @staticmethod
     def _save_1_crop(
-            crop : torch.Tensor,
-            mask : torch.Tensor | None,
-            path : str,
-        ) -> str:
+        crop: torch.Tensor,
+        mask: torch.Tensor | None,
+        path: str,
+    ) -> str:
         Image.fromarray(
-            obj=chw2hwc_uint8(crop, mask).detach().cpu().numpy(),
-            mode="RGB" if mask is None else "RGBA"
+            obj=chw2hwc_uint8(crop, mask).detach().cpu().numpy(), mode="RGB" if mask is None else "RGBA"
         ).save(path, compress_level=1)
         return path
 
     def save_crops(
-            self, 
-            outdir : str, 
-            basename : str | None=None, 
-            mask : bool=False, 
-            identifier : str | None=None,
-            wait : bool=False
-        ) -> list[str]:
+        self,
+        outdir: str,
+        basename: str | None = None,
+        mask: bool = False,
+        identifier: str | None = None,
+        wait: bool = False,
+    ) -> list[str]:
         """Save prediction crops."""
         if outdir is None or not os.path.exists(outdir) or not os.path.isdir(outdir):
             raise RuntimeError(f"Invalid outdir {outdir}, does not exist or is not a directory")
@@ -1112,14 +1117,17 @@ def save_crops(
         if identifier is None:
             identifier_field = ""
         else:
-            identifier_field = f'UUID_{identifier}'
-        
+            identifier_field = f"UUID_{identifier}"
+
         crops = self.crops
         if mask:
             crop_masks = self.crop_masks
         else:
             crop_masks = [None] * len(crops)
-        crop_paths = [os.path.join(outdir, f"crop_{basename}_CROPNUMBER_{i}_{identifier_field}{image_ext}") for i in range(len(crops))]
+        crop_paths = [
+            os.path.join(outdir, f"crop_{basename}_CROPNUMBER_{i}_{identifier_field}{image_ext}")
+            for i in range(len(crops))
+        ]
 
         for crop, _mask, path in zip(crops, crop_masks, crop_paths):
             if isinstance(_mask, torch.Tensor):
@@ -1129,7 +1137,7 @@ def save_crops(
             _executor.flush()
 
         return crop_paths
-    
+
     @property
     def json_data(self):
         """JSON-compatible dictionary with instance state data."""
@@ -1152,22 +1160,18 @@ def json_data(self):
             "image_height": self.image.shape[1],
             "mask_width": self.image.shape[2] if self.PREFER_POLYGONS else mdata.shape[2],
             "mask_height": self.image.shape[1] if self.PREFER_POLYGONS else mdata.shape[1],
-            "identifier": None
+            "identifier": None,
         }
 
     def serialize(
-            self, 
-            outpath: str, 
-            save_json: bool = True, 
-            save_pt: bool = False, 
-            identifier: str | None=None
-        ) -> None:
+        self, outpath: str, save_json: bool = True, save_pt: bool = False, identifier: str | None = None
+    ) -> None:
         """Serialize the `TensorPredictions` object to a .pt file and/or a .json file.
 
-        The .pt file contains an exact copy of the `TensorPredictions` object, while the .json file 
-        contains the data in a more human-readable format, which can be 
+        The .pt file contains an exact copy of the `TensorPredictions` object, while the .json file
+        contains the data in a more human-readable format, which can be
         deserialized into a `TensorPredictions` object using the 'load' function.
-        
+
         Args:
             outpath: The path to save the serialized data to. Defaults to None.
             save_json: Whether to save the .json file. Defaults to True. Recommended.
@@ -1176,17 +1180,21 @@ def serialize(
 
         """
         assert len(outpath) > 0, RuntimeError("Cannot serialize with empty outpath")
-        assert os.path.exists(os.path.dirname(outpath)), RuntimeError(f"Invalide outpath {outpath}, directory does not exist")
+        assert os.path.exists(os.path.dirname(outpath)), RuntimeError(
+            f"Invalide outpath {outpath}, directory does not exist"
+        )
 
-        # Check for file-extension on the outpath, it should have none - not really necessary anymore due to the check for directory above
+        # Check that the outpath doesn't have a file-extension
         outpath, ext = os.path.splitext(outpath)
         if ext != "" and len(ext) < 5:
-            logger.warning(f"serializer outpath ({outpath}) should not have a file-extension for 'TensorPredictions.serialize'!")
+            logger.warning(
+                f"serializer outpath ({outpath}) should not have a file-extension for 'TensorPredictions.serialize'!"
+            )
         else:
             outpath = f"{outpath}{ext}"
 
-        pt_path = f'{outpath}.pt'
-        json_path = f'{outpath}.json'
+        pt_path = f"{outpath}.pt"
+        json_path = f"{outpath}.json"
 
         if save_pt:
             if os.path.exists(pt_path):
@@ -1197,18 +1205,15 @@ def serialize(
             if os.path.exists(json_path):
                 logger.warning(f"JSON ({json_path}) already exists, overwriting!")
             json_data = self.json_data
-            json_data["identifier"] = identifier if identifier else self.image_path,
-            with open(json_path, 'w') as f:
+            json_data["identifier"] = (identifier if identifier else self.image_path,)
+            with open(json_path, "w") as f:
                 json.dump(json_data, f)
 
     @classmethod
-    def load(
-            cls, 
-            data: str | dict, 
-            device : DeviceLikeType | None=None, 
-            dtype : torch.types._dtype | None=None
-        ):
-        """Deserializes a TensorPredictions object from a .pt or .json file, or a dictionary. OBS: Mutates and returns the current object.
+    def load(cls, data: str | dict, device: DeviceLikeType | None = None, dtype: torch.types._dtype | None = None):
+        """Deserializes a TensorPredictions object from a .pt or .json file, or a dictionary.
+
+        OBS: Mutates and returns the current object.
 
         Args:
             data: The path to the file to load or a dictionary with the deserialized json data.
@@ -1242,7 +1247,9 @@ def load(
 
         empty_image = torch.zeros((3, data["image_height"], data["image_width"]), device=device, dtype=dtype) + 255  # type: ignore
         inst = cls(image=empty_image, device=device, dtype=dtype)
-        setattr(inst, "PREFER_POLYGONS", True) # Since we only store contours in the .json file, we prefer polygons on loading
+        setattr(
+            inst, "PREFER_POLYGONS", True
+        )  # Since we only store contours in the .json file, we prefer polygons on loading
 
         # Load constants
         for k, v in data.items():
@@ -1257,9 +1264,9 @@ def load(
             # Skip dynamically computed class property attributes
             if k in ["areas"]:
                 continue
-            # Skip the identifier 
+            # Skip the identifier
             if k in ["identifier", "image_height", "image_width"]:
-                continue                
+                continue
             # Catch attributes that don't need special treatment
             elif k in ["scales", "contours"]:
                 pass
@@ -1277,53 +1284,53 @@ def load(
         return inst
 
     def save(
-            self, 
-            output_directory: str, 
-            overview: bool | str=True,
-            crops: bool | str=True,
-            metadata: bool | str=True,
-            fast: bool=False,
-            mask_crops: bool=False,
-            identifier: str | None=None, 
-            basename: str | None=None,
-            wait: bool=False
-        ) -> str | None:
+        self,
+        output_directory: str,
+        overview: bool | str = True,
+        crops: bool | str = True,
+        metadata: bool | str = True,
+        fast: bool = False,
+        mask_crops: bool = False,
+        identifier: str | None = None,
+        basename: str | None = None,
+        wait: bool = False,
+    ) -> str | None:
         """Save the serialized prediction results, crops, and overview to the given output directory.
 
-        TODO: Add the identifier to the names of the files, 
+        TODO: Add the identifier to the names of the files,
         so that we can save multiple predictions for the same image or images with the same name.
 
         Args:
             output_directory: The directory to save the prediction results to.
-            overview: Whether to save the overview image. Defaults to True. 
+            overview: Whether to save the overview image. Defaults to True.
                 If a string is given, it is interpreted as a path to a directory to save the overview image to.
-            crops: Whether to save the crops. Defaults to True. 
+            crops: Whether to save the crops. Defaults to True.
                 If a string is given, it is interpreted as a path to a directory to save the crops to.
-            metadata: Whether to save the metadata. Defaults to True. 
+            metadata: Whether to save the metadata. Defaults to True.
                 If a string is given, it is interpreted as a path to a directory to save the metadata to.
-            fast: Whether to use the fast version of the overview image. Defaults to False. 
+            fast: Whether to use the fast version of the overview image. Defaults to False.
                 Saves the overview image at half the resolution.
             mask_crops: Whether to mask the crops. Defaults to False.
             identifier: An identifier for the serialized data. Defaults to None.
-            basename: The base name of the image. Defaults to None. 
+            basename: The base name of the image. Defaults to None.
                 If None, the base name is extracted from the image path, which must be set in this case.
-            wait: If true `save` blocks execution until results are finished saving, 
+            wait: If true `save` blocks execution until results are finished saving,
                 otherwise results will be saved asynchronously.
-        
+
         Returns:
             The path to the directory containing the serialized data.
             The crops and overview image(s) are also saved here by default.
             If the standard location is not used at all, the directory is not created and None is returned instead.
-        
+
         """
         if basename is None:
             if self.image_path is None:
                 raise ValueError("Unable to save prediction with unknown source file, when `basename` is not supplied.")
             # Get the base name of the image
             basename = os.path.splitext(os.path.basename(self.image_path))[0]
-        
+
         prediction_directory = os.path.join(output_directory, basename)
-        # Create the prediction directory if it does not exist and it is needed 
+        # Create the prediction directory if it does not exist and it is needed
         # (i.e. if we are saving crops, overview, or metadata to a standard location)
         prediction_directory_is_used = (overview is True) or (crops is True) or (metadata is True)
         if prediction_directory_is_used:
@@ -1345,7 +1352,7 @@ def save(
         # Save crops
         if crops:
             # Check if the crops path is overwritten and make sure the directory exists and is a directory
-            crop_directory = crops if isinstance(crops, str) else os.path.join(prediction_directory, "crops") 
+            crop_directory = crops if isinstance(crops, str) else os.path.join(prediction_directory, "crops")
             os.makedirs(crop_directory, exist_ok=True)
             assert os.path.isdir(crop_directory), RuntimeError(f"Invalid path for crops: {crop_directory}")
             self.save_crops(outdir=crop_directory, basename=basename, mask=mask_crops, identifier=identifier)
@@ -1356,7 +1363,7 @@ def save(
             metadata_directory = metadata if isinstance(metadata, str) else prediction_directory
             os.makedirs(metadata_directory, exist_ok=True)
             assert os.path.isdir(metadata_directory), RuntimeError(f"Invalid path for metadata: {metadata_directory}")
-            metadata_path = os.path.join(metadata_directory, f'metadata_{basename}_UUID_{identifier}')
+            metadata_path = os.path.join(metadata_directory, f"metadata_{basename}_UUID_{identifier}")
             self.serialize(outpath=metadata_path, identifier=identifier)
 
         if wait:
@@ -1364,20 +1371,22 @@ def save(
 
         return prediction_directory if prediction_directory_is_used else None
 
+
 def _process_batch(
-            image : torch.Tensor, 
-            offsets : list[tuple[tuple[int, int], tuple[int, int]]], 
-            tile_size : int, 
-            batch_start_idx : int, 
-            batch_size : int,
-            device : DeviceLikeType | None = None,
-            model : torch.nn.Module = None, # type: ignore # TODO: fixthis!
-            time : bool = False,
-            callback : str = "__call__",
-            **kwargs : Any # Swallow any extra arguments
-    ) -> tuple[torch.Tensor, Any, tuple[int, int, int] | None]:
-    start_batch_event = end_fetch_event = end_forward_event = \
-        end_batch_event = start_batch_event = current_device_stream = None
+    image: torch.Tensor,
+    offsets: list[tuple[tuple[int, int], tuple[int, int]]],
+    tile_size: int,
+    batch_start_idx: int,
+    batch_size: int,
+    device: DeviceLikeType | None = None,
+    model: torch.nn.Module = None,  # type: ignore # TODO: fixthis!
+    time: bool = False,
+    callback: str = "__call__",
+    **kwargs: Any,  # Swallow any extra arguments
+) -> tuple[torch.Tensor, Any, tuple[int, int, int] | None]:
+    start_batch_event = end_fetch_event = end_forward_event = end_batch_event = start_batch_event = (
+        current_device_stream
+    ) = None
     if time:
         start_batch_event = torch.cuda.Event(enable_timing=True)
         end_fetch_event = torch.cuda.Event(enable_timing=True)
@@ -1385,12 +1394,13 @@ def _process_batch(
         end_batch_event = torch.cuda.Event(enable_timing=True)
         current_device_stream = torch.cuda.current_stream(device=device)
         start_batch_event.record(current_device_stream)
-    
+
     # Get the offsets for the current batch and extract and stack the corresponding tiles
     batch = torch.stack([
-        image[:, o[0]: (o[0] + tile_size), o[1]: (o[1] + tile_size)] 
-        for (m, n), o in offsets[batch_start_idx:min((batch_start_idx + batch_size), len(offsets))]
+        image[:, o[0] : (o[0] + tile_size), o[1] : (o[1] + tile_size)]
+        for (m, n), o in offsets[batch_start_idx : min((batch_start_idx + batch_size), len(offsets))]
     ], dim=0)
+    
     if time:
         assert current_device_stream is not None and end_fetch_event is not None
         end_fetch_event.record(current_device_stream)
@@ -1398,7 +1408,7 @@ def _process_batch(
     # Forward pass the model on the batch tiles
     with torch.inference_mode():
         batch_outputs = getattr(model, callback)(batch)
-    
+
     if time:
         assert current_device_stream is not None and end_forward_event is not None
         end_forward_event.record(current_device_stream)
@@ -1416,12 +1426,13 @@ def _process_batch(
     else:
         return batch, batch_outputs, None
 
+
 class Predictor:
     """A flatbug predictor.
-    
+
     The flatbug is built to be used primarily via calling the instance itself
     (which is an alias for `Predictor.pyramid_predictions`):
-    
+
     ```
     model = Predictor(...)
     prediction = model(image)
@@ -1433,24 +1444,24 @@ class Predictor:
     RAM to avoid I/O), but also has export and visualization functionality.
     """
 
-    HYPERPARAMETERS : list[str] = CFG_PARAMS
+    HYPERPARAMETERS: list[str] = CFG_PARAMS
     """
     The available hyperparameters for the predictor. \\
     These can be set using the `set_hyperparameters` class method.
     """
 
     # Hyperparameters, set to None so they are visible in the class
-    MIN_MAX_OBJ_SIZE : tuple[int, int] = None  # type: ignore
+    MIN_MAX_OBJ_SIZE: tuple[int, int] = None  # type: ignore
     """
     Defines the minimum and maximum object size as seen in a single tile. \\
     Size is defined as the square root of the pixel area of the bounding box.
     """
-    MAX_MASK_SIZE : int = None  # type: ignore
+    MAX_MASK_SIZE: int = None  # type: ignore
     """
     Defines the maximum size of the segmentation masks. \\
     Only applies if PREFER_POLYGONS is False.
     """
-    SCORE_THRESHOLD : float = None  # type: ignore
+    SCORE_THRESHOLD: float = None  # type: ignore
     """
     The score threshold for the predictions. \\
     TODO: This should be called CONFIDENCE_THRESHOLD.
@@ -1459,46 +1470,46 @@ class Predictor:
     """
     The overlap (e.g. IOU) threshold used to determine if two instances are duplicates. \\
     """
-    MINIMUM_TILE_OVERLAP : int = None  # type: ignore
+    MINIMUM_TILE_OVERLAP: int = None  # type: ignore
     """
     The minimum - but not necessarily the maximum - overlap between tiles \\
     in a single layer of the pyramid. Increasing this value will increase \\
     the computation time, but may improve the detection of large instances. 
     """
-    EDGE_CASE_MARGIN : int = None  # type: ignore
+    EDGE_CASE_MARGIN: int = None  # type: ignore
     """
     The margin to add to the edge of the image to catch instances that are \\
     split between tiles. The margin is added to the edge of the image, such \\
     that instances on the true edge of the images are not removed.
     """
-    PREFER_POLYGONS : bool = None  # type: ignore
+    PREFER_POLYGONS: bool = None  # type: ignore
     """
     Whether to prefer representing the instance segmentation using polygons \\
     instead of masks. This is a much more compact representation, but cannot \\
     represent complex shapes (like holes in the mask), only concave polygons.
     """
-    EXPERIMENTAL_NMS_OPTIMIZATION : bool = None  # type: ignore
+    EXPERIMENTAL_NMS_OPTIMIZATION: bool = None  # type: ignore
     """
     Enables an experimental optimization for the NMS step. \\
     This optimization improves the performance of the NMS step when there are \\
     many instances in a large image and CUDA is available.
     """
-    OVERLAP_METRIC : str = None  # type: ignore
+    OVERLAP_METRIC: str = None  # type: ignore
     """
     Metric to use for NMS. One of "IOU" or "IOS", more might be added in the future.
     """
-    TIME : bool = None  # type: ignore
+    TIME: bool = None  # type: ignore
     """
     Whether to time the different parts of the prediction process. \\
     Enabling this will print a verbose output of the timing of the different \\
     parts of the prediction process.
     """
-    TILE_SIZE : int = None  # type: ignore
+    TILE_SIZE: int = None  # type: ignore
     """
     The size of the tiles to split the image into. \\
     This is defined by the model and should probably not be changed.
     """
-    BATCH_SIZE : int = None  # type: ignore
+    BATCH_SIZE: int = None  # type: ignore
     """
     The batch size to use for the prediction. \\
     This determines how many tiles are processed in parallel. \\
@@ -1509,20 +1520,20 @@ class Predictor:
     DEBUG = False
 
     def __init__(
-            self, 
-            model : str | pathlib.Path="flat_bug_M_v2.pt", 
-            cfg : dict | str | Path | None=None, 
-            device : str | torch.device | int | list[str | torch.device | int]=torch.device("cpu"), 
-            dtype : torch.types._dtype | str=torch.float32
-        ):
+        self,
+        model: str | pathlib.Path = "flat_bug_M_v2.pt",
+        cfg: dict | str | Path | None = None,
+        device: str | torch.device | int | list[str | torch.device | int] = torch.device("cpu"),
+        dtype: torch.types._dtype | str = torch.float32,
+    ):
         """Instantiate a flatbug predictor.
-        
+
         Args:
             model: Path to a local weight file, or the name of a weight file in the flatbug model zoo.
             cfg: A dictionary or a path to a YAML containing the flatbug config for this model instance.
             device: Which device to run the model on.
             dtype: Which dtype to run the model on.
-        
+
         """
         cfg = read_cfg(cfg, strict=True) if isinstance(cfg, (str, Path)) else (cfg or {})
         self.set_hyperparameters(**{**DEFAULT_CFG, **cfg})
@@ -1539,13 +1550,17 @@ def __init__(
             raise ValueError(f"Dtype '{dtype}' is not supported.")
         self._dtype = dtype
 
+        self._model: torch.nn.Module
         if isinstance(model, str):
             if not os.path.exists(model):
-                success = download_from_repository("models/" + "/".join(os.path.normpath(model).split(os.path.sep)), model, False)
+                success = download_from_repository(
+                    "models/" + "/".join(os.path.normpath(model).split(os.path.sep)), model, False
+                )
                 if not success:
                     raise FileNotFoundError(f"No such model or file: '{model}'")
-            
+
             yolo = YOLO(model, "segment", verbose=True)
+            assert isinstance(yolo.model, torch.nn.Module)
             self._model = yolo.model
             self._model.to(self._device, dtype=self._dtype)
             self._model.eval()
@@ -1574,13 +1589,8 @@ def set_hyperparameters(self, **kwargs):
             else:
                 raise ValueError(f"Unknown hyperparameter: {k}")
         return self
-    
-    def _detect_instances(
-            self,
-            image : torch.Tensor,
-            scale : float=1.0,
-            max_scale : bool = False
-        ) -> Prepared_Results:
+
+    def _detect_instances(self, image: torch.Tensor, scale: float = 1.0, max_scale: bool = False) -> Prepared_Results:
         TILE_SIZE = self.TILE_SIZE
         this_MIN_MAX_OBJ_SIZE = list(self.MIN_MAX_OBJ_SIZE)
         this_EDGE_CASE_MARGIN = self.EDGE_CASE_MARGIN
@@ -1595,7 +1605,7 @@ def _detect_instances(
             end_detect = torch.cuda.Event(enable_timing=True)
             main_stream = torch.cuda.current_stream(device=self._device)
             start_detect.record(main_stream)
-        
+
         orig_h, orig_w = image.shape[1:]
         w, h = orig_w, orig_h
         padded = False
@@ -1611,35 +1621,33 @@ def _detect_instances(
         if scale != 1:
             h, w = round(orig_h * scale / 4) * 4, round(orig_w * scale / 4) * 4
             real_scale = w / orig_w, h / orig_h
-            resize = transforms.Resize((h, w), antialias=True) 
+            resize = transforms.Resize((h, w), antialias=True)
             image = resize(image)
             h, w = image.shape[1:]
-        
+
         # If any of the sides are smaller than the TILE_SIZE, pad to TILE_SIZE
         if w < TILE_SIZE or h < TILE_SIZE:
             padded = True
             w_pad = max(0, TILE_SIZE - w) // 2
             h_pad = max(0, TILE_SIZE - h) // 2
             pad_lrtb = w_pad, w_pad + (w % 2 == 1), h_pad, h_pad + (h % 2 == 1)
-            image = torch.nn.functional.pad(image, pad_lrtb, mode="constant", value=0) # Pad with black
+            image = torch.nn.functional.pad(image, pad_lrtb, mode="constant", value=0)  # Pad with black
             h, w = image.shape[1:]
 
         offsets = calculate_tile_offsets(
-            image_size=(w, h),
-            tile_size=TILE_SIZE,
-            minimum_overlap=self.MINIMUM_TILE_OVERLAP
+            image_size=(w, h), tile_size=TILE_SIZE, minimum_overlap=self.MINIMUM_TILE_OVERLAP
         )
 
         hyperparams = {
-            "image" : image,
-            "batch_size" : self.BATCH_SIZE,
-            "tile_size" : TILE_SIZE,
-            "edge_case_margin" : this_EDGE_CASE_MARGIN,
-            "score_threshold" : self.SCORE_THRESHOLD,
-            "overlap_threshold" : self.OVERLAP_THRESHOLD,
-            "overlap_metric" : self.OVERLAP_METRIC,
-            "min_max_object_size" : this_MIN_MAX_OBJ_SIZE,
-            "time" : self.TIME
+            "image": image,
+            "batch_size": self.BATCH_SIZE,
+            "tile_size": TILE_SIZE,
+            "edge_case_margin": this_EDGE_CASE_MARGIN,
+            "score_threshold": self.SCORE_THRESHOLD,
+            "overlap_threshold": self.OVERLAP_THRESHOLD,
+            "overlap_metric": self.OVERLAP_METRIC,
+            "min_max_object_size": this_MIN_MAX_OBJ_SIZE,
+            "time": self.TIME,
         }
 
         if self.TIME:
@@ -1647,8 +1655,8 @@ def _detect_instances(
             start_event, end_event = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
             batch_times, fetch_times, forward_times, postprocess_times = [], [], [], []
             start_event.record(main_stream)
-        
-        postprocessed_results : list[Results] = [None for _ in range(len(offsets))]  # type: ignore
+
+        postprocessed_results: list[Results] = [None for _ in range(len(offsets))]  # type: ignore
         batches = 0
         with torch.no_grad():
             for batch_start_idx in range(0, len(offsets), self.BATCH_SIZE):
@@ -1666,18 +1674,18 @@ def _detect_instances(
                     postprocess_start.record(main_stream)
                 this_postprocessed_results = postprocess(
                     raw_results,
-                    imgs = batch,
-                    max_det = 1000,
-                    min_confidence = self.SCORE_THRESHOLD,
-                    overlap_threshold = self.OVERLAP_THRESHOLD,
-                    overlap_metric = self.OVERLAP_METRIC,
-                    nms = 3,
-                    valid_size_range = self.MIN_MAX_OBJ_SIZE,
-                    edge_margin = self.EDGE_CASE_MARGIN,
+                    imgs=batch,
+                    max_det=1000,
+                    min_confidence=self.SCORE_THRESHOLD,
+                    overlap_threshold=self.OVERLAP_THRESHOLD,
+                    overlap_metric=self.OVERLAP_METRIC,
+                    nms=3,
+                    valid_size_range=self.MIN_MAX_OBJ_SIZE,
+                    edge_margin=self.EDGE_CASE_MARGIN,
                 )
                 for batch_index in range(len(this_postprocessed_results)):
                     tr = Results(**this_postprocessed_results[batch_index])  # type: ignore
-                    tr.orig_img = None # Comment this line if we want debug output.
+                    tr.orig_img = None  # type: ignore # Comment this line if we want debug output.
                     postprocessed_results[batch_start_idx + batch_index] = tr
                 if self.TIME:
                     assert timing is not None
@@ -1685,13 +1693,13 @@ def _detect_instances(
                     fetch_times.append(timing[1])
                     forward_times.append(timing[2])
                     postprocess_end.record(main_stream)
-                    torch.cuda.synchronize(device = self._device)
+                    torch.cuda.synchronize(device=self._device)
                     postprocess_times.append(postprocess_start.elapsed_time(postprocess_end) / 1000)
-            
+
         if self.TIME:
             # Finish timing calculations
             end_event.record(main_stream)
-            torch.cuda.synchronize(device = self._device)
+            torch.cuda.synchronize(device=self._device)
             total_elapsed = start_event.elapsed_time(end_event) / 1000  # Convert to seconds
             fetch_time, forward_time, postprocess_time = sum(fetch_times), sum(forward_times), sum(postprocess_times)
             total_batch_time = sum(batch_times) + postprocess_time
@@ -1699,7 +1707,7 @@ def _detect_instances(
             fetch_prop, forward_prop, postprocess_prop = (
                 fetch_time / total_batch_time,
                 forward_time / total_batch_time,
-                postprocess_time / total_batch_time
+                postprocess_time / total_batch_time,
             )
 
         # ruff: disable[E501]
@@ -1724,28 +1732,34 @@ def _detect_instances(
         MASK_TO_IMG_RATIO = MASK_SIZE / torch.tensor(
             [TILE_SIZE, TILE_SIZE], dtype=torch.float32, device=self._device
         ).unsqueeze(0)
-        
+
         box_offsetters = torch.tensor(
-            [[o[1][0] - pad_lrtb[2], o[1][1] - pad_lrtb[0]] for o in offsets], 
-            dtype=torch.float32, device=self._device
+            [[o[1][0] - pad_lrtb[2], o[1][1] - pad_lrtb[0]] for o in offsets], dtype=torch.float32, device=self._device
         )
         mask_offsetters = torch.round(box_offsetters * MASK_TO_IMG_RATIO).long()
         new_mask_size = (
-            (mask_offsetters.max(dim=0).values + MASK_SIZE) - 
-            torch.tensor(pad_lrtb[1::2][::-1], dtype=torch.long, device=self._device) * MASK_TO_IMG_RATIO[0]
+            (mask_offsetters.max(dim=0).values + MASK_SIZE)
+            - torch.tensor(pad_lrtb[1::2][::-1], dtype=torch.long, device=self._device) * MASK_TO_IMG_RATIO[0]
         ).tolist()
-        orig_img = image[:, pad_lrtb[2]:(-pad_lrtb[3] if pad_lrtb[3] != 0 else None),
-                   pad_lrtb[0]:(-pad_lrtb[1] if pad_lrtb[1] != 0 else None)] if padded else image
+        orig_img = (
+            image[
+                :,
+                pad_lrtb[2] : (-pad_lrtb[3] if pad_lrtb[3] != 0 else None),
+                pad_lrtb[0] : (-pad_lrtb[1] if pad_lrtb[1] != 0 else None),
+            ]
+            if padded
+            else image
+        )
 
         merged_results = merge_tile_results(
-            results = postprocessed_results, 
-            orig_img = orig_img.permute(1, 2, 0), 
-            box_offsetters = box_offsetters.to(self._dtype),
-            mask_offsetters = mask_offsetters, 
-            new_shape = new_mask_size,
-            clamp_boxes = (h - sum(pad_lrtb[2:]), w - sum(pad_lrtb[:2])),
-            max_mask_size = self.MAX_MASK_SIZE, 
-            exclude_masks = self.PREFER_POLYGONS
+            results=postprocessed_results,
+            orig_img=orig_img.permute(1, 2, 0),
+            box_offsetters=box_offsetters.to(self._dtype),
+            mask_offsetters=mask_offsetters,
+            new_shape=new_mask_size,
+            clamp_boxes=(h - sum(pad_lrtb[2:]), w - sum(pad_lrtb[:2])),
+            max_mask_size=self.MAX_MASK_SIZE,
+            exclude_masks=self.PREFER_POLYGONS,
         )
 
         # ruff: disable[E501]
@@ -1769,38 +1783,39 @@ def _detect_instances(
             total_detect_time = start_detect.elapsed_time(end_detect) / 1000  # Convert to seconds
             pred_prop = total_elapsed / total_detect_time
             logger.info(
-                f'Prediction time: {total_elapsed:.3f}s/{pred_prop:>4.1%}'
-                f' (overhead: {overhead_prop:>4.1%}) |'
-                f' Fetch {fetch_prop:>4.1%} |'
-                f' Forward {forward_prop:>4.1%} |'
-                f' Postprocess {postprocess_prop:>4.1%} |'
-                f' Tiles {len(offsets)}'
+                f"Prediction time: {total_elapsed:.3f}s/{pred_prop:>4.1%}"
+                f" (overhead: {overhead_prop:>4.1%}) |"
+                f" Fetch {fetch_prop:>4.1%} |"
+                f" Forward {forward_prop:>4.1%} |"
+                f" Postprocess {postprocess_prop:>4.1%} |"
+                f" Tiles {len(offsets)}"
             )
             if hasattr(self, "total_detection_time"):
                 self.total_detection_time += total_detect_time
             if hasattr(self, "total_forward_time"):
                 self.total_forward_time += forward_time
         return Prepared_Results(
-            predictions = merged_results, 
-            scale = real_scale, 
-            device = self._device, 
-            dtype = self._dtype
+            predictions=merged_results,
+            scale=real_scale,
+            device=self._device,
+            dtype=self._dtype
         )
 
     def pyramid_predictions(
-            self, 
-            image : torch.Tensor | str, 
-            path : str | None=None,
-            scale_increment : float=2/3, 
-            scale_before : float | int=1, 
-            single_scale : bool=False
-        ) -> TensorPredictions:
+        self,
+        image: torch.Tensor | str,
+        path: str | None = None,
+        scale_increment: float = 2 / 3,
+        scale_before: float | int = 1,
+        single_scale: bool = False,
+    ) -> TensorPredictions:
         """Perform inference on an image at multiple scales and return the predictions.
-        
+
         Args:
             image: The image to run inference on. If a string is given, the image is read from the path.
                 If it is a `torch.Tensor`, the path must be provided.
-                We assume that floating point images are in the range [0, 1] and integer images are in the range [0, integer_type_max].
+                We assume that floating point images are in the range [0, 1]
+                and integer images are in the range [0, integer_type_max].
                 *(see https://github.com/pytorch/vision/blob/6d7851bd5e2bedc294e40e90532f0e375fcfee04/torchvision/transforms/_functional_tensor.py#L66)*
             path: The path to the image. Defaults to None. Must be provided if `image` is a `torch.Tensor`.
             scale_increment: The scale increment to use when resizing the image. Defaults to 2/3.
@@ -1818,8 +1833,8 @@ def pyramid_predictions(
         real_path = image if isinstance(image, str) else path
         if isinstance(image, str):
             tensor_image : torch.Tensor = decode_image(
-                input=image, 
-                mode=ImageReadMode.RGB, 
+                input=image,
+                mode=ImageReadMode.RGB,
                 apply_exif_orientation=True
             )
         elif isinstance(image, torch.Tensor):
@@ -1842,16 +1857,17 @@ def pyramid_predictions(
         if tensor_image.dtype in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]:
             transform_list.append(transforms.ConvertImageDtype(self._dtype))
 
-        # A border is always added now, to avoid edge-cases on the actual edge of the image. 
+        # A border is always added now, to avoid edge-cases on the actual edge of the image.
         # I.e. only detections on internal edges of tiles should be removed, not detections on the edge of the image.
         edge_case_margin_padding_multiplier = 2
-        padding_offset = torch.tensor(
-            (self.EDGE_CASE_MARGIN, self.EDGE_CASE_MARGIN), dtype=self._dtype
-        ) * edge_case_margin_padding_multiplier
+        padding_offset = (
+            torch.tensor((self.EDGE_CASE_MARGIN, self.EDGE_CASE_MARGIN), dtype=self._dtype)
+            * edge_case_margin_padding_multiplier
+        )
 
         if padding_offset.sum() > 0:
             padding_for_edge_cases = transforms.Pad(
-                padding=self.EDGE_CASE_MARGIN * edge_case_margin_padding_multiplier, 
+                padding=self.EDGE_CASE_MARGIN * edge_case_margin_padding_multiplier,
                 fill=0,
                 padding_mode='constant'
             )
@@ -1865,7 +1881,7 @@ def pyramid_predictions(
 
         assert len(transformed_image.shape) == 3, RuntimeError(
             f"transformed_image.shape {transformed_image.shape} != 3"
-        ) 
+        )
         assert transformed_image.shape[0] == 3, RuntimeError(
             f"transformed_image.shape[0] {transformed_image.shape[0]} != 3. "
             "The image is probably supplied in WxHxC instead of CxWxH, try image.permute(2, 1, 0) before passing it."
@@ -1894,34 +1910,38 @@ def pyramid_predictions(
 
         if self.TIME:
             self.total_detection_time, self.total_forward_time = 0, 0
-        all_preds = [self._detect_instances(transformed_image, scale=s, max_scale=s == min(scales)) for s in reversed(scales)]
+        all_preds = [
+            self._detect_instances(transformed_image, scale=s, max_scale=s == min(scales))
+            for s in reversed(scales)
+        ]
 
         if self.TIME:
             if self.total_detection_time > 0:
-                perc_forward = f'{self.total_forward_time / self.total_detection_time * 100:.3g}'
+                perc_forward = f"{self.total_forward_time / self.total_detection_time * 100:.3g}"
             else:
                 perc_forward = "N/A"
-            logger.info(
-                f'Total detection time: {self.total_detection_time:.3f}s'
-                f' ({perc_forward}% forward)'
+            logger.info(f"Total detection time: {self.total_detection_time:.3f}s ({perc_forward}% forward)")
+
+        all_preds = (
+            TensorPredictions(
+                predictions=all_preds,
+                image=tensor_image.to(self._device),
+                image_path=real_path,
+                dtype=self._dtype,
+                device=self._device,
+                time=self.TIME,
+                PREFER_POLYGONS=self.PREFER_POLYGONS,
+            )
+            .offset_scale_pad(
+                offset=-padding_offset,
+                scale=1 / scale_before,
+                pad=5,  # pad the boxes a bit to ensure they encapsulate the masks
+            )
+            .non_max_suppression(
+                overlap_threshold=self.OVERLAP_THRESHOLD,
+                metric=self.OVERLAP_METRIC,
+                group_first=self.EXPERIMENTAL_NMS_OPTIMIZATION,
             )
-
-        all_preds = TensorPredictions(
-            predictions     = all_preds,
-            image           = tensor_image.to(self._device),
-            image_path      = real_path,
-            dtype           = self._dtype,
-            device          = self._device,
-            time            = self.TIME,
-            PREFER_POLYGONS = self.PREFER_POLYGONS
-        ).offset_scale_pad(
-            offset = -padding_offset,
-            scale  = 1 / scale_before,
-            pad    = 5 # pad the boxes a bit to ensure they encapsulate the masks
-        ).non_max_suppression(
-            overlap_threshold = self.OVERLAP_THRESHOLD,
-            metric        = self.OVERLAP_METRIC,
-            group_first   = self.EXPERIMENTAL_NMS_OPTIMIZATION
         )
 
         if self.TIME:
@@ -1930,27 +1950,27 @@ def pyramid_predictions(
             torch.cuda.synchronize()
             total_pyramid_time = start_pyramid.elapsed_time(end_pyramid) / 1000
             logger.info(
-                f'Total pyramid time: {total_pyramid_time:.3f}s'
-                f' ({self.total_detection_time / total_pyramid_time * 100:.3g}% detection |'
-                f' {self.total_forward_time / total_pyramid_time * 100:.3g}% forward)'
+                f"Total pyramid time: {total_pyramid_time:.3f}s"
+                f" ({self.total_detection_time / total_pyramid_time * 100:.3g}% detection |"
+                f" {self.total_forward_time / total_pyramid_time * 100:.3g}% forward)"
             )
 
         return all_preds
 
     def __call__(
-            self, 
-            image : torch.Tensor | str, 
-            path : str | None=None, 
-            scale_increment : float=2/3, 
-            scale_before : float | int=1, 
-            single_scale : bool=False
-        ) -> TensorPredictions:
+        self,
+        image: torch.Tensor | str,
+        path: str | None = None,
+        scale_increment: float = 2 / 3,
+        scale_before: float | int = 1,
+        single_scale: bool = False,
+    ) -> TensorPredictions:
         """Perform inference on an image at multiple scales and return the predictions.
-        
+
         Args:
             image: The image to run inference on. If a string is given, the image is read from the path.
-                If it is a `torch.Tensor`, the path must be provided.
-                We assume that floating point images are in the range [0, 1] and integer images are in the range [0, integer_type_max].
+                If it is a `torch.Tensor`, the path must be provided. We assume that floating point images
+                are in the range [0, 1] and integer images are in the range [0, integer_type_max].
                 *(see https://github.com/pytorch/vision/blob/6d7851bd5e2bedc294e40e90532f0e375fcfee04/torchvision/transforms/_functional_tensor.py#L66)*
             path: The path to the image. Defaults to None. Must be provided if `image` is a `torch.Tensor`.
             scale_increment: The scale increment to use when resizing the image. Defaults to 2/3.
@@ -1963,4 +1983,4 @@ def __call__(
         """
         params = locals()
         params.pop("self", None)
-        return self.pyramid_predictions(**params)
\ No newline at end of file
+        return self.pyramid_predictions(**params)
diff --git a/src/flat_bug/trainers.py b/src/flat_bug/trainers.py
index 506e6b9..de09180 100644
--- a/src/flat_bug/trainers.py
+++ b/src/flat_bug/trainers.py
@@ -1,4 +1,5 @@
 """Custom modified YOLO segmentation training class and associated utilities."""
+
 import glob
 import json
 import os
@@ -18,23 +19,30 @@
 
 try:
     from ultralytics.nn.tasks import attempt_load_one_weight as _attempt_load_one_weight
+
     def _load_checkpoint(model):
         weights, ckpt = _attempt_load_one_weight(model)
         return weights, ckpt
 except ImportError:
     from ultralytics.nn.tasks import load_checkpoint as _load_checkpoint_raw  # ultralytics >= 8.4
+
     def _load_checkpoint(model):
         weights, ckpt = _load_checkpoint_raw(model)
         return weights, ckpt
+
+
 from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK, IterableSimpleNamespace
 
 try:
     from ultralytics.utils import yaml_load  # ultralytics < 8.4
 except ImportError:
     from ultralytics.utils import YAML as _YAML  # ultralytics >= 8.4
+
     def yaml_load(file):
         """Load a YAML file."""
         return _YAML.load(file)
+
+
 from ultralytics.utils.files import increment_path
 from ultralytics.utils.torch_utils import smart_inference_mode, torch_distributed_zero_first
 
@@ -42,11 +50,11 @@ def yaml_load(file):
 from flat_bug.datasets import FlatBugYOLODataset, FlatBugYOLOValidationDataset
 
 
-def remove_custom_fb_args(args : dict | IterableSimpleNamespace | Any) -> dict | IterableSimpleNamespace | Any:
+def remove_custom_fb_args(args: dict | IterableSimpleNamespace | Any) -> dict | IterableSimpleNamespace | Any:
     """Remove all custom flatbug key-value pairs from a dict or namespace.
-    
+
     All custom flatbug arguments must start with "_fb".
-    
+
     Returns:
         The dict or namespace without any custom flatbug key-value pairs.
 
@@ -62,11 +70,12 @@ def remove_custom_fb_args(args : dict | IterableSimpleNamespace | Any) -> dict |
 
     return args
 
-def extract_custom_fb_args(args : dict) -> dict:
+
+def extract_custom_fb_args(args: dict) -> dict:
     """Extract all custom flatbug arguments from a dictionary.
-    
+
     All custom flatbug arguments must start with "_fb".
-    
+
     Returns:
         The dictionary all, and only, custom flatbug key-value pairs.
 
@@ -78,11 +87,12 @@ def extract_custom_fb_args(args : dict) -> dict:
 
     return custom_fb_args
 
+
 @overload
-def data2labels(data : str) -> str: ...
+def data2labels(data: str) -> str: ...
 @overload
-def data2labels(data : Sequence[str]) -> list[str]: ...
-def data2labels(data : str | Sequence[str]) -> str | list[str]:
+def data2labels(data: Sequence[str]) -> list[str]: ...
+def data2labels(data: str | Sequence[str]) -> str | list[str]:
     """Infer label file(s) from image director[y/ies]."""
     if not isinstance(data, str):
         return [data2labels(d) for d in data]
@@ -90,9 +100,10 @@ def data2labels(data : str | Sequence[str]) -> str | list[str]:
     # Remove possible trailing directory separator
     if data[-1] == os.sep:
         data = data[:-1]
-    return data + f'{os.sep}instances_default.json'
+    return data + f"{os.sep}instances_default.json"
 
-def get_latest_weight(weight_dir : str | Path) -> str | None:
+
+def get_latest_weight(weight_dir: str | Path) -> str | None:
     """Get the most recently updated weights (files ending in ".pt") in a directory."""
     weights = glob.glob(f"{weight_dir}{os.sep}*.pt")
     if not weights:
@@ -100,42 +111,46 @@ def get_latest_weight(weight_dir : str | Path) -> str | None:
         return None
     return max(weights, key=os.path.getctime)
 
-def _custom_end_to_end_validation(self : "FlatBugSegmentationTrainer"):
-        if not self._do_custom_eval:
-            return
-        self._do_custom_eval = False
-        # Get image and label paths
-        train_data, val_data = self.data["train"], self.data["val"]
-        train_labels, val_labels = data2labels(train_data), data2labels(val_data)  # noqa: F841
-        train_paths, val_paths = self.training_image_paths, self.val_image_paths  # noqa: F841
-        if self._custom_num_images > -1 and self._custom_num_images < len(val_paths):
-            # Sample n images
-            # train_paths = random.sample(train_paths, self._custom_num_images)
-            val_paths = random.sample(val_paths, min(len(val_paths), self._custom_num_images))
-        val_pattern = '({})'.format(
-            "|".join([os.path.basename(f).replace(".", r"\.") for f in val_paths])
-        )
-        # Get latest model path
-        weight_dir = self.wdir
-        latest_weights = get_latest_weight(weight_dir)
-        # Construct end-to-end evaluation command
-        custom_eval_path = os.path.join(os.path.dirname(__file__), "..", "..", "scripts", "eval", "end_to_end_eval.sh")
-        # command = (
-        #     f'bash "{custom_eval_path}" -w "{latest_weights}" -d "{val_data}" -l "{val_labels}" '
-        #     f'-o "{self.save_dir}{os.sep}e2e_val{os.sep}{self.epoch}" -g "{self.args.device}" -p "{val_pattern}"'
-        # )
-        command = 'bash "{}" -w "{}" -d "{}" -l "{}" -o "{}" -g "{}" -p "{}"'.format(
-            custom_eval_path, latest_weights, val_data, val_labels, 
-            f"{self.save_dir}{os.sep}e2e_val{os.sep}{self.epochs}", 
-            self.args.device, val_pattern
-        )
-        logger.debug(f"Running custom end-to-end validation command: `{command}`")
-        # Run command
-        os.system(command)
 
-def findattr(o, name : str, filters : list=[lambda _ : True], exclude_prefix="_", label : str="object"):
+def _custom_end_to_end_validation(self: "FlatBugSegmentationTrainer"):
+    if not self._do_custom_eval:
+        return
+    self._do_custom_eval = False
+    # Get image and label paths
+    train_data, val_data = self.data["train"], self.data["val"]
+    train_labels, val_labels = data2labels(train_data), data2labels(val_data)  # noqa: F841
+    train_paths, val_paths = self.training_image_paths, self.val_image_paths  # noqa: F841
+    if self._custom_num_images > -1 and self._custom_num_images < len(val_paths):
+        # Sample n images
+        # train_paths = random.sample(train_paths, self._custom_num_images)
+        val_paths = random.sample(val_paths, min(len(val_paths), self._custom_num_images))
+    val_pattern = "({})".format("|".join([os.path.basename(f).replace(".", r"\.") for f in val_paths]))
+    # Get latest model path
+    weight_dir = self.wdir
+    latest_weights = get_latest_weight(weight_dir)
+    # Construct end-to-end evaluation command
+    custom_eval_path = os.path.join(os.path.dirname(__file__), "..", "..", "scripts", "eval", "end_to_end_eval.sh")
+    # command = (
+    #     f'bash "{custom_eval_path}" -w "{latest_weights}" -d "{val_data}" -l "{val_labels}" '
+    #     f'-o "{self.save_dir}{os.sep}e2e_val{os.sep}{self.epoch}" -g "{self.args.device}" -p "{val_pattern}"'
+    # )
+    command = 'bash "{}" -w "{}" -d "{}" -l "{}" -o "{}" -g "{}" -p "{}"'.format(
+        custom_eval_path,
+        latest_weights,
+        val_data,
+        val_labels,
+        f"{self.save_dir}{os.sep}e2e_val{os.sep}{self.epochs}",
+        self.args.device,
+        val_pattern,
+    )
+    logger.debug(f"Running custom end-to-end validation command: `{command}`")
+    # Run command
+    os.system(command)
+
+
+def findattr(o, name: str, filters: list = [lambda _: True], exclude_prefix="_", label: str = "object"):
     """Recursively extract certain attributes of an object or object's nested within the object's state.
-    
+
     TODO: This is very brittle and could easily result in a recursion loop and other errors.
 
     Returns:
@@ -152,19 +167,22 @@ def findattr(o, name : str, filters : list=[lambda _ : True], exclude_prefix="_"
             return {}
     values = {}
     for attr in attrs:
-        if (isinstance(attr, str) and attr.startswith(exclude_prefix)):
+        if isinstance(attr, str) and attr.startswith(exclude_prefix):
             continue
         new_label = f'{label}["{attr}"]' if isd else f"{label}.{attr}"
         val = o.get(attr) if isd else getattr(o, attr)
-        if name == attr and all(map(lambda f : f(val), filters)):
+        if name == attr and all(map(lambda f: f(val), filters)):
             values[new_label] = val
         else:
             values.update(findattr(val, name, filters=filters, exclude_prefix=exclude_prefix, label=new_label))
     return values
 
-def replaceattr(o : object, name : str, value, filters : list=[lambda _ : True], exclude_prefix="_", label : str="object"):
+
+def replaceattr(
+    o: object, name: str, value, filters: list = [lambda _: True], exclude_prefix="_", label: str = "object"
+):
     """Recursively replace certain attributes of an object or object's nested within the object's state.
-    
+
     TODO: This is very brittle and could easily result in a recursion loop and other errors.
 
     Returns:
@@ -181,23 +199,27 @@ def replaceattr(o : object, name : str, value, filters : list=[lambda _ : True],
         except Exception:
             return False
     for attr in attrs:
-        if (isinstance(attr, str) and attr.startswith(exclude_prefix)):
+        if isinstance(attr, str) and attr.startswith(exclude_prefix):
             continue
         new_label = f'{label}["{attr}"]' if isd else f"{label}.{attr}"
         val = o.get(attr) if isd else getattr(o, attr)
-        if name == attr and all(map(lambda f : f(val), filters)):
-            logger.debug(f'{new_label} : {val} ==> {value}')
+        if name == attr and all(map(lambda f: f(val), filters)):
+            logger.debug(f"{new_label} : {val} ==> {value}")
             if isd:
-                o.update({attr : value})
+                o.update({attr: value})
             else:
                 setattr(o, attr, value)
         else:
             replaceattr(
-                o.get(attr) if isd else getattr(o, attr), 
-                name, value, 
-                filters=filters, exclude_prefix=exclude_prefix, label=new_label
+                o.get(attr) if isd else getattr(o, attr),
+                name,
+                value,
+                filters=filters,
+                exclude_prefix=exclude_prefix,
+                label=new_label,
             )
 
+
 def apply_overrides_to_checkpoint(overrides):  # noqa: D103
     if not overrides.get("resume", False):
         return
@@ -211,20 +233,32 @@ def apply_overrides_to_checkpoint(overrides):  # noqa: D103
     if not os.path.exists(resume_model):
         raise FileNotFoundError(f"Resume checkpoint {resume_model} not found.")
     # Load original checkpoint
-    logger.debug(f'Loading checkpoint for resuming {resume_model} to `resume_ckpt`')
+    logger.debug(f"Loading checkpoint for resuming {resume_model} to `resume_ckpt`")
     resume_ckpt = torch.load(resume_model)
     logger.debug("Replacing values in `resume_ckpt`...")
     # Enforce overrides
     for k, v in overrides.items():
         if not k.startswith("fb_") and v is not None:
-            replaceattr(resume_ckpt, k, v, [lambda x : isinstance(x, (str, int, float)) or x is None], label="resume_ckpt")
+            replaceattr(
+                resume_ckpt, k, v, [lambda x: isinstance(x, (str, int, float)) or x is None], label="resume_ckpt"
+            )
     # Change save dir
     if "name" not in overrides:
-        overrides["name"] = (list(findattr(resume_ckpt, "name", [lambda x : isinstance(x, str)]).values()) or ["train"])[0]
+        overrides["name"] = (
+            list(findattr(resume_ckpt, "name", [lambda x : isinstance(x, str)]).values()) or ["train"]
+        )[0]
     if "project" not in overrides:
-        overrides["project"] = (list(findattr(resume_ckpt, "project", [lambda x : isinstance(x, str)]).values()) or ["runs/segment"])[0]
+        overrides["project"] = (
+            list(findattr(resume_ckpt, "project", [lambda x : isinstance(x, str)]).values()) or ["runs/segment"]
+        )[0]
     new_save_dir = increment_path(os.path.join(overrides["project"], overrides["name"]), False)
-    replaceattr(resume_ckpt, "save_dir", new_save_dir, [lambda x : isinstance(x, (str, int, float)) or x is None], label="resume_ckpt")
+    replaceattr(
+        resume_ckpt,
+        "save_dir",
+        new_save_dir,
+        [lambda x: isinstance(x, (str, int, float)) or x is None],
+        label="resume_ckpt",
+    )
     # Set epoch appropriately
     prior_epochs = resume_ckpt["train_results"]["epoch"]
     if len(prior_epochs) == 0 or max(prior_epochs) < 1:
@@ -234,10 +268,10 @@ def apply_overrides_to_checkpoint(overrides):  # noqa: D103
     tmp_resume_weight_dir = os.path.join(overrides["project"], "resume_weights")
     os.makedirs(tmp_resume_weight_dir, exist_ok=True)
     with NamedTemporaryFile(
-        delete=False, 
-        suffix=ckpt_ext, 
-        prefix="resume--" + "__".join(os.path.splitext(resume_model)[0].split(os.sep)) + "--", 
-        dir=tmp_resume_weight_dir
+        delete=False,
+        suffix=ckpt_ext,
+        prefix="resume--" + "__".join(os.path.splitext(resume_model)[0].split(os.sep)) + "--",
+        dir=tmp_resume_weight_dir,
     ) as tmp_model:
         torch.save(resume_ckpt, tmp_model)
     logger.debug(f"Saved altered checkpoint for resuming `resume_ckpt` as {tmp_model.name}")
@@ -250,33 +284,37 @@ def apply_overrides_to_checkpoint(overrides):  # noqa: D103
     # Return overrides for convenience, in fact this function mutates the original overrides object
     return overrides
 
+
 class FlatBugSegmentationTrainer(SegmentationTrainer):
     """Modified YOLO Segmentation trainer used for training flatbug."""
 
     def __init__(
-            self, 
-            cfg : IterableSimpleNamespace=DEFAULT_CFG, 
-            overrides : dict | None=None, 
-            _callbacks : Any=None, 
-            *args, 
-            **kwargs
-        ):
+        self,
+        cfg: IterableSimpleNamespace = DEFAULT_CFG,
+        overrides: dict | None = None,
+        _callbacks: Any = None,
+        *args,
+        **kwargs,
+    ):
         """Initialize a SegmentationTrainer object with given arguments."""
-        cfg = DEFAULT_CFG # In DDP mode, a CFG is created for each rank, but we always want the default one
+        cfg = DEFAULT_CFG  # In DDP mode, a CFG is created for each rank, but we always want the default one
         custom_fb_args = extract_custom_fb_args(overrides or {})
         self._max_instances = custom_fb_args["fb_max_instances"]
         self._max_images = custom_fb_args["fb_max_images"]
         self._exclude_datasets = custom_fb_args["fb_exclude_datasets"]
         self.custom_eval = custom_fb_args["fb_custom_eval"]
-        self._do_custom_eval = False # This is a dynamic signalling flag, not a hyperparameter
+        self._do_custom_eval = False  # This is a dynamic signalling flag, not a hyperparameter
         self._custom_num_images = custom_fb_args["fb_custom_eval_num_images"]
         assert self._custom_num_images != 0, (
-            'fb_custom_eval_num_images/custom_eval_num_images cannot be 0. '
-            'If you mean to disable custom eval set fb_custom_eval/custom_eval=False.'
+            "fb_custom_eval_num_images/custom_eval_num_images cannot be 0. "
+            "If you mean to disable custom eval set fb_custom_eval/custom_eval=False."
         )
-        assert self._max_instances != 0, 'fb_max_instances/max_instances cannot be 0.'
+        assert self._max_instances != 0, "fb_max_instances/max_instances cannot be 0."
         assert self._max_images != 0, "fb_max_images/max_images cannot be 0."
-        updated_overrides = remove_custom_fb_args(overrides or {}) # The custom arguments must be removed before calling super.__init___
+
+        # The custom arguments must be removed before calling super.__init___
+        updated_overrides = remove_custom_fb_args(overrides or {})
+
         # To use overrides we must apply these to the checkpoint file itself (only applies if we resume a training run)
         # otherwise the overrides are overwritten by the old training arguments stored within the checkpoint file
         apply_overrides_to_checkpoint(updated_overrides)
@@ -284,7 +322,8 @@ def __init__(
         if updated_overrides.get("resume", False):
             self.args.__dict__.update(updated_overrides)
         
-        self.args.__dict__.update(custom_fb_args) # But we need to add them back, otherwise they will be missing in DDP mode
+        # But we need to add them back, otherwise they will be missing in DDP mode
+        self.args.__dict__.update(custom_fb_args)
         if updated_overrides.get("resume", False):
             self.args.resume = True
         self.add_callback("on_train_epoch_start", FlatBugSegmentationTrainer.log_lr)
@@ -310,14 +349,15 @@ def setup_model(self) -> dict | None:  # noqa: D102
 
         model, weights = self.model, None
         ckpt = None
-        if str(model).endswith('.pt'):
+        if str(model).endswith(".pt"):
             if not os.path.exists(model):
                 # ultralytics < 8.4 doesn't auto-download in torch_safe_load, so do it explicitly
                 from ultralytics.utils.downloads import attempt_download_asset
+
                 model = attempt_download_asset(model)
             weights, ckpt = _load_checkpoint(model)
-            if ckpt is not None and hasattr(ckpt.get('model', None), 'yaml'):
-                cfg = ckpt['model'].yaml
+            if ckpt is not None and hasattr(ckpt.get("model", None), "yaml"):
+                cfg = ckpt["model"].yaml
             else:
                 cfg = weights.yaml
         else:
@@ -326,17 +366,14 @@ def setup_model(self) -> dict | None:  # noqa: D102
         if not self.args.resume:
             return None
         return ckpt
-    
+
     @property
     def exclude_pattern(self) -> str:  # noqa: D102
-        return f'^(?!({"|".join(self._exclude_datasets)}))' if self._exclude_datasets else ""
+        return f"^(?!({'|'.join(self._exclude_datasets)}))" if self._exclude_datasets else ""
 
     def build_dataset(  # noqa: D102
-            self, 
-            img_path : str, 
-            mode : str='train', 
-            batch : int | None=None
-        ) -> FlatBugYOLODataset | FlatBugYOLOValidationDataset:
+        self, img_path: str, mode: str = "train", batch: int | None = None
+    ) -> FlatBugYOLODataset | FlatBugYOLOValidationDataset:
         LOGGER.info(
             f"Building dataset with max instances ({self._max_instances}), "
             f"max images ({self._max_images}) and exclude pattern ({self.exclude_pattern})."
@@ -356,7 +393,7 @@ def build_dataset(  # noqa: D102
                 single_cls=self.args.single_cls or False,
                 max_instances=self._max_instances,
                 task="segment",
-                subset_args={"n" : self._max_images, "pattern" : self.exclude_pattern}
+                subset_args={"n": self._max_images, "pattern": self.exclude_pattern},
             )
         else:
             dataset = FlatBugYOLOValidationDataset(
@@ -373,18 +410,14 @@ def build_dataset(  # noqa: D102
                 single_cls=self.args.single_cls or False,
                 max_instances=np.inf,
                 task="segment",
-                subset_args={"n" : self._max_images, "pattern" : self.exclude_pattern}
+                subset_args={"n": self._max_images, "pattern": self.exclude_pattern},
             )
 
         return dataset
-    
+
     def get_dataloader(
-            self, 
-            dataset_path : str, 
-            batch_size : int | None=16, 
-            rank : int=0, 
-            mode : str="train"
-        ) -> InfiniteDataLoader:
+        self, dataset_path: str, batch_size: int | None = 16, rank: int = 0, mode: str = "train"
+    ) -> InfiniteDataLoader:
         """Construct and return dataloader."""
         assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}."
         if mode == "val":
@@ -462,7 +495,7 @@ def training_image_paths(self) -> list[str]:  # noqa: D102
                 "Accessing the training image paths is not possible, while training has not started."
             )
             raise e
-        
+
     @property
     def val_image_paths(self) -> list[str]:  # noqa: D102
         try:
@@ -478,20 +511,25 @@ def _reproducibility_setup(self):
         if RANK not in {-1, 0}:
             logger.warning("Reproducibility setup skipped for non-master rank.")
             return
+
         def log_data(self):
             with open(self.save_dir / "data_log.json", "w") as f:
                 json.dump(
                     obj={
-                        **{k : str(v) for k, v in self.data.items()}, 
-                        **{"train_images" : self.training_image_paths, "val_images" : self.val_image_paths}
-                    }, 
-                    fp=f
+                        **{k: str(v) for k, v in self.data.items()},
+                        **{"train_images": self.training_image_paths, "val_images": self.val_image_paths},
+                    },
+                    fp=f,
                 )
+
         self.add_callback("on_train_start", log_data)
 
     def get_validator(self) -> yolo.segment.SegmentationValidator:
         """Return an instance of SegmentationValidator for validation of YOLO model."""
         self.loss_names = "box_loss", "seg_loss", "cls_loss", "dfl_loss"
         return yolo.segment.SegmentationValidator(
-            self.test_loader, save_dir=self.save_dir, args=remove_custom_fb_args(copy(self.args)), _callbacks=self.callbacks
-        )
\ No newline at end of file
+            self.test_loader,
+            save_dir=self.save_dir,
+            args=remove_custom_fb_args(copy(self.args)),
+            _callbacks=self.callbacks,
+        )
diff --git a/src/flat_bug/yolo_helpers.py b/src/flat_bug/yolo_helpers.py
index 185ea18..110e760 100644
--- a/src/flat_bug/yolo_helpers.py
+++ b/src/flat_bug/yolo_helpers.py
@@ -1,4 +1,5 @@
 """Utilities for working with raw and processed YOLO outputs."""
+
 from argparse import Namespace
 from collections.abc import Callable, Sequence
 from functools import partial
@@ -16,20 +17,21 @@
 class Polygons(list[torch.Tensor]):
     """Convenience wrapper for a heterogenous list of torch.Tensor's used for storing polygon contours."""
 
-    def __getattr__(self, attr : str):
+    def __getattr__(self, attr: str):
         inner = getattr(torch.Tensor, attr)
         return partial(self._apply, func=inner)
 
-    def _apply(self, func : Callable[Concatenate[torch.Tensor, ...], torch.Tensor], *args, **kwargs):
+    def _apply(self, func: Callable[Concatenate[torch.Tensor, ...], torch.Tensor], *args, **kwargs):
         return Polygons(func(x, *args, **kwargs) for x in self)
-    
+
     def to_list(self):  # noqa: D102
         return [e for e in self]
 
+
 class ResultsWithTiles(Results):
     """Container for YOLO results with segmentation polygons and corresponding tile indices."""
 
-    def __init__(self, tiles : list[int] | torch.Tensor, polygons : list[torch.Tensor], *args, **kwargs):  # noqa: D107
+    def __init__(self, tiles: list[int] | torch.Tensor, polygons: list[torch.Tensor], *args, **kwargs):  # noqa: D107
         super().__init__(*args, **kwargs)
         self.tiles = torch.as_tensor(tiles)
         self.polygons = Polygons(polygons)
@@ -40,11 +42,8 @@ def __len__(self):
 
 
 def offset_box(
-        boxes : torch.Tensor, 
-        offset : torch.Tensor, 
-        max_x : int | float | None = None, 
-        max_y : int | float | None = None
-    ):
+    boxes: torch.Tensor, offset: torch.Tensor, max_x: int | float | None = None, max_y: int | float | None = None
+):
     """Move bounding boxes."""
     m = 4 / offset.shape[0]
     assert m // 1 == m, f"4 must be divisible by the number of offsets ({offset.shape[0]})"
@@ -55,23 +54,22 @@ def offset_box(
         boxes[:, [1, 3]] = boxes[:, [1, 3]].clamp(0, max_x - 1)
     return boxes
 
+
 def offset_mask(
-        mask : torch.Tensor, 
-        offset : torch.Tensor, 
-        new_shape : tuple[int, int] | list[int] | None=None, 
-        max_size=700
-    ):
+    mask: torch.Tensor, offset: torch.Tensor, new_shape: tuple[int, int] | list[int] | None = None, max_size=700
+):
     """Offset (move) binary masks in a new coordinate system."""
-    # Due to memory use, it is beneficial to restrict the maximum size of the masks. A 700x700 boolean tensor uses ~0.5 MB of memory
+    # Due to memory use, it is beneficial to restrict the maximum size of the masks.
+    # A 700x700 boolean tensor uses ~0.5 MB of memory
     n, h, w = mask.shape
     if new_shape is None:
         shape = n, h, w
     else:
         assert len(new_shape) == 2, f"new_shape must be a tuple or list of length 2, not {len(new_shape)}"
         shape = int(n), int(new_shape[0]), int(new_shape[1])
-    
+
     new_mask = torch.zeros(shape, dtype=torch.bool, device=mask.device)
-    
+
     # Calculate the possible clamped size of the mask (if it needs to be clamped)
     clamp_factor = (max(shape[1:]) / max_size) if max_size is not None else 1
     clamp_shape = [int(n), shape[1] / clamp_factor, shape[2] / clamp_factor]
@@ -86,7 +84,7 @@ def offset_mask(
             return torch.zeros(clamp_shape, dtype=torch.bool, device=mask.device)
 
     # Calculate the overlap of the mask with the new mask (in the new mask's coordinate system)
-    mask_overlap = [torch.empty((0, )), torch.empty((0, ))]
+    mask_overlap = [torch.empty((0,)), torch.empty((0,))]
     for i, (mask_d, new_mask_d, offset_d) in enumerate(zip([h, w], shape[1:], offset)):
         mask_overlap[i] = torch.arange(mask_d, device=mask.device) + offset_d
         mask_overlap[i] = mask_overlap[i][(mask_overlap[i] < new_mask_d) & (mask_overlap[i] >= 0)]
@@ -103,26 +101,28 @@ def offset_mask(
         (mask_overlap[1][0] - offset[1]):(mask_overlap[1][1] - offset[1])
     ]
 
-    # If the mask is larger than the maximum size, clamp it by downscaling it such that the largest dimension is max_size
+    # If the mask is larger than the maximum size,
+    # clamp it by downscaling it such that the largest dimension is max_size
     if clamp_factor > 1:
         # F.interpolate(
-        #   new_mask.float().unsqueeze(0), clamp_shape[1:], 
+        #   new_mask.float().unsqueeze(0), clamp_shape[1:],
         #   mode='bilinear', align_corners=False, antialias=True
         #   ).squeeze(0) > 0.25
         new_mask = resize_masks(new_mask, clamp_shape[1:])
-    
+
     return new_mask
 
+
 def merge_tile_results(
-        results : list[Results], 
-        orig_img : torch.Tensor | None=None,
-        box_offsetters : torch.Tensor | None=None,
-        mask_offsetters : torch.Tensor | None=None,
-        new_shape : tuple[int, int] | list[int] | None=None,
-        clamp_boxes : tuple[int | None, int | None] | list[int | None]=(None, None),
-        max_mask_size : int =700,
-        exclude_masks : bool=False
-    ) :
+    results: list[Results],
+    orig_img: torch.Tensor | None = None,
+    box_offsetters: torch.Tensor | None = None,
+    mask_offsetters: torch.Tensor | None = None,
+    new_shape: tuple[int, int] | list[int] | None = None,
+    clamp_boxes: tuple[int | None, int | None] | list[int | None] = (None, None),
+    max_mask_size: int = 700,
+    exclude_masks: bool = False,
+):
     """Merge results from multiple images into a single Results object, possibly with a new image."""
     assert results[0].boxes is not None and isinstance(results[0].boxes.data, torch.Tensor)
     _device = results[0].boxes.data.device
@@ -138,26 +138,34 @@ def merge_tile_results(
     mx, my = clamp_boxes
     path = results[0].path
     names = results[0].names
-    tile_indices = torch.concatenate([
-        torch.tensor([i] * (r.__len__() or 0), dtype=torch.long, device=_device)
-        for i, r in enumerate(results)
-    ])
-    boxes = torch.cat([
-        offset_box(torch.as_tensor(r.boxes.data), o.flip(0), mx, my)
-        for r, o in zip(results, box_offsetters) if r.boxes is not None
-    ])
+    tile_indices = torch.concatenate(
+        [
+            torch.tensor([i] * (r.__len__() or 0), dtype=torch.long, device=_device)
+            for i, r in enumerate(results)
+        ]
+    )
+    boxes = torch.cat(
+        [
+            offset_box(torch.as_tensor(r.boxes.data), o.flip(0), mx, my)
+            for r, o in zip(results, box_offsetters)
+            if r.boxes is not None
+        ]
+    )
     polygons = [
-        find_contours(resize_masks(mask, [256 * 3, 256 * 3]), True) * (1024 / 256) / 3 + o.flip(0).unsqueeze(0) 
+        find_contours(resize_masks(mask, [256 * 3, 256 * 3]), True) * (1024 / 256) / 3 + o.flip(0).unsqueeze(0)
         for r, o in zip(results, box_offsetters) if r.masks
         for mask in torch.as_tensor(r.masks.data)
     ]
     if exclude_masks:
         masks = torch.cat([torch.as_tensor(r.masks.data) for r in results if r.masks is not None])
     else:
-        masks = torch.cat([
-            offset_mask(torch.as_tensor(r.masks.data), o, new_shape, max_mask_size)
-            for r, o in zip(results, mask_offsetters) if r.masks is not None
-        ])
+        masks = torch.cat(
+            [
+                offset_mask(torch.as_tensor(r.masks.data), o, new_shape, max_mask_size)
+                for r, o in zip(results, mask_offsetters)
+                if r.masks is not None
+            ]
+        )
     if len(masks.shape) == 2:
         masks = masks.unsqueeze(0)
     if not all([r.probs is None for r in results]):
@@ -165,30 +173,28 @@ def merge_tile_results(
     if not all([r.keypoints is None for r in results]):
         raise NotImplementedError("'Keypoints' not implemented yet")
     return ResultsWithTiles(
-        tiles=tile_indices, 
-        orig_img=Namespace(shape=orig_img.shape), 
-        path=path, 
-        names=names, 
-        boxes=boxes, 
-        masks=masks, 
-        polygons=polygons, 
-        probs=None, 
-        keypoints=None
+        tiles=tile_indices,
+        orig_img=Namespace(shape=orig_img.shape),
+        path=path,
+        names=names,
+        boxes=boxes,
+        masks=masks,
+        polygons=polygons,
+        probs=None,
+        keypoints=None,
     )
 
-def stack_masks(
-        masks : list[Masks | torch.Tensor], 
-        orig_shape : tuple[int, int] | list[int] | None=None
-    ):
-    """Stacks a list of ultralytics.engine.results.Masks objects (or torch.Tensor) into a single ultralytics.engine.results.Masks object.
+
+def stack_masks(masks: list[Masks | torch.Tensor], orig_shape: tuple[int, int] | list[int] | None = None):
+    """Stacks a list of ultralytics `Masks` objects (or `torch.Tensor`) into a single `Masks` object.
 
     If the masks are not all the same size, they are resized to the largest size in the list.
 
     Args:
         masks: A list of ultralytics.engine.results.Masks objects (or torch.Tensor).
-        orig_shape: The original shape of the image. Defaults to None. 
+        orig_shape: The original shape of the image. Defaults to None.
             If None, the original shape is inferred from the first Masks object in the list if there is one.
-        antialias: A flag to indicate whether to use antialiasing when resizing the masks. Defaults to False.
+        antialias: Flag indicating whether to use antialiasing when resizing the masks. Defaults to False.
 
     Returns:
         A Masks object containing the stacked masks.
@@ -200,7 +206,9 @@ def stack_masks(
             orig_shape = m.orig_shape
             break
     tensor_masks = [torch.as_tensor(m.data) if isinstance(m, Masks) else m for m in masks]
-    assert all([isinstance(m, torch.Tensor) for m in tensor_masks]), f"'masks' must be a list of torch.Tensor, not {type(tensor_masks[0])}"
+    assert all([isinstance(m, torch.Tensor) for m in tensor_masks]), (
+        f"'masks' must be a list of torch.Tensor, not {type(tensor_masks[0])}"
+    )
     assert len(tensor_masks) != 0, f"'masks' ({tensor_masks}) must not be empty"
     _device = tensor_masks[0].device
 
@@ -223,10 +231,8 @@ def stack_masks(
 
     return Masks(new_masks, orig_shape=orig_shape)
 
-def crop_masks(
-        masks : torch.Tensor, 
-        boxes : torch.Tensor
-    ):
+
+def crop_masks(masks: torch.Tensor, boxes: torch.Tensor):
     """Crops masks to bounding boxes.
 
     Args:
@@ -244,13 +250,14 @@ def crop_masks(
 
     return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
 
+
 def process_mask(
-        protos : torch.Tensor, 
-        masks_in : torch.Tensor, 
-        bboxes : torch.Tensor, 
-        shape : tuple[int, int] | tuple[int, ...] | list[int], 
-        upsample : bool=False
-    ):
+    protos: torch.Tensor,
+    masks_in: torch.Tensor,
+    bboxes: torch.Tensor,
+    shape: tuple[int, int] | tuple[int, ...] | list[int],
+    upsample: bool = False,
+):
     """Apply masks to bounding boxes using the output of the mask head.
 
     Args:
@@ -258,7 +265,7 @@ def process_mask(
         masks_in: A tensor of shape [n, mask_dim], where n is the number of masks after NMS.
         bboxes: A tensor of shape [n, 4], where n is the number of masks after NMS.
         shape: A tuple of integers representing the size of the input image in the format (h, w).
-        upsample: A flag to indicate whether to upsample the mask to the original image size. Default is False.
+        upsample: Flag indicating whether to upsample the mask to the original image size. Default is False.
 
     Returns:
         A binary mask tensor of shape [n, h, w], where n is the number of masks after NMS, and h and w
@@ -268,8 +275,8 @@ def process_mask(
     c, mh, mw = protos.shape  # CHW
     ih, iw = shape
 
-    # CHW <- This line has been changed from the original implementation, 
-    # which had a superfluous type conversion which caused YOLOv8 to cast the masks to float32, 
+    # CHW <- This line has been changed from the original implementation,
+    # which had a superfluous type conversion which caused YOLOv8 to cast the masks to float32,
     # this change simply removes the type conversion enabling support for other data types
     masks = (masks_in.to(protos.dtype) @ protos.view(c, -1)).sigmoid().view(-1, mh, mw)
 
@@ -284,27 +291,27 @@ def process_mask(
     # masks = expand_bottom_right(masks)  # HW
 
     if upsample:
-        masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]  # CHW
-    
+        masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
+
     masks = masks.gt_(0.5).bool()
 
     return masks
 
-def expand_bottom_right(mask : torch.Tensor):
+
+def expand_bottom_right(mask: torch.Tensor):
     """Add an extra pixel above next to bottom/right edges of the region of 1s.
 
     Args:
         mask: A binary mask tensor of shape [h, w].
 
     Returns:
-        A binary mask tensor of shape [h, w], where an extra pixel is added above next to left/top edges of the region of 1s.
-
+        A binary mask tensor of shape [h, w],
+        where an extra pixel is added above next to left/top edges of the region of 1s.
     """
-    bottom_right_kernel = torch.tensor(
-        [
+    bottom_right_kernel = torch.tensor([
             [-1, -1, -1], 
-            [-1, -1, 1], 
-            [-1, 1, 1]
+            [-1, -1,  1], 
+            [-1,  1,  1]
         ], dtype=torch.float16, device=mask.device
     ).t()
     bottom_right = F.conv2d(
@@ -314,12 +321,12 @@ def expand_bottom_right(mask : torch.Tensor):
     ).squeeze(1).clamp(0)
     return mask + bottom_right
 
+
 ## These are taken from ultralytics to avoid unnecessary dependencies
 V = TypeVar("V", bound=torch.Tensor | np.ndarray)
-def clip_boxes(
-        boxes : V, 
-        shape : tuple[int, int]
-    ) -> V:
+
+
+def clip_boxes(boxes: V, shape: tuple[int, int]) -> V:
     """Clips bounding boxes to a specified shape (height, width).
 
     Args:
@@ -340,18 +347,19 @@ def clip_boxes(
         boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0])  # y1, y2
     return boxes
 
+
 def scale_boxes(
-        img1_shape : tuple[int, int], 
-        boxes : torch.Tensor, 
-        img0_shape : tuple[int, int], 
-        ratio_pad=None, 
-        padding : bool=True, 
-        xywh : bool=False
-    ):
+    img1_shape: tuple[int, int],
+    boxes: torch.Tensor,
+    img0_shape: tuple[int, int],
+    ratio_pad=None,
+    padding: bool = True,
+    xywh: bool = False,
+):
     """Rescales bounding boxes.
-    
+
     Bounding boxes are assumed to be in the "xyxy" format by default.
-    Reshapes bounding boxes from the shape of the image they were originally 
+    Reshapes bounding boxes from the shape of the image they were originally
         specified in (img1_shape) to the shape of a different image (img0_shape).
 
     Args:
@@ -387,6 +395,7 @@ def scale_boxes(
     boxes[..., :4] /= gain
     return clip_boxes(boxes, img0_shape)
 
+
 # Revised from ultralytics
 def _is_end2end_output(preds) -> bool:
     """Return True if preds is from an end2end model (e.g. YOLOv26).
@@ -406,16 +415,16 @@ def _is_end2end_output(preds) -> bool:
 
 
 def postprocess(
-        preds,
-        imgs : Sequence[torch.Tensor] | torch.Tensor,
-        max_det : int=300,
-        min_confidence : float=0,
-        overlap_threshold : float=0.1,
-        overlap_metric : str="IoU",
-        nms : int=0,
-        valid_size_range : tuple[int, int] | list[int] | None=None,
-        edge_margin : int | None=None
-    ) -> list[Results]:
+    preds,
+    imgs: Sequence[torch.Tensor] | torch.Tensor,
+    max_det: int = 300,
+    min_confidence: float = 0,
+    overlap_threshold: float = 0.1,
+    overlap_metric: str = "IoU",
+    nms: int = 0,
+    valid_size_range: tuple[int, int] | list[int] | None = None,
+    edge_margin: int | None = None,
+) -> list[Results]:
     """Postprocesses the predictions of the model.
 
     Args:
@@ -425,10 +434,12 @@ def postprocess(
         min_confidence: The minimum confidence of the predictions to return. Defaults to 0.
         overlap_threshold: The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to 0.1.
         overlap_metric: Overlap metric to use for NMS. Default is "IoU".
-        nms: The type of non-maximum suppression to use. Defaults to 0. 0 is no NMS, 1 is standard NMS, 2 is fancy NMS and 3 is mask NMS.
-        valid_size_range: The range of valid sizes for the bounding boxes in pixels. Defaults to None (no valid size range).
-        edge_margin: The minimum gap between the edge of the image and the bounding box in pixels for a prediction to be considered valid.
-            Defaults to None (no edge margin).
+        nms: The type of non-maximum suppression to use.
+            Defaults to 0. 0 is no NMS, 1 is standard NMS, 2 is fancy NMS and 3 is mask NMS.
+        valid_size_range: The range of valid sizes for the bounding boxes in pixels.
+            Defaults to None (no valid size range).
+        edge_margin: The minimum gap between the edge of the image and the bounding box in pixels
+            for a prediction to be considered valid. Defaults to None (no edge margin).
 
     Returns:
         A list of postprocessed predictions.
@@ -443,7 +454,7 @@ def postprocess(
         # YOLOv26 / end2end models: preds[0] = (dets[batch, n_dets, 4+1+1+32], protos[batch, 32, h, w])
         # Boxes are already in xyxy format and post-NMS; conf at dim 4, class at dim 5.
         p = preds[0][0].clone()  # (batch, n_dets, 38)
-        protos = preds[0][1]     # (batch, 32, h, w)
+        protos = preds[0][1]  # (batch, 32, h, w)
         if len(protos.shape) == 3:
             protos = protos.unsqueeze(0)
         # Filter zero-conf padding slots and apply min_confidence
@@ -468,16 +479,17 @@ def postprocess(
         assert isinstance(p, torch.Tensor)
         # Convert from xywh to xyxy
         p[:, :4, :] = torch.cat((
-                p[:, 0:2, :] - p[:, 2:4, :] / 2,  # x_min, y_min
-                p[:, 0:2, :] + p[:, 2:4, :] / 2   # x_max, y_max
-            ),
-            dim=1)
+            p[:, 0:2, :] - p[:, 2:4, :] / 2,  # x_min, y_min
+            p[:, 0:2, :] + p[:, 2:4, :] / 2   # x_max, y_max
+        ), dim=1)
         if min_confidence > 0:
             num_above_min_conf = (p[:, 4, :] > min_confidence).sum(dim=1)
             max_det = min(max_det, int(num_above_min_conf.max().item()))
         # Filter top-`max_det` predictions
         if max_det != 0:
-            p = p.gather(2, torch.argsort(p[:, 4, :], dim=1, descending=True)[:, :max_det].unsqueeze(1).expand(-1, p.size(1), -1))
+            p = p.gather(
+                2, torch.argsort(p[:, 4, :], dim=1, descending=True)[:, :max_det].unsqueeze(1).expand(-1, p.size(1), -1)
+            )
         # Change shape from (batch, xyxy + cls + masks, n) to (batch, n, xyxy + cls + masks)
         p = p.transpose(-2, -1)
         conf_mask = None
@@ -508,37 +520,40 @@ def postprocess(
             boxes = boxes[valid]
         # Remove predictions too close to the margin
         if edge_margin is not None and edge_margin > 0:
-            close_to_edge = (boxes[:, :2] < edge_margin).any(dim=1) | (boxes[:, 2:] > (tile_size - edge_margin)).any(dim=1)
+            close_to_edge = (
+                (boxes[:, :2] < edge_margin).any(dim=1) | 
+                (boxes[:, 2:] > (tile_size - edge_margin)).any(dim=1)
+            )
             pred = pred[~close_to_edge]
             boxes = boxes[~close_to_edge]
         # Deduplicate predictions
         if nms != 0:
             if nms == 1:
                 nms_ind = nms_boxes(
-                    boxes, pred[:, 4], 
+                    boxes, pred[:, 4],
                     overlap_threshold=overlap_threshold, overlap_fn=overlap_metric
                 )
             elif nms == 2:
                 nms_ind = fancy_nms(
-                    boxes, get_overlap_fn("box", overlap_metric), pred[:, 4], 
+                    boxes, get_overlap_fn("box", overlap_metric), pred[:, 4],
                     overlap_threshold=overlap_threshold, return_indices=True
                 )
             elif nms == 3:
                 # pred[:, -32:] - not sure this is correct for more than one class
                 masks = process_mask(
-                    protos[min(i, len(protos)-1)], 
-                    pred[:, -32:], 
+                    protos[min(i, len(protos)-1)],
+                    pred[:, -32:],
                     boxes,
                     (imgs[i].shape[-2], imgs[i].shape[-1]),
                     False
                 )
                 nms_ind = nms_masks(
-                    masks, pred[:, 4], 
-                    overlap_threshold=overlap_threshold, overlap_fn=overlap_metric, 
+                    masks, pred[:, 4],
+                    overlap_threshold=overlap_threshold, overlap_fn=overlap_metric,
                     return_indices=True, boxes=boxes / 4, group_first=False
                 )
-                # group_first is True, because nms_masks has vectorized IoU, 
-                # meaning that the overhead of doing connected-component clustering is 
+                # group_first is True, because nms_masks has vectorized IoU,
+                # meaning that the overhead of doing connected-component clustering is
                 # larger than the time-loss from redundant IoU calculations
                 masks = masks[nms_ind]
             else:
@@ -554,10 +569,10 @@ def postprocess(
         masks = masks[~too_small]
         pred[:, :4] = boxes
         results.append({
-            "orig_img" : Namespace(shape=imgs[i].permute(1,2,0).shape), 
-            "path" : "", 
-            "names" : ["insect"], 
-            "boxes" : pred[:, :6], 
-            "masks" : masks
+            "orig_img": Namespace(shape=imgs[i].permute(1, 2, 0).shape),
+            "path": "",
+            "names": ["insect"],
+            "boxes": pred[:, :6],
+            "masks": masks,
         })
-    return results
\ No newline at end of file
+    return results
diff --git a/tests/__init__.py b/tests/__init__.py
index e5b970e..6e03199 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1 +1 @@
-# noqa: D104
\ No newline at end of file
+# noqa: D104
diff --git a/tests/conftest.py b/tests/conftest.py
index 5ccb4b6..d7e8de0 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -7,8 +7,8 @@ def pytest_sessionfinish(session, exitstatus):
     """This hook runs after all tests have completed."""  # noqa: D401, D404
     script_path = os.path.join(os.path.dirname(__file__), "restore_assets.py")
     print(f"\n\nRunning Post-Test Cleanup:\n\t{script_path}\n")
-    
+
     try:
         subprocess.run(["python3", script_path], check=False)
     except Exception as e:
-        print(f"Post-test script failed to execute: {e}")
\ No newline at end of file
+        print(f"Post-test script failed to execute: {e}")
diff --git a/tests/generate_model_outputs.py b/tests/generate_model_outputs.py
index 23c86f7..701c51c 100644
--- a/tests/generate_model_outputs.py
+++ b/tests/generate_model_outputs.py
@@ -8,7 +8,7 @@
 from tests.test_predictor import ASSET_DIR, ASSET_NAME, TEST_MODEL_NAME, DummyModel
 
 # ruff: disable[E501]
-# Command I used: 
+# Command I used:
 #   python3 src/flat_bug/tests/generate_model_outputs.py --model model_snapshots/fb_2024-03-18_large_best.pt --image src/flat_bug/tests/assets/ALUS_Non-miteArachnids_Unknown_2020_11_03_4545.jpg --type both
 #
 # ruff: enable[E501]
@@ -16,7 +16,9 @@
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--model", type=str, default=TEST_MODEL_NAME, help="The model to test")
-    parser.add_argument("--image", type=str, default=os.path.join(ASSET_DIR, ASSET_NAME + ".jpg"), help="The image to test")
+    parser.add_argument(
+        "--image", type=str, default=os.path.join(ASSET_DIR, ASSET_NAME + ".jpg"), help="The image to test"
+    )
     parser.add_argument("--assets", type=str, default=ASSET_DIR, help="The directory to save the assets")
     parser.add_argument("--type", type=str, choices=["single_scale", "pyramid", "both"], required=True)
     parser.add_argument("--device", type=str, default="cuda:0", help="The device to use for inference")
@@ -24,8 +26,8 @@
 
     args = parser.parse_args()
 
-    image = read_image(args.image).to(torch.device(args.device), dtype=getattr(torch, args.dtype)) / 255.
-    test = DummyModel("single_scale", args.assets) # The type doesn't matter here
+    image = read_image(args.image).to(torch.device(args.device), dtype=getattr(torch, args.dtype)) / 255.0
+    test = DummyModel("single_scale", args.assets)  # The type doesn't matter here
     match args.type:
         case "single_scale":
             test.generate_single_scale_files(args.model, image)
@@ -36,4 +38,3 @@
             test.generate_pyramid_files(args.model, image, args.image)
         case _:
             raise ValueError(f"Invalid type {args.type}")
-    
\ No newline at end of file
diff --git a/tests/remote_lfs_fallback.py b/tests/remote_lfs_fallback.py
index 68abd12..d2c822b 100644
--- a/tests/remote_lfs_fallback.py
+++ b/tests/remote_lfs_fallback.py
@@ -10,8 +10,9 @@ def file_is_lfs_or_erda_pointer(file):  # noqa: D103
             return bool(re.search(r"git-lfs\.github\.com|ERDA Pointer", f.read()))
         except UnicodeDecodeError:
             return False
-    
-def check_file_with_remote_fallback(file, file_storage : str="https://anon.erda.au.dk/share_redirect/ecgKtuRWe5"):  # noqa: D103
+
+
+def check_file_with_remote_fallback(file, file_storage: str = "https://anon.erda.au.dk/share_redirect/ecgKtuRWe5"):  # noqa: D103
     if not os.path.exists(file) or file_is_lfs_or_erda_pointer(file):
         remote_uri = f"{file_storage}/{os.path.basename(file)}"
         try:
@@ -21,4 +22,4 @@ def check_file_with_remote_fallback(file, file_storage : str="https://anon.erda.
                 f"Failed to download test file {file} from remote file storage ({remote_uri})."
                 "\n\tPerhaps the file is not available."
             ) from e
-    return file
\ No newline at end of file
+    return file
diff --git a/tests/restore_assets.py b/tests/restore_assets.py
index 080e085..856c62f 100644
--- a/tests/restore_assets.py
+++ b/tests/restore_assets.py
@@ -15,4 +15,3 @@
         with open(asset, "w") as f:
             f.write("ERDA Pointer")
     print("-------------------------- Test assets restored! --------------------------")
-    
\ No newline at end of file
diff --git a/tests/test_augmentations.py b/tests/test_augmentations.py
index f404daf..8ef1887 100644
--- a/tests/test_augmentations.py
+++ b/tests/test_augmentations.py
@@ -1,4 +1,5 @@
 """Tests for flatbug augmentations (including their integration into the dataloader)."""
+
 import math
 import os
 from copy import deepcopy
@@ -29,7 +30,7 @@
     "min_size": 4,
     "imgsz": 1024,
     "use_segments": True,
-    "use_keypoints": False
+    "use_keypoints": False,
 }
 
 ASSET_DIR = os.path.join(os.path.dirname(__file__), "assets")
@@ -40,6 +41,7 @@
 check_file_with_remote_fallback(TEST_IMG)
 check_file_with_remote_fallback(TEST_LABEL)
 
+
 def generate_train_augmentation_pipeline(hyp):  # noqa: D103
     hyp = IterableSimpleNamespace(**hyp)
     return train_augmentation_pipeline(
@@ -48,18 +50,17 @@ def generate_train_augmentation_pipeline(hyp):  # noqa: D103
         max_instances=hyp.max_instances,
         min_size=hyp.min_size,
         use_segments=hyp.use_segments,
-        use_keypoints=hyp.use_keypoints
+        use_keypoints=hyp.use_keypoints,
     )
 
+
 def generate_validation_augmentation_pipeline(hyp):  # noqa: D103
     hyp = IterableSimpleNamespace(**hyp)
     return validation_augmentation_pipeline(
-        image_size=hyp.imgsz,
-        min_size=hyp.min_size,
-        use_segments=hyp.use_segments,
-        use_keypoints=hyp.use_keypoints
+        image_size=hyp.imgsz, min_size=hyp.min_size, use_segments=hyp.use_segments, use_keypoints=hyp.use_keypoints
     )
 
+
 def mock_verify_image_label(image_path, label_path):  # noqa: D103
     try:
         args = (image_path, label_path, "unit_test", False, 1, 0, 0)
@@ -70,7 +71,7 @@ def mock_verify_image_label(image_path, label_path):  # noqa: D103
     label = {
         "im_file": im_file,
         "shape": shape,
-        "cls": lb[:, 0:1],  
+        "cls": lb[:, 0:1],
         "bboxes": lb[:, 1:],
         "segments": segments,
         "keypoints": keypoint,
@@ -78,14 +79,15 @@ def mock_verify_image_label(image_path, label_path):  # noqa: D103
         "bbox_format": "xywh",
     }
     label["instances"] = Instances(
-        np.array(label["bboxes"]), 
+        np.array(label["bboxes"]),
         np.array(resample_segments(label["segments"])),
         label["keypoints"],
         bbox_format=label["bbox_format"],
-        normalized=label["normalized"]
+        normalized=label["normalized"],
     )
     return label
 
+
 def mock_yolo_base_dataset_load_image(image_path, imgsz, rect_mode=False):  # noqa: D103
     """Load an image from the given path and resize it if necessary.
 
@@ -102,16 +104,16 @@ def mock_yolo_base_dataset_load_image(image_path, imgsz, rect_mode=False):  # no
 
     """
     f = Path(image_path)
-    
+
     if not f.exists():
         raise FileNotFoundError(f"Image Not Found {f}")
 
     im = cv2.imread(str(f))  # Read image using OpenCV
     if im is None:
         raise FileNotFoundError(f"Image Not Found {f}")
-    
+
     h0, w0 = im.shape[:2]  # Original height and width
-    
+
     if rect_mode:  # Resize while maintaining aspect ratio
         r = imgsz / max(h0, w0)  # Ratio
         if r != 1:  # If sizes are not equal
@@ -119,10 +121,11 @@ def mock_yolo_base_dataset_load_image(image_path, imgsz, rect_mode=False):  # no
             im = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
     elif not (h0 == w0 == imgsz):  # Resize by stretching image to square imgsz
         im = cv2.resize(im, (imgsz, imgsz), interpolation=cv2.INTER_LINEAR)
-    
+
     return im, (h0, w0), im.shape[:2]
 
-def make_empty(obj : Any) -> Any:  # noqa: D103
+
+def make_empty(obj: Any) -> Any:  # noqa: D103
     if isinstance(obj, np.ndarray):
         obj = np.empty((0, *obj.shape[1:]), dtype=obj.dtype)
     elif isinstance(obj, torch.Tensor):
@@ -131,6 +134,7 @@ def make_empty(obj : Any) -> Any:  # noqa: D103
         obj = []
     return obj
 
+
 class TestMockYOLOHelpers:  # noqa: D101
     def test_mock_yolo_base_dataset_load_image(self):  # noqa: D102
         loaded_img, _, _ = mock_yolo_base_dataset_load_image(TEST_IMG, TEST_HYP["imgsz"])
@@ -142,7 +146,7 @@ def test_mock_yolo_base_dataset_load_image(self):  # noqa: D102
     def test_mock_verify_image_label(self):  # noqa: D102
         result = mock_verify_image_label(TEST_IMG, TEST_LABEL)
         assert isinstance(result, dict), f"Expected dict, got {type(result).__name__}"
-        correct : dict[str, type | None] = {
+        correct: dict[str, type | None] = {
             "im_file": str,
             "shape": tuple,
             "cls": np.ndarray,
@@ -151,7 +155,7 @@ def test_mock_verify_image_label(self):  # noqa: D102
             "keypoints": None,
             "normalized": bool,
             "bbox_format": str,
-            "instances": Instances
+            "instances": Instances,
         }
         for k, v in correct.items():
             assert k in result, f"Missing key '{k}' in result"
@@ -161,6 +165,7 @@ def test_mock_verify_image_label(self):  # noqa: D102
                 f"Invalid type for key '{k}'. Expected {v.__name__}, got {type(result[k]).__name__}"
             )
 
+
 class TestAugmentations:  # noqa: D101
     def test_generate_train_augmentation_pipeline(self):  # noqa: D102
         pipeline = generate_train_augmentation_pipeline(TEST_HYP)
@@ -196,11 +201,11 @@ def test_train_augmentation_pipeline(self):  # noqa: D102
                 continue
             if isinstance(v, Instances):
                 empty_pipeline_input[k] = Instances(
-                    make_empty(v.bboxes), 
+                    make_empty(v.bboxes),
                     make_empty(v.segments),
                     make_empty(v.keypoints),
                     bbox_format=empty_pipeline_input["bbox_format"],
-                    normalized=v.normalized
+                    normalized=v.normalized,
                 )
             else:
                 empty_pipeline_input[k] = make_empty(v)
@@ -234,7 +239,7 @@ def test_validation_augmentation_pipeline(self):  # noqa: D102
                     make_empty(v.segments),
                     make_empty(v.keypoints),
                     bbox_format=empty_pipeline_input["bbox_format"],
-                    normalized=v.normalized
+                    normalized=v.normalized,
                 )
             else:
                 empty_pipeline_input[k] = make_empty(v)
@@ -243,4 +248,4 @@ def test_validation_augmentation_pipeline(self):  # noqa: D102
         except Exception as e:
             e.add_note("Failed to execute validation augmentation pipeline on image without labels.")
             raise
-        assert isinstance(out, dict), "Invalid output of validation augmentation pipeline on image without labels."
\ No newline at end of file
+        assert isinstance(out, dict), "Invalid output of validation augmentation pipeline on image without labels."
diff --git a/tests/test_config.py b/tests/test_config.py
index d44e356..3b36729 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,4 +1,5 @@
 """Tests for flatbug config submodule."""
+
 import copy
 import os
 import tempfile
@@ -16,7 +17,7 @@
     "list of ints": [1, 2, 3],
     "tuple of ints": (1, 2, 3),
     "mixed list": [1, "string", True],
-    "list of mixed list and mixed tuple": [[1, "a"], ("b", 2)]
+    "list of mixed list and mixed tuple": [[1, "a"], ("b", 2)],
 }
 
 TEST_OBJECTS_TYPES_LIST_TUPLE_NOT_INTERCHANGEABLE = {
@@ -28,7 +29,7 @@
     "list of ints": [list, [int, int, int]],
     "tuple of ints": [tuple, [int, int, int]],
     "mixed list": [list, [int, str, bool]],
-    "list of mixed list and mixed tuple": [list, [[list, [int, str]], [tuple, [str, int]]]]
+    "list of mixed list and mixed tuple": [list, [[list, [int, str]], [tuple, [str, int]]]],
 }
 
 TEST_OBJECTS_TYPES_LIST_TUPLE_INTERCHANGEABLE = {
@@ -40,9 +41,10 @@
     "list of ints": [(tuple, list), [int, int, int]],
     "tuple of ints": [(tuple, list), [int, int, int]],
     "mixed list": [(tuple, list), [int, str, bool]],
-    "list of mixed list and mixed tuple": [(tuple, list), [[(tuple, list), [int, str]], [(tuple, list), [str, int]]]]
+    "list of mixed list and mixed tuple": [(tuple, list), [[(tuple, list), [int, str]], [(tuple, list), [str, int]]]],
 }
 
+
 def check_equals_recursive(obj1, obj2):  # noqa: D103
     if isinstance(obj1, (tuple, list)):
         if len(obj1) != len(obj2):
@@ -53,6 +55,7 @@ def check_equals_recursive(obj1, obj2):  # noqa: D103
         return True
     return obj1 == obj2
 
+
 class TestConfig:  # noqa: D101
     def test_check_types(self):  # noqa: D102
         for i, (key, obj) in enumerate(TEST_OBJECTS.items()):
@@ -60,7 +63,7 @@ def test_check_types(self):  # noqa: D102
             check_types(obj, expected_type, f"Object '{key}' ({i})")
             check_types(obj, TEST_OBJECTS_TYPES_LIST_TUPLE_NOT_INTERCHANGEABLE[key], f"Object '{key}' ({i})")
             check_types(obj, TEST_OBJECTS_TYPES_LIST_TUPLE_INTERCHANGEABLE[key], f"Object '{key}' ({i})")
-    
+
     def test_check_cfg_types(self):  # noqa: D102
         try:
             check_cfg_types(DEFAULT_CFG, strict=True)
@@ -76,13 +79,13 @@ def test_check_cfg_types(self):  # noqa: D102
             raise type(e)("Error raised when checking config with unknown key and strict=False:\n" + str(e))
 
     def test_get_type_def(self):  # noqa: D102
-        error_msg = \
-        """
+        error_msg = """
         Failed to generate the correct type definitions for the test objects with tuple_list_interchangeable={}.
         
         TEST_OBJECTS_TYPES should be a dictionary with:
             - keys: same as TEST_OBJECTS, 
-            - values: the type definitions of the corresponding values in TEST_OBJECTS, that pass the check_types function.
+            - values: the type definitions of the corresponding values in TEST_OBJECTS, 
+                that pass the check_types function.
         
         Either TEST_OBJECTS_TYPES is incorrect, 
             get_type_def is not generating the correct type definitions or test_check_types did not pass.
@@ -107,7 +110,7 @@ def test_get_type_def(self):  # noqa: D102
                 obj = TEST_OBJECTS[key]
                 type_def = get_type_def(obj, tuple_list_interchangeable=True)
                 # Check that the generated type definition is the same as the expected type definition
-                assert check_equals_recursive(type_def, expected_type), ( 
+                assert check_equals_recursive(type_def, expected_type), (
                     f"\nFailed on object:\n'{key}' => {obj}\n"
                     f"with generated type definition:\n{type_def}"
                     f"\nand expected type definition:\n{expected_type}"
@@ -145,8 +148,9 @@ def test_write_read_cfg(self):  # noqa: D102
             except Exception as e:
                 type(e)("Error raised when reading a config file with unknown key and strict=False:\n" + str(e))
         assert check_types(new_cfg, get_type_def(orig_cfg), "Reconstructed Config", strict=False), (
-            "Failed to reconstruct the original config with comparable types after writing and reading.")
+            "Failed to reconstruct the original config with comparable types after writing and reading."
+        )
         assert check_equals_recursive(orig_cfg, new_cfg), (
             "Failed to reconstruct the values of the original config after writing and reading. "
             "Although the types are comparable, the values are not equal."
-        )
\ No newline at end of file
+        )
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index b8bc75d..1296688 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1,4 +1,5 @@
 """Tests for the custom flatbug datasets used for training."""
+
 import glob
 import os
 import tempfile
@@ -40,7 +41,8 @@
 TEST_CFG = deepcopy(DEFAULT_CFG)
 setattr(TEST_CFG, "task", "segment")
 
-def mock_verify_image_label(image_path : str, label_path : str) -> dict:  # noqa: D103
+
+def mock_verify_image_label(image_path: str, label_path: str) -> dict:  # noqa: D103
     try:
         args = (image_path, label_path, "unit_test", False, 1, 0, 0)
         im_file, lb, shape, segments, keypoint, nm_f, nf_f, ne_f, nc_f, msg = verify_image_label(args)
@@ -50,7 +52,7 @@ def mock_verify_image_label(image_path : str, label_path : str) -> dict:  # noqa
     label = {
         "im_file": im_file,
         "shape": shape,
-        "cls": lb[:, 0:1],  
+        "cls": lb[:, 0:1],
         "bboxes": lb[:, 1:],
         "segments": segments,
         "keypoints": keypoint,
@@ -62,11 +64,12 @@ def mock_verify_image_label(image_path : str, label_path : str) -> dict:  # noqa
         np.array(resample_segments(label["segments"])),
         label["keypoints"],
         bbox_format=label["bbox_format"],
-        normalized=label["normalized"]
+        normalized=label["normalized"],
     )
     return label
 
-def create_train_dataset(args : IterableSimpleNamespace) -> FlatBugYOLODataset:  # noqa: D103
+
+def create_train_dataset(args: IterableSimpleNamespace) -> FlatBugYOLODataset:  # noqa: D103
     return FlatBugYOLODataset(
         data=ASSET_DATA,
         img_path=ASSET_DIR,
@@ -80,10 +83,11 @@ def create_train_dataset(args : IterableSimpleNamespace) -> FlatBugYOLODataset:
         single_cls=args.single_cls or False,
         max_instances=None,
         task="segment",
-        subset_args={"n" : 1, "pattern" : ASSET_NAME}
+        subset_args={"n": 1, "pattern": ASSET_NAME},
     )
 
-def create_validation_dataset(args : IterableSimpleNamespace) -> FlatBugYOLOValidationDataset:  # noqa: D103
+
+def create_validation_dataset(args: IterableSimpleNamespace) -> FlatBugYOLOValidationDataset:  # noqa: D103
     return FlatBugYOLOValidationDataset(
         data=ASSET_DATA,
         img_path=ASSET_DIR,
@@ -97,9 +101,10 @@ def create_validation_dataset(args : IterableSimpleNamespace) -> FlatBugYOLOVali
         single_cls=args.single_cls or False,
         max_instances=np.inf,
         task="segment",
-        subset_args={"n" : 1, "pattern" : ASSET_NAME}
+        subset_args={"n": 1, "pattern": ASSET_NAME},
     )
 
+
 def _test_plot_batch(batch, ni):  # noqa: D103
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as f:
         plot_images(
@@ -110,6 +115,7 @@ def _test_plot_batch(batch, ni):  # noqa: D103
             on_plot=os.remove,
         )
 
+
 @pytest.mark.filterwarnings("ignore:.*argument is set as true but no accelerator is found.*:UserWarning")
 class TestDataset:  # noqa: D101
     def test_train_dataset(self):  # noqa: D102
@@ -119,7 +125,7 @@ def test_train_dataset(self):  # noqa: D102
         dataloader_iter = dataloader.iterator
 
         for batch, _ in zip(dataloader_iter, range(1)):
-            assert batch["img"].shape[0] ==  BATCH_SIZE
+            assert batch["img"].shape[0] == BATCH_SIZE
 
         _test_plot_batch(batch, 0)
 
@@ -152,4 +158,4 @@ def teardown_class(cls):  # noqa: D102
         # Clean caches i.e. files ending with .cache or .cache.lock in the directory of this script
         cache_files = glob.glob(os.path.join(TEST_DIR, "*.cache*"))
         for cache_file in cache_files:
-            os.remove(cache_file)
\ No newline at end of file
+            os.remove(cache_file)
diff --git a/tests/test_predictor.py b/tests/test_predictor.py
index d44bc82..7989637 100644
--- a/tests/test_predictor.py
+++ b/tests/test_predictor.py
@@ -1,4 +1,5 @@
 """Tests for the public flatbug Predictor class."""
+
 import os
 import re
 import shutil
@@ -21,10 +22,7 @@
 ASSET_DIR = os.path.join(os.path.dirname(__file__), "assets")
 UUID = "XXXX"
 SERIALISED_TENSOR_PREDS = os.path.join(ASSET_DIR, f"metadata_{ASSET_NAME}_UUID_{UUID}.json")
-N_PREDICTIONS = {
-    "XXXX" : 11,
-    "ChangeThisTEMPORARY" : 10
-}
+N_PREDICTIONS = {"XXXX": 11, "ChangeThisTEMPORARY": 10}
 N_PREDICTIONS = N_PREDICTIONS.get(UUID, None)
 if N_PREDICTIONS is None:
     raise ValueError(f"Number of predictions for UUID {UUID} is not known")
@@ -40,9 +38,10 @@
     "EXPERIMENTAL_NMS_OPTIMIZATION": True,
     "TIME": False,
     "TILE_SIZE": 1024,
-    "BATCH_SIZE": 1
+    "BATCH_SIZE": 1,
 }
 
+
 class TestTensorPredictions:  # noqa: D101
     def test_load(self):  # noqa: D102
         tp = TensorPredictions.load(check_file_with_remote_fallback(SERIALISED_TENSOR_PREDS))
@@ -64,10 +63,10 @@ def test_save(self):  # noqa: D102
             n_crops = len(crops)
             # ###### DEBUG ######
             # [shutil.move(c, os.path.join(os.path.dirname(__file__), "assets", os.path.basename(c))) for c in crops]
-            # overview = glob(os.path.join(save_dir, "overview*"))[0] 
+            # overview = glob(os.path.join(save_dir, "overview*"))[0]
             # shutil.move(overview, os.path.join(os.path.dirname(__file__), "assets", os.path.basename(overview)))
             # ###################
-            assert n_crops == N_PREDICTIONS, ( 
+            assert n_crops == N_PREDICTIONS, (
                 f"Number of crops ({n_crops}) saved does not match the expected number of predictions ({N_PREDICTIONS})"
             )
             centroid_initial = [i.float().mean(dim=0).numpy() for i in tp.contours]
@@ -82,6 +81,7 @@ def test_save(self):  # noqa: D102
                 f"Centroid difference between initial and reloaded contours ({abs_diff}) is too large"
             )
 
+
 def cast_nested(obj, new_dtype):  # noqa: D103
     if not isinstance(obj, torch.Tensor):
         if hasattr(obj, "__iter__"):
@@ -89,15 +89,16 @@ def cast_nested(obj, new_dtype):  # noqa: D103
         return obj
     return obj.to(new_dtype)
 
+
 class DummyModel(torch.nn.Module):  # noqa: D101
-    def __init__(self, type : str, asset_dir : str):  # noqa: D107
+    def __init__(self, type: str, asset_dir: str):  # noqa: D107
         if type not in ["single_scale", "pyramid"]:
             raise ValueError(f"Invalid type {type}")
         self.type = type  # type: ignore
         self.asset_dir = asset_dir
         self.index = 1
 
-        self.save_counter = defaultdict(lambda : 0)
+        self.save_counter = defaultdict(lambda: 0)
 
     def to(self, *args, **kwargs):  # noqa: D102
         return self
@@ -107,7 +108,7 @@ def cpu(self):  # noqa: D102
 
     def cuda(self, *args, **kwargs):  # noqa: D102
         return self
-    
+
     def eval(self):  # noqa: D102
         return self
 
@@ -116,32 +117,35 @@ def train(self, mode=True):  # noqa: D102
 
     def __call__(self, image):  # noqa: D102
         try:
-            this_asset = os.path.join(self.asset_dir, f'{self.type}_tps_{self.index}.pt')
-            print(f'Processing asset {this_asset}')
+            this_asset = os.path.join(self.asset_dir, f"{self.type}_tps_{self.index}.pt")
+            print(f"Processing asset {this_asset}")
             check_file_with_remote_fallback(this_asset)
             out = cast_nested(torch.load(this_asset, map_location=image.device), image.dtype)
         except Exception as e:
             logger.error(
                 f'Failed to load test file "{self.type}_tps_{self.index}.pt" - '
-                'consider generating the test files with '
-                '`python3 src/flat_bug/tests/generate_model_outputs.py '
-                '--model model_snapshots/fb_2024-03-18_large_best.pt '
-                '--image src/flat_bug/tests/assets/ALUS_Non-miteArachnids_Unknown_2020_11_03_4545.jpg '
-                '--type both`'
+                "consider generating the test files with "
+                "`python3 src/flat_bug/tests/generate_model_outputs.py "
+                "--model model_snapshots/fb_2024-03-18_large_best.pt "
+                "--image src/flat_bug/tests/assets/ALUS_Non-miteArachnids_Unknown_2020_11_03_4545.jpg "
+                "--type both`"
             )
             raise e
         self.index += 1
         return out
-    
-    def hook_save_raw_output(self, model, label : str):  # noqa: D102
+
+    def hook_save_raw_output(self, model, label: str):  # noqa: D102
         ocall = model.__call__
 
         def call_wrapped(*args, **kwargs):
             output = ocall(*args, **kwargs)
             self.save_counter[label] += 1
-            torch.save(cast_nested(output, torch.device("cpu")), os.path.join(self.asset_dir, f'tps_{self.save_counter[label]}.pt'))
+            torch.save(
+                cast_nested(output, torch.device("cpu")),
+                os.path.join(self.asset_dir, f"tps_{self.save_counter[label]}.pt"),
+            )
             return output
-        
+
         model.__call__ = call_wrapped
 
     def generate_single_scale_files(self, weights, image):  # noqa: D102
@@ -152,13 +156,11 @@ def generate_single_scale_files(self, weights, image):  # noqa: D102
         model.total_detection_time = 0
         model.total_forward_time = 0
         output = model._detect_instances(
-            image,
-            scale=(model.TILE_SIZE / torch.tensor(image.shape[1:])).min().item(),
-            max_scale=False
+            image, scale=(model.TILE_SIZE / torch.tensor(image.shape[1:])).min().item(), max_scale=False
         )
         # Rename the files with the pattern "assets/tps_<NUMBER>.pt" to "assets/single_scale_tps_<NUMBER>.pt"
         [
-            shutil.move(f, os.path.join(self.asset_dir, re.sub(r'tps_', "single_scale_tps_", f))) 
+            shutil.move(f, os.path.join(self.asset_dir, re.sub(r"tps_", "single_scale_tps_", f)))
             for f in glob(os.path.join(self.asset_dir, "tps_*.pt"))
         ]
         # Create a file with the length of the output object as a reference
@@ -172,12 +174,11 @@ def generate_pyramid_files(self, weights, image, image_path):  # noqa: D102
         self.hook_save_raw_output(model._model, "pyramid")
         model.TIME = True
         output = model.pyramid_predictions(
-            image, image_path, 
-            scale_increment=1/2, scale_before=PYRAMID_SCALE_BEFORE, single_scale=False
+            image, image_path, scale_increment=1 / 2, scale_before=PYRAMID_SCALE_BEFORE, single_scale=False
         )
         # Rename the files with the pattern "assets/tps_<NUMBER>.pt" to "assets/pyramid_tps_<NUMBER>.pt"
         [
-            shutil.move(f, os.path.join(self.asset_dir, re.sub(r'tps_', "pyramid_tps_", f)))
+            shutil.move(f, os.path.join(self.asset_dir, re.sub(r"tps_", "pyramid_tps_", f)))
             for f in glob(os.path.join(self.asset_dir, "tps_*.pt"))
         ]
         # Create a file with the length of the output object as a reference
@@ -185,6 +186,7 @@ def generate_pyramid_files(self, weights, image, image_path):  # noqa: D102
         with open(os.path.join(self.asset_dir, "pyramid_output_length.txt"), "w") as f:
             f.write(str(len(output)))
 
+
 class TestPredictor:  # noqa: D101
     TOLERANCE = 0.1
 
@@ -193,34 +195,31 @@ def test_single_scale(self):  # noqa: D102
         predictor = Predictor(model=DummyModel("single_scale", ASSET_DIR), dtype=dtype, cfg=TEST_CFG)  # type: ignore
         image_path = os.path.join(ASSET_DIR, ASSET_NAME + ".jpg")
         check_file_with_remote_fallback(image_path)
-        image = read_image(image_path).to(torch.device("cpu"), dtype=dtype) / 255.
+        image = read_image(image_path).to(torch.device("cpu"), dtype=dtype) / 255.0
         output = predictor._detect_instances(
-            image, 
-            scale=(predictor.TILE_SIZE / torch.tensor(image.shape[1:])).min().item(),
-            max_scale=False
+            image, scale=(predictor.TILE_SIZE / torch.tensor(image.shape[1:])).min().item(), max_scale=False
         )
         output_length = len(output)
         with open(check_file_with_remote_fallback(os.path.join(ASSET_DIR, "single_scale_output_length.txt"))) as f:
             reference_length = int(f.read())
         # Check that the output length is within tolerance of the reference length
-        assert abs(1 - output_length/reference_length) < self.TOLERANCE, (
+        assert abs(1 - output_length / reference_length) < self.TOLERANCE, (
             f"Output length ({output_length}) does not match the reference length ({reference_length})"
         )
-    
+
     def test_pyramid(self):  # noqa: D102
         dtype = torch.float16
         predictor = Predictor(model=DummyModel("pyramid", ASSET_DIR), dtype=dtype, cfg=TEST_CFG)  # type: ignore
         image_path = os.path.join(ASSET_DIR, ASSET_NAME + ".jpg")
         check_file_with_remote_fallback(image_path)
-        image = read_image(image_path).to(torch.device("cpu"), dtype=dtype) / 255.
+        image = read_image(image_path).to(torch.device("cpu"), dtype=dtype) / 255.0
         output = predictor.pyramid_predictions(
-            image, image_path, 
-            scale_increment=1/2, scale_before=PYRAMID_SCALE_BEFORE, single_scale=False
+            image, image_path, scale_increment=1 / 2, scale_before=PYRAMID_SCALE_BEFORE, single_scale=False
         )
         output_length = len(output)
         with open(check_file_with_remote_fallback(os.path.join(ASSET_DIR, "pyramid_output_length.txt"))) as f:
             reference_length = int(f.read())
         # Check that the output length is within tolerance of the reference length
-        assert abs(1 - output_length/reference_length) < self.TOLERANCE, (
+        assert abs(1 - output_length / reference_length) < self.TOLERANCE, (
             f"Output length ({output_length}) does not match the reference length ({reference_length})"
-        )
\ No newline at end of file
+        )