integrate keypoints

Author: jveitchmichaelis

src/deepforest/augmentations.py

@@ -51,15 +51,17 @@ def get_available_augmentations() -> list[str]:
 
 def get_transform(
     augmentations: str | list[str] | dict[str, Any] | None = None,
+    task: str = "box",
 ) -> A.Compose:
-    """Create Albumentations transform for bounding boxes.
+    """Create Albumentations transform for boxes or keypoints.
 
     Args:
         augmentations: Augmentation configuration:
             - str: Single augmentation name
             - list: List of augmentation names
             - dict: Dict with names as keys and params as values
             - None: No augmentations
+        task: Task type - "box" for bounding boxes or "keypoint" for keypoints
 
     Returns:
         Composed albumentations transform
@@ -79,9 +81,13 @@ def get_transform(
         ...                              "HorizontalFlip": {"p": 0.5},
         ...                              "Downscale": {"scale_min": 0.25, "scale_max": 0.75}
         ...                          })
+
+        >>> # Keypoint augmentations
+        >>> transform = get_transform(augmentations=["HorizontalFlip"], task="keypoint")
     """
     transforms_list = []
     bbox_params = None
+    keypoint_params = None
 
     if augmentations is not None:
         augment_configs = _parse_augmentations(augmentations)
@@ -90,12 +96,19 @@ def get_transform(
             aug_transform = _create_augmentation(aug_name, aug_params)
             transforms_list.append(aug_transform)
 
-        bbox_params = A.BboxParams(format="pascal_voc", label_fields=["category_ids"])
+        if task == "box":
+            bbox_params = A.BboxParams(format="pascal_voc", label_fields=["labels"])
+        elif task == "keypoint":
+            keypoint_params = A.KeypointParams(format="xy", label_fields=["labels"])
+        else:
+            raise ValueError(f"Unsupported task: {task}. Must be 'box' or 'keypoint'.")
 
     # Always add ToTensorV2 at the end
     transforms_list.append(ToTensorV2())
 
-    return A.Compose(transforms_list, bbox_params=bbox_params)
+    return A.Compose(
+        transforms_list, bbox_params=bbox_params, keypoint_params=keypoint_params
+    )
 
 
 def _parse_augmentations(

src/deepforest/callbacks.py

@@ -151,12 +151,12 @@ def _log_last_predictions(self, trainer, pl_module):
         else:
             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
+        # 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]

src/deepforest/conf/bird.yaml

@@ -0,0 +1,10 @@
+# Ensure we inherit from default config + overlay these overrides.
+defaults:
+  - config
+  - _self_
+
+task: 'box'
+
+model:
+    name: 'weecology/deepforest-bird'
+    revision: 'main'

src/deepforest/conf/config.yaml

@@ -8,6 +8,7 @@ accelerator: auto
 batch_size: 1
 
 # Model Architecture
+task: 'box'
 architecture: 'retinanet'
 num_classes: 1
 nms_thresh: 0.05

src/deepforest/conf/keypoint.yaml

@@ -0,0 +1,32 @@
+# Config file for DeepForest keypoint detection tasks
+
+# Ensure we inherit from default config + overlay these overrides.
+defaults:
+  - config
+  - _self_
+
+# Task and Model Architecture
+task: 'keypoint'
+architecture: 'DeformableDetr'
+
+# Keypoint-specific parameters for Deformable DETR
+# point_cost: Relative weight of point distance in matching cost (default: 5.0)
+# point_loss_coefficient: Weight of point loss in total loss (default: 5.0)
+# point_loss_type: Type of loss for coordinates - "l1" (default) or "mse"
+point_cost: 5.0
+point_loss_coefficient: 5.0
+point_loss_type: 'l1'
+
+# For keypoint detection, start from pretrained Deformable DETR backbone
+model:
+    name: 'SenseTime/deformable-detr'
+    revision: 'main'
+
+# Transformer-based models often prefer lower learning rates
+train:
+    lr: 0.0001
+
+# Pixel distance threshold for keypoint matching (instead of IoU for boxes)
+# A prediction is considered correct if within this many pixels of ground truth
+validation:
+    pixel_distance_threshold: 10.0

src/deepforest/conf/livestock.yaml

@@ -0,0 +1,10 @@
+# Ensure we inherit from default config + overlay these overrides.
+defaults:
+  - box
+  - _self_
+
+task: 'box'
+
+model:
+    name: 'weecology/deepforest-livestock'
+    revision: 'main'

src/deepforest/conf/schema.py

@@ -68,6 +68,7 @@ class TrainConfig:
     fast_dev_run: bool = False
     preload_images: bool = False
     augmentations: list[str] | None = field(default_factory=lambda: ["HorizontalFlip"])
+    log_root: str = "logs"
 
 
 @dataclass
@@ -87,6 +88,8 @@ class ValidationConfig:
     val_accuracy_interval: int = 20
     lr_plateau_target: str = "val_loss"
     augmentations: list[str] | None = field(default_factory=lambda: [])
+    # Keypoint-specific validation (used when task="keypoint")
+    pixel_distance_threshold: float = 10.0
 
 
 @dataclass
@@ -130,10 +133,16 @@ class Config:
     accelerator: str = "auto"
     batch_size: int = 1
 
+    task: str = "box"
     architecture: str = "retinanet"
     num_classes: int = 1
     label_dict: dict[str, int] = field(default_factory=lambda: {"Tree": 0})
 
+    # Keypoint-specific parameters (used when task="keypoint")
+    point_cost: float = 5.0
+    point_loss_coefficient: float = 5.0
+    point_loss_type: str = "l1"
+
     nms_thresh: float = 0.05
     score_thresh: float = 0.1
     model: ModelConfig = field(default_factory=ModelConfig)

src/deepforest/conf/tree.yaml

@@ -0,0 +1,10 @@
+# Ensure we inherit from default config + overlay these overrides.
+defaults:
+  - box
+  - _self_
+
+task: 'box'
+
+model:
+    name: 'weecology/deepforest-tree'
+    revision: 'main'