add docs

luxonis · Dec 11, 2024 · bd7ff52 · bd7ff52
1 parent 6597687
commit bd7ff52
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 10 deletions.
diff --git a/luxonis_train/attached_modules/losses/README.md b/luxonis_train/attached_modules/losses/README.md
@@ -12,6 +12,8 @@ List of all the available loss functions.
 - [`AdaptiveDetectionLoss`](#adaptivedetectionloss)
 - [`EfficientKeypointBBoxLoss`](#efficientkeypointbboxloss)
 - [`FOMOLocalizationLoss`](#fomolocalizationLoss)
+- \[`PrecisionDFLDetectionLoss`\] (# precisiondfldetectionloss)
+- \[`PrecisionDFLSegmentationLoss`\] (# precisiondflsegmentationloss)
 
 ## `CrossEntropyLoss`
 
@@ -121,3 +123,29 @@ Adapted from [here](https://arxiv.org/abs/2108.07610).
 | Key             | Type    | Default value | Description                                     |
 | --------------- | ------- | ------------- | ----------------------------------------------- |
 | `object_weight` | `float` | `1000`        | Weight for the objects in the loss calculation. |
+
+## `PrecisionDFLDetectionLoss`
+
+Adapted from [here](https://arxiv.org/pdf/2207.02696.pdf) and [here](https://arxiv.org/pdf/2209.02976.pdf).
+
+**Parameters:**
+
+| Key                 | Type    | Default value | Description                                |
+| ------------------- | ------- | ------------- | ------------------------------------------ |
+| `tal_topk`          | `int`   | `10`          | Number of anchors considered in selection. |
+| `class_loss_weight` | `float` | `0.5`         | Weight for classification loss.            |
+| `bbox_loss_weight`  | `float` | `7.5`         | Weight for bbox loss.                      |
+| `dfl_loss_weigth`   | `float` | `1.5`         | Weight for DFL loss.                       |
+
+## `PrecisionDFLSegmentationLoss`
+
+Adapted from [here](https://arxiv.org/pdf/2207.02696.pdf) and [here](https://arxiv.org/pdf/2209.02976.pdf).
+
+**Parameters:**
+
+| Key                 | Type    | Default value | Description                                |
+| ------------------- | ------- | ------------- | ------------------------------------------ |
+| `tal_topk`          | `int`   | `10`          | Number of anchors considered in selection. |
+| `class_loss_weight` | `float` | `0.5`         | Weight for classification loss.            |
+| `bbox_loss_weight`  | `float` | `7.5`         | Weight for bbox and segmentation loss.     |
+| `dfl_loss_weigth`   | `float` | `1.5`         | Weight for DFL loss.                       |
diff --git a/luxonis_train/attached_modules/losses/precision_dfl_detection_loss.py b/luxonis_train/attached_modules/losses/precision_dfl_detection_loss.py
@@ -31,7 +31,6 @@ class PrecisionDFLDetectionLoss(
 
     def __init__(
         self,
-        reg_max: int = 16,
         tal_topk: int = 10,
         class_loss_weight: float = 0.5,
         bbox_loss_weight: float = 7.5,
@@ -43,8 +42,6 @@ def __init__(
         <https://arxiv.org/pdf/2209.02976.pdf>}.
         Code is adapted from U{https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/models}.
 
-        @type reg_max: int
-        @param reg_max: Maximum number of regression channels. Defaults to 16.
         @type tal_topk: int
         @param tal_topk: Number of anchors considered in selection. Defaults to 10.
         @type class_loss_weight: float
@@ -67,8 +64,8 @@ def __init__(
         self.assigner = TaskAlignedAssigner(
             n_classes=self.n_classes, topk=tal_topk, alpha=0.5, beta=6.0
         )
-        self.bbox_loss = CustomBboxLoss(reg_max)
-        self.proj = torch.arange(reg_max, dtype=torch.float)
+        self.bbox_loss = CustomBboxLoss(self.node.reg_max)
+        self.proj = torch.arange(self.node.reg_max, dtype=torch.float)
         self.bce = nn.BCEWithLogitsLoss(reduction="none")
 
     def prepare(

diff --git a/luxonis_train/attached_modules/losses/precision_dlf_segmentation_loss.py b/luxonis_train/attached_modules/losses/precision_dlf_segmentation_loss.py
@@ -29,7 +29,6 @@ class PrecisionDFLSegmentationLoss(PrecisionDFLDetectionLoss):
 
     def __init__(
         self,
-        reg_max: int = 16,
         tal_topk: int = 10,
         class_loss_weight: float = 0.5,
         bbox_loss_weight: float = 7.5,
@@ -41,8 +40,6 @@ def __init__(
         <https://arxiv.org/pdf/2209.02976.pdf>}.
         Code is adapted from U{https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/models}.
 
-        @type reg_max: int
-        @param reg_max: Maximum number of regression channels. Defaults to 16.
         @type tal_topk: int
         @param tal_topk: Number of anchors considered in selection. Defaults to 10.
         @type class_loss_weight: float
@@ -53,7 +50,6 @@ def __init__(
         @param dfl_loss_weight: Weight for DFL loss. Defaults to 1.5. For optimal results, multiply with accumulate_grad_batches.
         """
         super().__init__(
-            reg_max=reg_max,
             tal_topk=tal_topk,
             class_loss_weight=class_loss_weight,
             bbox_loss_weight=bbox_loss_weight,

diff --git a/luxonis_train/nodes/README.md b/luxonis_train/nodes/README.md
@@ -28,6 +28,8 @@ arbitrarily as long as the two nodes are compatible with each other. We've group
   - [`DDRNetSegmentationHead`](#ddrnetsegmentationhead)
   - [`DiscSubNetHead`](#discsubnet)
   - [`FOMOHead`](#fomohead)
+  - [`PrecisionBBoxHead`](#precisionbboxhead)
+  - [`PrecisionSegmentBBoxHead`](#precisionsegmentbboxhead)
     Every node takes these parameters:
 
 | Key                | Type          | Default value | Description                                                                 |
@@ -222,7 +224,7 @@ Adapted from [here](https://arxiv.org/pdf/2209.02976.pdf).
 
 | Key                  | Type    | Default value | Description                                                           |
 | -------------------- | ------- | ------------- | --------------------------------------------------------------------- |
-| `n_heads`            | `bool`  | `3`           | Number of output heads                                                |
+| `n_heads`            | `int`   | `3`           | Number of output heads                                                |
 | `conf_thres`         | `float` | `0.25`        | Confidence threshold for non-maxima-suppression (used for evaluation) |
 | `iou_thres`          | `float` | `0.45`        | `IoU` threshold for non-maxima-suppression (used for evaluation)      |
 | `max_det`            | `int`   | `300`         | Maximum number of detections retained after NMS                       |
@@ -272,3 +274,33 @@ Adapted from [here](https://arxiv.org/abs/2108.07610).
 | ----------------- | ----- | ------------- | ------------------------------------------------------- |
 | `num_conv_layers` | `int` | `3`           | Number of convolutional layers to use in the model.     |
 | `conv_channels`   | `int` | `16`          | Number of output channels for each convolutional layer. |
+
+## `PrecisionBBoxHead`
+
+Adapted from [here](https://arxiv.org/pdf/2207.02696.pdf) and [here](https://arxiv.org/pdf/2209.02976.pdf).
+
+**Parameters:**
+
+| Key          | Type    | Default value | Description                                                               |
+| ------------ | ------- | ------------- | ------------------------------------------------------------------------- |
+| `reg_max`    | `int`   | `16`          | Maximum number of regression channels                                     |
+| `n_heads`    | `int`   | `3`           | Number of output heads                                                    |
+| `conf_thres` | `float` | `0.25`        | Confidence threshold for non-maxima-suppression (used for evaluation)     |
+| `iou_thres`  | `float` | `0.45`        | IoU threshold for non-maxima-suppression (used for evaluation)            |
+| `max_det`    | `int`   | `300`         | Max number of detections for non-maxima-suppression (used for evaluation) |
+
+## `PrecisionSegmentBBoxHead`
+
+Adapted from [here](https://arxiv.org/pdf/2207.02696.pdf) and [here](https://arxiv.org/pdf/2209.02976.pdf).
+
+**Parameters:**
+
+| Key          | Type    | Default value | Description                                                                |
+| ------------ | ------- | ------------- | -------------------------------------------------------------------------- |
+| `reg_max`    | `int`   | `16`          | Maximum number of regression channels.                                     |
+| `n_heads`    | `int`   | `3`           | Number of output heads.                                                    |
+| `conf_thres` | `float` | `0.25`        | Confidence threshold for non-maxima-suppression (used for evaluation).     |
+| `iou_thres`  | `float` | `0.45`        | IoU threshold for non-maxima-suppression (used for evaluation).            |
+| `max_det`    | `int`   | `300`         | Max number of detections for non-maxima-suppression (used for evaluation). |
+| `n_masks`    | `int`   | `32`          | Number of of output instance segmentation masks at the output.             |
+| `n_proto`    | `int`   | `256`         | Number of prototypes generated from the prototype generator.               |
diff --git a/luxonis_train/nodes/heads/precision_bbox_head.py b/luxonis_train/nodes/heads/precision_bbox_head.py
@@ -47,6 +47,8 @@ def __init__(
         @param conf_thres: Confidence threshold for NMS.
         @type iou_thres: float
         @param iou_thres: IoU threshold for NMS.
+        @type max_det: int
+        @param max_det: Maximum number of detections retained after NMS.
         """
         super().__init__(**kwargs)
         self.reg_max = reg_max