Merge branch 'dev' into fix/update-nn-archive-generation

luxonis · Aug 14, 2024 · 7eaf631 · 7eaf631
2 parents 32f8098 + c116982
commit 7eaf631
Show file tree

Hide file tree

Showing 33 changed files with 554 additions and 555 deletions.
diff --git a/.gitignore b/.gitignore
@@ -147,3 +147,4 @@ models_venv/*
 
 # vscode settings
 .vscode
+tests/data
diff --git a/configs/classification_model.yaml b/configs/classification_model.yaml
diff --git a/configs/coco_model.yaml b/configs/coco_model.yaml
diff --git a/configs/coco_multitask_model.yaml b/configs/coco_multitask_model.yaml
diff --git a/configs/detection_model.yaml b/configs/detection_model.yaml
diff --git a/configs/example_export.yaml b/configs/example_export.yaml
diff --git a/configs/example_tuning.yaml b/configs/example_tuning.yaml
diff --git a/configs/keypoint_bbox_model.yaml b/configs/keypoint_bbox_model.yaml
diff --git a/configs/resnet_multitask_model.yaml b/configs/resnet_multitask_model.yaml
diff --git a/configs/segmentation_model.yaml b/configs/segmentation_model.yaml
diff --git a/luxonis_train/attached_modules/base_attached_module.py b/luxonis_train/attached_modules/base_attached_module.py
@@ -133,6 +133,10 @@ def get_label(
         @rtype: tuple[Tensor, LabelType]
         @return: Extracted label and its type.
         """
+        if label_type is None:
+            if len(self.required_labels) == 1:
+                label_type = self.required_labels[0]
+
         if label_type is not None:
             task_name = self.node.get_task_name(label_type)
             if task_name not in labels:

diff --git a/luxonis_train/attached_modules/losses/implicit_keypoint_bbox_loss.py b/luxonis_train/attached_modules/losses/implicit_keypoint_bbox_loss.py
@@ -167,7 +167,7 @@ def prepare(
         targets[:, self.box_offset + 2 :: 3] = kpts[:, 3::3]  # insert kp y coordinates
         targets[:, self.box_offset + 3 :: 3] = kpts[:, 4::3]  # insert kp visibility
 
-        n_targets = len(targets)
+        n_targets = targets.shape[0]
 
         class_targets: list[Tensor] = []
         box_targets: list[Tensor] = []

diff --git a/luxonis_train/attached_modules/losses/smooth_bce_with_logits.py b/luxonis_train/attached_modules/losses/smooth_bce_with_logits.py
@@ -67,5 +67,7 @@ def forward(self, predictions: list[Tensor], target: Tensor) -> Tensor:
             self.negative_smooth_const,
             device=prediction.device,
         )
-        smoothed_target[torch.arange(len(target)), target] = self.positive_smooth_const
+        smoothed_target[
+            torch.arange(target.shape[0]), target
+        ] = self.positive_smooth_const
         return self.criterion.forward(prediction, smoothed_target)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -147,3 +147,4 @@ models_venv/*

		# vscode settings
		.vscode
		tests/data