add dinov3

Author: jveitchmichaelis

.gitignore

@@ -20,4 +20,5 @@ tests/__pycache__
 tests/data/*
 .vscode/
 *ipynb_checkpoints/
-docs/user_guide/deepforestr.md
+docs/user_guide/deepforestr.md
+.env

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "h5py",
     "huggingface_hub>=0.25.0",
     "hydra-core",
+    "geopandas>=1.0.0",
     "matplotlib",
     "numpy<2.0",
     "omegaconf",
@@ -56,11 +57,28 @@ dependencies = [
     "supervision",
     "tensorboard",
     "timm",
-    "torch>=2.2.0,<2.3.0",
-    "torchvision>=0.17.0,<0.18.0",
+    "torch>=2.7.0",
+    "torchvision>=0.17.0",
     "tqdm",
-    "transformers",
+    "transformers>=4.56",
     "xmltodict",
+    "transformers",
+    "timm>=1.0.15",
+    "faster-coco-eval>=1.6.7",
+    "comet-ml>=3.51.0",
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+torchvision = [
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
 ]
 
 [project.urls]

src/deepforest/callbacks.py

@@ -5,19 +5,26 @@
 """
 
 import glob
+import os
+import warnings
+from pathlib import Path
 
+import matplotlib.pyplot as plt
 import numpy as np
 import supervision as sv
+import torch
+from PIL import Image
 from pytorch_lightning import Callback
 
-from deepforest import visualize
+from deepforest import utilities, visualize
+from deepforest.datasets.training import BoxDataset
 
 
-class images_callback(Callback):
+class ImagesCallback(Callback):
     """Log evaluation images during training.
 
     Args:
-        savedir: Directory to save predicted images
+        save_dir: Directory to save predicted images
         n: Number of images to process
         every_n_epochs: Run interval in epochs
         select_random: Whether to select random images
@@ -26,61 +33,190 @@ class images_callback(Callback):
     """
 
     def __init__(
-        self, savedir, n=2, every_n_epochs=5, select_random=False, color=None, thickness=1
+        self,
+        save_dir,
+        prediction_samples=2,
+        dataset_samples=5,
+        every_n_epochs=5,
+        select_random=False,
+        color=None,
+        thickness=1,
     ):
-        self.savedir = savedir
-        self.n = n
+        self.savedir = save_dir
+        self.prediction_samples = prediction_samples
+        self.dataset_samples = dataset_samples
         self.color = color
         self.thickness = thickness
         self.select_random = select_random
         self.every_n_epochs = every_n_epochs
 
-    def log_images(self, pl_module):
-        """Log images to the logger."""
+    def on_train_start(self, trainer, pl_module):
+        """Log sample images from training and validation datasets at training
+        start."""
+
+        if trainer.fast_dev_run:
+            return
+
+        self.trainer = trainer
+        self.pl_module = pl_module
+
+        # Training samples
+        self.pl_module.print("Logging training dataset samples.")
+        train_ds = self.trainer.train_dataloader.dataset
+        self._log_dataset_sample(train_ds, split="train")
+
+        # Validation samples
+        if self.trainer.val_dataloaders:
+            self.pl_module.print("Logging validation dataset samples.")
+            val_ds = self.trainer.val_dataloaders.dataset
+            self._log_dataset_sample(val_ds, split="validation")
+
+    def on_validation_end(self, trainer, pl_module):
+        """Run callback at validation end."""
+        if trainer.sanity_checking or trainer.fast_dev_run:
+            return
+
+        if trainer.current_epoch % self.every_n_epochs == 0:
+            pl_module.print("Running image callback")
+            self._log_last_predictions(trainer, pl_module)
+
+    def _log_dataset_sample(self, dataset: BoxDataset, split: str):
+        """Log random samples from a DeepForest BoxDataset."""
+
+        if self.dataset_samples == 0:
+            return
+
+        out_dir = os.path.join(self.savedir, split + "_sample")
+        os.makedirs(out_dir, exist_ok=True)
+        n_samples = min(self.dataset_samples, len(dataset))
+        sample_indices = torch.randperm(len(dataset))[:n_samples]
+
+        sample_data = [dataset[idx] for idx in sample_indices]
+        sample_images = [data[0] for data in sample_data]
+        sample_targets = [data[1] for data in sample_data]
+        sample_paths = [data[2] for data in sample_data]
+
+        for image, target, path in zip(
+            sample_images, sample_targets, sample_paths, strict=False
+        ):
+            image_annotations = target.copy()
+            image_annotations = utilities.format_geometry(image_annotations, scores=False)
+            image_annotations.root_dir = dataset.root_dir
+            image_annotations["image_path"] = path
+
+            # Plot transformed image
+            basename = Path(path).stem
+            image = (255 * image.cpu().numpy().transpose((1, 2, 0))).astype(np.uint8)
+            fig = visualize.plot_annotations(
+                image=image,
+                annotations=image_annotations,
+                savedir=out_dir,
+                basename=basename,
+                thickness=self.thickness,
+                show=False,
+            )
+            plt.close(fig)
+
+            self._log_to_all(
+                image=os.path.join(out_dir, basename + ".png"),
+                trainer=self.trainer,
+                tag=f"{split} dataset sample",
+            )
+
+    def _log_last_predictions(self, trainer, pl_module):
+        """Log sample of predictions + targets from last validation."""
+        if self.prediction_samples == 0:
+            return
+
+        out_dir = os.path.join(self.savedir, "predictions")
+        os.makedirs(out_dir, exist_ok=True)
         df = pl_module.predictions
 
+        # Add root_dir to the dataframe
+        if "root_dir" not in df.columns:
+            df["root_dir"] = trainer.val_dataloaders.dataset.root_dir
+
         # Limit to n images, potentially randomly selected
         if self.select_random:
-            selected_images = np.random.choice(df.image_path.unique(), self.n)
+            selected_images = np.random.choice(
+                df.image_path.unique(), self.prediction_samples
+            )
         else:
-            selected_images = df.image_path.unique()[: self.n]
-        df = df[df.image_path.isin(selected_images)]
+            selected_images = df.image_path.unique()[: self.prediction_samples]
+
+            # Ensure color is correctly assigned
+            if self.color is None:
+                num_classes = len(df["label"].unique())
+                results_color = sv.ColorPalette.from_matplotlib("viridis", num_classes)
+            else:
+                results_color = self.color
+
+        for image_name in selected_images:
+            pred_df = df[df.image_path == image_name]
+            targets = utilities.format_geometry(
+                pl_module.targets[image_name], scores=False
+            )
 
-        # Add root_dir to the dataframe
-        if "root_dir" not in df.columns:
-            df["root_dir"] = pl_module.config.validation.root_dir
+            # Assume that validation images are un-augmented
+            fig = visualize.plot_results(
+                results=pred_df,
+                ground_truth=targets,
+                savedir=out_dir,
+                results_color=results_color,
+                thickness=self.thickness,
+                show=False,
+            )
+            plt.close(fig)
 
-        # Ensure color is correctly assigned
-        if self.color is None:
-            num_classes = len(df["label"].unique())
-            results_color = sv.ColorPalette.from_matplotlib("viridis", num_classes)
-        else:
-            results_color = self.color
-
-        # Plot results
-        visualize.plot_results(
-            results=df,
-            savedir=self.savedir,
-            results_color=results_color,
-            thickness=self.thickness,
-        )
+        saved_plots = glob.glob(f"{out_dir}/*.png")
+        for saved_plot in saved_plots:
+            self._log_to_all(image=saved_plot, trainer=trainer, tag="prediction sample")
+
+    def _log_to_all(self, image: str, trainer, tag):
+        """Log to all connected loggers.
 
+        Since Comet will pickup image logs to Tensorboard by default, we
+        add a check to log images preferentially to Tensorboard if both
+        are enabled.
+        """
         try:
-            saved_plots = glob.glob(f"{self.savedir}/*.png")
-            for x in saved_plots:
-                pl_module.logger.experiment.log_image(x)
-        except Exception as e:
-            print(
-                "Could not find comet logger in lightning module, "
-                f"skipping upload, images were saved to {self.savedir}, "
-                f"error was raised {e}"
+            img = np.array(Image.open(image).convert("RGB"))
+
+            loggers = [lg for lg in trainer.loggers if hasattr(lg, "experiment")]
+
+            tb = next((lg for lg in loggers if hasattr(lg.experiment, "add_image")), None)
+            if tb is not None:
+                tb.experiment.add_image(
+                    tag=f"{tag}/{os.path.basename(image)}",
+                    img_tensor=img,
+                    global_step=trainer.global_step,
+                    dataformats="HWC",
+                )
+                return
+
+            comet = next(
+                (lg.experiment for lg in loggers if hasattr(lg.experiment, "log_image")),
+                None,
             )
+            if comet is not None:
+                comet.experiment.log_image(
+                    img,
+                    name=tag,
+                    step=trainer.global_step,
+                    metadata={
+                        "image_name": os.path.basename(image),
+                        "context": tag,
+                        "step": trainer.global_step,
+                    },
+                )
 
-    def on_validation_end(self, trainer, pl_module):
-        """Run callback at validation end."""
-        if trainer.sanity_checking:
-            return
+        except Exception as e:
+            warnings.warn(f"Tried to log {image} exception raised: {e}", stacklevel=2)
 
-        if trainer.current_epoch % self.every_n_epochs == 0:
-            print("Running image callback")
-            self.log_images(pl_module)
+
+class images_callback(ImagesCallback):
+    def __init__(self, savedir, **kwargs):
+        warnings.warn(
+            "Please use ImagesCallback instead.", DeprecationWarning, stacklevel=2
+        )
+        super().__init__(save_dir=savedir, **kwargs)

src/deepforest/conf/config.yaml

@@ -30,6 +30,9 @@ rgb_dir:
 path_to_rgb:
 
 train:
+    # Sanity check annotations on dataset load
+    check_annotations: False
+    log_root: logs
     csv_file:
     root_dir:
 

src/deepforest/conf/dinov3.yaml

@@ -0,0 +1,17 @@
+# RetinaNet Base Configuration - Shared parameters for all folds
+defaults:
+  - config
+  - _self_
+
+model:
+  name: "facebook/dinov3-vitl16-pretrain-sat493m"
+  revision: 'main'
+
+train:
+  epochs: 75
+  lr: 0.01
+  scheduler:
+    type: cosine
+    params:
+      T_max: 75
+      eta_min: 0.0001

src/deepforest/conf/schema.py

@@ -62,12 +62,14 @@ class TrainConfig:
 
     csv_file: str | None = MISSING
     root_dir: str | None = MISSING
+    log_root: str = "logs"
     lr: float = 0.001
     scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
     epochs: int = 1
     fast_dev_run: bool = False
     preload_images: bool = False
     augmentations: list[str] | None = field(default_factory=lambda: ["HorizontalFlip"])
+    check_annotations: bool = False
 
 
 @dataclass