From 7eb43727e7c1fee73044f60fd5bf13bed94f4a0f Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Mon, 2 Jun 2025 15:12:54 +0200
Subject: [PATCH 01/24] Added png support. Added folder in configs. Added Early
 Stopping but not tested yet.

---
 .../5_cm/finetuning_5_cm_TreeAI_full.toml     | 19 ++++
 src/deepforest_finetuning/config/_config.py   |  1 +
 .../prediction/_prediction_dataset.py         | 23 ++++-
 .../training/_finetuning.py                   | 93 ++++++++++++++++---
 4 files changed, 123 insertions(+), 13 deletions(-)
 create mode 100644 configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml

diff --git a/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml b/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
new file mode 100644
index 0000000..2451120
--- /dev/null
+++ b/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
@@ -0,0 +1,19 @@
+base_dir = "/mnt/daten/TreeAI/finetuning/for_finetuning/5cm"
+tmp_dir = "./tmp"
+patch_size = 640
+patch_overlap = 0.2
+image_folder = "images"
+train_annotation_files = [
+    "annotations",
+]
+test_annotation_files = [
+    "test_annotations",
+]
+epochs = 20
+seeds = [0, 1, 2, 3, 4]
+learning_rate = 0.0001
+checkpoint_dir = "checkpoints"
+early_stopping_patience = 2  # Set to the number of epochs to wait before stopping, or remove to disable early stopping
+
+[prediction_export]
+output_folder = "predictions"
diff --git a/src/deepforest_finetuning/config/_config.py b/src/deepforest_finetuning/config/_config.py
index 4061719..7fc727c 100644
--- a/src/deepforest_finetuning/config/_config.py
+++ b/src/deepforest_finetuning/config/_config.py
@@ -96,6 +96,7 @@ class TrainingConfig:  # pylint: disable=too-many-instance-attributes
     precision: str = "16-mixed"
     float32_matmul_precision: str = "medium"
     log_dir: str = "./logs"
+    early_stopping_patience: Optional[int] = None
 
 
 @dataclasses.dataclass
diff --git a/src/deepforest_finetuning/prediction/_prediction_dataset.py b/src/deepforest_finetuning/prediction/_prediction_dataset.py
index 3fac452..7693231 100644
--- a/src/deepforest_finetuning/prediction/_prediction_dataset.py
+++ b/src/deepforest_finetuning/prediction/_prediction_dataset.py
@@ -8,6 +8,7 @@
 import numpy as np
 import numpy.typing as npt
 from PIL import Image
+import os
 from tifffile import imread
 from torch.utils.data import Dataset
 
@@ -42,8 +43,26 @@ def __getitem__(self, idx: int) -> npt.NDArray:
         Returns: Image data.
         """
         img_path = self.image_files[idx]
-        image_array = np.array(imread(img_path))[:, :, :3].astype(np.uint8)
-
+        # Check file extension to use appropriate method for reading the image
+        file_ext = os.path.splitext(img_path)[1].lower()
+        
+        if file_ext in ['.tif', '.tiff']:
+            # Use tifffile for TIFF images
+            image_array = np.array(imread(img_path))
+            # Ensure we only take the first 3 channels if there are more
+            if image_array.ndim >= 3 and image_array.shape[2] > 3:
+                image_array = image_array[:, :, :3]
+        else:
+            # Use PIL for other image formats (PNG, JPG, etc.)
+            image = Image.open(img_path)
+            image_array = np.array(image)
+            # Ensure we only take the first 3 channels if there are more
+            if image_array.ndim >= 3 and image_array.shape[2] > 3:
+                image_array = image_array[:, :, :3]
+        
+        # Ensure consistent dtype
+        image_array = image_array.astype(np.uint8)
+        
         if self.resize_images_to is not None:
             image = Image.fromarray(image_array)
             image = image.resize((self.resize_images_to, self.resize_images_to))
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 2459dc3..435216e 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -15,7 +15,7 @@
 from deepforest import utilities, preprocess
 from deepforest import main as deepforest_main
 from pytorch_lightning import seed_everything, Trainer, LightningModule
-from pytorch_lightning.callbacks import Callback, ModelCheckpoint
+from pytorch_lightning.callbacks import Callback, ModelCheckpoint, EarlyStopping
 from pytorch_lightning.loggers import CSVLogger
 from pytorch_lightning.utilities.seed import isolate_rng
 import numpy as np
@@ -43,12 +43,11 @@ def get_transform(augment: bool, seed: Optional[int] = None):
         transform = A.Compose(
             [A.HorizontalFlip(p=0.5), ToTensorV2()],
             bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
-            seed=seed,
         )
 
     else:
         transform = A.Compose(
-            [ToTensorV2()], bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]), seed=seed
+            [ToTensorV2()], bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"])
         )
 
     return transform
@@ -107,6 +106,37 @@ def split_images_into_patches(
     return annotations_path
 
 
+def _process_annotation_paths(base_dir: Path, annotation_files: List[str]) -> List[str]:
+    """
+    Process annotation file paths, handling both individual files and directories.
+
+    Args:
+        base_dir: Base directory for relative paths.
+        annotation_files: List of annotation file paths or directories containing annotation files.
+
+    Returns:
+        List of processed annotation file paths.
+    """
+    processed_paths = []
+
+    for file_path in annotation_files:
+        path = base_dir / file_path
+        if path.is_dir():
+            # If it's a directory, collect all JSON files inside
+            json_files = list(path.glob("*.json"))
+            # Convert paths to strings relative to base_dir
+            rel_paths = [str(js_file.relative_to(base_dir)) for js_file in json_files]
+            processed_paths.extend(rel_paths)
+            print(f"INFO: Found {len(json_files)} JSON files in directory {path}.")
+        else:
+            # If it's a file, add it directly
+            processed_paths.append(file_path)
+            print(f"INFO: Using annotation file {file_path}.")
+
+
+    return processed_paths
+
+
 class EvaluationCallBack(Callback):
     """
     Callback that evaluates the model after each training epoch.
@@ -136,9 +166,12 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
         pl_module.trainer = trainer
         # trainer_state = copy.deepcopy(trainer.state)
         # evaluate on training and test set
+        processed_train_files = _process_annotation_paths(self._base_dir, self._config.train_annotation_files)
+        processed_test_files = _process_annotation_paths(self._base_dir, self._config.test_annotation_files)
+        
         for prefix, annotation_files in [
-            ("train", self._config.train_annotation_files),
-            ("test", self._config.test_annotation_files),
+            ("train", processed_train_files),
+            ("test", processed_test_files),
         ]:
             image_files = []
             annotations = []
@@ -201,9 +234,19 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
     preprocessed_image_folders = {}
     preprocessed_annotation_files = {}
 
-    splitting_configs = [("train", config.train_annotation_files)]
+    extracted_train_annotation_files = config.train_annotation_files
+
+    # Process annotation file paths for train and pretraining (if available)
+    processed_train_files = _process_annotation_paths(base_dir, config.train_annotation_files)
+    
+    print(processed_train_files)
+
+    splitting_configs = [("train", processed_train_files)]
     if config.pretrain_annotation_files is not None and len(config.pretrain_annotation_files) > 0:
-        splitting_configs.append(("pretraining", config.pretrain_annotation_files))
+        processed_pretrain_files = _process_annotation_paths(base_dir, config.pretrain_annotation_files)
+        splitting_configs.append(("pretraining", processed_pretrain_files))
+
+    print(f"INFO: Found {len(splitting_configs)} annotation configurations to process.")
 
     for prefix, annotation_files in splitting_configs:
         annotations = []
@@ -225,6 +268,9 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
         )
 
     print("\nStarting training ...")
+    
+    if config.early_stopping_patience is not None:
+        print(f"Early stopping enabled with patience of {config.early_stopping_patience} epochs")
 
     for seed in config.seeds:
         # set seeds for reproducibility
@@ -249,16 +295,31 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
             model.config["train"]["epochs"] = config.epochs
             model.config["save-snapshot"] = False
 
-            if "pretrain" in annotation_files:
-                model.config["train"]["csv_file"] = preprocessed_annotation_files["pretrain"]
-                model.config["train"]["root_dir"] = preprocessed_image_folders["pretrain"]
+            if "pretraining" in preprocessed_annotation_files:
+                model.config["train"]["csv_file"] = preprocessed_annotation_files["pretraining"]
+                model.config["train"]["root_dir"] = preprocessed_image_folders["pretraining"]
                 logger = CSVLogger(config.log_dir, name=f"{config.epochs}_epochs_seed_{seed}_pretraining")
+                
+                # Add pretraining callbacks
+                pretraining_callbacks = []
+                if config.early_stopping:
+                    pretraining_callbacks.append(
+                        EarlyStopping(
+                            monitor=config.early_stopping_monitor,
+                            min_delta=config.early_stopping_min_delta,
+                            patience=config.early_stopping_patience,
+                            verbose=True,
+                            mode=config.early_stopping_mode,
+                        )
+                    )
+                
                 model.create_trainer(
                     precision=config.precision if torch.cuda.is_available() else 32,
                     log_every_n_steps=1,
                     benchmark=False,
                     deterministic=True,
                     logger=logger,
+                    callbacks=pretraining_callbacks if pretraining_callbacks else None,
                 )
                 model.trainer.fit(model)
 
@@ -269,7 +330,6 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
 
             callbacks: List[Callback] = [EvaluationCallBack(config, seed)]
             if config.checkpoint_dir is not None:
-
                 callbacks.append(
                     ModelCheckpoint(
                         dirpath=base_dir / config.checkpoint_dir,
@@ -279,6 +339,17 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
                         enable_version_counter=False,
                     )
                 )
+                
+            # Add early stopping callback if patience is set
+            if config.early_stopping_patience is not None:
+                callbacks.append(
+                    EarlyStopping(
+                        monitor="val_loss",
+                        patience=config.early_stopping_patience,
+                        verbose=True,
+                        mode="min",
+                    )
+                )
 
             model.create_trainer(
                 precision=config.precision if torch.cuda.is_available() else 32,

From fc31d4e605a078b10ba16f90d888eb0cbc8e9a3d Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Wed, 4 Jun 2025 11:15:08 +0200
Subject: [PATCH 02/24] finished Early Stopping

---
 src/deepforest_finetuning/evaluation/_evaluate.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/deepforest_finetuning/evaluation/_evaluate.py b/src/deepforest_finetuning/evaluation/_evaluate.py
index fc62c35..b85d9f8 100644
--- a/src/deepforest_finetuning/evaluation/_evaluate.py
+++ b/src/deepforest_finetuning/evaluation/_evaluate.py
@@ -4,7 +4,7 @@
 
 from pathlib import Path
 import warnings
-from typing import Union
+from typing import Dict, Union
 
 from deepforest.evaluate import evaluate_boxes
 import pandas as pd
@@ -12,7 +12,7 @@
 
 def evaluate(
     predictions: pd.DataFrame, annotations: pd.DataFrame, iou_threshold: float, output_file: Union[str, Path]
-) -> None:
+) -> Dict[str, float]:
     """
     Evaluates a model's predictions and stores the evaluation metrics as CSV file.
 
@@ -22,6 +22,9 @@ def evaluate(
         iou_threshold: Threshold for the IoU between predicted and target bounding boxes at which predicted bounding
             boxes are counted as true positives.
         output_file: Path of the CSV file in which to store the evaluation metrics.
+        
+    Returns:
+        Dictionary containing the evaluation metrics (precision, recall, f1).
     """
 
     # ignore deprecated warnings from pandas raised by deepforest.IoU (line 113: iou_df = pd.concat(iou_df))
@@ -48,3 +51,10 @@ def evaluate(
     Path(output_file).parent.mkdir(exist_ok=True, parents=True)
 
     df.to_csv(output_file, index=False)
+    
+    # Return metrics dictionary for use with Lightning logger
+    return {
+        "precision": results["precision"],
+        "recall": results["recall"],
+        "f1": results["f1"]
+    }

From 7d9aed606235e042770781a40334ef83034b6a54 Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Wed, 4 Jun 2025 11:15:25 +0200
Subject: [PATCH 03/24] forgot to stage one file. Early Stopping finished.

---
 .../training/_finetuning.py                   | 42 +++++++++++++++----
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 435216e..8ecdba0 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -161,10 +161,14 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
             pl_module: Model to be evaluated.
         """
 
-        trainer = copy.deepcopy(trainer)
-        pl_module = copy.deepcopy(pl_module)
-        pl_module.trainer = trainer
-        # trainer_state = copy.deepcopy(trainer.state)
+        # Create a copy of the model for evaluation to avoid affecting the training
+        eval_model = copy.deepcopy(pl_module)
+        eval_trainer = copy.deepcopy(trainer)
+        eval_model.trainer = eval_trainer
+        
+        # Keep a reference to the original trainer for logging metrics
+        original_trainer = trainer
+        
         # evaluate on training and test set
         processed_train_files = _process_annotation_paths(self._base_dir, self._config.train_annotation_files)
         processed_test_files = _process_annotation_paths(self._base_dir, self._config.test_annotation_files)
@@ -194,7 +198,7 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
             export_config.output_file_name = f"{prefix}_predictions_seed_{self._seed}.csv"
 
             run_prediction(
-                pl_module,
+                eval_model,
                 image_files=image_files,
                 predict_tile=True,
                 patch_size=self._config.patch_size,
@@ -213,12 +217,33 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
                 / f"{trainer.current_epoch + 1}_epochs"
                 / f"{prefix}_metrics_seed_{self._seed}.csv"
             )
-            evaluate(
+            metrics = evaluate(
                 prediction,
                 pd.concat(annotations),
                 self._config.iou_threshold,
                 metrics_file,
             )
+            
+            # Log metrics to Lightning logger for each prefix (train/test)
+            # This makes them available for callbacks like EarlyStopping
+            for metric_name, metric_value in metrics.items():
+                # Log with trainer
+                original_trainer.logger.log_metrics(
+                    {f"{prefix}_{metric_name}": metric_value}, 
+                    step=original_trainer.current_epoch
+                )
+                
+                # For test metrics, also log them as validation metrics for EarlyStopping
+                if prefix == "test":
+                    if metric_name == "f1":
+                        # Log F1 as val_f1 for early stopping (maximize)
+                        # Access callback_metrics directly on the original trainer
+                        original_trainer.callback_metrics[f"val_{metric_name}"] = torch.tensor(metric_value)
+                    
+                    # Always log inverse F1 as val_loss for compatibility with default early stopping
+                    if metric_name == "f1":
+                        # For loss, use 1-F1 as val_loss (minimize is better)
+                        original_trainer.callback_metrics["val_loss"] = torch.tensor(1.0 - metric_value)
 
 
 def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-many-statements
@@ -342,12 +367,13 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
                 
             # Add early stopping callback if patience is set
             if config.early_stopping_patience is not None:
+                # We provide both val_loss and val_f1 - let's use val_f1 for early stopping (higher is better)
                 callbacks.append(
                     EarlyStopping(
-                        monitor="val_loss",
+                        monitor="val_f1",  # Use F1 score for early stopping
                         patience=config.early_stopping_patience,
                         verbose=True,
-                        mode="min",
+                        mode="max",  # Higher F1 is better
                     )
                 )
 

From e29bb57c2633a39b94898f2aa36e24b9581c46b4 Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Wed, 4 Jun 2025 11:34:37 +0200
Subject: [PATCH 04/24] added copilot readme

---
 README.md | 268 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 268 insertions(+)
 create mode 100644 README.md

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..65f6251
--- /dev/null
+++ b/README.md
@@ -0,0 +1,268 @@
+# DeepForest Finetuning
+
+A Python package for fine-tuning the [DeepForest](https://github.com/weecology/DeepForest) model on custom data for tree detection in aerial imagery. This project provides a streamlined workflow for preprocessing training data, fine-tuning the DeepForest model, making predictions, and evaluating results.
+
+## Table of Contents
+- [Overview](#overview)
+- [Installation](#installation)
+  - [Using Conda](#using-conda)
+  - [Using Docker](#using-docker)
+- [Project Structure](#project-structure)
+- [Workflow](#workflow)
+  - [1. Data Preprocessing](#1-data-preprocessing)
+  - [2. Model Fine-tuning](#2-model-fine-tuning)
+  - [3. Making Predictions](#3-making-predictions)
+  - [4. Evaluating Results](#4-evaluating-results)
+- [Configuration Files](#configuration-files)
+- [Example Usage](#example-usage)
+- [Advanced Features](#advanced-features)
+- [License](#license)
+
+## Overview
+
+DeepForest is a deep learning model designed for detecting trees in aerial RGB imagery. This package extends DeepForest by providing a comprehensive framework to fine-tune the model on your own datasets. Key features include:
+
+- Data preprocessing for various input formats
+- Automatic label projection from 3D point clouds to 2D orthophotos
+- Image rescaling and patch generation
+- Model fine-tuning with multiple random seeds for robust evaluation
+- Prediction on new images with customizable tiling
+- Evaluation metrics calculation (precision, recall, F1 score)
+- Support for experiment tracking with [Weights & Biases](https://wandb.ai/)
+
+## Installation
+
+### Using Conda
+
+1. Clone this repository:
+   ```bash
+   git clone https://github.com/yourusername/deepforest-finetuning.git
+   cd deepforest-finetuning
+   ```
+
+2. Create and activate a conda environment from the provided environment.yml file:
+   ```bash
+   conda env create -f environment.yml
+   conda activate deepforest-env
+   ```
+
+3. Install the package in development mode:
+   ```bash
+   pip install -e .
+   ```
+
+### Using Docker
+
+A Dockerfile is provided for containerized usage:
+
+```bash
+docker build -t deepforest-finetuning .
+docker run --gpus all -it -v /path/to/your/data:/data deepforest-finetuning /bin/bash
+```
+
+## Project Structure
+
+```
+deepforest-finetuning/
+├── configs/              # Configuration files for different workflows
+│   ├── baseline/         # Configurations for the baseline model
+│   ├── finetuning/       # Configurations for fine-tuning
+│   └── preprocessing/    # Configurations for data preprocessing
+├── scripts/              # Executable scripts for main workflows
+│   ├── evaluate.py       # Script for model evaluation
+│   ├── finetuning.py     # Script for fine-tuning
+│   ├── prediction.py     # Script for making predictions
+│   └── preprocessing.py  # Script for data preprocessing
+└── src/                  # Source code
+    └── deepforest_finetuning/
+        ├── config/       # Configuration dataclasses
+        ├── evaluation/   # Model evaluation logic
+        ├── prediction/   # Prediction logic
+        ├── preprocessing/# Data preprocessing logic
+        ├── training/     # Model training and fine-tuning logic
+        └── utils/        # Utility functions
+```
+
+## Workflow
+
+### 1. Data Preprocessing
+
+The package supports multiple preprocessing steps:
+
+#### a. Projecting Labels from Point Clouds
+
+If you have 3D point cloud data with tree positions, you can project them to 2D bounding boxes:
+
+```bash
+python scripts/preprocessing.py configs/preprocessing/project_point_cloud_labels.toml
+```
+
+Required configuration:
+```toml
+base_dir = "/path/to/your/data"
+point_cloud_paths = ["pointcloud1.las", "pointcloud2.las"]
+image_paths = ["image1.tif", "image2.tif"]
+label_json_output_paths = ["labels1.json", "labels2.json"]
+```
+
+#### b. Preprocessing Manually Corrected Labels
+
+Convert manually created or corrected labels to the format required by DeepForest:
+
+```bash
+python scripts/preprocessing.py configs/preprocessing/preprocess_manually_corrected_labels.toml
+```
+
+#### c. Filtering Labels
+
+Filter out unwanted labels based on size, position, etc.:
+
+```bash
+python scripts/preprocessing.py configs/preprocessing/filter_labels.toml
+```
+
+#### d. Image Rescaling
+
+Rescale images to different resolutions:
+
+```bash
+python scripts/preprocessing.py configs/preprocessing/rescale_images.toml
+```
+
+### 2. Model Fine-tuning
+
+Fine-tune the DeepForest model on your custom dataset:
+
+```bash
+python scripts/finetuning.py configs/finetuning/finetuning_5_cm_manual_labeling_small.toml
+```
+
+Example configuration:
+```toml
+base_dir = "/path/to/your/data"
+tmp_dir = "./tmp"
+patch_size = 640
+patch_overlap = 0.2
+image_folder = "images"
+train_annotation_files = ["annotations"]
+test_annotation_files = ["test_annotations"]
+epochs = 20
+seeds = [0, 1, 2, 3, 4]
+learning_rate = 0.0001
+checkpoint_dir = "checkpoints"
+early_stopping_patience = 2
+```
+
+This will:
+1. Split images into patches
+2. Create training and test datasets
+3. Fine-tune the model for the specified number of epochs
+4. Run with multiple random seeds for robust evaluation
+5. Save checkpoints and logs
+
+### 3. Making Predictions
+
+Make predictions with the fine-tuned model:
+
+```bash
+python scripts/prediction.py configs/finetuning/predict_finetuned_5_cm.toml
+```
+
+Example configuration:
+```toml
+checkpoint_path = "/path/to/checkpoint.pt"
+image_files = ["/path/to/image.tif"]
+predict_tile = true
+patch_size = 1000
+patch_overlap = 0.2
+
+[prediction_export]
+output_folder = "/path/to/predictions"
+output_file_name = "predictions.csv"
+```
+
+### 4. Evaluating Results
+
+Evaluate predictions against ground truth:
+
+```bash
+python scripts/evaluate.py configs/evaluation/evaluate_finetuned_5_cm.toml
+```
+
+Example configuration:
+```toml
+prediction_file = "/path/to/predictions.csv"
+label_file = "/path/to/ground_truth.csv"
+iou_threshold = 0.4
+output_file = "/path/to/evaluation_results.csv"
+```
+
+## Configuration Files
+
+All workflows are configured using TOML files:
+
+- **Preprocessing configs**: Define data paths and parameters for preprocessing steps
+- **Training configs**: Specify training hyperparameters, data paths, and evaluation settings
+- **Prediction configs**: Set model checkpoint, input images, and output format
+- **Evaluation configs**: Define prediction and ground truth file paths, and evaluation metrics
+
+## Example Usage
+
+### Complete Fine-tuning Pipeline Example
+
+1. Preprocess your data (e.g., project point cloud labels, rescale images)
+2. Fine-tune the model:
+   ```bash
+   python scripts/finetuning.py configs/finetuning/my_finetuning_config.toml
+   ```
+3. Make predictions with the fine-tuned model:
+   ```bash
+   python scripts/prediction.py configs/prediction/my_prediction_config.toml
+   ```
+4. Evaluate the results:
+   ```bash
+   python scripts/evaluate.py configs/evaluation/my_evaluation_config.toml
+   ```
+
+### Using Different Image Resolutions
+
+The package supports working with images at various resolutions. The config directories contain subdirectories for different resolutions (e.g., `2_5_cm`, `5_cm`, `7_5_cm`, `10_cm`).
+
+## Advanced Features
+
+### Experiment Tracking with Weights & Biases
+
+The fine-tuning process supports integration with Weights & Biases for experiment tracking:
+
+```toml
+# Add to your fine-tuning config
+use_wandb = true
+wandb_project = "deepforest-finetuning"
+wandb_entity = "your-wandb-username"
+```
+
+This will log metrics, hyperparameters, and validation results to your W&B account.
+
+### Multiple Random Seeds
+
+To ensure robust evaluation, you can run fine-tuning with multiple random seeds:
+
+```toml
+seeds = [0, 1, 2, 3, 4]
+```
+
+This will train separate models with different weight initializations and report the average performance.
+
+### Early Stopping
+
+To prevent overfitting, you can enable early stopping:
+
+```toml
+early_stopping_patience = 2
+```
+
+This will stop training if the validation performance doesn't improve for the specified number of epochs.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.

From 13db21a7ed8962bb42e9cdaa5fc7c7476f0281cf Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Wed, 4 Jun 2025 12:01:00 +0200
Subject: [PATCH 05/24] black. mypy.

---
 scripts/preprocessing.py                      |  10 +-
 .../evaluation/_evaluate.py                   |  21 ++-
 .../prediction/_prediction.py                 |  12 +-
 .../prediction/_prediction_dataset.py         |   8 +-
 .../preprocessing/_filter_labels.py           |  25 +++-
 .../_preprocess_manually_corrected_labels.py  |  58 ++++++--
 .../_project_point_cloud_labels.py            | 106 ++++++++++----
 .../preprocessing/_rescale_images.py          |  53 +++++--
 .../training/_finetuning.py                   | 131 ++++++++++++------
 .../utils/_annotations_to_coco.py             |  22 ++-
 .../utils/_rescale_coco_json.py               |  24 +++-
 11 files changed, 345 insertions(+), 125 deletions(-)

diff --git a/scripts/preprocessing.py b/scripts/preprocessing.py
index b4fe769..c256ce7 100644
--- a/scripts/preprocessing.py
+++ b/scripts/preprocessing.py
@@ -20,7 +20,9 @@
 from deepforest_finetuning.utils import load_config
 
 
-def preprocessing_step(config_path: str, config_type: Type, script_function: Callable[[Any], None]):
+def preprocessing_step(
+    config_path: str, config_type: Type, script_function: Callable[[Any], None]
+):
     """
     Loads the specified config file, parses it based on the given configuration type and then calls the script function
     with the given configuration.
@@ -43,7 +45,11 @@ def preprocessing_step(config_path: str, config_type: Type, script_function: Cal
             preprocess_manually_corrected_labels,
             ManuallyCorrectedLabelPreprocessingConfig,
         ),
-        ("project_point_cloud_labels", project_point_cloud_labels, PointCloudLabelProjectionConfig),
+        (
+            "project_point_cloud_labels",
+            project_point_cloud_labels,
+            PointCloudLabelProjectionConfig,
+        ),
         ("rescale_images", rescale_images, ImageRescalingConfig),
     ]
     fire_dict = {}
diff --git a/src/deepforest_finetuning/evaluation/_evaluate.py b/src/deepforest_finetuning/evaluation/_evaluate.py
index b85d9f8..f044459 100644
--- a/src/deepforest_finetuning/evaluation/_evaluate.py
+++ b/src/deepforest_finetuning/evaluation/_evaluate.py
@@ -11,7 +11,10 @@
 
 
 def evaluate(
-    predictions: pd.DataFrame, annotations: pd.DataFrame, iou_threshold: float, output_file: Union[str, Path]
+    predictions: pd.DataFrame,
+    annotations: pd.DataFrame,
+    iou_threshold: float,
+    output_file: Union[str, Path],
 ) -> Dict[str, float]:
     """
     Evaluates a model's predictions and stores the evaluation metrics as CSV file.
@@ -22,7 +25,7 @@ def evaluate(
         iou_threshold: Threshold for the IoU between predicted and target bounding boxes at which predicted bounding
             boxes are counted as true positives.
         output_file: Path of the CSV file in which to store the evaluation metrics.
-        
+
     Returns:
         Dictionary containing the evaluation metrics (precision, recall, f1).
     """
@@ -38,9 +41,15 @@ def evaluate(
 
     results["precision"] = results.pop("box_precision")
     results["recall"] = results.pop("box_recall")
-    results["f1"] = 2 * (results["precision"] * results["recall"]) / (results["precision"] + results["recall"])
+    results["f1"] = (
+        2
+        * (results["precision"] * results["recall"])
+        / (results["precision"] + results["recall"])
+    )
 
-    print(f"Precision:\t{results['precision']}\nRecall:\t\t{results['recall']}\nF1:\t\t{results['f1']}")
+    print(
+        f"Precision:\t{results['precision']}\nRecall:\t\t{results['recall']}\nF1:\t\t{results['f1']}"
+    )
 
     metrics = []
     metrics.append({"metric": "precision", "score": results["precision"]})
@@ -51,10 +60,10 @@ def evaluate(
     Path(output_file).parent.mkdir(exist_ok=True, parents=True)
 
     df.to_csv(output_file, index=False)
-    
+
     # Return metrics dictionary for use with Lightning logger
     return {
         "precision": results["precision"],
         "recall": results["recall"],
-        "f1": results["f1"]
+        "f1": results["f1"],
     }
diff --git a/src/deepforest_finetuning/prediction/_prediction.py b/src/deepforest_finetuning/prediction/_prediction.py
index de697cd..ed43e77 100644
--- a/src/deepforest_finetuning/prediction/_prediction.py
+++ b/src/deepforest_finetuning/prediction/_prediction.py
@@ -37,9 +37,13 @@ def prediction(
 
     if predict_tile:
         if patch_size is None:
-            raise ValueError("Patch size must be specified when predict_tile is set to True.")
+            raise ValueError(
+                "Patch size must be specified when predict_tile is set to True."
+            )
         if patch_overlap is None:
-            raise ValueError("Patch overlap must be specified when predict_tile is set to True.")
+            raise ValueError(
+                "Patch overlap must be specified when predict_tile is set to True."
+            )
 
     print("\nLoading dataset and model ...")
 
@@ -70,7 +74,9 @@ def prediction(
 
         export_labels(
             pred,
-            export_path=(Path(export_config.output_folder) / image_name).with_suffix(".csv"),
+            export_path=(Path(export_config.output_folder) / image_name).with_suffix(
+                ".csv"
+            ),
             column_order=export_config.column_order,
             index_as_label_suffix=export_config.index_as_label_suffix,
             sort_by=export_config.sort_by,
diff --git a/src/deepforest_finetuning/prediction/_prediction_dataset.py b/src/deepforest_finetuning/prediction/_prediction_dataset.py
index 7693231..d864b0c 100644
--- a/src/deepforest_finetuning/prediction/_prediction_dataset.py
+++ b/src/deepforest_finetuning/prediction/_prediction_dataset.py
@@ -45,8 +45,8 @@ def __getitem__(self, idx: int) -> npt.NDArray:
         img_path = self.image_files[idx]
         # Check file extension to use appropriate method for reading the image
         file_ext = os.path.splitext(img_path)[1].lower()
-        
-        if file_ext in ['.tif', '.tiff']:
+
+        if file_ext in [".tif", ".tiff"]:
             # Use tifffile for TIFF images
             image_array = np.array(imread(img_path))
             # Ensure we only take the first 3 channels if there are more
@@ -59,10 +59,10 @@ def __getitem__(self, idx: int) -> npt.NDArray:
             # Ensure we only take the first 3 channels if there are more
             if image_array.ndim >= 3 and image_array.shape[2] > 3:
                 image_array = image_array[:, :, :3]
-        
+
         # Ensure consistent dtype
         image_array = image_array.astype(np.uint8)
-        
+
         if self.resize_images_to is not None:
             image = Image.fromarray(image_array)
             image = image.resize((self.resize_images_to, self.resize_images_to))
diff --git a/src/deepforest_finetuning/preprocessing/_filter_labels.py b/src/deepforest_finetuning/preprocessing/_filter_labels.py
index a758bd6..dec32eb 100644
--- a/src/deepforest_finetuning/preprocessing/_filter_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_filter_labels.py
@@ -14,7 +14,9 @@
 from deepforest_finetuning.config import LabelFilteringConfig
 
 
-def filter_bounding_boxed_with_size_based_nms(coco_json: Dict[str, Any], iou_threshold: float) -> Dict[str, Any]:
+def filter_bounding_boxed_with_size_based_nms(
+    coco_json: Dict[str, Any], iou_threshold: float
+) -> Dict[str, Any]:
     """
     Applies non-maximum suppression to the bounding boxes, using the box sizes as scores.
 
@@ -38,7 +40,12 @@ def filter_bounding_boxed_with_size_based_nms(coco_json: Dict[str, Any], iou_thr
     for annotation in coco_json["annotations"]:
         bounding_box = annotation["bbox"]
         bounding_boxes.append(
-            [bounding_box[0], bounding_box[1], bounding_box[0] + bounding_box[2], bounding_box[1] + bounding_box[3]]
+            [
+                bounding_box[0],
+                bounding_box[1],
+                bounding_box[0] + bounding_box[2],
+                bounding_box[1] + bounding_box[3],
+            ]
         )
         bounding_box_sizes.append(bounding_box[2] * bounding_box[3])
 
@@ -58,7 +65,11 @@ def filter_labels(config: LabelFilteringConfig):
     base_dir = Path(config.base_dir)
     label_folder = base_dir / config.input_label_folder
 
-    subfolders = [file for file in os.listdir(label_folder) if os.path.isdir(os.path.join(label_folder, file))]
+    subfolders = [
+        file
+        for file in os.listdir(label_folder)
+        if os.path.isdir(os.path.join(label_folder, file))
+    ]
 
     for subfolder in subfolders:
         for file in os.listdir(label_folder / subfolder):
@@ -68,9 +79,13 @@ def filter_labels(config: LabelFilteringConfig):
                     coco_json = json.load(f)
 
                 assert len(coco_json["images"]) == 1
-                coco_json = filter_bounding_boxed_with_size_based_nms(coco_json, iou_threshold=config.iou_threshold)
+                coco_json = filter_bounding_boxed_with_size_based_nms(
+                    coco_json, iou_threshold=config.iou_threshold
+                )
 
-                output_file_path = base_dir / config.output_label_folder / subfolder / file
+                output_file_path = (
+                    base_dir / config.output_label_folder / subfolder / file
+                )
                 output_file_path.parent.mkdir(exist_ok=True, parents=True)
                 with open(output_file_path, "w", encoding="utf-8") as f:
                     json.dump(coco_json, f, indent=4)
diff --git a/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py b/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py
index dd3e5aa..e6ac704 100644
--- a/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py
@@ -12,7 +12,11 @@
 import rasterio
 
 from deepforest_finetuning.config import ManuallyCorrectedLabelPreprocessingConfig
-from deepforest_finetuning.utils import annotations_to_coco, get_image_size_from_pascal_voc, rescale_coco_json
+from deepforest_finetuning.utils import (
+    annotations_to_coco,
+    get_image_size_from_pascal_voc,
+    rescale_coco_json,
+)
 
 
 def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, too-many-nested-blocks, too-many-statements
@@ -53,19 +57,27 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
 
             target_image_path = base_dir / config.input_image_folder / image_path
 
-            coco_json = annotations_to_coco(annotations, image_width, image_height, capture_date=capture_date)
+            coco_json = annotations_to_coco(
+                annotations, image_width, image_height, capture_date=capture_date
+            )
             coco_json = rescale_coco_json(
-                coco_json, target_image_path, source_image_shape=np.array([image_height, image_width])
+                coco_json,
+                target_image_path,
+                source_image_shape=np.array([image_height, image_width]),
             )
 
-            output_file_path = (output_label_folder / f"{Path(image_path).stem}_coco").with_suffix(".json")
+            output_file_path = (
+                output_label_folder / f"{Path(image_path).stem}_coco"
+            ).with_suffix(".json")
             with open(output_file_path, "w", encoding="utf-8") as f:
                 json.dump(coco_json, f, indent=4)
 
     # s1_p1_ext_mc contains labels for a larger tile
     # we crop the part for which we also have fully manually and automatically created labels
     if "s1_p1_ext_mc_coco.json" in os.listdir(output_label_folder):
-        with open(output_label_folder / "s1_p1_ext_mc_coco.json", "r", encoding="utf-8") as f:
+        with open(
+            output_label_folder / "s1_p1_ext_mc_coco.json", "r", encoding="utf-8"
+        ) as f:
             coco_json = json.load(f)
 
         source_image_path = base_dir / config.input_image_folder / "s1_p1_ext_mc.tif"
@@ -74,7 +86,9 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
         with rasterio.open(source_image_path) as image:
             transform = image.transform
 
-            src_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
+            src_pixel_size = np.abs(
+                np.array([transform[0], transform[4]], dtype=np.float64)
+            )
 
             src_top_left = np.array([transform.c, transform.f], dtype=np.float64)
             src_bottom_left = src_top_left.copy()
@@ -87,7 +101,9 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
             target_width = image.width
             target_height = image.height
 
-            target_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
+            target_pixel_size = np.abs(
+                np.array([transform[0], transform[4]], dtype=np.float64)
+            )
 
             target_top_left = np.array([transform.c, transform.f], dtype=np.float64)
             target_bottom_left = target_top_left.copy()
@@ -98,9 +114,15 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
         clipped_annotations = []
         for annotation in coco_json["annotations"]:
             x_min_meter = src_top_left[0] + annotation["bbox"][0] * src_pixel_size[0]
-            x_max_meter = src_top_left[0] + (annotation["bbox"][0] + annotation["bbox"][2]) * src_pixel_size[0]
+            x_max_meter = (
+                src_top_left[0]
+                + (annotation["bbox"][0] + annotation["bbox"][2]) * src_pixel_size[0]
+            )
             y_min_meter = src_top_left[1] - annotation["bbox"][1] * src_pixel_size[1]
-            y_max_meter = src_top_left[1] - (annotation["bbox"][1] + annotation["bbox"][3]) * src_pixel_size[1]
+            y_max_meter = (
+                src_top_left[1]
+                - (annotation["bbox"][1] + annotation["bbox"][3]) * src_pixel_size[1]
+            )
 
             if (
                 x_max_meter < target_top_left[0]
@@ -110,10 +132,20 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
             ):
                 continue
 
-            clipped_x_min = max(int((x_min_meter - target_top_left[0]) / target_pixel_size[0]), 0)
-            clipped_x_max = min(int((x_max_meter - target_top_left[0]) / target_pixel_size[0]), target_width)
-            clipped_y_min = max(int((target_top_left[1] - y_min_meter) / target_pixel_size[1]), 0)
-            clipped_y_max = min(int((target_top_left[1] - y_max_meter) / target_pixel_size[1]), target_height)
+            clipped_x_min = max(
+                int((x_min_meter - target_top_left[0]) / target_pixel_size[0]), 0
+            )
+            clipped_x_max = min(
+                int((x_max_meter - target_top_left[0]) / target_pixel_size[0]),
+                target_width,
+            )
+            clipped_y_min = max(
+                int((target_top_left[1] - y_min_meter) / target_pixel_size[1]), 0
+            )
+            clipped_y_max = min(
+                int((target_top_left[1] - y_max_meter) / target_pixel_size[1]),
+                target_height,
+            )
 
             clipped_annotation = deepcopy(annotation)
             clipped_annotation["bbox"] = [
diff --git a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
index a709d36..0cac1a8 100644
--- a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
@@ -9,7 +9,11 @@
 import numpy as np
 from pointtorch import read
 from pointtorch.operations.numpy import make_labels_consecutive
-from pointtree.operations import cloth_simulation_filtering, create_digital_terrain_model, distance_to_dtm
+from pointtree.operations import (
+    cloth_simulation_filtering,
+    create_digital_terrain_model,
+    distance_to_dtm,
+)
 import rasterio
 from rasterio.transform import from_origin
 from skimage.filters.rank import modal
@@ -28,11 +32,15 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
     if len(config.point_cloud_paths) != len(config.image_paths):
         raise ValueError("point_cloud_paths and image_paths must have the same length.")
     if len(config.point_cloud_paths) != len(config.label_json_output_paths):
-        raise ValueError("point_cloud_paths and label_json_output_paths must have the same length.")
-    if config.label_image_output_paths is not None and len(config.label_json_output_paths) != len(
-        config.label_image_output_paths
-    ):
-        raise ValueError("label_json_output_paths and label_image_output_paths must have the same length.")
+        raise ValueError(
+            "point_cloud_paths and label_json_output_paths must have the same length."
+        )
+    if config.label_image_output_paths is not None and len(
+        config.label_json_output_paths
+    ) != len(config.label_image_output_paths):
+        raise ValueError(
+            "label_json_output_paths and label_image_output_paths must have the same length."
+        )
 
     base_dir = Path(config.base_dir)
 
@@ -48,11 +56,15 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
             image_height = image.height
             crs = image.crs
 
-            pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
+            pixel_size = np.abs(
+                np.array([transform[0], transform[4]], dtype=np.float64)
+            )
             image_width_meter = image_width * pixel_size[1]
             image_height_meter = image_height * pixel_size[0]
 
-            tile_upper_left_corner = np.array([transform.c, transform.f], dtype=np.float64)
+            tile_upper_left_corner = np.array(
+                [transform.c, transform.f], dtype=np.float64
+            )
             tile_lower_left_corner = tile_upper_left_corner.copy()
             tile_lower_left_corner[1] -= image.height * pixel_size[1]
 
@@ -63,7 +75,9 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
 
         point_cloud = read(base_dir / point_cloud_path)
 
-        print(f"Loaded point cloud {Path(point_cloud_path).name} with {len(point_cloud)} points.")
+        print(
+            f"Loaded point cloud {Path(point_cloud_path).name} with {len(point_cloud)} points."
+        )
 
         point_cloud["instance_id_prediction"] = make_labels_consecutive(
             point_cloud["instance_id_prediction"].to_numpy(), ignore_id=0
@@ -71,27 +85,33 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
 
         if (
             point_cloud["instance_id_prediction"].min() != 0
-            or point_cloud["instance_id_prediction"].max() != len(point_cloud["instance_id_prediction"].unique()) - 1
+            or point_cloud["instance_id_prediction"].max()
+            != len(point_cloud["instance_id_prediction"].unique()) - 1
         ):
-            raise ValueError("The predicted instance IDs must be continuous, starting from zero.")
+            raise ValueError(
+                "The predicted instance IDs must be continuous, starting from zero."
+            )
 
-        pixel_indices = np.floor((point_cloud.xyz()[:, :2] - tile_lower_left_corner) / config.grid_resolution).astype(
-            np.int64
-        )
+        pixel_indices = np.floor(
+            (point_cloud.xyz()[:, :2] - tile_lower_left_corner) / config.grid_resolution
+        ).astype(np.int64)
 
-        label_image_shape = np.ceil(np.array([image_height_meter, image_width_meter]) / config.grid_resolution).astype(
-            np.int64
-        )
+        label_image_shape = np.ceil(
+            np.array([image_height_meter, image_width_meter]) / config.grid_resolution
+        ).astype(np.int64)
         label_image = np.zeros(label_image_shape, dtype=np.int64)
 
         valid_mask = np.logical_and(
-            (pixel_indices >= 0).all(axis=-1), (pixel_indices < np.flip(label_image_shape)).all(axis=-1)
+            (pixel_indices >= 0).all(axis=-1),
+            (pixel_indices < np.flip(label_image_shape)).all(axis=-1),
         )
 
         pixel_indices = pixel_indices[valid_mask]
         valid_point_cloud = point_cloud[valid_mask]
 
-        unique_pixel_indices, inverse_indices = np.unique(pixel_indices, axis=0, return_inverse=True)
+        unique_pixel_indices, inverse_indices = np.unique(
+            pixel_indices, axis=0, return_inverse=True
+        )
 
         terrain_classification = cloth_simulation_filtering(
             valid_point_cloud.xyz(),
@@ -116,18 +136,26 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
         )
 
         max_height, max_indices = scatter_max(
-            torch.from_numpy(dist_to_dtm), torch.from_numpy(inverse_indices).long(), dim=-1
+            torch.from_numpy(dist_to_dtm),
+            torch.from_numpy(inverse_indices).long(),
+            dim=-1,
         )
 
-        instance_ids = valid_point_cloud["instance_id_prediction"].to_numpy()[max_indices.cpu().numpy()]
+        instance_ids = valid_point_cloud["instance_id_prediction"].to_numpy()[
+            max_indices.cpu().numpy()
+        ]
 
         instance_ids[max_height.cpu().numpy() < config.min_tree_height] = 0
 
-        label_image[unique_pixel_indices[:, 1], unique_pixel_indices[:, 0]] = instance_ids
+        label_image[unique_pixel_indices[:, 1], unique_pixel_indices[:, 0]] = (
+            instance_ids
+        )
         # image coordinate system starts in upper left corner and not in lower left
         label_image = np.flip(label_image, axis=0)
 
-        label_image = modal(label_image.astype(np.uint16), footprint=np.ones((3, 3), dtype=np.uint16))
+        label_image = modal(
+            label_image.astype(np.uint16), footprint=np.ones((3, 3), dtype=np.uint16)
+        )
 
         transform = from_origin(
             west=tile_upper_left_corner[0],
@@ -175,14 +203,24 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
             height_meter = height * pixel_size[1]
 
             if width_meter < config.min_bounding_box_width:
-                print(f"Skipping bounding box {instance_id} with width of {np.round(width_meter, 2)} m.")
+                print(
+                    f"Skipping bounding box {instance_id} with width of {np.round(width_meter, 2)} m."
+                )
                 continue
             if height_meter < config.min_bounding_box_width:
-                print(f"Skipping bounding box {instance_id} with width of {np.round(height_meter, 2)} m.")
+                print(
+                    f"Skipping bounding box {instance_id} with width of {np.round(height_meter, 2)} m."
+                )
                 continue
 
             bounding_box = [x_min, y_min, width, height]
-            annotation = {"id": next_id, "image_id": 0, "category_id": 0, "iscrowd": 0, "bbox": bounding_box}
+            annotation = {
+                "id": next_id,
+                "image_id": 0,
+                "category_id": 0,
+                "iscrowd": 0,
+                "bbox": bounding_box,
+            }
             annotation["segmentation"] = coco_bbox_to_polygon(bounding_box)
             annotations.append(annotation)
             next_id += 1
@@ -190,14 +228,24 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
         # in the image files used in our dataset, the capture date is encoded in the file name
         if len(image_path.stem) >= 8 and (image_path.stem[:8]).isnumeric():
             date_prefix = Path(image_path).stem[:8]
-            image_capture_date = f"{date_prefix[:4]}-{date_prefix[4:6]}-{date_prefix[6:]}"
+            image_capture_date = (
+                f"{date_prefix[:4]}-{date_prefix[4:6]}-{date_prefix[6:]}"
+            )
         else:
             image_capture_date = ""
 
         coco_json = {
-            "info": {"year": "2024", "version": "1.0.0", "date_created": datetime.today().strftime("%Y-%m-%d")},
+            "info": {
+                "year": "2024",
+                "version": "1.0.0",
+                "date_created": datetime.today().strftime("%Y-%m-%d"),
+            },
             "licenses": [
-                {"id": 0, "name": "Attribution License", "url": "https://creativecommons.org/licenses/by/4.0/"}
+                {
+                    "id": 0,
+                    "name": "Attribution License",
+                    "url": "https://creativecommons.org/licenses/by/4.0/",
+                }
             ],
             "images": [
                 {
diff --git a/src/deepforest_finetuning/preprocessing/_rescale_images.py b/src/deepforest_finetuning/preprocessing/_rescale_images.py
index 8daab5d..54cc8f9 100644
--- a/src/deepforest_finetuning/preprocessing/_rescale_images.py
+++ b/src/deepforest_finetuning/preprocessing/_rescale_images.py
@@ -20,14 +20,22 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
     """
 
     if len(config.output_folders) != len(config.target_resolutions):
-        raise ValueError("The number of output folders and target resolutions must be the same.")
+        raise ValueError(
+            "The number of output folders and target resolutions must be the same."
+        )
 
     base_dir = Path(config.base_dir)
 
-    for output_folder, target_resolution in zip(config.output_folders, config.target_resolutions):
+    for output_folder, target_resolution in zip(
+        config.output_folders, config.target_resolutions
+    ):
         if isinstance(config.input_images, str):
             input_folder = base_dir / config.input_images
-            image_files = [input_folder / file for file in os.listdir(input_folder) if file.endswith(".tif")]
+            image_files = [
+                input_folder / file
+                for file in os.listdir(input_folder)
+                if file.endswith(".tif")
+            ]
         else:
             image_files = [base_dir / file for file in config.input_images]
 
@@ -39,11 +47,20 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
 
             with rasterio.open(original_image_path) as src:
                 transform, width, height = calculate_default_transform(
-                    src.crs, src.crs, src.width, src.height, *src.bounds, resolution=target_resolution
+                    src.crs,
+                    src.crs,
+                    src.width,
+                    src.height,
+                    *src.bounds,
+                    resolution=target_resolution,
                 )
 
-                input_pixel_size = np.abs(np.array([src.transform[0], src.transform[4]], dtype=np.float64))
-                target_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
+                input_pixel_size = np.abs(
+                    np.array([src.transform[0], src.transform[4]], dtype=np.float64)
+                )
+                target_pixel_size = np.abs(
+                    np.array([transform[0], transform[4]], dtype=np.float64)
+                )
 
                 if width > src.width or height > src.height:
                     raise ValueError(
@@ -52,7 +69,9 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
                     )
 
                 kwargs = src.meta.copy()
-                kwargs.update({"transform": transform, "width": width, "height": height})
+                kwargs.update(
+                    {"transform": transform, "width": width, "height": height}
+                )
 
                 with rasterio.open(target_image_path, "w", **kwargs) as dst:
                     for i in range(1, src.count + 1):  # reproject each channel
@@ -69,10 +88,16 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
             for folder in config.input_label_folders:
                 input_label_folder = base_dir / folder
 
-                label_output_folder = base_dir / output_folder.replace("images", Path(folder).stem)
+                label_output_folder = base_dir / output_folder.replace(
+                    "images", Path(folder).stem
+                )
                 label_output_folder.mkdir(exist_ok=True, parents=True)
 
-                label_subfolders = [x for x in os.listdir(input_label_folder) if os.path.isdir(input_label_folder / x)]
+                label_subfolders = [
+                    x
+                    for x in os.listdir(input_label_folder)
+                    if os.path.isdir(input_label_folder / x)
+                ]
                 for label_subfolder in label_subfolders:
                     label_file_name = f"{original_image_path.stem}_coco.json"
                     label_file = input_label_folder / label_subfolder / label_file_name
@@ -80,13 +105,19 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
                     if not label_file.exists():
                         continue
 
-                    target_label_path = label_output_folder / label_subfolder / label_file_name
+                    target_label_path = (
+                        label_output_folder / label_subfolder / label_file_name
+                    )
                     target_label_path.parent.mkdir(exist_ok=True, parents=True)
 
                     with open(label_file, "r", encoding="utf-8") as f:
                         coco_json = json.load(f)
 
-                    coco_json = rescale_coco_json(coco_json, target_image_path, source_image_path=original_image_path)
+                    coco_json = rescale_coco_json(
+                        coco_json,
+                        target_image_path,
+                        source_image_path=original_image_path,
+                    )
 
                     with open(target_label_path, "w", encoding="utf-8") as f:
                         json.dump(coco_json, f, indent=4)
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 8ecdba0..6169e36 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -42,12 +42,17 @@ def get_transform(augment: bool, seed: Optional[int] = None):
     if augment:
         transform = A.Compose(
             [A.HorizontalFlip(p=0.5), ToTensorV2()],
-            bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
+            bbox_params=A.BboxParams(
+                format="pascal_voc", label_fields=["category_ids"]
+            ),
         )
 
     else:
         transform = A.Compose(
-            [ToTensorV2()], bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"])
+            [ToTensorV2()],
+            bbox_params=A.BboxParams(
+                format="pascal_voc", label_fields=["category_ids"]
+            ),
         )
 
     return transform
@@ -85,7 +90,9 @@ def split_images_into_patches(
         output_dir = Path(output_dir)
 
     for image_file in annotations["image_path"].unique():
-        image_annotations = annotations.loc[annotations["image_path"] == image_file].copy()
+        image_annotations = annotations.loc[
+            annotations["image_path"] == image_file
+        ].copy()
         image_annotations["label"] = "Tree"
 
         _ = preprocess.split_raster(
@@ -133,7 +140,6 @@ def _process_annotation_paths(base_dir: Path, annotation_files: List[str]) -> Li
             processed_paths.append(file_path)
             print(f"INFO: Using annotation file {file_path}.")
 
-
     return processed_paths
 
 
@@ -165,14 +171,18 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
         eval_model = copy.deepcopy(pl_module)
         eval_trainer = copy.deepcopy(trainer)
         eval_model.trainer = eval_trainer
-        
+
         # Keep a reference to the original trainer for logging metrics
         original_trainer = trainer
-        
+
         # evaluate on training and test set
-        processed_train_files = _process_annotation_paths(self._base_dir, self._config.train_annotation_files)
-        processed_test_files = _process_annotation_paths(self._base_dir, self._config.test_annotation_files)
-        
+        processed_train_files = _process_annotation_paths(
+            self._base_dir, self._config.train_annotation_files
+        )
+        processed_test_files = _process_annotation_paths(
+            self._base_dir, self._config.test_annotation_files
+        )
+
         for prefix, annotation_files in [
             ("train", processed_train_files),
             ("test", processed_test_files),
@@ -180,7 +190,9 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
             image_files = []
             annotations = []
             for file_path in annotation_files:
-                current_annotations = utilities.read_file(str(self._base_dir / file_path), label="Tree")
+                current_annotations = utilities.read_file(
+                    str(self._base_dir / file_path), label="Tree"
+                )
                 annotations.append(current_annotations)
 
                 image_files.extend(
@@ -193,9 +205,13 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
 
             export_config = copy.deepcopy(self._config.prediction_export)
             export_config.output_folder = str(
-                self._base_dir / export_config.output_folder / f"{trainer.current_epoch + 1}_epochs"
+                self._base_dir
+                / export_config.output_folder
+                / f"{trainer.current_epoch + 1}_epochs"
+            )
+            export_config.output_file_name = (
+                f"{prefix}_predictions_seed_{self._seed}.csv"
             )
-            export_config.output_file_name = f"{prefix}_predictions_seed_{self._seed}.csv"
 
             run_prediction(
                 eval_model,
@@ -207,7 +223,11 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
             )
 
             prediction = utilities.read_file(
-                str(self._base_dir / export_config.output_folder / f"{prefix}_predictions_seed_{self._seed}.csv"),
+                str(
+                    self._base_dir
+                    / export_config.output_folder
+                    / f"{prefix}_predictions_seed_{self._seed}.csv"
+                ),
                 label="Tree",
             )
 
@@ -223,30 +243,37 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
                 self._config.iou_threshold,
                 metrics_file,
             )
-            
+
             # Log metrics to Lightning logger for each prefix (train/test)
             # This makes them available for callbacks like EarlyStopping
             for metric_name, metric_value in metrics.items():
-                # Log with trainer
-                original_trainer.logger.log_metrics(
-                    {f"{prefix}_{metric_name}": metric_value}, 
-                    step=original_trainer.current_epoch
-                )
-                
+                # Log with trainer if logger is available
+                if original_trainer.logger is not None:
+                    original_trainer.logger.log_metrics(
+                        {f"{prefix}_{metric_name}": metric_value},
+                        step=original_trainer.current_epoch,
+                    )
+
                 # For test metrics, also log them as validation metrics for EarlyStopping
                 if prefix == "test":
                     if metric_name == "f1":
                         # Log F1 as val_f1 for early stopping (maximize)
                         # Access callback_metrics directly on the original trainer
-                        original_trainer.callback_metrics[f"val_{metric_name}"] = torch.tensor(metric_value)
-                    
+                        original_trainer.callback_metrics[f"val_{metric_name}"] = (
+                            torch.tensor(metric_value)
+                        )
+
                     # Always log inverse F1 as val_loss for compatibility with default early stopping
                     if metric_name == "f1":
                         # For loss, use 1-F1 as val_loss (minimize is better)
-                        original_trainer.callback_metrics["val_loss"] = torch.tensor(1.0 - metric_value)
+                        original_trainer.callback_metrics["val_loss"] = torch.tensor(
+                            1.0 - metric_value
+                        )
 
 
-def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-many-statements
+def finetuning(
+    config: TrainingConfig,
+):  # pylint: disable=too-many-locals, too-many-statements
     """Fine-tunes the DeepForest model."""
 
     torch.set_float32_matmul_precision(config.float32_matmul_precision)
@@ -262,13 +289,20 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
     extracted_train_annotation_files = config.train_annotation_files
 
     # Process annotation file paths for train and pretraining (if available)
-    processed_train_files = _process_annotation_paths(base_dir, config.train_annotation_files)
-    
+    processed_train_files = _process_annotation_paths(
+        base_dir, config.train_annotation_files
+    )
+
     print(processed_train_files)
 
     splitting_configs = [("train", processed_train_files)]
-    if config.pretrain_annotation_files is not None and len(config.pretrain_annotation_files) > 0:
-        processed_pretrain_files = _process_annotation_paths(base_dir, config.pretrain_annotation_files)
+    if (
+        config.pretrain_annotation_files is not None
+        and len(config.pretrain_annotation_files) > 0
+    ):
+        processed_pretrain_files = _process_annotation_paths(
+            base_dir, config.pretrain_annotation_files
+        )
         splitting_configs.append(("pretraining", processed_pretrain_files))
 
     print(f"INFO: Found {len(splitting_configs)} annotation configurations to process.")
@@ -293,9 +327,11 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
         )
 
     print("\nStarting training ...")
-    
+
     if config.early_stopping_patience is not None:
-        print(f"Early stopping enabled with patience of {config.early_stopping_patience} epochs")
+        print(
+            f"Early stopping enabled with patience of {config.early_stopping_patience} epochs"
+        )
 
     for seed in config.seeds:
         # set seeds for reproducibility
@@ -305,7 +341,9 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
             print(f"INFO: Training for {config.epochs} epochs with seed {seed}...")
 
             # load model
-            model = deepforest_main.deepforest(transforms=partial(get_transform, seed=seed))
+            model = deepforest_main.deepforest(
+                transforms=partial(get_transform, seed=seed)
+            )
             model.use_release()
 
             # copy config to avoid overwriting
@@ -321,23 +359,30 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
             model.config["save-snapshot"] = False
 
             if "pretraining" in preprocessed_annotation_files:
-                model.config["train"]["csv_file"] = preprocessed_annotation_files["pretraining"]
-                model.config["train"]["root_dir"] = preprocessed_image_folders["pretraining"]
-                logger = CSVLogger(config.log_dir, name=f"{config.epochs}_epochs_seed_{seed}_pretraining")
-                
+                model.config["train"]["csv_file"] = preprocessed_annotation_files[
+                    "pretraining"
+                ]
+                model.config["train"]["root_dir"] = preprocessed_image_folders[
+                    "pretraining"
+                ]
+                logger = CSVLogger(
+                    config.log_dir,
+                    name=f"{config.epochs}_epochs_seed_{seed}_pretraining",
+                )
+
                 # Add pretraining callbacks
                 pretraining_callbacks = []
-                if config.early_stopping:
+                if config.early_stopping_patience is not None:
                     pretraining_callbacks.append(
                         EarlyStopping(
-                            monitor=config.early_stopping_monitor,
-                            min_delta=config.early_stopping_min_delta,
+                            monitor="f1",  # Use F1 score for early stopping
+                            min_delta=0.0,  # Minimum change to qualify as an improvement
                             patience=config.early_stopping_patience,
                             verbose=True,
-                            mode=config.early_stopping_mode,
+                            mode="max",  # Higher F1 is better
                         )
                     )
-                
+
                 model.create_trainer(
                     precision=config.precision if torch.cuda.is_available() else 32,
                     log_every_n_steps=1,
@@ -351,7 +396,9 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
             model.config["train"]["lr"] = current_config.learning_rate
             model.config["train"]["csv_file"] = preprocessed_annotation_files["train"]
             model.config["train"]["root_dir"] = preprocessed_image_folders["train"]
-            logger = CSVLogger(config.log_dir, name=f"{config.epochs}_epochs_seed_{seed}")
+            logger = CSVLogger(
+                config.log_dir, name=f"{config.epochs}_epochs_seed_{seed}"
+            )
 
             callbacks: List[Callback] = [EvaluationCallBack(config, seed)]
             if config.checkpoint_dir is not None:
@@ -364,7 +411,7 @@ def finetuning(config: TrainingConfig):  # pylint: disable=too-many-locals, too-
                         enable_version_counter=False,
                     )
                 )
-                
+
             # Add early stopping callback if patience is set
             if config.early_stopping_patience is not None:
                 # We provide both val_loss and val_f1 - let's use val_f1 for early stopping (higher is better)
diff --git a/src/deepforest_finetuning/utils/_annotations_to_coco.py b/src/deepforest_finetuning/utils/_annotations_to_coco.py
index 84365a9..aee0273 100644
--- a/src/deepforest_finetuning/utils/_annotations_to_coco.py
+++ b/src/deepforest_finetuning/utils/_annotations_to_coco.py
@@ -11,7 +11,10 @@
 
 
 def annotations_to_coco(
-    annotations: pd.DataFrame, image_width: int, image_height: int, capture_date: Optional[str] = None
+    annotations: pd.DataFrame,
+    image_width: int,
+    image_height: int,
+    capture_date: Optional[str] = None,
 ) -> Dict[str, Any]:
     """
     Converts DeepForest annotations into COCO format.
@@ -72,8 +75,18 @@ def annotations_to_coco(
         next_id += 1
 
     coco_json = {
-        "info": {"year": "2024", "version": "1.0.0", "date_created": datetime.today().strftime("%Y-%m-%d")},
-        "licenses": [{"id": 0, "name": "Attribution License", "url": "https://creativecommons.org/licenses/by/4.0/"}],
+        "info": {
+            "year": "2024",
+            "version": "1.0.0",
+            "date_created": datetime.today().strftime("%Y-%m-%d"),
+        },
+        "licenses": [
+            {
+                "id": 0,
+                "name": "Attribution License",
+                "url": "https://creativecommons.org/licenses/by/4.0/",
+            }
+        ],
         "images": [
             {
                 "id": 0,
@@ -85,7 +98,8 @@ def annotations_to_coco(
         ],
         "annotations": coco_annotations,
         "categories": [
-            {"id": idx, "name": category, "supercategory": category} for category, idx in category_to_id.items()
+            {"id": idx, "name": category, "supercategory": category}
+            for category, idx in category_to_id.items()
         ],
     }
 
diff --git a/src/deepforest_finetuning/utils/_rescale_coco_json.py b/src/deepforest_finetuning/utils/_rescale_coco_json.py
index 48ca6bd..bf84ca9 100644
--- a/src/deepforest_finetuning/utils/_rescale_coco_json.py
+++ b/src/deepforest_finetuning/utils/_rescale_coco_json.py
@@ -39,7 +39,9 @@ def rescale_coco_json(
     """
 
     if source_image_path is None and source_image_shape is None:
-        raise ValueError("Either source_image_path or source_image_shape must not be None.")
+        raise ValueError(
+            "Either source_image_path or source_image_shape must not be None."
+        )
 
     assert len(coco_json["images"]) == 1
     coco_json = deepcopy(coco_json)
@@ -50,22 +52,32 @@ def rescale_coco_json(
         target_height = image.height
         target_image_shape = np.array([target_height, target_width], dtype=np.int32)
 
-        target_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
+        target_pixel_size = np.abs(
+            np.array([transform[0], transform[4]], dtype=np.float64)
+        )
 
     if source_image_path is not None:
         with rasterio.open(source_image_path) as image:
             transform = image.transform
-            source_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
+            source_pixel_size = np.abs(
+                np.array([transform[0], transform[4]], dtype=np.float64)
+            )
     else:
         source_image_shape = cast(npt.NDArray, source_image_shape)
-        source_pixel_size = (target_image_shape / source_image_shape) * target_pixel_size
+        source_pixel_size = (
+            target_image_shape / source_image_shape
+        ) * target_pixel_size
         source_pixel_size = np.flip(source_pixel_size)
 
     annotations = []
     for annotation in coco_json["annotations"]:
         bounding_box = np.array(annotation["bbox"])
-        bounding_box[[0, 2]] = bounding_box[[0, 2]] * source_pixel_size[0] / target_pixel_size[0]
-        bounding_box[[1, 3]] = bounding_box[[1, 3]] * source_pixel_size[1] / target_pixel_size[1]
+        bounding_box[[0, 2]] = (
+            bounding_box[[0, 2]] * source_pixel_size[0] / target_pixel_size[0]
+        )
+        bounding_box[[1, 3]] = (
+            bounding_box[[1, 3]] * source_pixel_size[1] / target_pixel_size[1]
+        )
         annotation["bbox"] = bounding_box.astype(int).tolist()
         annotation["segmentation"] = coco_bbox_to_polygon(annotation["bbox"])
         annotations.append(annotation)

From 62dfee592142b7980f87dc1171dfbbc1c0e41c78 Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Wed, 4 Jun 2025 13:50:52 +0200
Subject: [PATCH 06/24] get_transform is now without seed because transform
 itself is too

---
 .../prediction/_prediction_dataset.py                    | 2 +-
 .../preprocessing/_project_point_cloud_labels.py         | 2 +-
 src/deepforest_finetuning/training/_finetuning.py        | 9 ++++-----
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/deepforest_finetuning/prediction/_prediction_dataset.py b/src/deepforest_finetuning/prediction/_prediction_dataset.py
index d864b0c..9b792cb 100644
--- a/src/deepforest_finetuning/prediction/_prediction_dataset.py
+++ b/src/deepforest_finetuning/prediction/_prediction_dataset.py
@@ -4,11 +4,11 @@
 
 from pathlib import Path
 from typing import List, Optional
+import os
 
 import numpy as np
 import numpy.typing as npt
 from PIL import Image
-import os
 from tifffile import imread
 from torch.utils.data import Dataset
 
diff --git a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
index 0cac1a8..5e4aa63 100644
--- a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
@@ -18,7 +18,7 @@
 from rasterio.transform import from_origin
 from skimage.filters.rank import modal
 import torch
-from torch_scatter import scatter_max
+from torch_scatter import scatter_max  # pylint: disable=import-error
 
 from deepforest_finetuning.utils import coco_bbox_to_polygon
 from deepforest_finetuning.config import PointCloudLabelProjectionConfig
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 6169e36..a005f67 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -27,13 +27,12 @@
 from deepforest_finetuning.prediction import prediction as run_prediction
 
 
-def get_transform(augment: bool, seed: Optional[int] = None):
+def get_transform(augment: bool):
     """
     Albumentations transformation of bounding boxes.
 
     Args:
         augment: Whether to apply data augmentations.
-        seed: Random seed for data augmentations to ensure reproducibility. Defaults to :code:`None`.
 
     Returns:
         Transforms.
@@ -158,7 +157,9 @@ def __init__(self, config: TrainingConfig, seed: int):
         self._base_dir = Path(config.base_dir)
         self._seed = seed
 
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+    def on_train_epoch_end(
+        self, trainer: Trainer, pl_module: LightningModule
+    ):  # pylint: disable=too-many-locals
         """
         Hook that evaluates the model after each training epoch.
 
@@ -286,8 +287,6 @@ def finetuning(
     preprocessed_image_folders = {}
     preprocessed_annotation_files = {}
 
-    extracted_train_annotation_files = config.train_annotation_files
-
     # Process annotation file paths for train and pretraining (if available)
     processed_train_files = _process_annotation_paths(
         base_dir, config.train_annotation_files

From 9d6d7ddb19f7d3982ac86e6186e49a1803757342 Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Wed, 4 Jun 2025 14:11:44 +0200
Subject: [PATCH 07/24] made everything ready for CI/CD

---
 scripts/preprocessing.py                      |  4 +-
 .../evaluation/_evaluate.py                   | 10 +--
 .../prediction/_prediction.py                 | 12 +--
 .../preprocessing/_filter_labels.py           | 18 +---
 .../_preprocess_manually_corrected_labels.py  | 38 ++-------
 .../_project_point_cloud_labels.py            | 73 +++++-----------
 .../preprocessing/_rescale_images.py          | 40 ++-------
 .../training/_finetuning.py                   | 85 +++++--------------
 .../utils/_annotations_to_coco.py             |  3 +-
 .../utils/_rescale_coco_json.py               | 24 ++----
 10 files changed, 79 insertions(+), 228 deletions(-)

diff --git a/scripts/preprocessing.py b/scripts/preprocessing.py
index c256ce7..12575a4 100644
--- a/scripts/preprocessing.py
+++ b/scripts/preprocessing.py
@@ -20,9 +20,7 @@
 from deepforest_finetuning.utils import load_config
 
 
-def preprocessing_step(
-    config_path: str, config_type: Type, script_function: Callable[[Any], None]
-):
+def preprocessing_step(config_path: str, config_type: Type, script_function: Callable[[Any], None]):
     """
     Loads the specified config file, parses it based on the given configuration type and then calls the script function
     with the given configuration.
diff --git a/src/deepforest_finetuning/evaluation/_evaluate.py b/src/deepforest_finetuning/evaluation/_evaluate.py
index f044459..29eeb38 100644
--- a/src/deepforest_finetuning/evaluation/_evaluate.py
+++ b/src/deepforest_finetuning/evaluation/_evaluate.py
@@ -41,15 +41,9 @@ def evaluate(
 
     results["precision"] = results.pop("box_precision")
     results["recall"] = results.pop("box_recall")
-    results["f1"] = (
-        2
-        * (results["precision"] * results["recall"])
-        / (results["precision"] + results["recall"])
-    )
+    results["f1"] = 2 * (results["precision"] * results["recall"]) / (results["precision"] + results["recall"])
 
-    print(
-        f"Precision:\t{results['precision']}\nRecall:\t\t{results['recall']}\nF1:\t\t{results['f1']}"
-    )
+    print(f"Precision:\t{results['precision']}\nRecall:\t\t{results['recall']}\nF1:\t\t{results['f1']}")
 
     metrics = []
     metrics.append({"metric": "precision", "score": results["precision"]})
diff --git a/src/deepforest_finetuning/prediction/_prediction.py b/src/deepforest_finetuning/prediction/_prediction.py
index ed43e77..de697cd 100644
--- a/src/deepforest_finetuning/prediction/_prediction.py
+++ b/src/deepforest_finetuning/prediction/_prediction.py
@@ -37,13 +37,9 @@ def prediction(
 
     if predict_tile:
         if patch_size is None:
-            raise ValueError(
-                "Patch size must be specified when predict_tile is set to True."
-            )
+            raise ValueError("Patch size must be specified when predict_tile is set to True.")
         if patch_overlap is None:
-            raise ValueError(
-                "Patch overlap must be specified when predict_tile is set to True."
-            )
+            raise ValueError("Patch overlap must be specified when predict_tile is set to True.")
 
     print("\nLoading dataset and model ...")
 
@@ -74,9 +70,7 @@ def prediction(
 
         export_labels(
             pred,
-            export_path=(Path(export_config.output_folder) / image_name).with_suffix(
-                ".csv"
-            ),
+            export_path=(Path(export_config.output_folder) / image_name).with_suffix(".csv"),
             column_order=export_config.column_order,
             index_as_label_suffix=export_config.index_as_label_suffix,
             sort_by=export_config.sort_by,
diff --git a/src/deepforest_finetuning/preprocessing/_filter_labels.py b/src/deepforest_finetuning/preprocessing/_filter_labels.py
index dec32eb..88ea6c3 100644
--- a/src/deepforest_finetuning/preprocessing/_filter_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_filter_labels.py
@@ -14,9 +14,7 @@
 from deepforest_finetuning.config import LabelFilteringConfig
 
 
-def filter_bounding_boxed_with_size_based_nms(
-    coco_json: Dict[str, Any], iou_threshold: float
-) -> Dict[str, Any]:
+def filter_bounding_boxed_with_size_based_nms(coco_json: Dict[str, Any], iou_threshold: float) -> Dict[str, Any]:
     """
     Applies non-maximum suppression to the bounding boxes, using the box sizes as scores.
 
@@ -65,11 +63,7 @@ def filter_labels(config: LabelFilteringConfig):
     base_dir = Path(config.base_dir)
     label_folder = base_dir / config.input_label_folder
 
-    subfolders = [
-        file
-        for file in os.listdir(label_folder)
-        if os.path.isdir(os.path.join(label_folder, file))
-    ]
+    subfolders = [file for file in os.listdir(label_folder) if os.path.isdir(os.path.join(label_folder, file))]
 
     for subfolder in subfolders:
         for file in os.listdir(label_folder / subfolder):
@@ -79,13 +73,9 @@ def filter_labels(config: LabelFilteringConfig):
                     coco_json = json.load(f)
 
                 assert len(coco_json["images"]) == 1
-                coco_json = filter_bounding_boxed_with_size_based_nms(
-                    coco_json, iou_threshold=config.iou_threshold
-                )
+                coco_json = filter_bounding_boxed_with_size_based_nms(coco_json, iou_threshold=config.iou_threshold)
 
-                output_file_path = (
-                    base_dir / config.output_label_folder / subfolder / file
-                )
+                output_file_path = base_dir / config.output_label_folder / subfolder / file
                 output_file_path.parent.mkdir(exist_ok=True, parents=True)
                 with open(output_file_path, "w", encoding="utf-8") as f:
                     json.dump(coco_json, f, indent=4)
diff --git a/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py b/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py
index e6ac704..877eb11 100644
--- a/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_preprocess_manually_corrected_labels.py
@@ -57,27 +57,21 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
 
             target_image_path = base_dir / config.input_image_folder / image_path
 
-            coco_json = annotations_to_coco(
-                annotations, image_width, image_height, capture_date=capture_date
-            )
+            coco_json = annotations_to_coco(annotations, image_width, image_height, capture_date=capture_date)
             coco_json = rescale_coco_json(
                 coco_json,
                 target_image_path,
                 source_image_shape=np.array([image_height, image_width]),
             )
 
-            output_file_path = (
-                output_label_folder / f"{Path(image_path).stem}_coco"
-            ).with_suffix(".json")
+            output_file_path = (output_label_folder / f"{Path(image_path).stem}_coco").with_suffix(".json")
             with open(output_file_path, "w", encoding="utf-8") as f:
                 json.dump(coco_json, f, indent=4)
 
     # s1_p1_ext_mc contains labels for a larger tile
     # we crop the part for which we also have fully manually and automatically created labels
     if "s1_p1_ext_mc_coco.json" in os.listdir(output_label_folder):
-        with open(
-            output_label_folder / "s1_p1_ext_mc_coco.json", "r", encoding="utf-8"
-        ) as f:
+        with open(output_label_folder / "s1_p1_ext_mc_coco.json", "r", encoding="utf-8") as f:
             coco_json = json.load(f)
 
         source_image_path = base_dir / config.input_image_folder / "s1_p1_ext_mc.tif"
@@ -86,9 +80,7 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
         with rasterio.open(source_image_path) as image:
             transform = image.transform
 
-            src_pixel_size = np.abs(
-                np.array([transform[0], transform[4]], dtype=np.float64)
-            )
+            src_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
 
             src_top_left = np.array([transform.c, transform.f], dtype=np.float64)
             src_bottom_left = src_top_left.copy()
@@ -101,9 +93,7 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
             target_width = image.width
             target_height = image.height
 
-            target_pixel_size = np.abs(
-                np.array([transform[0], transform[4]], dtype=np.float64)
-            )
+            target_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
 
             target_top_left = np.array([transform.c, transform.f], dtype=np.float64)
             target_bottom_left = target_top_left.copy()
@@ -114,15 +104,9 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
         clipped_annotations = []
         for annotation in coco_json["annotations"]:
             x_min_meter = src_top_left[0] + annotation["bbox"][0] * src_pixel_size[0]
-            x_max_meter = (
-                src_top_left[0]
-                + (annotation["bbox"][0] + annotation["bbox"][2]) * src_pixel_size[0]
-            )
+            x_max_meter = src_top_left[0] + (annotation["bbox"][0] + annotation["bbox"][2]) * src_pixel_size[0]
             y_min_meter = src_top_left[1] - annotation["bbox"][1] * src_pixel_size[1]
-            y_max_meter = (
-                src_top_left[1]
-                - (annotation["bbox"][1] + annotation["bbox"][3]) * src_pixel_size[1]
-            )
+            y_max_meter = src_top_left[1] - (annotation["bbox"][1] + annotation["bbox"][3]) * src_pixel_size[1]
 
             if (
                 x_max_meter < target_top_left[0]
@@ -132,16 +116,12 @@ def preprocess_manually_corrected_labels(  # pylint: disable=too-many-locals, to
             ):
                 continue
 
-            clipped_x_min = max(
-                int((x_min_meter - target_top_left[0]) / target_pixel_size[0]), 0
-            )
+            clipped_x_min = max(int((x_min_meter - target_top_left[0]) / target_pixel_size[0]), 0)
             clipped_x_max = min(
                 int((x_max_meter - target_top_left[0]) / target_pixel_size[0]),
                 target_width,
             )
-            clipped_y_min = max(
-                int((target_top_left[1] - y_min_meter) / target_pixel_size[1]), 0
-            )
+            clipped_y_min = max(int((target_top_left[1] - y_min_meter) / target_pixel_size[1]), 0)
             clipped_y_max = min(
                 int((target_top_left[1] - y_max_meter) / target_pixel_size[1]),
                 target_height,
diff --git a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
index 5e4aa63..10895ff 100644
--- a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
@@ -32,15 +32,11 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
     if len(config.point_cloud_paths) != len(config.image_paths):
         raise ValueError("point_cloud_paths and image_paths must have the same length.")
     if len(config.point_cloud_paths) != len(config.label_json_output_paths):
-        raise ValueError(
-            "point_cloud_paths and label_json_output_paths must have the same length."
-        )
-    if config.label_image_output_paths is not None and len(
-        config.label_json_output_paths
-    ) != len(config.label_image_output_paths):
-        raise ValueError(
-            "label_json_output_paths and label_image_output_paths must have the same length."
-        )
+        raise ValueError("point_cloud_paths and label_json_output_paths must have the same length.")
+    if config.label_image_output_paths is not None and len(config.label_json_output_paths) != len(
+        config.label_image_output_paths
+    ):
+        raise ValueError("label_json_output_paths and label_image_output_paths must have the same length.")
 
     base_dir = Path(config.base_dir)
 
@@ -56,15 +52,11 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
             image_height = image.height
             crs = image.crs
 
-            pixel_size = np.abs(
-                np.array([transform[0], transform[4]], dtype=np.float64)
-            )
+            pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
             image_width_meter = image_width * pixel_size[1]
             image_height_meter = image_height * pixel_size[0]
 
-            tile_upper_left_corner = np.array(
-                [transform.c, transform.f], dtype=np.float64
-            )
+            tile_upper_left_corner = np.array([transform.c, transform.f], dtype=np.float64)
             tile_lower_left_corner = tile_upper_left_corner.copy()
             tile_lower_left_corner[1] -= image.height * pixel_size[1]
 
@@ -75,9 +67,7 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
 
         point_cloud = read(base_dir / point_cloud_path)
 
-        print(
-            f"Loaded point cloud {Path(point_cloud_path).name} with {len(point_cloud)} points."
-        )
+        print(f"Loaded point cloud {Path(point_cloud_path).name} with {len(point_cloud)} points.")
 
         point_cloud["instance_id_prediction"] = make_labels_consecutive(
             point_cloud["instance_id_prediction"].to_numpy(), ignore_id=0
@@ -85,20 +75,17 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
 
         if (
             point_cloud["instance_id_prediction"].min() != 0
-            or point_cloud["instance_id_prediction"].max()
-            != len(point_cloud["instance_id_prediction"].unique()) - 1
+            or point_cloud["instance_id_prediction"].max() != len(point_cloud["instance_id_prediction"].unique()) - 1
         ):
-            raise ValueError(
-                "The predicted instance IDs must be continuous, starting from zero."
-            )
+            raise ValueError("The predicted instance IDs must be continuous, starting from zero.")
 
-        pixel_indices = np.floor(
-            (point_cloud.xyz()[:, :2] - tile_lower_left_corner) / config.grid_resolution
-        ).astype(np.int64)
+        pixel_indices = np.floor((point_cloud.xyz()[:, :2] - tile_lower_left_corner) / config.grid_resolution).astype(
+            np.int64
+        )
 
-        label_image_shape = np.ceil(
-            np.array([image_height_meter, image_width_meter]) / config.grid_resolution
-        ).astype(np.int64)
+        label_image_shape = np.ceil(np.array([image_height_meter, image_width_meter]) / config.grid_resolution).astype(
+            np.int64
+        )
         label_image = np.zeros(label_image_shape, dtype=np.int64)
 
         valid_mask = np.logical_and(
@@ -109,9 +96,7 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
         pixel_indices = pixel_indices[valid_mask]
         valid_point_cloud = point_cloud[valid_mask]
 
-        unique_pixel_indices, inverse_indices = np.unique(
-            pixel_indices, axis=0, return_inverse=True
-        )
+        unique_pixel_indices, inverse_indices = np.unique(pixel_indices, axis=0, return_inverse=True)
 
         terrain_classification = cloth_simulation_filtering(
             valid_point_cloud.xyz(),
@@ -141,21 +126,15 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
             dim=-1,
         )
 
-        instance_ids = valid_point_cloud["instance_id_prediction"].to_numpy()[
-            max_indices.cpu().numpy()
-        ]
+        instance_ids = valid_point_cloud["instance_id_prediction"].to_numpy()[max_indices.cpu().numpy()]
 
         instance_ids[max_height.cpu().numpy() < config.min_tree_height] = 0
 
-        label_image[unique_pixel_indices[:, 1], unique_pixel_indices[:, 0]] = (
-            instance_ids
-        )
+        label_image[unique_pixel_indices[:, 1], unique_pixel_indices[:, 0]] = instance_ids
         # image coordinate system starts in upper left corner and not in lower left
         label_image = np.flip(label_image, axis=0)
 
-        label_image = modal(
-            label_image.astype(np.uint16), footprint=np.ones((3, 3), dtype=np.uint16)
-        )
+        label_image = modal(label_image.astype(np.uint16), footprint=np.ones((3, 3), dtype=np.uint16))
 
         transform = from_origin(
             west=tile_upper_left_corner[0],
@@ -203,14 +182,10 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
             height_meter = height * pixel_size[1]
 
             if width_meter < config.min_bounding_box_width:
-                print(
-                    f"Skipping bounding box {instance_id} with width of {np.round(width_meter, 2)} m."
-                )
+                print(f"Skipping bounding box {instance_id} with width of {np.round(width_meter, 2)} m.")
                 continue
             if height_meter < config.min_bounding_box_width:
-                print(
-                    f"Skipping bounding box {instance_id} with width of {np.round(height_meter, 2)} m."
-                )
+                print(f"Skipping bounding box {instance_id} with width of {np.round(height_meter, 2)} m.")
                 continue
 
             bounding_box = [x_min, y_min, width, height]
@@ -228,9 +203,7 @@ def project_point_cloud_labels(  # pylint: disable=too-many-locals, too-many-sta
         # in the image files used in our dataset, the capture date is encoded in the file name
         if len(image_path.stem) >= 8 and (image_path.stem[:8]).isnumeric():
             date_prefix = Path(image_path).stem[:8]
-            image_capture_date = (
-                f"{date_prefix[:4]}-{date_prefix[4:6]}-{date_prefix[6:]}"
-            )
+            image_capture_date = f"{date_prefix[:4]}-{date_prefix[4:6]}-{date_prefix[6:]}"
         else:
             image_capture_date = ""
 
diff --git a/src/deepforest_finetuning/preprocessing/_rescale_images.py b/src/deepforest_finetuning/preprocessing/_rescale_images.py
index 54cc8f9..5b9a7fd 100644
--- a/src/deepforest_finetuning/preprocessing/_rescale_images.py
+++ b/src/deepforest_finetuning/preprocessing/_rescale_images.py
@@ -20,22 +20,14 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
     """
 
     if len(config.output_folders) != len(config.target_resolutions):
-        raise ValueError(
-            "The number of output folders and target resolutions must be the same."
-        )
+        raise ValueError("The number of output folders and target resolutions must be the same.")
 
     base_dir = Path(config.base_dir)
 
-    for output_folder, target_resolution in zip(
-        config.output_folders, config.target_resolutions
-    ):
+    for output_folder, target_resolution in zip(config.output_folders, config.target_resolutions):
         if isinstance(config.input_images, str):
             input_folder = base_dir / config.input_images
-            image_files = [
-                input_folder / file
-                for file in os.listdir(input_folder)
-                if file.endswith(".tif")
-            ]
+            image_files = [input_folder / file for file in os.listdir(input_folder) if file.endswith(".tif")]
         else:
             image_files = [base_dir / file for file in config.input_images]
 
@@ -55,12 +47,8 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
                     resolution=target_resolution,
                 )
 
-                input_pixel_size = np.abs(
-                    np.array([src.transform[0], src.transform[4]], dtype=np.float64)
-                )
-                target_pixel_size = np.abs(
-                    np.array([transform[0], transform[4]], dtype=np.float64)
-                )
+                input_pixel_size = np.abs(np.array([src.transform[0], src.transform[4]], dtype=np.float64))
+                target_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
 
                 if width > src.width or height > src.height:
                     raise ValueError(
@@ -69,9 +57,7 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
                     )
 
                 kwargs = src.meta.copy()
-                kwargs.update(
-                    {"transform": transform, "width": width, "height": height}
-                )
+                kwargs.update({"transform": transform, "width": width, "height": height})
 
                 with rasterio.open(target_image_path, "w", **kwargs) as dst:
                     for i in range(1, src.count + 1):  # reproject each channel
@@ -88,16 +74,10 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
             for folder in config.input_label_folders:
                 input_label_folder = base_dir / folder
 
-                label_output_folder = base_dir / output_folder.replace(
-                    "images", Path(folder).stem
-                )
+                label_output_folder = base_dir / output_folder.replace("images", Path(folder).stem)
                 label_output_folder.mkdir(exist_ok=True, parents=True)
 
-                label_subfolders = [
-                    x
-                    for x in os.listdir(input_label_folder)
-                    if os.path.isdir(input_label_folder / x)
-                ]
+                label_subfolders = [x for x in os.listdir(input_label_folder) if os.path.isdir(input_label_folder / x)]
                 for label_subfolder in label_subfolders:
                     label_file_name = f"{original_image_path.stem}_coco.json"
                     label_file = input_label_folder / label_subfolder / label_file_name
@@ -105,9 +85,7 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
                     if not label_file.exists():
                         continue
 
-                    target_label_path = (
-                        label_output_folder / label_subfolder / label_file_name
-                    )
+                    target_label_path = label_output_folder / label_subfolder / label_file_name
                     target_label_path.parent.mkdir(exist_ok=True, parents=True)
 
                     with open(label_file, "r", encoding="utf-8") as f:
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index a005f67..6249263 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -7,7 +7,7 @@
 import os
 from pathlib import Path
 import shutil
-from typing import List, Optional, Union
+from typing import List, Union
 import uuid
 
 import albumentations as A
@@ -41,17 +41,13 @@ def get_transform(augment: bool):
     if augment:
         transform = A.Compose(
             [A.HorizontalFlip(p=0.5), ToTensorV2()],
-            bbox_params=A.BboxParams(
-                format="pascal_voc", label_fields=["category_ids"]
-            ),
+            bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
         )
 
     else:
         transform = A.Compose(
             [ToTensorV2()],
-            bbox_params=A.BboxParams(
-                format="pascal_voc", label_fields=["category_ids"]
-            ),
+            bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
         )
 
     return transform
@@ -89,9 +85,7 @@ def split_images_into_patches(
         output_dir = Path(output_dir)
 
     for image_file in annotations["image_path"].unique():
-        image_annotations = annotations.loc[
-            annotations["image_path"] == image_file
-        ].copy()
+        image_annotations = annotations.loc[annotations["image_path"] == image_file].copy()
         image_annotations["label"] = "Tree"
 
         _ = preprocess.split_raster(
@@ -157,9 +151,7 @@ def __init__(self, config: TrainingConfig, seed: int):
         self._base_dir = Path(config.base_dir)
         self._seed = seed
 
-    def on_train_epoch_end(
-        self, trainer: Trainer, pl_module: LightningModule
-    ):  # pylint: disable=too-many-locals
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):  # pylint: disable=too-many-locals
         """
         Hook that evaluates the model after each training epoch.
 
@@ -177,12 +169,8 @@ def on_train_epoch_end(
         original_trainer = trainer
 
         # evaluate on training and test set
-        processed_train_files = _process_annotation_paths(
-            self._base_dir, self._config.train_annotation_files
-        )
-        processed_test_files = _process_annotation_paths(
-            self._base_dir, self._config.test_annotation_files
-        )
+        processed_train_files = _process_annotation_paths(self._base_dir, self._config.train_annotation_files)
+        processed_test_files = _process_annotation_paths(self._base_dir, self._config.test_annotation_files)
 
         for prefix, annotation_files in [
             ("train", processed_train_files),
@@ -191,9 +179,7 @@ def on_train_epoch_end(
             image_files = []
             annotations = []
             for file_path in annotation_files:
-                current_annotations = utilities.read_file(
-                    str(self._base_dir / file_path), label="Tree"
-                )
+                current_annotations = utilities.read_file(str(self._base_dir / file_path), label="Tree")
                 annotations.append(current_annotations)
 
                 image_files.extend(
@@ -206,13 +192,9 @@ def on_train_epoch_end(
 
             export_config = copy.deepcopy(self._config.prediction_export)
             export_config.output_folder = str(
-                self._base_dir
-                / export_config.output_folder
-                / f"{trainer.current_epoch + 1}_epochs"
-            )
-            export_config.output_file_name = (
-                f"{prefix}_predictions_seed_{self._seed}.csv"
+                self._base_dir / export_config.output_folder / f"{trainer.current_epoch + 1}_epochs"
             )
+            export_config.output_file_name = f"{prefix}_predictions_seed_{self._seed}.csv"
 
             run_prediction(
                 eval_model,
@@ -224,11 +206,7 @@ def on_train_epoch_end(
             )
 
             prediction = utilities.read_file(
-                str(
-                    self._base_dir
-                    / export_config.output_folder
-                    / f"{prefix}_predictions_seed_{self._seed}.csv"
-                ),
+                str(self._base_dir / export_config.output_folder / f"{prefix}_predictions_seed_{self._seed}.csv"),
                 label="Tree",
             )
 
@@ -260,16 +238,12 @@ def on_train_epoch_end(
                     if metric_name == "f1":
                         # Log F1 as val_f1 for early stopping (maximize)
                         # Access callback_metrics directly on the original trainer
-                        original_trainer.callback_metrics[f"val_{metric_name}"] = (
-                            torch.tensor(metric_value)
-                        )
+                        original_trainer.callback_metrics[f"val_{metric_name}"] = torch.tensor(metric_value)
 
                     # Always log inverse F1 as val_loss for compatibility with default early stopping
                     if metric_name == "f1":
                         # For loss, use 1-F1 as val_loss (minimize is better)
-                        original_trainer.callback_metrics["val_loss"] = torch.tensor(
-                            1.0 - metric_value
-                        )
+                        original_trainer.callback_metrics["val_loss"] = torch.tensor(1.0 - metric_value)
 
 
 def finetuning(
@@ -288,20 +262,13 @@ def finetuning(
     preprocessed_annotation_files = {}
 
     # Process annotation file paths for train and pretraining (if available)
-    processed_train_files = _process_annotation_paths(
-        base_dir, config.train_annotation_files
-    )
+    processed_train_files = _process_annotation_paths(base_dir, config.train_annotation_files)
 
     print(processed_train_files)
 
     splitting_configs = [("train", processed_train_files)]
-    if (
-        config.pretrain_annotation_files is not None
-        and len(config.pretrain_annotation_files) > 0
-    ):
-        processed_pretrain_files = _process_annotation_paths(
-            base_dir, config.pretrain_annotation_files
-        )
+    if config.pretrain_annotation_files is not None and len(config.pretrain_annotation_files) > 0:
+        processed_pretrain_files = _process_annotation_paths(base_dir, config.pretrain_annotation_files)
         splitting_configs.append(("pretraining", processed_pretrain_files))
 
     print(f"INFO: Found {len(splitting_configs)} annotation configurations to process.")
@@ -328,9 +295,7 @@ def finetuning(
     print("\nStarting training ...")
 
     if config.early_stopping_patience is not None:
-        print(
-            f"Early stopping enabled with patience of {config.early_stopping_patience} epochs"
-        )
+        print(f"Early stopping enabled with patience of {config.early_stopping_patience} epochs")
 
     for seed in config.seeds:
         # set seeds for reproducibility
@@ -340,9 +305,7 @@ def finetuning(
             print(f"INFO: Training for {config.epochs} epochs with seed {seed}...")
 
             # load model
-            model = deepforest_main.deepforest(
-                transforms=partial(get_transform, seed=seed)
-            )
+            model = deepforest_main.deepforest(transforms=partial(get_transform))
             model.use_release()
 
             # copy config to avoid overwriting
@@ -358,12 +321,8 @@ def finetuning(
             model.config["save-snapshot"] = False
 
             if "pretraining" in preprocessed_annotation_files:
-                model.config["train"]["csv_file"] = preprocessed_annotation_files[
-                    "pretraining"
-                ]
-                model.config["train"]["root_dir"] = preprocessed_image_folders[
-                    "pretraining"
-                ]
+                model.config["train"]["csv_file"] = preprocessed_annotation_files["pretraining"]
+                model.config["train"]["root_dir"] = preprocessed_image_folders["pretraining"]
                 logger = CSVLogger(
                     config.log_dir,
                     name=f"{config.epochs}_epochs_seed_{seed}_pretraining",
@@ -395,9 +354,7 @@ def finetuning(
             model.config["train"]["lr"] = current_config.learning_rate
             model.config["train"]["csv_file"] = preprocessed_annotation_files["train"]
             model.config["train"]["root_dir"] = preprocessed_image_folders["train"]
-            logger = CSVLogger(
-                config.log_dir, name=f"{config.epochs}_epochs_seed_{seed}"
-            )
+            logger = CSVLogger(config.log_dir, name=f"{config.epochs}_epochs_seed_{seed}")
 
             callbacks: List[Callback] = [EvaluationCallBack(config, seed)]
             if config.checkpoint_dir is not None:
diff --git a/src/deepforest_finetuning/utils/_annotations_to_coco.py b/src/deepforest_finetuning/utils/_annotations_to_coco.py
index aee0273..0250fae 100644
--- a/src/deepforest_finetuning/utils/_annotations_to_coco.py
+++ b/src/deepforest_finetuning/utils/_annotations_to_coco.py
@@ -98,8 +98,7 @@ def annotations_to_coco(
         ],
         "annotations": coco_annotations,
         "categories": [
-            {"id": idx, "name": category, "supercategory": category}
-            for category, idx in category_to_id.items()
+            {"id": idx, "name": category, "supercategory": category} for category, idx in category_to_id.items()
         ],
     }
 
diff --git a/src/deepforest_finetuning/utils/_rescale_coco_json.py b/src/deepforest_finetuning/utils/_rescale_coco_json.py
index bf84ca9..48ca6bd 100644
--- a/src/deepforest_finetuning/utils/_rescale_coco_json.py
+++ b/src/deepforest_finetuning/utils/_rescale_coco_json.py
@@ -39,9 +39,7 @@ def rescale_coco_json(
     """
 
     if source_image_path is None and source_image_shape is None:
-        raise ValueError(
-            "Either source_image_path or source_image_shape must not be None."
-        )
+        raise ValueError("Either source_image_path or source_image_shape must not be None.")
 
     assert len(coco_json["images"]) == 1
     coco_json = deepcopy(coco_json)
@@ -52,32 +50,22 @@ def rescale_coco_json(
         target_height = image.height
         target_image_shape = np.array([target_height, target_width], dtype=np.int32)
 
-        target_pixel_size = np.abs(
-            np.array([transform[0], transform[4]], dtype=np.float64)
-        )
+        target_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
 
     if source_image_path is not None:
         with rasterio.open(source_image_path) as image:
             transform = image.transform
-            source_pixel_size = np.abs(
-                np.array([transform[0], transform[4]], dtype=np.float64)
-            )
+            source_pixel_size = np.abs(np.array([transform[0], transform[4]], dtype=np.float64))
     else:
         source_image_shape = cast(npt.NDArray, source_image_shape)
-        source_pixel_size = (
-            target_image_shape / source_image_shape
-        ) * target_pixel_size
+        source_pixel_size = (target_image_shape / source_image_shape) * target_pixel_size
         source_pixel_size = np.flip(source_pixel_size)
 
     annotations = []
     for annotation in coco_json["annotations"]:
         bounding_box = np.array(annotation["bbox"])
-        bounding_box[[0, 2]] = (
-            bounding_box[[0, 2]] * source_pixel_size[0] / target_pixel_size[0]
-        )
-        bounding_box[[1, 3]] = (
-            bounding_box[[1, 3]] * source_pixel_size[1] / target_pixel_size[1]
-        )
+        bounding_box[[0, 2]] = bounding_box[[0, 2]] * source_pixel_size[0] / target_pixel_size[0]
+        bounding_box[[1, 3]] = bounding_box[[1, 3]] * source_pixel_size[1] / target_pixel_size[1]
         annotation["bbox"] = bounding_box.astype(int).tolist()
         annotation["segmentation"] = coco_bbox_to_polygon(annotation["bbox"])
         annotations.append(annotation)

From 82def70767d4d6d0f382b2cd412c294a52c252c5 Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Wed, 4 Jun 2025 14:56:50 +0200
Subject: [PATCH 08/24] rerun CI/CD

---
 src/deepforest_finetuning/evaluation/_evaluate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/deepforest_finetuning/evaluation/_evaluate.py b/src/deepforest_finetuning/evaluation/_evaluate.py
index 29eeb38..0b156a9 100644
--- a/src/deepforest_finetuning/evaluation/_evaluate.py
+++ b/src/deepforest_finetuning/evaluation/_evaluate.py
@@ -30,7 +30,7 @@ def evaluate(
         Dictionary containing the evaluation metrics (precision, recall, f1).
     """
 
-    # ignore deprecated warnings from pandas raised by deepforest.IoU (line 113: iou_df = pd.concat(iou_df))
+    # ignore deprecated pandas warnings raised by deepforest.IoU (line 113: iou_df = pd.concat(iou_df))
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         results = evaluate_boxes(

From c568b2faf983785d4737741e57b95d93f7c8aa0a Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Thu, 5 Jun 2025 12:09:50 +0200
Subject: [PATCH 09/24] Added different target metrics and save_top_k parameter
 to early stopping

---
 README.md                                      | 12 ++++++++----
 .../5_cm/finetuning_5_cm_TreeAI_full.toml      |  3 +++
 src/deepforest_finetuning/config/_config.py    |  2 ++
 .../training/_finetuning.py                    | 18 ++++++++++++------
 4 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 65f6251..899aef8 100644
--- a/README.md
+++ b/README.md
@@ -151,6 +151,8 @@ seeds = [0, 1, 2, 3, 4]
 learning_rate = 0.0001
 checkpoint_dir = "checkpoints"
 early_stopping_patience = 2
+save_top_k = 1
+target_metric = "val_f1"
 ```
 
 This will:
@@ -253,15 +255,17 @@ seeds = [0, 1, 2, 3, 4]
 
 This will train separate models with different weight initializations and report the average performance.
 
-### Early Stopping
+### Early Stopping and Model Checkpointing
 
-To prevent overfitting, you can enable early stopping:
+To prevent overfitting, you can enable early stopping and control model checkpointing:
 
 ```toml
-early_stopping_patience = 2
+early_stopping_patience = 2  # Stop training if performance doesn't improve for this many epochs
+save_top_k = 1  # Save the top k best models based on the target metric
+target_metric = "val_f1"  # Metric to monitor for early stopping and checkpointing
 ```
 
-This will stop training if the validation performance doesn't improve for the specified number of epochs.
+The `mode` (min/max) is automatically inferred from the metric name. Metrics containing "loss" use "min" mode (lower is better), all others use "max" mode (higher is better).
 
 ## License
 
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml b/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
index 2451120..256f662 100644
--- a/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
+++ b/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
@@ -14,6 +14,9 @@ seeds = [0, 1, 2, 3, 4]
 learning_rate = 0.0001
 checkpoint_dir = "checkpoints"
 early_stopping_patience = 2  # Set to the number of epochs to wait before stopping, or remove to disable early stopping
+save_top_k = 1  # Save the top k best models based on target_metric
+target_metric = "val_f1"  # Metric to monitor for early stopping and model checkpointing.
+# Mode min/max is inferred from the metric name.
 
 [prediction_export]
 output_folder = "predictions"
diff --git a/src/deepforest_finetuning/config/_config.py b/src/deepforest_finetuning/config/_config.py
index 7fc727c..cf44a3a 100644
--- a/src/deepforest_finetuning/config/_config.py
+++ b/src/deepforest_finetuning/config/_config.py
@@ -97,6 +97,8 @@ class TrainingConfig:  # pylint: disable=too-many-instance-attributes
     float32_matmul_precision: str = "medium"
     log_dir: str = "./logs"
     early_stopping_patience: Optional[int] = None
+    save_top_k: int = 1  # Save top k best models based on target_metric
+    target_metric: str = "val_f1"  # Metric to monitor for early stopping and model checkpointing
 
 
 @dataclasses.dataclass
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 6249263..461894f 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -333,11 +333,14 @@ def finetuning(
                 if config.early_stopping_patience is not None:
                     pretraining_callbacks.append(
                         EarlyStopping(
-                            monitor="f1",  # Use F1 score for early stopping
+                            monitor=config.target_metric.replace(
+                                "val_", ""
+                            ),  # Adjust for pretraining where we don't use "val_" prefix
                             min_delta=0.0,  # Minimum change to qualify as an improvement
                             patience=config.early_stopping_patience,
                             verbose=True,
-                            mode="max",  # Higher F1 is better
+                            # Automatically infer mode from the metric name
+                            mode="min" if "loss" in config.target_metric else "max",
                         )
                     )
 
@@ -362,21 +365,24 @@ def finetuning(
                     ModelCheckpoint(
                         dirpath=base_dir / config.checkpoint_dir,
                         filename="{epoch}_" + f"seed={seed}",
-                        save_top_k=-1,
+                        monitor=config.target_metric,
+                        save_top_k=config.save_top_k,
                         every_n_epochs=1,
                         enable_version_counter=False,
+                        # Automatically infer mode from the metric name
+                        mode="min" if "loss" in config.target_metric else "max",
                     )
                 )
 
             # Add early stopping callback if patience is set
             if config.early_stopping_patience is not None:
-                # We provide both val_loss and val_f1 - let's use val_f1 for early stopping (higher is better)
                 callbacks.append(
                     EarlyStopping(
-                        monitor="val_f1",  # Use F1 score for early stopping
+                        monitor=config.target_metric,  # Use configured metric for early stopping
                         patience=config.early_stopping_patience,
                         verbose=True,
-                        mode="max",  # Higher F1 is better
+                        # Automatically infer mode from the metric name
+                        mode="min" if "loss" in config.target_metric else "max",
                     )
                 )
 

From 44e38f9521c6b1d5aeae85e2bac1c89738a6eb2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anton=20Eichst=C3=A4dt?= <antoleeic05@gmail.com>
Date: Thu, 5 Jun 2025 12:27:41 +0200
Subject: [PATCH 10/24] Update
 src/deepforest_finetuning/training/_finetuning.py

Co-authored-by: Josafat-Mattias Burmeister <33292321+josafatburmeister@users.noreply.github.com>
---
 src/deepforest_finetuning/training/_finetuning.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 461894f..97c8eb1 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -160,7 +160,7 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):  # p
             pl_module: Model to be evaluated.
         """
 
-        # Create a copy of the model for evaluation to avoid affecting the training
+        # Create a copy of the model for evaluation to avoid affecting the random state of the training
         eval_model = copy.deepcopy(pl_module)
         eval_trainer = copy.deepcopy(trainer)
         eval_model.trainer = eval_trainer

From 615cb41cff776c78028119b99a565fef7af2d154 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anton=20Eichst=C3=A4dt?= <antoleeic05@gmail.com>
Date: Thu, 5 Jun 2025 12:28:25 +0200
Subject: [PATCH 11/24] Update
 src/deepforest_finetuning/training/_finetuning.py

Co-authored-by: Josafat-Mattias Burmeister <33292321+josafatburmeister@users.noreply.github.com>
---
 src/deepforest_finetuning/training/_finetuning.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 97c8eb1..f11a6dd 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -125,7 +125,7 @@ def _process_annotation_paths(base_dir: Path, annotation_files: List[str]) -> Li
             # If it's a directory, collect all JSON files inside
             json_files = list(path.glob("*.json"))
             # Convert paths to strings relative to base_dir
-            rel_paths = [str(js_file.relative_to(base_dir)) for js_file in json_files]
+            rel_paths = [str(json_file.relative_to(base_dir)) for json_file in json_files]
             processed_paths.extend(rel_paths)
             print(f"INFO: Found {len(json_files)} JSON files in directory {path}.")
         else:

From 3e5ec93b122717791909549eb031532ba0d42e52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anton=20Eichst=C3=A4dt?= <antoleeic05@gmail.com>
Date: Thu, 5 Jun 2025 12:34:02 +0200
Subject: [PATCH 12/24] Update README.md

Co-authored-by: Josafat-Mattias Burmeister <33292321+josafatburmeister@users.noreply.github.com>
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 899aef8..24ccc89 100644
--- a/README.md
+++ b/README.md
@@ -20,7 +20,7 @@ A Python package for fine-tuning the [DeepForest](https://github.com/weecology/D
 
 ## Overview
 
-DeepForest is a deep learning model designed for detecting trees in aerial RGB imagery. This package extends DeepForest by providing a comprehensive framework to fine-tune the model on your own datasets. Key features include:
+DeepForest is a deep learning model designed for detecting trees in aerial RGB imagery. This package extends DeepForest by providing a framework to fine-tune the model on your own datasets. Key features include:
 
 - Data preprocessing for various input formats
 - Automatic label projection from 3D point clouds to 2D orthophotos

From f04b37dffd99771b4aa81cea4b643232ef011358 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anton=20Eichst=C3=A4dt?= <antoleeic05@gmail.com>
Date: Thu, 5 Jun 2025 12:34:15 +0200
Subject: [PATCH 13/24] Update README.md

Co-authored-by: Josafat-Mattias Burmeister <33292321+josafatburmeister@users.noreply.github.com>
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 24ccc89..23a18a4 100644
--- a/README.md
+++ b/README.md
@@ -24,7 +24,7 @@ DeepForest is a deep learning model designed for detecting trees in aerial RGB i
 
 - Data preprocessing for various input formats
 - Automatic label projection from 3D point clouds to 2D orthophotos
-- Image rescaling and patch generation
+- Image rescaling and tiling
 - Model fine-tuning with multiple random seeds for robust evaluation
 - Prediction on new images with customizable tiling
 - Evaluation metrics calculation (precision, recall, F1 score)

From ddb6cc9b7f4310f49ef67334a62cb808a1e210ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anton=20Eichst=C3=A4dt?= <antoleeic05@gmail.com>
Date: Thu, 5 Jun 2025 12:36:52 +0200
Subject: [PATCH 14/24] Update README.md

Co-authored-by: Josafat-Mattias Burmeister <33292321+josafatburmeister@users.noreply.github.com>
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 23a18a4..e899397 100644
--- a/README.md
+++ b/README.md
@@ -91,7 +91,7 @@ The package supports multiple preprocessing steps:
 
 #### a. Projecting Labels from Point Clouds
 
-If you have 3D point cloud data with tree positions, you can project them to 2D bounding boxes:
+If you have 3D point cloud data with pointwise tree instance labels, you can project them to 2D bounding boxes:
 
 ```bash
 python scripts/preprocessing.py configs/preprocessing/project_point_cloud_labels.toml

From 4f842817c389ef970c61ee2ade7b7f89aa9eb403 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anton=20Eichst=C3=A4dt?= <antoleeic05@gmail.com>
Date: Thu, 5 Jun 2025 12:37:47 +0200
Subject: [PATCH 15/24] Update README.md

Co-authored-by: Josafat-Mattias Burmeister <33292321+josafatburmeister@users.noreply.github.com>
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index e899397..58de78e 100644
--- a/README.md
+++ b/README.md
@@ -115,7 +115,7 @@ python scripts/preprocessing.py configs/preprocessing/preprocess_manually_correc
 
 #### c. Filtering Labels
 
-Filter out unwanted labels based on size, position, etc.:
+Filter out unwanted labels based on overlap and size:
 
 ```bash
 python scripts/preprocessing.py configs/preprocessing/filter_labels.toml

From bb5a68ca9846cd5be3dcd968830012cdcda61d5a Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Thu, 5 Jun 2025 17:05:13 +0200
Subject: [PATCH 16/24] eveyrthing fixed from PR review

---
 README.md                                     | 36 ++----------
 .../5_cm/finetuning_5_cm_TreeAI_full.toml     |  2 +-
 .../prediction/_prediction_dataset.py         | 13 ++---
 .../_project_point_cloud_labels.py            |  2 +-
 .../training/_finetuning.py                   | 55 +++++++------------
 5 files changed, 30 insertions(+), 78 deletions(-)

diff --git a/README.md b/README.md
index 58de78e..20fa9e5 100644
--- a/README.md
+++ b/README.md
@@ -57,7 +57,7 @@ A Dockerfile is provided for containerized usage:
 
 ```bash
 docker build -t deepforest-finetuning .
-docker run --gpus all -it -v /path/to/your/data:/data deepforest-finetuning /bin/bash
+docker run --gpus all -it -v /path/to/your/data:/workspace/data/deepforest-finetuning
 ```
 
 ## Project Structure
@@ -65,9 +65,6 @@ docker run --gpus all -it -v /path/to/your/data:/data deepforest-finetuning /bin
 ```
 deepforest-finetuning/
 ├── configs/              # Configuration files for different workflows
-│   ├── baseline/         # Configurations for the baseline model
-│   ├── finetuning/       # Configurations for fine-tuning
-│   └── preprocessing/    # Configurations for data preprocessing
 ├── scripts/              # Executable scripts for main workflows
 │   ├── evaluate.py       # Script for model evaluation
 │   ├── finetuning.py     # Script for fine-tuning
@@ -105,15 +102,7 @@ image_paths = ["image1.tif", "image2.tif"]
 label_json_output_paths = ["labels1.json", "labels2.json"]
 ```
 
-#### b. Preprocessing Manually Corrected Labels
-
-Convert manually created or corrected labels to the format required by DeepForest:
-
-```bash
-python scripts/preprocessing.py configs/preprocessing/preprocess_manually_corrected_labels.toml
-```
-
-#### c. Filtering Labels
+#### b. Filtering Labels
 
 Filter out unwanted labels based on overlap and size:
 
@@ -121,7 +110,7 @@ Filter out unwanted labels based on overlap and size:
 python scripts/preprocessing.py configs/preprocessing/filter_labels.toml
 ```
 
-#### d. Image Rescaling
+#### c. Image Rescaling
 
 Rescale images to different resolutions:
 
@@ -226,24 +215,7 @@ All workflows are configured using TOML files:
    python scripts/evaluate.py configs/evaluation/my_evaluation_config.toml
    ```
 
-### Using Different Image Resolutions
-
-The package supports working with images at various resolutions. The config directories contain subdirectories for different resolutions (e.g., `2_5_cm`, `5_cm`, `7_5_cm`, `10_cm`).
-
-## Advanced Features
-
-### Experiment Tracking with Weights & Biases
-
-The fine-tuning process supports integration with Weights & Biases for experiment tracking:
-
-```toml
-# Add to your fine-tuning config
-use_wandb = true
-wandb_project = "deepforest-finetuning"
-wandb_entity = "your-wandb-username"
-```
-
-This will log metrics, hyperparameters, and validation results to your W&B account.
+## Other Features
 
 ### Multiple Random Seeds
 
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml b/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
index 256f662..e8419f4 100644
--- a/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
+++ b/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
@@ -1,7 +1,7 @@
 base_dir = "/mnt/daten/TreeAI/finetuning/for_finetuning/5cm"
 tmp_dir = "./tmp"
 patch_size = 640
-patch_overlap = 0.2
+patch_overlap = 0.1
 image_folder = "images"
 train_annotation_files = [
     "annotations",
diff --git a/src/deepforest_finetuning/prediction/_prediction_dataset.py b/src/deepforest_finetuning/prediction/_prediction_dataset.py
index 9b792cb..83bd3fe 100644
--- a/src/deepforest_finetuning/prediction/_prediction_dataset.py
+++ b/src/deepforest_finetuning/prediction/_prediction_dataset.py
@@ -43,22 +43,19 @@ def __getitem__(self, idx: int) -> npt.NDArray:
         Returns: Image data.
         """
         img_path = self.image_files[idx]
-        # Check file extension to use appropriate method for reading the image
-        file_ext = os.path.splitext(img_path)[1].lower()
+        file_ext = Path(img_path).suffix.lower()
 
         if file_ext in [".tif", ".tiff"]:
             # Use tifffile for TIFF images
             image_array = np.array(imread(img_path))
-            # Ensure we only take the first 3 channels if there are more
-            if image_array.ndim >= 3 and image_array.shape[2] > 3:
-                image_array = image_array[:, :, :3]
         else:
             # Use PIL for other image formats (PNG, JPG, etc.)
             image = Image.open(img_path)
             image_array = np.array(image)
-            # Ensure we only take the first 3 channels if there are more
-            if image_array.ndim >= 3 and image_array.shape[2] > 3:
-                image_array = image_array[:, :, :3]
+
+        # Ensure we only take the first 3 channels if there are more
+        if image_array.ndim >= 3 and image_array.shape[2] > 3:
+            image_array = image_array[:, :, :3]
 
         # Ensure consistent dtype
         image_array = image_array.astype(np.uint8)
diff --git a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
index 10895ff..77adb23 100644
--- a/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
+++ b/src/deepforest_finetuning/preprocessing/_project_point_cloud_labels.py
@@ -18,7 +18,7 @@
 from rasterio.transform import from_origin
 from skimage.filters.rank import modal
 import torch
-from torch_scatter import scatter_max  # pylint: disable=import-error
+from torch_scatter import scatter_max
 
 from deepforest_finetuning.utils import coco_bbox_to_polygon
 from deepforest_finetuning.config import PointCloudLabelProjectionConfig
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index f11a6dd..42e1043 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -106,7 +106,7 @@ def split_images_into_patches(
     return annotations_path
 
 
-def _process_annotation_paths(base_dir: Path, annotation_files: List[str]) -> List[str]:
+def _collect_annotation_paths(base_dir: Path, annotation_files: List[str]) -> List[str]:
     """
     Process annotation file paths, handling both individual files and directories.
 
@@ -117,23 +117,20 @@ def _process_annotation_paths(base_dir: Path, annotation_files: List[str]) -> Li
     Returns:
         List of processed annotation file paths.
     """
-    processed_paths = []
+    annotation_paths = []
 
     for file_path in annotation_files:
         path = base_dir / file_path
         if path.is_dir():
-            # If it's a directory, collect all JSON files inside
             json_files = list(path.glob("*.json"))
-            # Convert paths to strings relative to base_dir
-            rel_paths = [str(json_file.relative_to(base_dir)) for json_file in json_files]
-            processed_paths.extend(rel_paths)
+            relative_paths = [str(js_file.relative_to(base_dir)) for js_file in json_files]
+            annotation_paths.extend(relative_paths)
             print(f"INFO: Found {len(json_files)} JSON files in directory {path}.")
         else:
-            # If it's a file, add it directly
-            processed_paths.append(file_path)
+            annotation_paths.append(file_path)
             print(f"INFO: Using annotation file {file_path}.")
 
-    return processed_paths
+    return annotation_paths
 
 
 class EvaluationCallBack(Callback):
@@ -165,12 +162,9 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):  # p
         eval_trainer = copy.deepcopy(trainer)
         eval_model.trainer = eval_trainer
 
-        # Keep a reference to the original trainer for logging metrics
-        original_trainer = trainer
-
         # evaluate on training and test set
-        processed_train_files = _process_annotation_paths(self._base_dir, self._config.train_annotation_files)
-        processed_test_files = _process_annotation_paths(self._base_dir, self._config.test_annotation_files)
+        processed_train_files = _collect_annotation_paths(self._base_dir, self._config.train_annotation_files)
+        processed_test_files = _collect_annotation_paths(self._base_dir, self._config.test_annotation_files)
 
         for prefix, annotation_files in [
             ("train", processed_train_files),
@@ -227,23 +221,16 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):  # p
             # This makes them available for callbacks like EarlyStopping
             for metric_name, metric_value in metrics.items():
                 # Log with trainer if logger is available
-                if original_trainer.logger is not None:
-                    original_trainer.logger.log_metrics(
+                if trainer.logger is not None:
+                    trainer.logger.log_metrics(
                         {f"{prefix}_{metric_name}": metric_value},
-                        step=original_trainer.current_epoch,
+                        step=trainer.current_epoch,
                     )
 
-                # For test metrics, also log them as validation metrics for EarlyStopping
-                if prefix == "test":
-                    if metric_name == "f1":
-                        # Log F1 as val_f1 for early stopping (maximize)
-                        # Access callback_metrics directly on the original trainer
-                        original_trainer.callback_metrics[f"val_{metric_name}"] = torch.tensor(metric_value)
-
-                    # Always log inverse F1 as val_loss for compatibility with default early stopping
-                    if metric_name == "f1":
-                        # For loss, use 1-F1 as val_loss (minimize is better)
-                        original_trainer.callback_metrics["val_loss"] = torch.tensor(1.0 - metric_value)
+                if metric_name == "f1":
+                    # Log F1 as val_f1 for early stopping (maximize)
+                    # Access callback_metrics directly on the original trainer
+                    trainer.callback_metrics[f"val_{metric_name}"] = torch.tensor(metric_value)
 
 
 def finetuning(
@@ -262,16 +249,12 @@ def finetuning(
     preprocessed_annotation_files = {}
 
     # Process annotation file paths for train and pretraining (if available)
-    processed_train_files = _process_annotation_paths(base_dir, config.train_annotation_files)
+    train_annotation_files = _collect_annotation_paths(base_dir, config.train_annotation_files)
 
-    print(processed_train_files)
-
-    splitting_configs = [("train", processed_train_files)]
+    splitting_configs = [("train", train_annotation_files)]
     if config.pretrain_annotation_files is not None and len(config.pretrain_annotation_files) > 0:
-        processed_pretrain_files = _process_annotation_paths(base_dir, config.pretrain_annotation_files)
-        splitting_configs.append(("pretraining", processed_pretrain_files))
-
-    print(f"INFO: Found {len(splitting_configs)} annotation configurations to process.")
+        pretrain_annotation_files = _collect_annotation_paths(base_dir, config.pretrain_annotation_files)
+        splitting_configs.append(("pretraining", pretrain_annotation_files))
 
     for prefix, annotation_files in splitting_configs:
         annotations = []

From 0f115a5892c0e5c2fcef37abca20a1e76ab0f0af Mon Sep 17 00:00:00 2001
From: Anton <Anton.Eichstaedt@student.hpi.uni.potsdam.de>
Date: Thu, 5 Jun 2025 18:09:36 +0200
Subject: [PATCH 17/24] unused import fix

---
 src/deepforest_finetuning/prediction/_prediction_dataset.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/deepforest_finetuning/prediction/_prediction_dataset.py b/src/deepforest_finetuning/prediction/_prediction_dataset.py
index 83bd3fe..5ae0a6a 100644
--- a/src/deepforest_finetuning/prediction/_prediction_dataset.py
+++ b/src/deepforest_finetuning/prediction/_prediction_dataset.py
@@ -4,7 +4,6 @@
 
 from pathlib import Path
 from typing import List, Optional
-import os
 
 import numpy as np
 import numpy.typing as npt

From 462f3f96685e09314976b78d187d0e20920cab05 Mon Sep 17 00:00:00 2001
From: Josafat-Mattias Burmeister <burmeister@uni-potsdam.de>
Date: Sat, 30 Aug 2025 12:10:31 +0200
Subject: [PATCH 18/24] update Readme

---
 README.md | 100 +++++++++++++-----------------------------------------
 1 file changed, 24 insertions(+), 76 deletions(-)

diff --git a/README.md b/README.md
index 74f49ca..abc1ce7 100644
--- a/README.md
+++ b/README.md
@@ -1,26 +1,6 @@
 ## Fine-Tuning DeepForest for Forest Tree Detection in High-Resolution UAV Imagery
 
-A Python package for fine-tuning the [DeepForest](https://github.com/weecology/DeepForest) model on custom data for tree detection in aerial imagery. This project provides a streamlined workflow for preprocessing training data, fine-tuning the DeepForest model, making predictions, and evaluating results.
-
-## Table of Contents
-- [Overview](#overview)
-- [Installation](#installation)
-  - [Using Conda](#using-conda)
-  - [Using Docker](#using-docker)
-- [Project Structure](#project-structure)
-- [Workflow](#workflow)
-  - [1. Data Preprocessing](#1-data-preprocessing)
-  - [2. Model Fine-tuning](#2-model-fine-tuning)
-  - [3. Making Predictions](#3-making-predictions)
-  - [4. Evaluating Results](#4-evaluating-results)
-- [Configuration Files](#configuration-files)
-- [Example Usage](#example-usage)
-- [Advanced Features](#advanced-features)
-- [License](#license)
-
-## Overview
-
-DeepForest is a deep learning model designed for detecting trees in aerial RGB imagery. This package extends DeepForest by providing a framework to fine-tune the model on your own datasets. Key features include:
+A Python package for fine-tuning the [DeepForest](https://github.com/weecology/DeepForest) model on custom data. DeepForest is a deep learning model for detecting trees in aerial RGB imagery. This package extends DeepForest by providing a workflow to fine-tune the model on your own datasets. Key features include:
 
 - Data preprocessing for various input formats
 - Automatic label projection from 3D point clouds to 2D orthophotos
@@ -28,7 +8,6 @@ DeepForest is a deep learning model designed for detecting trees in aerial RGB i
 - Model fine-tuning with multiple random seeds for robust evaluation
 - Prediction on new images with customizable tiling
 - Evaluation metrics calculation (precision, recall, F1 score)
-- Support for experiment tracking with [Weights & Biases](https://wandb.ai/)
 
 ## Installation
 
@@ -57,27 +36,7 @@ A Dockerfile is provided for containerized usage:
 
 ```bash
 docker build -t deepforest-finetuning .
-docker run --gpus all -it -v /path/to/your/data:/workspace/data/deepforest-finetuning
-```
-
-## Project Structure
-
-```
-deepforest-finetuning/
-├── configs/              # Configuration files for different workflows
-├── scripts/              # Executable scripts for main workflows
-│   ├── evaluate.py       # Script for model evaluation
-│   ├── finetuning.py     # Script for fine-tuning
-│   ├── prediction.py     # Script for making predictions
-│   └── preprocessing.py  # Script for data preprocessing
-└── src/                  # Source code
-    └── deepforest_finetuning/
-        ├── config/       # Configuration dataclasses
-        ├── evaluation/   # Model evaluation logic
-        ├── prediction/   # Prediction logic
-        ├── preprocessing/# Data preprocessing logic
-        ├── training/     # Model training and fine-tuning logic
-        └── utils/        # Utility functions
+docker run --gpus all --rm -it -v /path/to/your/data:/workspace/data/ deepforest-finetuning
 ```
 
 ## Workflow
@@ -88,7 +47,7 @@ The package supports multiple preprocessing steps:
 
 #### a. Projecting Labels from Point Clouds
 
-If you have 3D point cloud data with pointwise tree instance labels, you can project them to 2D bounding boxes:
+If you have 3D point cloud data with pointwise tree instance labels in addition to 2D aerial images, you can project them to 2D bounding boxes:
 
 ```bash
 python scripts/preprocessing.py configs/preprocessing/project_point_cloud_labels.toml
@@ -104,20 +63,37 @@ label_json_output_paths = ["labels1.json", "labels2.json"]
 
 #### b. Filtering Labels
 
-Filter out unwanted labels based on overlap and size:
+Filters labels using non-maximum suppression based on overlap and size:
 
 ```bash
 python scripts/preprocessing.py configs/preprocessing/filter_labels.toml
 ```
 
+Required configuration:
+```toml
+base_dir = "/path/to/your/data"
+input_label_folder = "labels"
+output_label_folder = "labels_filtered"
+iou_threshold = 0.5
+```
+
 #### c. Image Rescaling
 
-Rescale images to different resolutions:
+Rescale images and corresponding labels to different resolutions:
 
 ```bash
 python scripts/preprocessing.py configs/preprocessing/rescale_images.toml
 ```
 
+Required configuration:
+```toml
+base_dir = "/path/to/your/data"
+input_images = ["image1.tif", "image2.tif"]
+input_label_folders = ["labels"]
+output_folders = ["rescaled_2_5_cm", "rescaled_5_cm"]
+target_resolutions = [0.025, 0.05]
+```
+
 ### 2. Model Fine-tuning
 
 Fine-tune the DeepForest model on your custom dataset:
@@ -145,8 +121,7 @@ target_metric = "val_f1"
 ```
 
 This will:
-1. Split images into patches
-2. Create training and test datasets
+1. Split images into patches and load training and test datasets
 3. Fine-tune the model for the specified number of epochs
 4. Run with multiple random seeds for robust evaluation
 5. Save checkpoints and logs
@@ -190,30 +165,7 @@ output_file = "/path/to/evaluation_results.csv"
 
 ## Configuration Files
 
-All workflows are configured using TOML files:
-
-- **Preprocessing configs**: Define data paths and parameters for preprocessing steps
-- **Training configs**: Specify training hyperparameters, data paths, and evaluation settings
-- **Prediction configs**: Set model checkpoint, input images, and output format
-- **Evaluation configs**: Define prediction and ground truth file paths, and evaluation metrics
-
-## Example Usage
-
-### Complete Fine-tuning Pipeline Example
-
-1. Preprocess your data (e.g., project point cloud labels, rescale images)
-2. Fine-tune the model:
-   ```bash
-   python scripts/finetuning.py configs/finetuning/my_finetuning_config.toml
-   ```
-3. Make predictions with the fine-tuned model:
-   ```bash
-   python scripts/prediction.py configs/prediction/my_prediction_config.toml
-   ```
-4. Evaluate the results:
-   ```bash
-   python scripts/evaluate.py configs/evaluation/my_evaluation_config.toml
-   ```
+All workflows are configured using TOML files. Example configurations are provided in the `configs` folder.
 
 ## Other Features
 
@@ -239,10 +191,6 @@ target_metric = "val_f1"  # Metric to monitor for early stopping and checkpointi
 
 The `mode` (min/max) is automatically inferred from the metric name. Metrics containing "loss" use "min" mode (lower is better), all others use "max" mode (higher is better).
 
-## License
-
-This project is licensed under the MIT License - see the LICENSE file for details.
-
 ### How to Cite
 
 If you use our code, please consider citing our paper:

From f1358bee846d21f9cd884902e446c5db36aae408 Mon Sep 17 00:00:00 2001
From: Josafat-Mattias Burmeister <burmeister@uni-potsdam.de>
Date: Sat, 30 Aug 2025 12:38:47 +0200
Subject: [PATCH 19/24] clean up code

---
 README.md                                     |  4 +-
 .../evaluate_without_finetuning_10_cm.toml    |  0
 .../evaluate_without_finetuning_2_5_cm.toml   |  0
 .../evaluate_without_finetuning_5_cm.toml     |  0
 .../evaluate_without_finetuning_7_5_cm.toml   |  0
 .../predict_without_finetuning_10_cm.toml     |  0
 .../predict_without_finetuning_2_5_cm.toml    |  0
 .../predict_without_finetuning_5_cm.toml      |  0
 .../predict_without_finetuning_7_5_cm.toml    |  0
 ...netuning_10_cm_automatic_labeling_ext.toml |  0
 ...tuning_10_cm_automatic_labeling_small.toml |  0
 ...inetuning_10_cm_manual_correction_ext.toml |  0
 ...etuning_10_cm_manual_correction_small.toml |  0
 .../finetuning_10_cm_manual_labeling_ext.toml |  0
 ...inetuning_10_cm_manual_labeling_small.toml |  0
 ...etuning_2_5_cm_automatic_labeling_ext.toml |  0
 ...uning_2_5_cm_automatic_labeling_small.toml |  0
 ...netuning_2_5_cm_manual_correction_ext.toml |  0
 ...tuning_2_5_cm_manual_correction_small.toml |  0
 ...finetuning_2_5_cm_manual_labeling_ext.toml |  0
 ...netuning_2_5_cm_manual_labeling_small.toml |  0
 ...inetuning_5_cm_automatic_labeling_ext.toml |  0
 ...etuning_5_cm_automatic_labeling_small.toml |  0
 ...finetuning_5_cm_manual_correction_ext.toml |  0
 ...netuning_5_cm_manual_correction_small.toml |  0
 .../finetuning_5_cm_manual_labeling_ext.toml  |  0
 ...finetuning_5_cm_manual_labeling_small.toml |  0
 ...etuning_7_5_cm_automatic_labeling_ext.toml |  0
 ...uning_7_5_cm_automatic_labeling_small.toml |  0
 ...netuning_7_5_cm_manual_correction_ext.toml |  0
 ...tuning_7_5_cm_manual_correction_small.toml |  0
 ...finetuning_7_5_cm_manual_labeling_ext.toml |  0
 ...netuning_7_5_cm_manual_labeling_small.toml |  0
 .../finetuning/run_finetuning.bat             |  0
 .../preprocessing/filter_labels.toml          |  0
 .../preprocess_manually_corrected_labels.toml |  0
 .../project_point_cloud_labels.toml           |  0
 .../preprocessing/rescale_images.toml         |  0
 .../finetuning_5_cm_TreeAI_full.toml          |  2 +-
 src/deepforest_finetuning/config/_config.py   |  4 +-
 .../prediction/_prediction_dataset.py         |  1 -
 .../training/_finetuning.py                   | 40 ++++++++-----------
 42 files changed, 21 insertions(+), 30 deletions(-)
 rename configs/{ => 3d-geoinfo-2025}/baseline/evaluate_without_finetuning_10_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/baseline/evaluate_without_finetuning_2_5_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/baseline/evaluate_without_finetuning_5_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/baseline/evaluate_without_finetuning_7_5_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/baseline/predict_without_finetuning_10_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/baseline/predict_without_finetuning_2_5_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/baseline/predict_without_finetuning_5_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/baseline/predict_without_finetuning_7_5_cm.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/10_cm/finetuning_10_cm_automatic_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/10_cm/finetuning_10_cm_automatic_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/10_cm/finetuning_10_cm_manual_correction_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/10_cm/finetuning_10_cm_manual_correction_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/10_cm/finetuning_10_cm_manual_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/10_cm/finetuning_10_cm_manual_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/5_cm/finetuning_5_cm_automatic_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/5_cm/finetuning_5_cm_automatic_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/5_cm/finetuning_5_cm_manual_correction_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/5_cm/finetuning_5_cm_manual_correction_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/5_cm/finetuning_5_cm_manual_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/5_cm/finetuning_5_cm_manual_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_ext.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_small.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/finetuning/run_finetuning.bat (100%)
 rename configs/{ => 3d-geoinfo-2025}/preprocessing/filter_labels.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/preprocessing/preprocess_manually_corrected_labels.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/preprocessing/project_point_cloud_labels.toml (100%)
 rename configs/{ => 3d-geoinfo-2025}/preprocessing/rescale_images.toml (100%)
 rename configs/{finetuning/5_cm => tree-ai-2025}/finetuning_5_cm_TreeAI_full.toml (87%)

diff --git a/README.md b/README.md
index abc1ce7..3298eaf 100644
--- a/README.md
+++ b/README.md
@@ -117,7 +117,7 @@ learning_rate = 0.0001
 checkpoint_dir = "checkpoints"
 early_stopping_patience = 2
 save_top_k = 1
-target_metric = "val_f1"
+target_metric = "test_f1"
 ```
 
 This will:
@@ -186,7 +186,7 @@ To prevent overfitting, you can enable early stopping and control model checkpoi
 ```toml
 early_stopping_patience = 2  # Stop training if performance doesn't improve for this many epochs
 save_top_k = 1  # Save the top k best models based on the target metric
-target_metric = "val_f1"  # Metric to monitor for early stopping and checkpointing
+target_metric = "test_f1"  # Metric to monitor for early stopping and checkpointing
 ```
 
 The `mode` (min/max) is automatically inferred from the metric name. Metrics containing "loss" use "min" mode (lower is better), all others use "max" mode (higher is better).
diff --git a/configs/baseline/evaluate_without_finetuning_10_cm.toml b/configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_10_cm.toml
similarity index 100%
rename from configs/baseline/evaluate_without_finetuning_10_cm.toml
rename to configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_10_cm.toml
diff --git a/configs/baseline/evaluate_without_finetuning_2_5_cm.toml b/configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_2_5_cm.toml
similarity index 100%
rename from configs/baseline/evaluate_without_finetuning_2_5_cm.toml
rename to configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_2_5_cm.toml
diff --git a/configs/baseline/evaluate_without_finetuning_5_cm.toml b/configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_5_cm.toml
similarity index 100%
rename from configs/baseline/evaluate_without_finetuning_5_cm.toml
rename to configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_5_cm.toml
diff --git a/configs/baseline/evaluate_without_finetuning_7_5_cm.toml b/configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_7_5_cm.toml
similarity index 100%
rename from configs/baseline/evaluate_without_finetuning_7_5_cm.toml
rename to configs/3d-geoinfo-2025/baseline/evaluate_without_finetuning_7_5_cm.toml
diff --git a/configs/baseline/predict_without_finetuning_10_cm.toml b/configs/3d-geoinfo-2025/baseline/predict_without_finetuning_10_cm.toml
similarity index 100%
rename from configs/baseline/predict_without_finetuning_10_cm.toml
rename to configs/3d-geoinfo-2025/baseline/predict_without_finetuning_10_cm.toml
diff --git a/configs/baseline/predict_without_finetuning_2_5_cm.toml b/configs/3d-geoinfo-2025/baseline/predict_without_finetuning_2_5_cm.toml
similarity index 100%
rename from configs/baseline/predict_without_finetuning_2_5_cm.toml
rename to configs/3d-geoinfo-2025/baseline/predict_without_finetuning_2_5_cm.toml
diff --git a/configs/baseline/predict_without_finetuning_5_cm.toml b/configs/3d-geoinfo-2025/baseline/predict_without_finetuning_5_cm.toml
similarity index 100%
rename from configs/baseline/predict_without_finetuning_5_cm.toml
rename to configs/3d-geoinfo-2025/baseline/predict_without_finetuning_5_cm.toml
diff --git a/configs/baseline/predict_without_finetuning_7_5_cm.toml b/configs/3d-geoinfo-2025/baseline/predict_without_finetuning_7_5_cm.toml
similarity index 100%
rename from configs/baseline/predict_without_finetuning_7_5_cm.toml
rename to configs/3d-geoinfo-2025/baseline/predict_without_finetuning_7_5_cm.toml
diff --git a/configs/finetuning/10_cm/finetuning_10_cm_automatic_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_automatic_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/10_cm/finetuning_10_cm_automatic_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_automatic_labeling_ext.toml
diff --git a/configs/finetuning/10_cm/finetuning_10_cm_automatic_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_automatic_labeling_small.toml
similarity index 100%
rename from configs/finetuning/10_cm/finetuning_10_cm_automatic_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_automatic_labeling_small.toml
diff --git a/configs/finetuning/10_cm/finetuning_10_cm_manual_correction_ext.toml b/configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_correction_ext.toml
similarity index 100%
rename from configs/finetuning/10_cm/finetuning_10_cm_manual_correction_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_correction_ext.toml
diff --git a/configs/finetuning/10_cm/finetuning_10_cm_manual_correction_small.toml b/configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_correction_small.toml
similarity index 100%
rename from configs/finetuning/10_cm/finetuning_10_cm_manual_correction_small.toml
rename to configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_correction_small.toml
diff --git a/configs/finetuning/10_cm/finetuning_10_cm_manual_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/10_cm/finetuning_10_cm_manual_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_labeling_ext.toml
diff --git a/configs/finetuning/10_cm/finetuning_10_cm_manual_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_labeling_small.toml
similarity index 100%
rename from configs/finetuning/10_cm/finetuning_10_cm_manual_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/10_cm/finetuning_10_cm_manual_labeling_small.toml
diff --git a/configs/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_ext.toml
diff --git a/configs/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_small.toml
similarity index 100%
rename from configs/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_automatic_labeling_small.toml
diff --git a/configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_ext.toml b/configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_ext.toml
similarity index 100%
rename from configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_ext.toml
diff --git a/configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_small.toml b/configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_small.toml
similarity index 100%
rename from configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_small.toml
rename to configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_correction_small.toml
diff --git a/configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_ext.toml
diff --git a/configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_small.toml
similarity index 100%
rename from configs/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/2_5_cm/finetuning_2_5_cm_manual_labeling_small.toml
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_automatic_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_automatic_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/5_cm/finetuning_5_cm_automatic_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_automatic_labeling_ext.toml
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_automatic_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_automatic_labeling_small.toml
similarity index 100%
rename from configs/finetuning/5_cm/finetuning_5_cm_automatic_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_automatic_labeling_small.toml
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_manual_correction_ext.toml b/configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_correction_ext.toml
similarity index 100%
rename from configs/finetuning/5_cm/finetuning_5_cm_manual_correction_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_correction_ext.toml
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_manual_correction_small.toml b/configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_correction_small.toml
similarity index 100%
rename from configs/finetuning/5_cm/finetuning_5_cm_manual_correction_small.toml
rename to configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_correction_small.toml
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_manual_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/5_cm/finetuning_5_cm_manual_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_labeling_ext.toml
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_manual_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_labeling_small.toml
similarity index 100%
rename from configs/finetuning/5_cm/finetuning_5_cm_manual_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/5_cm/finetuning_5_cm_manual_labeling_small.toml
diff --git a/configs/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_ext.toml
diff --git a/configs/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_small.toml
similarity index 100%
rename from configs/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_automatic_labeling_small.toml
diff --git a/configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_ext.toml b/configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_ext.toml
similarity index 100%
rename from configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_ext.toml
diff --git a/configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_small.toml b/configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_small.toml
similarity index 100%
rename from configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_small.toml
rename to configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_correction_small.toml
diff --git a/configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_ext.toml b/configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_ext.toml
similarity index 100%
rename from configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_ext.toml
rename to configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_ext.toml
diff --git a/configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_small.toml b/configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_small.toml
similarity index 100%
rename from configs/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_small.toml
rename to configs/3d-geoinfo-2025/finetuning/7_5_cm/finetuning_7_5_cm_manual_labeling_small.toml
diff --git a/configs/finetuning/run_finetuning.bat b/configs/3d-geoinfo-2025/finetuning/run_finetuning.bat
similarity index 100%
rename from configs/finetuning/run_finetuning.bat
rename to configs/3d-geoinfo-2025/finetuning/run_finetuning.bat
diff --git a/configs/preprocessing/filter_labels.toml b/configs/3d-geoinfo-2025/preprocessing/filter_labels.toml
similarity index 100%
rename from configs/preprocessing/filter_labels.toml
rename to configs/3d-geoinfo-2025/preprocessing/filter_labels.toml
diff --git a/configs/preprocessing/preprocess_manually_corrected_labels.toml b/configs/3d-geoinfo-2025/preprocessing/preprocess_manually_corrected_labels.toml
similarity index 100%
rename from configs/preprocessing/preprocess_manually_corrected_labels.toml
rename to configs/3d-geoinfo-2025/preprocessing/preprocess_manually_corrected_labels.toml
diff --git a/configs/preprocessing/project_point_cloud_labels.toml b/configs/3d-geoinfo-2025/preprocessing/project_point_cloud_labels.toml
similarity index 100%
rename from configs/preprocessing/project_point_cloud_labels.toml
rename to configs/3d-geoinfo-2025/preprocessing/project_point_cloud_labels.toml
diff --git a/configs/preprocessing/rescale_images.toml b/configs/3d-geoinfo-2025/preprocessing/rescale_images.toml
similarity index 100%
rename from configs/preprocessing/rescale_images.toml
rename to configs/3d-geoinfo-2025/preprocessing/rescale_images.toml
diff --git a/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml b/configs/tree-ai-2025/finetuning_5_cm_TreeAI_full.toml
similarity index 87%
rename from configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
rename to configs/tree-ai-2025/finetuning_5_cm_TreeAI_full.toml
index e8419f4..873a55f 100644
--- a/configs/finetuning/5_cm/finetuning_5_cm_TreeAI_full.toml
+++ b/configs/tree-ai-2025/finetuning_5_cm_TreeAI_full.toml
@@ -15,7 +15,7 @@ learning_rate = 0.0001
 checkpoint_dir = "checkpoints"
 early_stopping_patience = 2  # Set to the number of epochs to wait before stopping, or remove to disable early stopping
 save_top_k = 1  # Save the top k best models based on target_metric
-target_metric = "val_f1"  # Metric to monitor for early stopping and model checkpointing.
+target_metric = "test_f1"  # Metric to monitor for early stopping and model checkpointing.
 # Mode min/max is inferred from the metric name.
 
 [prediction_export]
diff --git a/src/deepforest_finetuning/config/_config.py b/src/deepforest_finetuning/config/_config.py
index cf44a3a..68fe43a 100644
--- a/src/deepforest_finetuning/config/_config.py
+++ b/src/deepforest_finetuning/config/_config.py
@@ -97,8 +97,8 @@ class TrainingConfig:  # pylint: disable=too-many-instance-attributes
     float32_matmul_precision: str = "medium"
     log_dir: str = "./logs"
     early_stopping_patience: Optional[int] = None
-    save_top_k: int = 1  # Save top k best models based on target_metric
-    target_metric: str = "val_f1"  # Metric to monitor for early stopping and model checkpointing
+    save_top_k: int = 1
+    target_metric: str = "test_f1"
 
 
 @dataclasses.dataclass
diff --git a/src/deepforest_finetuning/prediction/_prediction_dataset.py b/src/deepforest_finetuning/prediction/_prediction_dataset.py
index 5ae0a6a..4236a6d 100644
--- a/src/deepforest_finetuning/prediction/_prediction_dataset.py
+++ b/src/deepforest_finetuning/prediction/_prediction_dataset.py
@@ -56,7 +56,6 @@ def __getitem__(self, idx: int) -> npt.NDArray:
         if image_array.ndim >= 3 and image_array.shape[2] > 3:
             image_array = image_array[:, :, :3]
 
-        # Ensure consistent dtype
         image_array = image_array.astype(np.uint8)
 
         if self.resize_images_to is not None:
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 42e1043..f3c4340 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -7,7 +7,7 @@
 import os
 from pathlib import Path
 import shutil
-from typing import List, Union
+from typing import List, Optional, Union
 import uuid
 
 import albumentations as A
@@ -27,7 +27,7 @@
 from deepforest_finetuning.prediction import prediction as run_prediction
 
 
-def get_transform(augment: bool):
+def get_transform(augment: bool, seed: Optional[int] = None):
     """
     Albumentations transformation of bounding boxes.
 
@@ -42,12 +42,14 @@ def get_transform(augment: bool):
         transform = A.Compose(
             [A.HorizontalFlip(p=0.5), ToTensorV2()],
             bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
+            seed=seed
         )
 
     else:
         transform = A.Compose(
             [ToTensorV2()],
             bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
+            seed=seed
         )
 
     return transform
@@ -123,7 +125,7 @@ def _collect_annotation_paths(base_dir: Path, annotation_files: List[str]) -> Li
         path = base_dir / file_path
         if path.is_dir():
             json_files = list(path.glob("*.json"))
-            relative_paths = [str(js_file.relative_to(base_dir)) for js_file in json_files]
+            relative_paths = [str(json_file.relative_to(base_dir)) for json_file in json_files]
             annotation_paths.extend(relative_paths)
             print(f"INFO: Found {len(json_files)} JSON files in directory {path}.")
         else:
@@ -159,16 +161,15 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):  # p
 
         # Create a copy of the model for evaluation to avoid affecting the random state of the training
         eval_model = copy.deepcopy(pl_module)
-        eval_trainer = copy.deepcopy(trainer)
-        eval_model.trainer = eval_trainer
+        eval_model.trainer = copy.deepcopy(trainer)
 
         # evaluate on training and test set
-        processed_train_files = _collect_annotation_paths(self._base_dir, self._config.train_annotation_files)
-        processed_test_files = _collect_annotation_paths(self._base_dir, self._config.test_annotation_files)
+        train_annotation_files = _collect_annotation_paths(self._base_dir, self._config.train_annotation_files)
+        test_annotation_files = _collect_annotation_paths(self._base_dir, self._config.test_annotation_files)
 
         for prefix, annotation_files in [
-            ("train", processed_train_files),
-            ("test", processed_test_files),
+            ("train", train_annotation_files),
+            ("test", test_annotation_files),
         ]:
             image_files = []
             annotations = []
@@ -220,7 +221,6 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):  # p
             # Log metrics to Lightning logger for each prefix (train/test)
             # This makes them available for callbacks like EarlyStopping
             for metric_name, metric_value in metrics.items():
-                # Log with trainer if logger is available
                 if trainer.logger is not None:
                     trainer.logger.log_metrics(
                         {f"{prefix}_{metric_name}": metric_value},
@@ -228,9 +228,8 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):  # p
                     )
 
                 if metric_name == "f1":
-                    # Log F1 as val_f1 for early stopping (maximize)
                     # Access callback_metrics directly on the original trainer
-                    trainer.callback_metrics[f"val_{metric_name}"] = torch.tensor(metric_value)
+                    trainer.callback_metrics[f"{prefix}_{metric_name}"] = torch.tensor(metric_value)
 
 
 def finetuning(
@@ -248,7 +247,6 @@ def finetuning(
     preprocessed_image_folders = {}
     preprocessed_annotation_files = {}
 
-    # Process annotation file paths for train and pretraining (if available)
     train_annotation_files = _collect_annotation_paths(base_dir, config.train_annotation_files)
 
     splitting_configs = [("train", train_annotation_files)]
@@ -288,7 +286,7 @@ def finetuning(
             print(f"INFO: Training for {config.epochs} epochs with seed {seed}...")
 
             # load model
-            model = deepforest_main.deepforest(transforms=partial(get_transform))
+            model = deepforest_main.deepforest(transforms=partial(get_transform, seed=seed))
             model.use_release()
 
             # copy config to avoid overwriting
@@ -311,18 +309,14 @@ def finetuning(
                     name=f"{config.epochs}_epochs_seed_{seed}_pretraining",
                 )
 
-                # Add pretraining callbacks
                 pretraining_callbacks = []
                 if config.early_stopping_patience is not None:
                     pretraining_callbacks.append(
                         EarlyStopping(
-                            monitor=config.target_metric.replace(
-                                "val_", ""
-                            ),  # Adjust for pretraining where we don't use "val_" prefix
-                            min_delta=0.0,  # Minimum change to qualify as an improvement
+                            monitor=config.target_metric,
+                            min_delta=0.0,
                             patience=config.early_stopping_patience,
                             verbose=True,
-                            # Automatically infer mode from the metric name
                             mode="min" if "loss" in config.target_metric else "max",
                         )
                     )
@@ -352,19 +346,17 @@ def finetuning(
                         save_top_k=config.save_top_k,
                         every_n_epochs=1,
                         enable_version_counter=False,
-                        # Automatically infer mode from the metric name
                         mode="min" if "loss" in config.target_metric else "max",
                     )
                 )
 
-            # Add early stopping callback if patience is set
             if config.early_stopping_patience is not None:
                 callbacks.append(
                     EarlyStopping(
-                        monitor=config.target_metric,  # Use configured metric for early stopping
+                        monitor=config.target_metric,
+                        min_delta=0.0,
                         patience=config.early_stopping_patience,
                         verbose=True,
-                        # Automatically infer mode from the metric name
                         mode="min" if "loss" in config.target_metric else "max",
                     )
                 )

From 9ded35714685ab0dd1fef9da3fdf95dce802d2e7 Mon Sep 17 00:00:00 2001
From: Josafat-Mattias Burmeister <burmeister@uni-potsdam.de>
Date: Sat, 30 Aug 2025 17:31:43 +0200
Subject: [PATCH 20/24] update code formatting

---
 src/deepforest_finetuning/training/_finetuning.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index f3c4340..7657e35 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -42,14 +42,12 @@ def get_transform(augment: bool, seed: Optional[int] = None):
         transform = A.Compose(
             [A.HorizontalFlip(p=0.5), ToTensorV2()],
             bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
-            seed=seed
+            seed=seed,
         )
 
     else:
         transform = A.Compose(
-            [ToTensorV2()],
-            bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
-            seed=seed
+            [ToTensorV2()], bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]), seed=seed
         )
 
     return transform

From 6fe14aa1da2519cc3cfe75747ac4175ce148bf23 Mon Sep 17 00:00:00 2001
From: Josafat-Mattias Burmeister <josafat-mattias.burm@hpi.de>
Date: Tue, 2 Sep 2025 20:04:41 +0200
Subject: [PATCH 21/24] update code to latest deepforest changes

---
 src/deepforest_finetuning/config/_config.py   |  4 +--
 .../prediction/_prediction.py                 |  2 --
 .../preprocessing/_rescale_images.py          |  1 +
 .../training/_finetuning.py                   | 35 ++++++++-----------
 4 files changed, 17 insertions(+), 25 deletions(-)

diff --git a/src/deepforest_finetuning/config/_config.py b/src/deepforest_finetuning/config/_config.py
index 68fe43a..599c8ff 100644
--- a/src/deepforest_finetuning/config/_config.py
+++ b/src/deepforest_finetuning/config/_config.py
@@ -84,8 +84,8 @@ class TrainingConfig:  # pylint: disable=too-many-instance-attributes
     patch_overlap: float
     learning_rate: float
     tmp_dir: str
-    train_annotation_files: List[str]
-    test_annotation_files: List[str]
+    train_annotation_files: Union[List[str], str]
+    test_annotation_files: Union[List[str], str]
     prediction_export: ExportConfig
     checkpoint_dir: Optional[str] = None
     iou_threshold: float = 0.5
diff --git a/src/deepforest_finetuning/prediction/_prediction.py b/src/deepforest_finetuning/prediction/_prediction.py
index de697cd..afd5c43 100644
--- a/src/deepforest_finetuning/prediction/_prediction.py
+++ b/src/deepforest_finetuning/prediction/_prediction.py
@@ -55,14 +55,12 @@ def prediction(
         if predict_tile:
             pred = model.predict_tile(
                 image=tree_dataset[img_idx].astype(np.float32),
-                return_plot=False,
                 patch_size=patch_size,
                 patch_overlap=patch_overlap,
             )
         else:
             pred = model.predict_image(
                 image=tree_dataset[img_idx].astype(np.float32),
-                return_plot=False,
             )
         image_name = tree_dataset.__getname__(img_idx)
         pred["image_path"] = image_name
diff --git a/src/deepforest_finetuning/preprocessing/_rescale_images.py b/src/deepforest_finetuning/preprocessing/_rescale_images.py
index 5b9a7fd..33ab7e9 100644
--- a/src/deepforest_finetuning/preprocessing/_rescale_images.py
+++ b/src/deepforest_finetuning/preprocessing/_rescale_images.py
@@ -78,6 +78,7 @@ def rescale_images(config: ImageRescalingConfig):  # pylint: disable=too-many-lo
                 label_output_folder.mkdir(exist_ok=True, parents=True)
 
                 label_subfolders = [x for x in os.listdir(input_label_folder) if os.path.isdir(input_label_folder / x)]
+                label_subfolders.append(".")
                 for label_subfolder in label_subfolders:
                     label_file_name = f"{original_image_path.stem}_coco.json"
                     label_file = input_label_folder / label_subfolder / label_file_name
diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 7657e35..36f31bb 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -27,28 +27,19 @@
 from deepforest_finetuning.prediction import prediction as run_prediction
 
 
-def get_transform(augment: bool, seed: Optional[int] = None):
+def get_transform(seed: Optional[int] = None):
     """
     Albumentations transformation of bounding boxes.
 
-    Args:
-        augment: Whether to apply data augmentations.
-
     Returns:
         Transforms.
     """
 
-    if augment:
-        transform = A.Compose(
-            [A.HorizontalFlip(p=0.5), ToTensorV2()],
-            bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
-            seed=seed,
-        )
-
-    else:
-        transform = A.Compose(
-            [ToTensorV2()], bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]), seed=seed
-        )
+    transform = A.Compose(
+        [A.HorizontalFlip(p=0.5), ToTensorV2()],
+        bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
+        seed=seed,
+    )
 
     return transform
 
@@ -106,7 +97,7 @@ def split_images_into_patches(
     return annotations_path
 
 
-def _collect_annotation_paths(base_dir: Path, annotation_files: List[str]) -> List[str]:
+def _collect_annotation_paths(base_dir: Path, annotation_files: Union[List[str], str]) -> List[str]:
     """
     Process annotation file paths, handling both individual files and directories.
 
@@ -119,6 +110,9 @@ def _collect_annotation_paths(base_dir: Path, annotation_files: List[str]) -> Li
     """
     annotation_paths = []
 
+    if isinstance(annotation_files, str):
+        annotation_files = [annotation_files]
+
     for file_path in annotation_files:
         path = base_dir / file_path
         if path.is_dir():
@@ -284,7 +278,7 @@ def finetuning(
             print(f"INFO: Training for {config.epochs} epochs with seed {seed}...")
 
             # load model
-            model = deepforest_main.deepforest(transforms=partial(get_transform, seed=seed))
+            model = deepforest_main.deepforest(transforms=get_transform(seed=seed))
             model.use_release()
 
             # copy config to avoid overwriting
@@ -292,12 +286,11 @@ def finetuning(
 
             # configure model
             if current_config.pretrain_learning_rate is None:
-                model.config["train"]["lr"] = current_config.learning_rate
+                model.config.train.lr = current_config.learning_rate
             else:
-                model.config["train"]["lr"] = current_config.pretrain_learning_rate
+                model.config.train.lr = current_config.pretrain_learning_rate
 
-            model.config["train"]["epochs"] = config.epochs
-            model.config["save-snapshot"] = False
+            model.config.train.epochs = config.epochs
 
             if "pretraining" in preprocessed_annotation_files:
                 model.config["train"]["csv_file"] = preprocessed_annotation_files["pretraining"]

From 08f2b17b3d9162713243c11d8b6ee81be3fdcbcc Mon Sep 17 00:00:00 2001
From: Josafat-Mattias Burmeister <josafat-mattias.burm@hpi.de>
Date: Tue, 2 Sep 2025 20:04:57 +0200
Subject: [PATCH 22/24] update Readme

---
 README.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/README.md b/README.md
index 3298eaf..752467a 100644
--- a/README.md
+++ b/README.md
@@ -88,8 +88,11 @@ python scripts/preprocessing.py configs/preprocessing/rescale_images.toml
 Required configuration:
 ```toml
 base_dir = "/path/to/your/data"
+# input_images can either be a list of individual file paths or string specifying a folder path
 input_images = ["image1.tif", "image2.tif"]
+# if no labels are available, input_label_folders can be left empty
 input_label_folders = ["labels"]
+# there must be one output folder for each target resolution
 output_folders = ["rescaled_2_5_cm", "rescaled_5_cm"]
 target_resolutions = [0.025, 0.05]
 ```

From f6db2fc3204e10f7ef4462207a9c4ab509ecfbcf Mon Sep 17 00:00:00 2001
From: Josafat-Mattias Burmeister <josafat-mattias.burm@hpi.de>
Date: Tue, 2 Sep 2025 20:21:45 +0200
Subject: [PATCH 23/24] update Readme

---
 README.md | 36 +++++++++++++++++++++++++++++++++---
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 752467a..77be1d4 100644
--- a/README.md
+++ b/README.md
@@ -13,19 +13,49 @@ A Python package for fine-tuning the [DeepForest](https://github.com/weecology/D
 
 ### Using Conda
 
-1. Clone this repository:
+1. Make sure you have conda installed. If conda is installed, `conda --version` should output the conda version.
+
+2. Clone this repository:
    ```bash
    git clone https://github.com/yourusername/deepforest-finetuning.git
    cd deepforest-finetuning
    ```
 
-2. Create and activate a conda environment from the provided environment.yml file:
+3. Create and activate a conda environment from the provided environment.yml file:
    ```bash
    conda env create -f environment.yml
    conda activate deepforest-env
    ```
 
-3. Install the package in development mode:
+4. Install the package in development mode:
+   ```bash
+   pip install -e .
+   ```
+
+### Using pip
+
+1. Make sure that Python3 and pip are installed.
+
+2. Clone this repository:
+   ```bash
+   git clone https://github.com/yourusername/deepforest-finetuning.git
+   cd deepforest-finetuning
+   ```
+
+3. Install the [pointtorch](https://ai4trees.github.io/pointtorch/v0.2.0/) package and its dependencies (`${TORCH}` should be replaced by the PyTorch version and `${CUDA}` by `cpu`, `cu126`, etc., depending on the PyTorch installation):
+   ```bash
+   pip install torch torchvision --index-url https://download.pytorch.org/whl/${CUDA}
+   pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
+   pip install pointtorch
+   ``` 
+
+4. Install the [DeepForest](https://deepforest.readthedocs.io/en/v1.5.0/getting_started/install.html) package:
+
+   ```bash
+   pip install "git+https://github.com/weecology/DeepForest.git"
+   ```
+
+5. Install the package in development mode:
    ```bash
    pip install -e .
    ```

From e853e34d72d24863bd9fa05e43e8b1139a9c9aa9 Mon Sep 17 00:00:00 2001
From: Josafat-Mattias Burmeister <josafat-mattias.burm@hpi.de>
Date: Tue, 2 Sep 2025 20:33:21 +0200
Subject: [PATCH 24/24]  fix pylint error

---
 src/deepforest_finetuning/training/_finetuning.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/deepforest_finetuning/training/_finetuning.py b/src/deepforest_finetuning/training/_finetuning.py
index 36f31bb..9c72725 100644
--- a/src/deepforest_finetuning/training/_finetuning.py
+++ b/src/deepforest_finetuning/training/_finetuning.py
@@ -3,7 +3,6 @@
 __all__ = ["split_images_into_patches", "finetuning"]
 
 import copy
-from functools import partial
 import os
 from pathlib import Path
 import shutil