From bd7ff525ea0a18e60e05c43713c0273c4dff6f1f Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Wed, 11 Dec 2024 09:58:21 +0100 Subject: [PATCH] add docs --- .../attached_modules/losses/README.md | 28 +++++++++++++++ .../losses/precision_dfl_detection_loss.py | 7 ++-- .../losses/precision_dlf_segmentation_loss.py | 4 --- luxonis_train/nodes/README.md | 34 ++++++++++++++++++- .../nodes/heads/precision_bbox_head.py | 2 ++ 5 files changed, 65 insertions(+), 10 deletions(-) diff --git a/luxonis_train/attached_modules/losses/README.md b/luxonis_train/attached_modules/losses/README.md index 38f8b42f..a2f07106 100644 --- 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 index fb6b559f..cb80b105 100644 --- 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__( }. 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 index af777a80..27f05809 100644 --- 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__( }. 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 index 31f1f6c2..92046745 100644 --- 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 index c2c1893a..27c2fb9f 100644 --- 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