From 241c74a976657b28c7379243d545a11038358a90 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Fri, 6 Dec 2024 18:54:03 +0100 Subject: [PATCH 01/57] Added necessary requirements. Added integration with pytorch-metric-learning (loss). Added custom reID/embedding metrics. Implemented a test to verify trainability, exportability and inference. Removed GhostFaceNetsV2 from the backbone tests as it only generates embeddings instead of the usual features. --- .../attached_modules/losses/__init__.py | 2 + .../attached_modules/losses/pml_loss.py | 119 ++++ .../attached_modules/metrics/__init__.py | 3 + .../attached_modules/metrics/pml_metrics.py | 248 ++++++++ .../attached_modules/visualizers/__init__.py | 2 + .../visualizers/embeddings_visualizer.py | 95 ++++ luxonis_train/loaders/utils.py | 1 + luxonis_train/nodes/backbones/__init__.py | 2 + luxonis_train/nodes/backbones/ghostfacenet.py | 534 ++++++++++++++++++ requirements.txt | 2 + tests/configs/reid.yaml | 60 ++ tests/integration/test_detection.py | 4 +- tests/integration/test_reid.py | 91 +++ tests/integration/test_segmentation.py | 4 +- 14 files changed, 1165 insertions(+), 2 deletions(-) create mode 100644 luxonis_train/attached_modules/losses/pml_loss.py create mode 100644 luxonis_train/attached_modules/metrics/pml_metrics.py create mode 100644 luxonis_train/attached_modules/visualizers/embeddings_visualizer.py create mode 100644 luxonis_train/nodes/backbones/ghostfacenet.py create mode 100644 tests/configs/reid.yaml create mode 100644 tests/integration/test_reid.py diff --git a/luxonis_train/attached_modules/losses/__init__.py b/luxonis_train/attached_modules/losses/__init__.py index ff0bafc8..b320fada 100644 --- a/luxonis_train/attached_modules/losses/__init__.py +++ b/luxonis_train/attached_modules/losses/__init__.py @@ -7,6 +7,7 @@ from .ohem_bce_with_logits import OHEMBCEWithLogitsLoss from .ohem_cross_entropy import OHEMCrossEntropyLoss from .ohem_loss import OHEMLoss +from .pml_loss import MetricLearningLoss from .reconstruction_segmentation_loss import ReconstructionSegmentationLoss from .sigmoid_focal_loss import SigmoidFocalLoss from .smooth_bce_with_logits import SmoothBCEWithLogitsLoss @@ -26,4 +27,5 @@ "OHEMCrossEntropyLoss", "OHEMBCEWithLogitsLoss", "FOMOLocalizationLoss", + "MetricLearningLoss", ] diff --git a/luxonis_train/attached_modules/losses/pml_loss.py b/luxonis_train/attached_modules/losses/pml_loss.py new file mode 100644 index 00000000..aacd667b --- /dev/null +++ b/luxonis_train/attached_modules/losses/pml_loss.py @@ -0,0 +1,119 @@ +import warnings + +from pytorch_metric_learning.losses import ( + AngularLoss, + ArcFaceLoss, + CircleLoss, + ContrastiveLoss, + CosFaceLoss, + CrossBatchMemory, + DynamicSoftMarginLoss, + FastAPLoss, + GeneralizedLiftedStructureLoss, + HistogramLoss, + InstanceLoss, + IntraPairVarianceLoss, + LargeMarginSoftmaxLoss, + LiftedStructureLoss, + ManifoldLoss, + MarginLoss, + MultiSimilarityLoss, + NCALoss, + NormalizedSoftmaxLoss, + NPairsLoss, + NTXentLoss, + P2SGradLoss, + PNPLoss, + ProxyAnchorLoss, + ProxyNCALoss, + RankedListLoss, + SignalToNoiseRatioContrastiveLoss, + SoftTripleLoss, + SphereFaceLoss, + SubCenterArcFaceLoss, + SupConLoss, + TripletMarginLoss, + TupletMarginLoss, +) +from torch import Tensor + +from .base_loss import BaseLoss + +# Dictionary mapping string keys to loss classes +loss_dict = { + "AngularLoss": AngularLoss, + "ArcFaceLoss": ArcFaceLoss, + "CircleLoss": CircleLoss, + "ContrastiveLoss": ContrastiveLoss, + "CosFaceLoss": CosFaceLoss, + "DynamicSoftMarginLoss": DynamicSoftMarginLoss, + "FastAPLoss": FastAPLoss, + "GeneralizedLiftedStructureLoss": GeneralizedLiftedStructureLoss, + "InstanceLoss": InstanceLoss, + "HistogramLoss": HistogramLoss, + "IntraPairVarianceLoss": IntraPairVarianceLoss, + "LargeMarginSoftmaxLoss": LargeMarginSoftmaxLoss, + "LiftedStructureLoss": LiftedStructureLoss, + "ManifoldLoss": ManifoldLoss, + "MarginLoss": MarginLoss, + "MultiSimilarityLoss": MultiSimilarityLoss, + "NCALoss": NCALoss, + "NormalizedSoftmaxLoss": NormalizedSoftmaxLoss, + "NPairsLoss": NPairsLoss, + "NTXentLoss": NTXentLoss, + "P2SGradLoss": P2SGradLoss, + "PNPLoss": PNPLoss, + "ProxyAnchorLoss": ProxyAnchorLoss, + "ProxyNCALoss": ProxyNCALoss, + "RankedListLoss": RankedListLoss, + "SignalToNoiseRatioContrastiveLoss": SignalToNoiseRatioContrastiveLoss, + "SoftTripleLoss": SoftTripleLoss, + "SphereFaceLoss": SphereFaceLoss, + "SubCenterArcFaceLoss": SubCenterArcFaceLoss, + "SupConLoss": SupConLoss, + "TripletMarginLoss": TripletMarginLoss, + "TupletMarginLoss": TupletMarginLoss, +} + + +class MetricLearningLoss(BaseLoss): + def __init__( + self, + loss_name: str, + embedding_size: int = 512, + cross_batch_memory_size=0, + loss_kwargs: dict | None = None, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + if loss_kwargs is None: + loss_kwargs = {} + self.loss_func = loss_dict[loss_name]( + **loss_kwargs + ) # Instantiate the loss object + if cross_batch_memory_size > 0: + if loss_name in CrossBatchMemory.supported_losses(): + self.loss_func = CrossBatchMemory( + self.loss_func, embedding_size=embedding_size + ) + else: + # Warn that cross_batch_memory_size is ignored + warnings.warn( + f"Cross batch memory is not supported for {loss_name}. Ignoring cross_batch_memory_size" + ) + + # self.miner_func = miner_func + + def prepare(self, inputs, labels): + embeddings = inputs["features"][0] + + IDs = labels["id"][0][:, 0] + return embeddings, IDs + + def forward(self, inputs: Tensor, target: Tensor): + # miner_output = self.miner_func(inputs, target) + + loss = self.loss_func(inputs, target) + + return loss diff --git a/luxonis_train/attached_modules/metrics/__init__.py b/luxonis_train/attached_modules/metrics/__init__.py index b1dc40ea..c43f32b4 100644 --- a/luxonis_train/attached_modules/metrics/__init__.py +++ b/luxonis_train/attached_modules/metrics/__init__.py @@ -2,6 +2,7 @@ from .mean_average_precision import MeanAveragePrecision from .mean_average_precision_keypoints import MeanAveragePrecisionKeypoints from .object_keypoint_similarity import ObjectKeypointSimilarity +from .pml_metrics import ClosestIsPositiveAccuracy, MedianDistances from .torchmetrics import Accuracy, F1Score, JaccardIndex, Precision, Recall __all__ = [ @@ -14,4 +15,6 @@ "ObjectKeypointSimilarity", "Precision", "Recall", + "ClosestIsPositiveAccuracy", + "MedianDistances", ] diff --git a/luxonis_train/attached_modules/metrics/pml_metrics.py b/luxonis_train/attached_modules/metrics/pml_metrics.py new file mode 100644 index 00000000..b280742d --- /dev/null +++ b/luxonis_train/attached_modules/metrics/pml_metrics.py @@ -0,0 +1,248 @@ +import torch +from torch import Tensor + +from .base_metric import BaseMetric + +# Converted from https://omoindrot.github.io/triplet-loss#offline-and-online-triplet-mining +# to PyTorch from TensorFlow + + +def _pairwise_distances(embeddings, squared=False): + """Compute the 2D matrix of distances between all the embeddings. + + Args: + embeddings: tensor of shape (batch_size, embed_dim) + squared: Boolean. If true, output is the pairwise squared euclidean distance matrix. + If false, output is the pairwise euclidean distance matrix. + + Returns: + pairwise_distances: tensor of shape (batch_size, batch_size) + """ + # Get the dot product between all embeddings + # shape (batch_size, batch_size) + dot_product = torch.matmul(embeddings, embeddings.t()) + + # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`. + # This also provides more numerical stability (the diagonal of the result will be exactly 0). + # shape (batch_size,) + square_norm = torch.diag(dot_product) + + # Compute the pairwise distance matrix as we have: + # ||a - b||^2 = ||a||^2 - 2 + ||b||^2 + # shape (batch_size, batch_size) + distances = ( + square_norm.unsqueeze(0) - 2.0 * dot_product + square_norm.unsqueeze(1) + ) + + # Because of computation errors, some distances might be negative so we put everything >= 0.0 + distances = torch.max(distances, torch.tensor(0.0)) + + if not squared: + # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal) + # we need to add a small epsilon where distances == 0.0 + mask = (distances == 0.0).float() + distances = distances + mask * 1e-16 + + distances = torch.sqrt(distances) + + # Correct the epsilon added: set the distances on the mask to be exactly 0.0 + distances = distances * (1.0 - mask) + + return distances + + +def _get_anchor_positive_triplet_mask(labels): + indices_equal = torch.eye( + labels.shape[0], dtype=torch.uint8, device=labels.device + ) + indices_not_equal = ~indices_equal + labels_equal = labels.unsqueeze(0) == labels.unsqueeze(1) + mask = indices_not_equal & labels_equal + return mask + + +class ClosestIsPositiveAccuracy(BaseMetric): + def __init__(self, cross_batch_memory_size=0, **kwargs): + super().__init__(**kwargs) + self.cross_batch_memory_size = cross_batch_memory_size + self.add_state("cross_batch_memory", default=[], dist_reduce_fx="cat") + self.add_state( + "correct_predictions", + default=torch.tensor(0), + dist_reduce_fx="sum", + ) + self.add_state( + "total_predictions", default=torch.tensor(0), dist_reduce_fx="sum" + ) + + def prepare(self, inputs, labels): + embeddings = inputs["features"][0] + IDs = labels["id"][0][:, 0] + return embeddings, IDs + + def update(self, inputs: Tensor, target: Tensor): + embeddings, labels = inputs, target + + if self.cross_batch_memory_size > 0: + # Append embedding and labels to the memory + self.cross_batch_memory.extend(list(zip(embeddings, labels))) + + # If the memory is full, remove the oldest elements + if len(self.cross_batch_memory) > self.cross_batch_memory_size: + self.cross_batch_memory = self.cross_batch_memory[ + -self.cross_batch_memory_size : + ] + + # If the memory is not full, return + if len(self.cross_batch_memory) < self.cross_batch_memory_size: + return + + # Get the embeddings and labels from the memory + embeddings, labels = zip(*self.cross_batch_memory) + embeddings = torch.stack(embeddings) + labels = torch.stack(labels) + + # print(f"Calculating accuracy for {len(embeddings)} embeddings") + + # Get the pairwise distances between all embeddings + pairwise_distances = _pairwise_distances(embeddings) + + # Set diagonal to infinity so that the closest embedding is not the same embedding + pairwise_distances.fill_diagonal_(float("inf")) + + # Find the closest embedding for each query embedding + closest_indices = torch.argmin(pairwise_distances, dim=1) + + # Get the labels of the closest embeddings + closest_labels = labels[closest_indices] + + # Filter out embeddings that don't have both positive and negative examples + positive_mask = _get_anchor_positive_triplet_mask(labels) + num_positives = positive_mask.sum(dim=1) + has_at_least_one_positive_and_negative = (num_positives > 0) & ( + num_positives < len(labels) + ) + + # Filter embeddings, labels, and closest indices based on valid indices + filtered_labels = labels[has_at_least_one_positive_and_negative] + filtered_closest_labels = closest_labels[ + has_at_least_one_positive_and_negative + ] + + # Calculate the number of correct predictions where the closest is positive + correct_predictions = ( + filtered_labels == filtered_closest_labels + ).sum() + + # Update the metric state + self.correct_predictions += correct_predictions + self.total_predictions += len(filtered_labels) + + def compute(self): + return self.correct_predictions / self.total_predictions + + +class MedianDistances(BaseMetric): + def __init__(self, cross_batch_memory_size=0, **kwargs): + super().__init__(**kwargs) + self.cross_batch_memory_size = cross_batch_memory_size + self.add_state("cross_batch_memory", default=[], dist_reduce_fx="cat") + self.add_state("all_distances", default=[], dist_reduce_fx="cat") + self.add_state("closest_distances", default=[], dist_reduce_fx="cat") + self.add_state("positive_distances", default=[], dist_reduce_fx="cat") + self.add_state( + "closest_vs_positive_distances", default=[], dist_reduce_fx="cat" + ) + + def prepare(self, inputs, labels): + embeddings = inputs["features"][0] + IDs = labels["id"][0][:, 0] + return embeddings, IDs + + def update(self, inputs: Tensor, target: Tensor): + embeddings, labels = inputs, target + + if self.cross_batch_memory_size > 0: + # Append embedding and labels to the memory + self.cross_batch_memory.extend(list(zip(embeddings, labels))) + + # If the memory is full, remove the oldest elements + if len(self.cross_batch_memory) > self.cross_batch_memory_size: + self.cross_batch_memory = self.cross_batch_memory[ + -self.cross_batch_memory_size : + ] + + # If the memory is not full, return + if len(self.cross_batch_memory) < self.cross_batch_memory_size: + return + + # Get the embeddings and labels from the memory + embeddings, labels = zip(*self.cross_batch_memory) + embeddings = torch.stack(embeddings) + labels = torch.stack(labels) + + # Get the pairwise distances between all embeddings + pairwise_distances = _pairwise_distances(embeddings) + # Append only upper triangular part of the matrix + self.all_distances.append( + pairwise_distances[ + torch.triu(torch.ones_like(pairwise_distances), diagonal=1) + == 1 + ].flatten() + ) + + # Set diagonal to infinity so that the closest embedding is not the same embedding + pairwise_distances.fill_diagonal_(float("inf")) + + # Get the closest distance for each query embedding + closest_distances, _ = torch.min(pairwise_distances, dim=1) + self.closest_distances.append(closest_distances) + + # Get the positive mask and convert it to boolean + positive_mask = _get_anchor_positive_triplet_mask(labels).bool() + + only_positive_distances = pairwise_distances.clone() + only_positive_distances[~positive_mask] = float("inf") + + closest_positive_distances, _ = torch.min( + only_positive_distances, dim=1 + ) + + non_inf_mask = closest_positive_distances != float("inf") + difference = closest_positive_distances - closest_distances + difference = difference[non_inf_mask] + + # Update the metric state + self.closest_vs_positive_distances.append(difference) + self.positive_distances.append( + closest_positive_distances[non_inf_mask] + ) + + def compute(self): + if len(self.all_distances) == 0: + # Return NaN tensor if no distances were calculated + return { + "MedianDistance": torch.tensor(float("nan")), + "MedianClosestDistance": torch.tensor(float("nan")), + "MedianClosestPositiveDistance": torch.tensor(float("nan")), + "MedianClosestVsClosestPositiveDistance": torch.tensor( + float("nan") + ), + } + + all_distances = torch.cat(self.all_distances) + closest_distances = torch.cat(self.closest_distances) + positive_distances = torch.cat(self.positive_distances) + closest_vs_positive_distances = torch.cat( + self.closest_vs_positive_distances + ) + + # Return medians + return { + "MedianDistance": torch.median(all_distances), + "MedianClosestDistance": torch.median(closest_distances), + "MedianClosestPositiveDistance": torch.median(positive_distances), + "MedianClosestVsClosestPositiveDistance": torch.median( + closest_vs_positive_distances + ), + } diff --git a/luxonis_train/attached_modules/visualizers/__init__.py b/luxonis_train/attached_modules/visualizers/__init__.py index 50b90471..69ecc3c4 100644 --- a/luxonis_train/attached_modules/visualizers/__init__.py +++ b/luxonis_train/attached_modules/visualizers/__init__.py @@ -1,6 +1,7 @@ from .base_visualizer import BaseVisualizer from .bbox_visualizer import BBoxVisualizer from .classification_visualizer import ClassificationVisualizer +from .embeddings_visualizer import EmbeddingsVisualizer from .keypoint_visualizer import KeypointVisualizer from .multi_visualizer import MultiVisualizer from .segmentation_visualizer import SegmentationVisualizer @@ -23,6 +24,7 @@ "KeypointVisualizer", "MultiVisualizer", "SegmentationVisualizer", + "EmbeddingsVisualizer", "combine_visualizations", "draw_bounding_box_labels", "draw_keypoint_labels", diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py new file mode 100644 index 00000000..b5fb5f0e --- /dev/null +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -0,0 +1,95 @@ +import logging + +from matplotlib import pyplot as plt +from sklearn.manifold import TSNE +from torch import Tensor + +from luxonis_train.utils import Labels, Packet + +from .base_visualizer import BaseVisualizer +from .utils import ( + figure_to_torch, +) + +logger = logging.getLogger(__name__) +log_disable = False + + +class EmbeddingsVisualizer(BaseVisualizer[Tensor, Tensor]): + # supported_tasks: list[TaskType] = [TaskType.LABEL] + + def __init__( + self, + **kwargs, + ): + """Visualizer for embedding tasks like reID.""" + super().__init__(**kwargs) + + def prepare( + self, inputs: Packet[Tensor], labels: Labels | None + ) -> tuple[Tensor, Tensor]: + embeddings = inputs["features"][0] + IDs = labels["id"][0] + return embeddings, IDs + + def forward( + self, + label_canvas: Tensor, + prediction_canvas: Tensor, + embeddings: Tensor, + IDs: Tensor | None, + **kwargs, + ) -> Tensor: + """Creates a visualization of the embeddings. + + @type label_canvas: Tensor + @param label_canvas: The canvas to draw the labels on. + @type prediction_canvas: Tensor + @param prediction_canvas: The canvas to draw the predictions on. + @type embeddings: Tensor + @param embeddings: The embeddings to visualize. + @type IDs: Tensor + @param IDs: The IDs to visualize. + @rtype: Tensor + @return: An embedding space projection. + """ + + # Embeddings: [B, D], D = e.g. 512 + # IDs: [B, 1], corresponding to the embeddings + + # Convert embeddings to numpy array + embeddings_np = embeddings.detach().cpu().numpy() + + # Perplexity must be less than the number of samples + perplexity = min(30, embeddings_np.shape[0] - 1) + + # Reduce dimensionality to 2D using t-SNE + tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity) + embeddings_2d = tsne.fit_transform(embeddings_np) + + # Plot the embeddings + fig, ax = plt.subplots(figsize=(10, 10)) + scatter = ax.scatter( + embeddings_2d[:, 0], + embeddings_2d[:, 1], + c=IDs.detach().cpu().numpy(), + cmap="viridis", + s=5, + ) + fig.colorbar(scatter, ax=ax) + ax.set_title("Embeddings Visualization") + ax.set_xlabel("Dimension 1") + ax.set_ylabel("Dimension 2") + + # Convert figure to tensor + image_tensor = figure_to_torch( + fig, width=label_canvas.shape[3], height=label_canvas.shape[2] + ) + + # Close the figure to free memory + plt.close(fig) + + # Add fake batch dimension + image_tensor = image_tensor.unsqueeze(0) + + return image_tensor diff --git a/luxonis_train/loaders/utils.py b/luxonis_train/loaders/utils.py index b030e218..2782500e 100644 --- a/luxonis_train/loaders/utils.py +++ b/luxonis_train/loaders/utils.py @@ -38,6 +38,7 @@ def collate_fn( TaskType.CLASSIFICATION, TaskType.SEGMENTATION, TaskType.ARRAY, + TaskType.LABEL, ]: out_labels[task] = torch.stack(annos, 0), task_type diff --git a/luxonis_train/nodes/backbones/__init__.py b/luxonis_train/nodes/backbones/__init__.py index cc621625..f5319981 100644 --- a/luxonis_train/nodes/backbones/__init__.py +++ b/luxonis_train/nodes/backbones/__init__.py @@ -2,6 +2,7 @@ from .ddrnet import DDRNet from .efficientnet import EfficientNet from .efficientrep import EfficientRep +from .ghostfacenet import GhostFaceNetsV2 from .micronet import MicroNet from .mobilenetv2 import MobileNetV2 from .mobileone import MobileOne @@ -22,4 +23,5 @@ "ResNet", "DDRNet", "RecSubNet", + "GhostFaceNetsV2", ] diff --git a/luxonis_train/nodes/backbones/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet.py new file mode 100644 index 00000000..b4b17758 --- /dev/null +++ b/luxonis_train/nodes/backbones/ghostfacenet.py @@ -0,0 +1,534 @@ +# Original source: https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py + + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from luxonis_train.nodes.base_node import BaseNode + + +def _make_divisible(v, divisor, min_value=None): + """This function is taken from the original tf repo. + + It ensures that all layers have a channel number that is divisible by 8 + It can be seen here: + https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py + """ + if min_value is None: + min_value = divisor + new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_v < 0.9 * v: + new_v += divisor + return new_v + + +def hard_sigmoid(x, inplace: bool = False): + if inplace: + return x.add_(3.0).clamp_(0.0, 6.0).div_(6.0) + else: + return F.relu6(x + 3.0) / 6.0 + + +class SqueezeExcite(nn.Module): + def __init__( + self, + in_chs, + se_ratio=0.25, + reduced_base_chs=None, + act_layer=nn.PReLU, + gate_fn=hard_sigmoid, + divisor=4, + **_, + ): + super(SqueezeExcite, self).__init__() + self.gate_fn = gate_fn + reduced_chs = _make_divisible( + (reduced_base_chs or in_chs) * se_ratio, divisor + ) + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True) + self.act1 = act_layer() + self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True) + + def forward(self, x): + x_se = self.avg_pool(x) + x_se = self.conv_reduce(x_se) + x_se = self.act1(x_se) + x_se = self.conv_expand(x_se) + x = x * self.gate_fn(x_se) + return x + + +class ConvBnAct(nn.Module): + def __init__( + self, in_chs, out_chs, kernel_size, stride=1, act_layer=nn.PReLU + ): + super(ConvBnAct, self).__init__() + self.conv = nn.Conv2d( + in_chs, out_chs, kernel_size, stride, kernel_size // 2, bias=False + ) + self.bn1 = nn.BatchNorm2d(out_chs) + self.act1 = act_layer() + + def forward(self, x): + x = self.conv(x) + x = self.bn1(x) + x = self.act1(x) + return x + + +class ModifiedGDC(nn.Module): + def __init__( + self, image_size, in_chs, num_classes, dropout, emb=512 + ): # dropout implementation is in the original code but not in the paper + super(ModifiedGDC, self).__init__() + + if image_size % 32 == 0: + self.conv_dw = nn.Conv2d( + in_chs, + in_chs, + kernel_size=(image_size // 32), + groups=in_chs, + bias=False, + ) + else: + self.conv_dw = nn.Conv2d( + in_chs, + in_chs, + kernel_size=(image_size // 32 + 1), + groups=in_chs, + bias=False, + ) + self.bn1 = nn.BatchNorm2d(in_chs) + self.dropout = nn.Dropout(dropout) + + self.conv = nn.Conv2d(in_chs, emb, kernel_size=1, bias=False) + self.bn2 = nn.BatchNorm1d(emb) + self.linear = ( + nn.Linear(emb, num_classes) if num_classes else nn.Identity() + ) + + def forward(self, inps): + x = inps + x = self.conv_dw(x) + x = self.bn1(x) + x = self.dropout(x) + # # Add spots to the features + # x = torch.cat([x, spots.view(spots.size(0), -1, 1, 1)], dim=1) + x = self.conv(x) + x = x.view(x.size(0), -1) # Flatten + x = self.bn2(x) + x = self.linear(x) + return x + + +class GhostModuleV2(nn.Module): + def __init__( + self, + inp, + oup, + kernel_size=1, + ratio=2, + dw_size=3, + stride=1, + prelu=True, + mode=None, + args=None, + ): + super(GhostModuleV2, self).__init__() + self.mode = mode + self.gate_fn = nn.Sigmoid() + + if self.mode in ["original"]: + self.oup = oup + init_channels = math.ceil(oup / ratio) + new_channels = init_channels * (ratio - 1) + self.primary_conv = nn.Sequential( + nn.Conv2d( + inp, + init_channels, + kernel_size, + stride, + kernel_size // 2, + bias=False, + ), + nn.BatchNorm2d(init_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + self.cheap_operation = nn.Sequential( + nn.Conv2d( + init_channels, + new_channels, + dw_size, + 1, + dw_size // 2, + groups=init_channels, + bias=False, + ), + nn.BatchNorm2d(new_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + elif self.mode in ["attn"]: # DFC + self.oup = oup + init_channels = math.ceil(oup / ratio) + new_channels = init_channels * (ratio - 1) + self.primary_conv = nn.Sequential( + nn.Conv2d( + inp, + init_channels, + kernel_size, + stride, + kernel_size // 2, + bias=False, + ), + nn.BatchNorm2d(init_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + self.cheap_operation = nn.Sequential( + nn.Conv2d( + init_channels, + new_channels, + dw_size, + 1, + dw_size // 2, + groups=init_channels, + bias=False, + ), + nn.BatchNorm2d(new_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + self.short_conv = nn.Sequential( + nn.Conv2d( + inp, oup, kernel_size, stride, kernel_size // 2, bias=False + ), + nn.BatchNorm2d(oup), + nn.Conv2d( + oup, + oup, + kernel_size=(1, 5), + stride=1, + padding=(0, 2), + groups=oup, + bias=False, + ), + nn.BatchNorm2d(oup), + nn.Conv2d( + oup, + oup, + kernel_size=(5, 1), + stride=1, + padding=(2, 0), + groups=oup, + bias=False, + ), + nn.BatchNorm2d(oup), + ) + + def forward(self, x): + if self.mode in ["original"]: + x1 = self.primary_conv(x) + x2 = self.cheap_operation(x1) + out = torch.cat([x1, x2], dim=1) + return out[:, : self.oup, :, :] + elif self.mode in ["attn"]: + res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2)) + x1 = self.primary_conv(x) + x2 = self.cheap_operation(x1) + out = torch.cat([x1, x2], dim=1) + return out[:, : self.oup, :, :] * F.interpolate( + self.gate_fn(res), + size=(out.shape[-2], out.shape[-1]), + mode="nearest", + ) + + +class GhostBottleneckV2(nn.Module): + def __init__( + self, + in_chs, + mid_chs, + out_chs, + dw_kernel_size=3, + stride=1, + act_layer=nn.PReLU, + se_ratio=0.0, + layer_id=None, + args=None, + ): + super(GhostBottleneckV2, self).__init__() + has_se = se_ratio is not None and se_ratio > 0.0 + self.stride = stride + + # Point-wise expansion + if layer_id <= 1: + self.ghost1 = GhostModuleV2( + in_chs, mid_chs, prelu=True, mode="original", args=args + ) + else: + self.ghost1 = GhostModuleV2( + in_chs, mid_chs, prelu=True, mode="attn", args=args + ) + + # Depth-wise convolution + if self.stride > 1: + self.conv_dw = nn.Conv2d( + mid_chs, + mid_chs, + dw_kernel_size, + stride=stride, + padding=(dw_kernel_size - 1) // 2, + groups=mid_chs, + bias=False, + ) + self.bn_dw = nn.BatchNorm2d(mid_chs) + + # Squeeze-and-excitation + if has_se: + self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio) + else: + self.se = None + + self.ghost2 = GhostModuleV2( + mid_chs, out_chs, prelu=False, mode="original", args=args + ) + + # shortcut + if in_chs == out_chs and self.stride == 1: + self.shortcut = nn.Sequential() + else: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_chs, + in_chs, + dw_kernel_size, + stride=stride, + padding=(dw_kernel_size - 1) // 2, + groups=in_chs, + bias=False, + ), + nn.BatchNorm2d(in_chs), + nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(out_chs), + ) + + def forward(self, x): + residual = x + x = self.ghost1(x) + if self.stride > 1: + x = self.conv_dw(x) + x = self.bn_dw(x) + if self.se is not None: + x = self.se(x) + x = self.ghost2(x) + x += self.shortcut(residual) + return x + + +# NODES.register_module() +class GhostFaceNetsV2(BaseNode[torch.Tensor, list[torch.Tensor]]): + def unwrap(self, inputs): + return [inputs[0]["features"][0]] + + def wrap(self, outputs): + return {"features": [outputs]} + + def set_export_mode(self, mode: bool = True): + self.export_mode = mode + self.train(not mode) + + def __init__( + self, + cfgs=None, + embedding_size=512, + num_classes=0, + width=1.0, + dropout=0.2, + block=GhostBottleneckV2, + add_pointwise_conv=False, + bn_momentum=0.9, + bn_epsilon=1e-5, + init_kaiming=True, + block_args=None, + *args, + **kwargs, + ): + # kwargs['_tasks'] = {TaskType.LABEL: 'features'} + super().__init__(*args, **kwargs) + + inp_shape = kwargs["input_shapes"][0]["features"][0] + # spots_shape = kwargs['input_shapes'][0]['features'][1] + + image_size = inp_shape[2] + channels = inp_shape[1] + if cfgs is None: + self.cfgs = [ + # k, t, c, SE, s + [[3, 16, 16, 0, 1]], + [[3, 48, 24, 0, 2]], + [[3, 72, 24, 0, 1]], + [[5, 72, 40, 0.25, 2]], + [[5, 120, 40, 0.25, 1]], + [[3, 240, 80, 0, 2]], + [ + [3, 200, 80, 0, 1], + [3, 184, 80, 0, 1], + [3, 184, 80, 0, 1], + [3, 480, 112, 0.25, 1], + [3, 672, 112, 0.25, 1], + ], + [[5, 672, 160, 0.25, 2]], + [ + [5, 960, 160, 0, 1], + [5, 960, 160, 0.25, 1], + [5, 960, 160, 0, 1], + [5, 960, 160, 0.25, 1], + ], + ] + else: + self.cfgs = cfgs + + # building first layer + output_channel = _make_divisible(16 * width, 4) + self.conv_stem = nn.Conv2d( + channels, output_channel, 3, 2, 1, bias=False + ) + self.bn1 = nn.BatchNorm2d(output_channel) + self.act1 = nn.PReLU() + input_channel = output_channel + + # building inverted residual blocks + stages = [] + layer_id = 0 + for cfg in self.cfgs: + layers = [] + for k, exp_size, c, se_ratio, s in cfg: + output_channel = _make_divisible(c * width, 4) + hidden_channel = _make_divisible(exp_size * width, 4) + if block == GhostBottleneckV2: + layers.append( + block( + input_channel, + hidden_channel, + output_channel, + k, + s, + se_ratio=se_ratio, + layer_id=layer_id, + args=block_args, + ) + ) + input_channel = output_channel + layer_id += 1 + stages.append(nn.Sequential(*layers)) + + output_channel = _make_divisible(exp_size * width, 4) + stages.append( + nn.Sequential(ConvBnAct(input_channel, output_channel, 1)) + ) + + self.blocks = nn.Sequential(*stages) + + # building last several layers + pointwise_conv = [] + if add_pointwise_conv: + pointwise_conv.append( + nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=True) + ) + pointwise_conv.append(nn.BatchNorm2d(output_channel)) + pointwise_conv.append(nn.PReLU()) + else: + pointwise_conv.append(nn.Sequential()) + + self.pointwise_conv = nn.Sequential(*pointwise_conv) + self.classifier = ModifiedGDC( + image_size, output_channel, num_classes, dropout, embedding_size + ) + + # Initialize weights + for m in self.modules(): + if init_kaiming: + if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) + negative_slope = 0.25 # Default value for PReLU in PyTorch, change it if you use custom value + m.weight.data.normal_( + 0, math.sqrt(2.0 / (fan_in * (1 + negative_slope**2))) + ) + if isinstance(m, nn.BatchNorm2d): + m.momentum, m.eps = bn_momentum, bn_epsilon + + def forward(self, inps): + x = inps[0] + x = self.conv_stem(x) + x = self.bn1(x) + x = self.act1(x) + x = self.blocks(x) + x = self.pointwise_conv(x) + x = self.classifier(x) + return x + + # @property + # def task(self) -> str: + # return "label" + + # @property + # def tasks(self) -> dict: + # return [TaskType.LABEL] + + +if __name__ == "__main__": + W, H = 256, 256 + model = GhostFaceNetsV2(image_size=W) + model.eval() # Set the model to evaluation mode + + # Create a dummy input tensor of the appropriate size + x = torch.randn(1, 3, H, W) + + # Export the model + onnx_path = "ghostfacenet.onnx" + torch.onnx.export( + model, # model being run + x, # model input (or a tuple for multiple inputs) + onnx_path, # where to save the model (can be a file or file-like object) + export_params=True, # store the trained parameter weights inside the model file + opset_version=12, # the ONNX version to export the model to + do_constant_folding=True, # whether to execute constant folding for optimization + input_names=["input"], # the model's input names + output_names=["output"], # the model's output names + # dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes + # 'output' : {0 : 'batch_size'}} + ) + import os + + import numpy as np + import onnx + import onnxsim + + # logger.info("Simplifying ONNX model...") + model_onnx = onnx.load(onnx_path) + onnx_model, check = onnxsim.simplify(model_onnx) + if not check: + raise RuntimeError("Onnx simplify failed.") + onnx.save(onnx_model, onnx_path) + + # Add calibration data + dir = "shared_with_container/calibration_data/" + for file in os.listdir(dir): + os.remove(dir + file) + for i in range(20): + np_array = np.random.rand(1, 3, H, W).astype(np.float32) + np.save(f"{dir}{i:02d}.npy", np_array) + np_array.tofile(f"{dir}{i:02d}.raw") + + # Test backpropagation on the model + # Create a dummy target tensor of the appropriate size + Y = model(x) + target = torch.randn(1, 512) + loss_fn = torch.nn.MSELoss() + loss = loss_fn(Y, target) + model.zero_grad() + loss.backward() + print("Backpropagation test successful") diff --git a/requirements.txt b/requirements.txt index 5d0fcb28..94badc1c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,3 +18,5 @@ mlflow>=2.10.0 psutil>=5.0.0 tabulate>=0.9.0 grad-cam>=1.5.4 +pytorch_metric_learning>=2.7.0 +scikit-learn>=1.5.0 \ No newline at end of file diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml new file mode 100644 index 00000000..d9c0ec11 --- /dev/null +++ b/tests/configs/reid.yaml @@ -0,0 +1,60 @@ +loader: + name: CustomReIDLoader + +model: + name: reid_test + nodes: + - name: GhostFaceNetsV2 + input_sources: + - image + params: + embedding_size: &embedding_size 512 + + losses: + - name: MetricLearningLoss + params: + loss_name: SupConLoss + embedding_size: *embedding_size + cross_batch_memory_size: &memory_size 200 + attached_to: GhostFaceNetsV2 + + metrics: + - name: ClosestIsPositiveAccuracy + params: + cross_batch_memory_size: *memory_size + attached_to: GhostFaceNetsV2 + is_main_metric: True + - name: MedianDistances + params: + cross_batch_memory_size: *memory_size + attached_to: GhostFaceNetsV2 + is_main_metric: False + + visualizers: + - name: EmbeddingsVisualizer + attached_to: GhostFaceNetsV2 + +trainer: + preprocessing: + train_image_size: [256, 256] + + batch_size: 16 + epochs: 10 + n_workers: 0 + validation_interval: 10 + + callbacks: + - name: ExportOnTrainEnd + + optimizer: + name: Adam + params: + lr: 0.01 + +tracker: + project_name: reid_example + is_tensorboard: True + +exporter: + onnx: + opset_version: 11 \ No newline at end of file diff --git a/tests/integration/test_detection.py b/tests/integration/test_detection.py index 45e83f0a..060e84e2 100644 --- a/tests/integration/test_detection.py +++ b/tests/integration/test_detection.py @@ -103,7 +103,9 @@ def train_and_test( assert value > 0.8, f"{name} = {value} (expected > 0.8)" -@pytest.mark.parametrize("backbone", BACKBONES) +@pytest.mark.parametrize( + "backbone", [b for b in BACKBONES if b != "GhostFaceNetsV2"] +) def test_backbones( backbone: str, config: dict[str, Any], diff --git a/tests/integration/test_reid.py b/tests/integration/test_reid.py new file mode 100644 index 00000000..9ed4e867 --- /dev/null +++ b/tests/integration/test_reid.py @@ -0,0 +1,91 @@ +import shutil +from pathlib import Path +from typing import Any + +import pytest +import torch + +from luxonis_train.core import LuxonisModel +from luxonis_train.enums import TaskType +from luxonis_train.loaders import BaseLoaderTorch + +from .multi_input_modules import * + +INFER_PATH = Path("tests/integration/infer-save-directory") +ONNX_PATH = Path("tests/integration/_model.onnx") +STUDY_PATH = Path("study_local.db") + + +class CustomReIDLoader(BaseLoaderTorch): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @property + def input_shapes(self): + return { + "image": torch.Size([3, 256, 256]), + "id": torch.Size([1]), + } + + def __getitem__(self, _): # pragma: no cover + # Fake data + image = torch.rand(3, 256, 256, dtype=torch.float32) + inputs = { + "image": image, + } + + # Fake labels + id = torch.randint(0, 1000, (1,), dtype=torch.int64) + labels = { + "id": (id, TaskType.LABEL), + } + + return inputs, labels + + def __len__(self): + return 10 + + def get_classes(self) -> dict[TaskType, list[str]]: + return {TaskType.LABEL: ["id"]} + + +@pytest.fixture +def infer_path() -> Path: + if INFER_PATH.exists(): + shutil.rmtree(INFER_PATH) + INFER_PATH.mkdir() + return INFER_PATH + + +@pytest.fixture +def opts(test_output_dir: Path) -> dict[str, Any]: + return { + "trainer.epochs": 1, + "trainer.batch_size": 2, + "trainer.validation_interval": 1, + "trainer.callbacks": "[]", + "tracker.save_directory": str(test_output_dir), + "tuner.n_trials": 4, + } + + +@pytest.fixture(scope="function", autouse=True) +def clear_files(): + yield + STUDY_PATH.unlink(missing_ok=True) + ONNX_PATH.unlink(missing_ok=True) + + +def test_reid(opts: dict[str, Any], infer_path: Path): + config_file = "tests/configs/reid.yaml" + model = LuxonisModel(config_file, opts) + model.train() + model.test(view="val") + + assert not ONNX_PATH.exists() + model.export(str(ONNX_PATH)) + assert ONNX_PATH.exists() + + assert len(list(infer_path.iterdir())) == 0 + model.infer(view="val", save_dir=infer_path) + assert infer_path.exists() diff --git a/tests/integration/test_segmentation.py b/tests/integration/test_segmentation.py index a8b4df91..4ab4478a 100644 --- a/tests/integration/test_segmentation.py +++ b/tests/integration/test_segmentation.py @@ -123,7 +123,9 @@ def train_and_test( assert value > 0.8, f"{name} = {value} (expected > 0.8)" -@pytest.mark.parametrize("backbone", BACKBONES) +@pytest.mark.parametrize( + "backbone", [b for b in BACKBONES if b != "GhostFaceNetsV2"] +) def test_backbones( backbone: str, config: dict[str, Any], From be4c2d23f3039496c0e220f5efea7e6f9256fbdc Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Fri, 6 Dec 2024 19:05:22 +0100 Subject: [PATCH 02/57] Add detailed docstring for GhostFaceNetsV2 backbone class --- luxonis_train/nodes/backbones/ghostfacenet.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/luxonis_train/nodes/backbones/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet.py index b4b17758..c242633f 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet.py @@ -356,6 +356,44 @@ def __init__( *args, **kwargs, ): + """GhostFaceNetsV2 backbone. + + GhostFaceNetsV2 is a convolutional neural network architecture focused on face recognition, but it is + adaptable to generic embedding tasks. It is based on the GhostNet architecture and uses Ghost BottleneckV2 blocks. + + Source: U{https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py} + + @license: U{MIT License + } + + @see: U{GhostFaceNets: Lightweight Face Recognition Model From Cheap Operations + } + + @type cfgs: list[list[list[int]]] | None + @param cfgs: List of Ghost BottleneckV2 configurations. Defaults to None, which uses the original GhostFaceNetsV2 configuration. + @type embedding_size: int + @param embedding_size: Size of the embedding. Defaults to 512. + @type num_classes: int + @param num_classes: Number of classes. Defaults to 0, which makes the network output the raw embeddings. Otherwise it can be used to + add another linear layer to the network, which is useful for training using ArcFace or similar classification-based losses that + require the user to drop the last layer of the network. + @type width: float + @param width: Width multiplier. Increases complexity and number of parameters. Defaults to 1.0. + @type dropout: float + @param dropout: Dropout rate. Defaults to 0.2. + @type block: nn.Module + @param block: Ghost BottleneckV2 block. Defaults to GhostBottleneckV2. + @type add_pointwise_conv: bool + @param add_pointwise_conv: If True, adds a pointwise convolution layer at the end of the network. Defaults to False. + @type bn_momentum: float + @param bn_momentum: Batch normalization momentum. Defaults to 0.9. + @type bn_epsilon: float + @param bn_epsilon: Batch normalization epsilon. Defaults to 1e-5. + @type init_kaiming: bool + @param init_kaiming: If True, initializes the weights using the Kaiming initialization. Defaults to True. + @type block_args: dict + @param block_args: Arguments to pass to the block. Defaults to None. + """ # kwargs['_tasks'] = {TaskType.LABEL: 'features'} super().__init__(*args, **kwargs) From c5c4f16463efc41f20898a47e65573740513b0c4 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Fri, 6 Dec 2024 19:08:01 +0100 Subject: [PATCH 03/57] fix: update docstring for pairwise_distances function in pml_metrics.py --- .../attached_modules/metrics/pml_metrics.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/luxonis_train/attached_modules/metrics/pml_metrics.py b/luxonis_train/attached_modules/metrics/pml_metrics.py index b280742d..fdd66a41 100644 --- a/luxonis_train/attached_modules/metrics/pml_metrics.py +++ b/luxonis_train/attached_modules/metrics/pml_metrics.py @@ -10,13 +10,15 @@ def _pairwise_distances(embeddings, squared=False): """Compute the 2D matrix of distances between all the embeddings. - Args: - embeddings: tensor of shape (batch_size, embed_dim) - squared: Boolean. If true, output is the pairwise squared euclidean distance matrix. - If false, output is the pairwise euclidean distance matrix. - - Returns: - pairwise_distances: tensor of shape (batch_size, batch_size) + @param embeddings: tensor of shape (batch_size, embed_dim) + @type embeddings: torch.Tensor + @param squared: If true, output is the pairwise squared euclidean + distance matrix. If false, output is the pairwise euclidean + distance matrix. + @type squared: bool + @return: pairwise_distances: tensor of shape (batch_size, + batch_size) + @rtype: torch.Tensor """ # Get the dot product between all embeddings # shape (batch_size, batch_size) From c360f15bc953405cf1ea15c6787ced733e479e41 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Fri, 6 Dec 2024 19:12:30 +0100 Subject: [PATCH 04/57] Fixed type errors --- .../attached_modules/losses/pml_loss.py | 3 + .../attached_modules/metrics/pml_metrics.py | 8 +++ .../visualizers/embeddings_visualizer.py | 25 +++++-- luxonis_train/nodes/backbones/ghostfacenet.py | 65 +------------------ 4 files changed, 31 insertions(+), 70 deletions(-) diff --git a/luxonis_train/attached_modules/losses/pml_loss.py b/luxonis_train/attached_modules/losses/pml_loss.py index aacd667b..1727d091 100644 --- a/luxonis_train/attached_modules/losses/pml_loss.py +++ b/luxonis_train/attached_modules/losses/pml_loss.py @@ -108,6 +108,9 @@ def __init__( def prepare(self, inputs, labels): embeddings = inputs["features"][0] + assert ( + labels is not None and "id" in labels + ), "ID labels are required for metric learning losses" IDs = labels["id"][0][:, 0] return embeddings, IDs diff --git a/luxonis_train/attached_modules/metrics/pml_metrics.py b/luxonis_train/attached_modules/metrics/pml_metrics.py index fdd66a41..a6d4effa 100644 --- a/luxonis_train/attached_modules/metrics/pml_metrics.py +++ b/luxonis_train/attached_modules/metrics/pml_metrics.py @@ -79,6 +79,10 @@ def __init__(self, cross_batch_memory_size=0, **kwargs): def prepare(self, inputs, labels): embeddings = inputs["features"][0] + + assert ( + labels is not None and "id" in labels + ), "ID labels are required for metric learning losses" IDs = labels["id"][0][:, 0] return embeddings, IDs @@ -158,6 +162,10 @@ def __init__(self, cross_batch_memory_size=0, **kwargs): def prepare(self, inputs, labels): embeddings = inputs["features"][0] + + assert ( + labels is not None and "id" in labels + ), "ID labels are required for metric learning losses" IDs = labels["id"][0][:, 0] return embeddings, IDs diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py index b5fb5f0e..d1096bfa 100644 --- a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -29,6 +29,10 @@ def prepare( self, inputs: Packet[Tensor], labels: Labels | None ) -> tuple[Tensor, Tensor]: embeddings = inputs["features"][0] + + assert ( + labels is not None and "id" in labels + ), "ID labels are required for metric learning losses" IDs = labels["id"][0] return embeddings, IDs @@ -69,13 +73,20 @@ def forward( # Plot the embeddings fig, ax = plt.subplots(figsize=(10, 10)) - scatter = ax.scatter( - embeddings_2d[:, 0], - embeddings_2d[:, 1], - c=IDs.detach().cpu().numpy(), - cmap="viridis", - s=5, - ) + if IDs is not None: + scatter = ax.scatter( + embeddings_2d[:, 0], + embeddings_2d[:, 1], + c=IDs.detach().cpu().numpy(), + cmap="viridis", + s=5, + ) + else: + scatter = ax.scatter( + embeddings_2d[:, 0], + embeddings_2d[:, 1], + s=5, + ) fig.colorbar(scatter, ax=ax) ax.set_title("Embeddings Visualization") ax.set_xlabel("Dimension 1") diff --git a/luxonis_train/nodes/backbones/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet.py index c242633f..9641596d 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet.py @@ -263,6 +263,8 @@ def __init__( has_se = se_ratio is not None and se_ratio > 0.0 self.stride = stride + assert layer_id is not None, "Layer ID must be explicitly provided" + # Point-wise expansion if layer_id <= 1: self.ghost1 = GhostModuleV2( @@ -507,66 +509,3 @@ def forward(self, inps): x = self.pointwise_conv(x) x = self.classifier(x) return x - - # @property - # def task(self) -> str: - # return "label" - - # @property - # def tasks(self) -> dict: - # return [TaskType.LABEL] - - -if __name__ == "__main__": - W, H = 256, 256 - model = GhostFaceNetsV2(image_size=W) - model.eval() # Set the model to evaluation mode - - # Create a dummy input tensor of the appropriate size - x = torch.randn(1, 3, H, W) - - # Export the model - onnx_path = "ghostfacenet.onnx" - torch.onnx.export( - model, # model being run - x, # model input (or a tuple for multiple inputs) - onnx_path, # where to save the model (can be a file or file-like object) - export_params=True, # store the trained parameter weights inside the model file - opset_version=12, # the ONNX version to export the model to - do_constant_folding=True, # whether to execute constant folding for optimization - input_names=["input"], # the model's input names - output_names=["output"], # the model's output names - # dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes - # 'output' : {0 : 'batch_size'}} - ) - import os - - import numpy as np - import onnx - import onnxsim - - # logger.info("Simplifying ONNX model...") - model_onnx = onnx.load(onnx_path) - onnx_model, check = onnxsim.simplify(model_onnx) - if not check: - raise RuntimeError("Onnx simplify failed.") - onnx.save(onnx_model, onnx_path) - - # Add calibration data - dir = "shared_with_container/calibration_data/" - for file in os.listdir(dir): - os.remove(dir + file) - for i in range(20): - np_array = np.random.rand(1, 3, H, W).astype(np.float32) - np.save(f"{dir}{i:02d}.npy", np_array) - np_array.tofile(f"{dir}{i:02d}.raw") - - # Test backpropagation on the model - # Create a dummy target tensor of the appropriate size - Y = model(x) - target = torch.randn(1, 512) - loss_fn = torch.nn.MSELoss() - loss = loss_fn(Y, target) - model.zero_grad() - loss.backward() - print("Backpropagation test successful") From 6eda12af3fc0a591b95adf498434f28aa6863c09 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Mon, 16 Dec 2024 11:51:08 +0100 Subject: [PATCH 05/57] Implemented improvements and suggestions. Separated GFN into class, blocks and variants. Added tests for all supported pytorch metric learning losses. --- .../attached_modules/losses/__init__.py | 4 +- .../attached_modules/losses/pml_loss.py | 183 +++---- .../attached_modules/metrics/pml_metrics.py | 125 +++-- .../visualizers/embeddings_visualizer.py | 16 +- luxonis_train/nodes/backbones/__init__.py | 2 +- luxonis_train/nodes/backbones/ghostfacenet.py | 511 ------------------ .../nodes/backbones/ghostfacenet/__init__.py | 3 + .../nodes/backbones/ghostfacenet/blocks.py | 256 +++++++++ .../backbones/ghostfacenet/ghostfacenet.py | 159 ++++++ .../nodes/backbones/ghostfacenet/variants.py | 214 ++++++++ .../nodes/backbones/micronet/blocks.py | 23 +- tests/configs/reid.yaml | 2 +- tests/integration/test_reid.py | 28 +- 13 files changed, 837 insertions(+), 689 deletions(-) delete mode 100644 luxonis_train/nodes/backbones/ghostfacenet.py create mode 100644 luxonis_train/nodes/backbones/ghostfacenet/__init__.py create mode 100644 luxonis_train/nodes/backbones/ghostfacenet/blocks.py create mode 100644 luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py create mode 100644 luxonis_train/nodes/backbones/ghostfacenet/variants.py diff --git a/luxonis_train/attached_modules/losses/__init__.py b/luxonis_train/attached_modules/losses/__init__.py index b320fada..2d0c77e1 100644 --- a/luxonis_train/attached_modules/losses/__init__.py +++ b/luxonis_train/attached_modules/losses/__init__.py @@ -7,7 +7,7 @@ from .ohem_bce_with_logits import OHEMBCEWithLogitsLoss from .ohem_cross_entropy import OHEMCrossEntropyLoss from .ohem_loss import OHEMLoss -from .pml_loss import MetricLearningLoss +from .pml_loss import EmbeddingLossWrapper from .reconstruction_segmentation_loss import ReconstructionSegmentationLoss from .sigmoid_focal_loss import SigmoidFocalLoss from .smooth_bce_with_logits import SmoothBCEWithLogitsLoss @@ -27,5 +27,5 @@ "OHEMCrossEntropyLoss", "OHEMBCEWithLogitsLoss", "FOMOLocalizationLoss", - "MetricLearningLoss", + "EmbeddingLossWrapper", ] diff --git a/luxonis_train/attached_modules/losses/pml_loss.py b/luxonis_train/attached_modules/losses/pml_loss.py index 1727d091..959a5f68 100644 --- a/luxonis_train/attached_modules/losses/pml_loss.py +++ b/luxonis_train/attached_modules/losses/pml_loss.py @@ -1,87 +1,67 @@ -import warnings - -from pytorch_metric_learning.losses import ( - AngularLoss, - ArcFaceLoss, - CircleLoss, - ContrastiveLoss, - CosFaceLoss, - CrossBatchMemory, - DynamicSoftMarginLoss, - FastAPLoss, - GeneralizedLiftedStructureLoss, - HistogramLoss, - InstanceLoss, - IntraPairVarianceLoss, - LargeMarginSoftmaxLoss, - LiftedStructureLoss, - ManifoldLoss, - MarginLoss, - MultiSimilarityLoss, - NCALoss, - NormalizedSoftmaxLoss, - NPairsLoss, - NTXentLoss, - P2SGradLoss, - PNPLoss, - ProxyAnchorLoss, - ProxyNCALoss, - RankedListLoss, - SignalToNoiseRatioContrastiveLoss, - SoftTripleLoss, - SphereFaceLoss, - SubCenterArcFaceLoss, - SupConLoss, - TripletMarginLoss, - TupletMarginLoss, -) +import logging + +import pytorch_metric_learning.losses as pml_losses +from pytorch_metric_learning.losses import CrossBatchMemory from torch import Tensor from .base_loss import BaseLoss -# Dictionary mapping string keys to loss classes -loss_dict = { - "AngularLoss": AngularLoss, - "ArcFaceLoss": ArcFaceLoss, - "CircleLoss": CircleLoss, - "ContrastiveLoss": ContrastiveLoss, - "CosFaceLoss": CosFaceLoss, - "DynamicSoftMarginLoss": DynamicSoftMarginLoss, - "FastAPLoss": FastAPLoss, - "GeneralizedLiftedStructureLoss": GeneralizedLiftedStructureLoss, - "InstanceLoss": InstanceLoss, - "HistogramLoss": HistogramLoss, - "IntraPairVarianceLoss": IntraPairVarianceLoss, - "LargeMarginSoftmaxLoss": LargeMarginSoftmaxLoss, - "LiftedStructureLoss": LiftedStructureLoss, - "ManifoldLoss": ManifoldLoss, - "MarginLoss": MarginLoss, - "MultiSimilarityLoss": MultiSimilarityLoss, - "NCALoss": NCALoss, - "NormalizedSoftmaxLoss": NormalizedSoftmaxLoss, - "NPairsLoss": NPairsLoss, - "NTXentLoss": NTXentLoss, - "P2SGradLoss": P2SGradLoss, - "PNPLoss": PNPLoss, - "ProxyAnchorLoss": ProxyAnchorLoss, - "ProxyNCALoss": ProxyNCALoss, - "RankedListLoss": RankedListLoss, - "SignalToNoiseRatioContrastiveLoss": SignalToNoiseRatioContrastiveLoss, - "SoftTripleLoss": SoftTripleLoss, - "SphereFaceLoss": SphereFaceLoss, - "SubCenterArcFaceLoss": SubCenterArcFaceLoss, - "SupConLoss": SupConLoss, - "TripletMarginLoss": TripletMarginLoss, - "TupletMarginLoss": TupletMarginLoss, -} - - -class MetricLearningLoss(BaseLoss): +logger = logging.getLogger(__name__) + +ALL_EMBEDDING_LOSSES = [ + "AngularLoss", + "ArcFaceLoss", + "CircleLoss", + "ContrastiveLoss", + "CosFaceLoss", + "DynamicSoftMarginLoss", + "FastAPLoss", + "HistogramLoss", + "InstanceLoss", + "IntraPairVarianceLoss", + "LargeMarginSoftmaxLoss", + "GeneralizedLiftedStructureLoss", + "LiftedStructureLoss", + "MarginLoss", + "MultiSimilarityLoss", + "NPairsLoss", + "NCALoss", + "NormalizedSoftmaxLoss", + "NTXentLoss", + "PNPLoss", + "ProxyAnchorLoss", + "ProxyNCALoss", + "RankedListLoss", + "SignalToNoiseRatioContrastiveLoss", + "SoftTripleLoss", + "SphereFaceLoss", + "SubCenterArcFaceLoss", + "SupConLoss", + "ThresholdConsistentMarginLoss", + "TripletMarginLoss", + "TupletMarginLoss", +] + +CLASS_EMBEDDING_LOSSES = [ + "ArcFaceLoss", + "CosFaceLoss", + "LargeMarginSoftmaxLoss", + "NormalizedSoftmaxLoss", + "ProxyAnchorLoss", + "ProxyNCALoss", + "SoftTripleLoss", + "SphereFaceLoss", + "SubCenterArcFaceLoss", +] + + +class EmbeddingLossWrapper(BaseLoss): def __init__( self, loss_name: str, embedding_size: int = 512, cross_batch_memory_size=0, + num_classes: int = 0, loss_kwargs: dict | None = None, *args, **kwargs, @@ -89,34 +69,51 @@ def __init__( super().__init__(*args, **kwargs) if loss_kwargs is None: loss_kwargs = {} - self.loss_func = loss_dict[loss_name]( - **loss_kwargs - ) # Instantiate the loss object + + try: + loss_cls = getattr(pml_losses, loss_name) + except AttributeError as e: + raise ValueError( + f"Loss {loss_name} not found in pytorch_metric_learning" + ) from e + + if loss_name in CLASS_EMBEDDING_LOSSES: + if num_classes < 0: + raise ValueError( + f"Loss {loss_name} requires num_classes to be set to a positive value" + ) + loss_kwargs["num_classes"] = num_classes + loss_kwargs["embedding_size"] = embedding_size + + # If we wanted to support these losses, we would need to add a separate optimizer for them. + # They may be useful in some scenarios, so leaving this here for future reference. + raise ValueError( + f"Loss {loss_name} requires its own optimizer, and that is not currently supported." + ) + + self.loss_func = loss_cls(**loss_kwargs) + if cross_batch_memory_size > 0: if loss_name in CrossBatchMemory.supported_losses(): self.loss_func = CrossBatchMemory( self.loss_func, embedding_size=embedding_size ) else: - # Warn that cross_batch_memory_size is ignored - warnings.warn( - f"Cross batch memory is not supported for {loss_name}. Ignoring cross_batch_memory_size" + logger.warning( + f"Cross batch memory is not supported for {loss_name}. Ignoring cross_batch_memory_size." ) - # self.miner_func = miner_func - - def prepare(self, inputs, labels): - embeddings = inputs["features"][0] + def prepare( + self, inputs: dict[str, list[Tensor]], labels: dict[str, list[Tensor]] + ) -> tuple[Tensor, Tensor]: + embeddings = self.get_input_tensors(inputs, "features")[0] - assert ( - labels is not None and "id" in labels - ), "ID labels are required for metric learning losses" - IDs = labels["id"][0][:, 0] - return embeddings, IDs + if labels is None or "id" not in labels: + raise ValueError("Labels must contain 'id' key") - def forward(self, inputs: Tensor, target: Tensor): - # miner_output = self.miner_func(inputs, target) + ids = labels["id"][0][:, 0] + return embeddings, ids + def forward(self, inputs: Tensor, target: Tensor) -> Tensor: loss = self.loss_func(inputs, target) - return loss diff --git a/luxonis_train/attached_modules/metrics/pml_metrics.py b/luxonis_train/attached_modules/metrics/pml_metrics.py index a6d4effa..ad8b0d88 100644 --- a/luxonis_train/attached_modules/metrics/pml_metrics.py +++ b/luxonis_train/attached_modules/metrics/pml_metrics.py @@ -7,62 +7,6 @@ # to PyTorch from TensorFlow -def _pairwise_distances(embeddings, squared=False): - """Compute the 2D matrix of distances between all the embeddings. - - @param embeddings: tensor of shape (batch_size, embed_dim) - @type embeddings: torch.Tensor - @param squared: If true, output is the pairwise squared euclidean - distance matrix. If false, output is the pairwise euclidean - distance matrix. - @type squared: bool - @return: pairwise_distances: tensor of shape (batch_size, - batch_size) - @rtype: torch.Tensor - """ - # Get the dot product between all embeddings - # shape (batch_size, batch_size) - dot_product = torch.matmul(embeddings, embeddings.t()) - - # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`. - # This also provides more numerical stability (the diagonal of the result will be exactly 0). - # shape (batch_size,) - square_norm = torch.diag(dot_product) - - # Compute the pairwise distance matrix as we have: - # ||a - b||^2 = ||a||^2 - 2 + ||b||^2 - # shape (batch_size, batch_size) - distances = ( - square_norm.unsqueeze(0) - 2.0 * dot_product + square_norm.unsqueeze(1) - ) - - # Because of computation errors, some distances might be negative so we put everything >= 0.0 - distances = torch.max(distances, torch.tensor(0.0)) - - if not squared: - # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal) - # we need to add a small epsilon where distances == 0.0 - mask = (distances == 0.0).float() - distances = distances + mask * 1e-16 - - distances = torch.sqrt(distances) - - # Correct the epsilon added: set the distances on the mask to be exactly 0.0 - distances = distances * (1.0 - mask) - - return distances - - -def _get_anchor_positive_triplet_mask(labels): - indices_equal = torch.eye( - labels.shape[0], dtype=torch.uint8, device=labels.device - ) - indices_not_equal = ~indices_equal - labels_equal = labels.unsqueeze(0) == labels.unsqueeze(1) - mask = indices_not_equal & labels_equal - return mask - - class ClosestIsPositiveAccuracy(BaseMetric): def __init__(self, cross_batch_memory_size=0, **kwargs): super().__init__(**kwargs) @@ -83,8 +27,8 @@ def prepare(self, inputs, labels): assert ( labels is not None and "id" in labels ), "ID labels are required for metric learning losses" - IDs = labels["id"][0][:, 0] - return embeddings, IDs + ids = labels["id"][0][:, 0] + return embeddings, ids def update(self, inputs: Tensor, target: Tensor): embeddings, labels = inputs, target @@ -166,8 +110,8 @@ def prepare(self, inputs, labels): assert ( labels is not None and "id" in labels ), "ID labels are required for metric learning losses" - IDs = labels["id"][0][:, 0] - return embeddings, IDs + ids = labels["id"][0][:, 0] + return embeddings, ids def update(self, inputs: Tensor, target: Tensor): embeddings, labels = inputs, target @@ -211,13 +155,18 @@ def update(self, inputs: Tensor, target: Tensor): # Get the positive mask and convert it to boolean positive_mask = _get_anchor_positive_triplet_mask(labels).bool() + # Filter out distances to negative elements w.r.t. each query embedding only_positive_distances = pairwise_distances.clone() only_positive_distances[~positive_mask] = float("inf") + # From the positive distances, get the closest positive distance for each query embedding closest_positive_distances, _ = torch.min( only_positive_distances, dim=1 ) + # Calculate the difference between the closest distance (any) and closest positive distances + # - this tells us how much closer should the closest positive be in order for the embedding + # to be considered correct non_inf_mask = closest_positive_distances != float("inf") difference = closest_positive_distances - closest_distances difference = difference[non_inf_mask] @@ -256,3 +205,59 @@ def compute(self): closest_vs_positive_distances ), } + + +def _pairwise_distances(embeddings, squared=False): + """Compute the 2D matrix of distances between all the embeddings. + + @param embeddings: tensor of shape (batch_size, embed_dim) + @type embeddings: torch.Tensor + @param squared: If true, output is the pairwise squared euclidean + distance matrix. If false, output is the pairwise euclidean + distance matrix. + @type squared: bool + @return: pairwise_distances: tensor of shape (batch_size, + batch_size) + @rtype: torch.Tensor + """ + # Get the dot product between all embeddings + # shape (batch_size, batch_size) + dot_product = torch.matmul(embeddings, embeddings.t()) + + # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`. + # This also provides more numerical stability (the diagonal of the result will be exactly 0). + # shape (batch_size,) + square_norm = torch.diag(dot_product) + + # Compute the pairwise distance matrix as we have: + # ||a - b||^2 = ||a||^2 - 2 + ||b||^2 + # shape (batch_size, batch_size) + distances = ( + square_norm.unsqueeze(0) - 2.0 * dot_product + square_norm.unsqueeze(1) + ) + + # Because of computation errors, some distances might be negative so we put everything >= 0.0 + distances = torch.max(distances, torch.tensor(0.0)) + + if not squared: + # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal) + # we need to add a small epsilon where distances == 0.0 + mask = (distances == 0.0).float() + distances = distances + mask * 1e-16 + + distances = torch.sqrt(distances) + + # Correct the epsilon added: set the distances on the mask to be exactly 0.0 + distances = distances * (1.0 - mask) + + return distances + + +def _get_anchor_positive_triplet_mask(labels): + indices_equal = torch.eye( + labels.shape[0], dtype=torch.uint8, device=labels.device + ) + indices_not_equal = ~indices_equal + labels_equal = labels.unsqueeze(0) == labels.unsqueeze(1) + mask = indices_not_equal & labels_equal + return mask diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py index d1096bfa..d8e5c940 100644 --- a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -33,15 +33,15 @@ def prepare( assert ( labels is not None and "id" in labels ), "ID labels are required for metric learning losses" - IDs = labels["id"][0] - return embeddings, IDs + ids = labels["id"][0] + return embeddings, ids def forward( self, label_canvas: Tensor, prediction_canvas: Tensor, embeddings: Tensor, - IDs: Tensor | None, + ids: Tensor | None, **kwargs, ) -> Tensor: """Creates a visualization of the embeddings. @@ -52,14 +52,14 @@ def forward( @param prediction_canvas: The canvas to draw the predictions on. @type embeddings: Tensor @param embeddings: The embeddings to visualize. - @type IDs: Tensor - @param IDs: The IDs to visualize. + @type ids: Tensor + @param ids: The ids to visualize. @rtype: Tensor @return: An embedding space projection. """ # Embeddings: [B, D], D = e.g. 512 - # IDs: [B, 1], corresponding to the embeddings + # ids: [B, 1], corresponding to the embeddings # Convert embeddings to numpy array embeddings_np = embeddings.detach().cpu().numpy() @@ -73,11 +73,11 @@ def forward( # Plot the embeddings fig, ax = plt.subplots(figsize=(10, 10)) - if IDs is not None: + if ids is not None: scatter = ax.scatter( embeddings_2d[:, 0], embeddings_2d[:, 1], - c=IDs.detach().cpu().numpy(), + c=ids.detach().cpu().numpy(), cmap="viridis", s=5, ) diff --git a/luxonis_train/nodes/backbones/__init__.py b/luxonis_train/nodes/backbones/__init__.py index f5319981..da063a5e 100644 --- a/luxonis_train/nodes/backbones/__init__.py +++ b/luxonis_train/nodes/backbones/__init__.py @@ -2,7 +2,7 @@ from .ddrnet import DDRNet from .efficientnet import EfficientNet from .efficientrep import EfficientRep -from .ghostfacenet import GhostFaceNetsV2 +from .ghostfacenet.ghostfacenet import GhostFaceNetsV2 from .micronet import MicroNet from .mobilenetv2 import MobileNetV2 from .mobileone import MobileOne diff --git a/luxonis_train/nodes/backbones/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet.py deleted file mode 100644 index 9641596d..00000000 --- a/luxonis_train/nodes/backbones/ghostfacenet.py +++ /dev/null @@ -1,511 +0,0 @@ -# Original source: https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py - - -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from luxonis_train.nodes.base_node import BaseNode - - -def _make_divisible(v, divisor, min_value=None): - """This function is taken from the original tf repo. - - It ensures that all layers have a channel number that is divisible by 8 - It can be seen here: - https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py - """ - if min_value is None: - min_value = divisor - new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) - # Make sure that round down does not go down by more than 10%. - if new_v < 0.9 * v: - new_v += divisor - return new_v - - -def hard_sigmoid(x, inplace: bool = False): - if inplace: - return x.add_(3.0).clamp_(0.0, 6.0).div_(6.0) - else: - return F.relu6(x + 3.0) / 6.0 - - -class SqueezeExcite(nn.Module): - def __init__( - self, - in_chs, - se_ratio=0.25, - reduced_base_chs=None, - act_layer=nn.PReLU, - gate_fn=hard_sigmoid, - divisor=4, - **_, - ): - super(SqueezeExcite, self).__init__() - self.gate_fn = gate_fn - reduced_chs = _make_divisible( - (reduced_base_chs or in_chs) * se_ratio, divisor - ) - self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True) - self.act1 = act_layer() - self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True) - - def forward(self, x): - x_se = self.avg_pool(x) - x_se = self.conv_reduce(x_se) - x_se = self.act1(x_se) - x_se = self.conv_expand(x_se) - x = x * self.gate_fn(x_se) - return x - - -class ConvBnAct(nn.Module): - def __init__( - self, in_chs, out_chs, kernel_size, stride=1, act_layer=nn.PReLU - ): - super(ConvBnAct, self).__init__() - self.conv = nn.Conv2d( - in_chs, out_chs, kernel_size, stride, kernel_size // 2, bias=False - ) - self.bn1 = nn.BatchNorm2d(out_chs) - self.act1 = act_layer() - - def forward(self, x): - x = self.conv(x) - x = self.bn1(x) - x = self.act1(x) - return x - - -class ModifiedGDC(nn.Module): - def __init__( - self, image_size, in_chs, num_classes, dropout, emb=512 - ): # dropout implementation is in the original code but not in the paper - super(ModifiedGDC, self).__init__() - - if image_size % 32 == 0: - self.conv_dw = nn.Conv2d( - in_chs, - in_chs, - kernel_size=(image_size // 32), - groups=in_chs, - bias=False, - ) - else: - self.conv_dw = nn.Conv2d( - in_chs, - in_chs, - kernel_size=(image_size // 32 + 1), - groups=in_chs, - bias=False, - ) - self.bn1 = nn.BatchNorm2d(in_chs) - self.dropout = nn.Dropout(dropout) - - self.conv = nn.Conv2d(in_chs, emb, kernel_size=1, bias=False) - self.bn2 = nn.BatchNorm1d(emb) - self.linear = ( - nn.Linear(emb, num_classes) if num_classes else nn.Identity() - ) - - def forward(self, inps): - x = inps - x = self.conv_dw(x) - x = self.bn1(x) - x = self.dropout(x) - # # Add spots to the features - # x = torch.cat([x, spots.view(spots.size(0), -1, 1, 1)], dim=1) - x = self.conv(x) - x = x.view(x.size(0), -1) # Flatten - x = self.bn2(x) - x = self.linear(x) - return x - - -class GhostModuleV2(nn.Module): - def __init__( - self, - inp, - oup, - kernel_size=1, - ratio=2, - dw_size=3, - stride=1, - prelu=True, - mode=None, - args=None, - ): - super(GhostModuleV2, self).__init__() - self.mode = mode - self.gate_fn = nn.Sigmoid() - - if self.mode in ["original"]: - self.oup = oup - init_channels = math.ceil(oup / ratio) - new_channels = init_channels * (ratio - 1) - self.primary_conv = nn.Sequential( - nn.Conv2d( - inp, - init_channels, - kernel_size, - stride, - kernel_size // 2, - bias=False, - ), - nn.BatchNorm2d(init_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - self.cheap_operation = nn.Sequential( - nn.Conv2d( - init_channels, - new_channels, - dw_size, - 1, - dw_size // 2, - groups=init_channels, - bias=False, - ), - nn.BatchNorm2d(new_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - elif self.mode in ["attn"]: # DFC - self.oup = oup - init_channels = math.ceil(oup / ratio) - new_channels = init_channels * (ratio - 1) - self.primary_conv = nn.Sequential( - nn.Conv2d( - inp, - init_channels, - kernel_size, - stride, - kernel_size // 2, - bias=False, - ), - nn.BatchNorm2d(init_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - self.cheap_operation = nn.Sequential( - nn.Conv2d( - init_channels, - new_channels, - dw_size, - 1, - dw_size // 2, - groups=init_channels, - bias=False, - ), - nn.BatchNorm2d(new_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - self.short_conv = nn.Sequential( - nn.Conv2d( - inp, oup, kernel_size, stride, kernel_size // 2, bias=False - ), - nn.BatchNorm2d(oup), - nn.Conv2d( - oup, - oup, - kernel_size=(1, 5), - stride=1, - padding=(0, 2), - groups=oup, - bias=False, - ), - nn.BatchNorm2d(oup), - nn.Conv2d( - oup, - oup, - kernel_size=(5, 1), - stride=1, - padding=(2, 0), - groups=oup, - bias=False, - ), - nn.BatchNorm2d(oup), - ) - - def forward(self, x): - if self.mode in ["original"]: - x1 = self.primary_conv(x) - x2 = self.cheap_operation(x1) - out = torch.cat([x1, x2], dim=1) - return out[:, : self.oup, :, :] - elif self.mode in ["attn"]: - res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2)) - x1 = self.primary_conv(x) - x2 = self.cheap_operation(x1) - out = torch.cat([x1, x2], dim=1) - return out[:, : self.oup, :, :] * F.interpolate( - self.gate_fn(res), - size=(out.shape[-2], out.shape[-1]), - mode="nearest", - ) - - -class GhostBottleneckV2(nn.Module): - def __init__( - self, - in_chs, - mid_chs, - out_chs, - dw_kernel_size=3, - stride=1, - act_layer=nn.PReLU, - se_ratio=0.0, - layer_id=None, - args=None, - ): - super(GhostBottleneckV2, self).__init__() - has_se = se_ratio is not None and se_ratio > 0.0 - self.stride = stride - - assert layer_id is not None, "Layer ID must be explicitly provided" - - # Point-wise expansion - if layer_id <= 1: - self.ghost1 = GhostModuleV2( - in_chs, mid_chs, prelu=True, mode="original", args=args - ) - else: - self.ghost1 = GhostModuleV2( - in_chs, mid_chs, prelu=True, mode="attn", args=args - ) - - # Depth-wise convolution - if self.stride > 1: - self.conv_dw = nn.Conv2d( - mid_chs, - mid_chs, - dw_kernel_size, - stride=stride, - padding=(dw_kernel_size - 1) // 2, - groups=mid_chs, - bias=False, - ) - self.bn_dw = nn.BatchNorm2d(mid_chs) - - # Squeeze-and-excitation - if has_se: - self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio) - else: - self.se = None - - self.ghost2 = GhostModuleV2( - mid_chs, out_chs, prelu=False, mode="original", args=args - ) - - # shortcut - if in_chs == out_chs and self.stride == 1: - self.shortcut = nn.Sequential() - else: - self.shortcut = nn.Sequential( - nn.Conv2d( - in_chs, - in_chs, - dw_kernel_size, - stride=stride, - padding=(dw_kernel_size - 1) // 2, - groups=in_chs, - bias=False, - ), - nn.BatchNorm2d(in_chs), - nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False), - nn.BatchNorm2d(out_chs), - ) - - def forward(self, x): - residual = x - x = self.ghost1(x) - if self.stride > 1: - x = self.conv_dw(x) - x = self.bn_dw(x) - if self.se is not None: - x = self.se(x) - x = self.ghost2(x) - x += self.shortcut(residual) - return x - - -# NODES.register_module() -class GhostFaceNetsV2(BaseNode[torch.Tensor, list[torch.Tensor]]): - def unwrap(self, inputs): - return [inputs[0]["features"][0]] - - def wrap(self, outputs): - return {"features": [outputs]} - - def set_export_mode(self, mode: bool = True): - self.export_mode = mode - self.train(not mode) - - def __init__( - self, - cfgs=None, - embedding_size=512, - num_classes=0, - width=1.0, - dropout=0.2, - block=GhostBottleneckV2, - add_pointwise_conv=False, - bn_momentum=0.9, - bn_epsilon=1e-5, - init_kaiming=True, - block_args=None, - *args, - **kwargs, - ): - """GhostFaceNetsV2 backbone. - - GhostFaceNetsV2 is a convolutional neural network architecture focused on face recognition, but it is - adaptable to generic embedding tasks. It is based on the GhostNet architecture and uses Ghost BottleneckV2 blocks. - - Source: U{https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py} - - @license: U{MIT License - } - - @see: U{GhostFaceNets: Lightweight Face Recognition Model From Cheap Operations - } - - @type cfgs: list[list[list[int]]] | None - @param cfgs: List of Ghost BottleneckV2 configurations. Defaults to None, which uses the original GhostFaceNetsV2 configuration. - @type embedding_size: int - @param embedding_size: Size of the embedding. Defaults to 512. - @type num_classes: int - @param num_classes: Number of classes. Defaults to 0, which makes the network output the raw embeddings. Otherwise it can be used to - add another linear layer to the network, which is useful for training using ArcFace or similar classification-based losses that - require the user to drop the last layer of the network. - @type width: float - @param width: Width multiplier. Increases complexity and number of parameters. Defaults to 1.0. - @type dropout: float - @param dropout: Dropout rate. Defaults to 0.2. - @type block: nn.Module - @param block: Ghost BottleneckV2 block. Defaults to GhostBottleneckV2. - @type add_pointwise_conv: bool - @param add_pointwise_conv: If True, adds a pointwise convolution layer at the end of the network. Defaults to False. - @type bn_momentum: float - @param bn_momentum: Batch normalization momentum. Defaults to 0.9. - @type bn_epsilon: float - @param bn_epsilon: Batch normalization epsilon. Defaults to 1e-5. - @type init_kaiming: bool - @param init_kaiming: If True, initializes the weights using the Kaiming initialization. Defaults to True. - @type block_args: dict - @param block_args: Arguments to pass to the block. Defaults to None. - """ - # kwargs['_tasks'] = {TaskType.LABEL: 'features'} - super().__init__(*args, **kwargs) - - inp_shape = kwargs["input_shapes"][0]["features"][0] - # spots_shape = kwargs['input_shapes'][0]['features'][1] - - image_size = inp_shape[2] - channels = inp_shape[1] - if cfgs is None: - self.cfgs = [ - # k, t, c, SE, s - [[3, 16, 16, 0, 1]], - [[3, 48, 24, 0, 2]], - [[3, 72, 24, 0, 1]], - [[5, 72, 40, 0.25, 2]], - [[5, 120, 40, 0.25, 1]], - [[3, 240, 80, 0, 2]], - [ - [3, 200, 80, 0, 1], - [3, 184, 80, 0, 1], - [3, 184, 80, 0, 1], - [3, 480, 112, 0.25, 1], - [3, 672, 112, 0.25, 1], - ], - [[5, 672, 160, 0.25, 2]], - [ - [5, 960, 160, 0, 1], - [5, 960, 160, 0.25, 1], - [5, 960, 160, 0, 1], - [5, 960, 160, 0.25, 1], - ], - ] - else: - self.cfgs = cfgs - - # building first layer - output_channel = _make_divisible(16 * width, 4) - self.conv_stem = nn.Conv2d( - channels, output_channel, 3, 2, 1, bias=False - ) - self.bn1 = nn.BatchNorm2d(output_channel) - self.act1 = nn.PReLU() - input_channel = output_channel - - # building inverted residual blocks - stages = [] - layer_id = 0 - for cfg in self.cfgs: - layers = [] - for k, exp_size, c, se_ratio, s in cfg: - output_channel = _make_divisible(c * width, 4) - hidden_channel = _make_divisible(exp_size * width, 4) - if block == GhostBottleneckV2: - layers.append( - block( - input_channel, - hidden_channel, - output_channel, - k, - s, - se_ratio=se_ratio, - layer_id=layer_id, - args=block_args, - ) - ) - input_channel = output_channel - layer_id += 1 - stages.append(nn.Sequential(*layers)) - - output_channel = _make_divisible(exp_size * width, 4) - stages.append( - nn.Sequential(ConvBnAct(input_channel, output_channel, 1)) - ) - - self.blocks = nn.Sequential(*stages) - - # building last several layers - pointwise_conv = [] - if add_pointwise_conv: - pointwise_conv.append( - nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=True) - ) - pointwise_conv.append(nn.BatchNorm2d(output_channel)) - pointwise_conv.append(nn.PReLU()) - else: - pointwise_conv.append(nn.Sequential()) - - self.pointwise_conv = nn.Sequential(*pointwise_conv) - self.classifier = ModifiedGDC( - image_size, output_channel, num_classes, dropout, embedding_size - ) - - # Initialize weights - for m in self.modules(): - if init_kaiming: - if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): - fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) - negative_slope = 0.25 # Default value for PReLU in PyTorch, change it if you use custom value - m.weight.data.normal_( - 0, math.sqrt(2.0 / (fan_in * (1 + negative_slope**2))) - ) - if isinstance(m, nn.BatchNorm2d): - m.momentum, m.eps = bn_momentum, bn_epsilon - - def forward(self, inps): - x = inps[0] - x = self.conv_stem(x) - x = self.bn1(x) - x = self.act1(x) - x = self.blocks(x) - x = self.pointwise_conv(x) - x = self.classifier(x) - return x diff --git a/luxonis_train/nodes/backbones/ghostfacenet/__init__.py b/luxonis_train/nodes/backbones/ghostfacenet/__init__.py new file mode 100644 index 00000000..85ed4447 --- /dev/null +++ b/luxonis_train/nodes/backbones/ghostfacenet/__init__.py @@ -0,0 +1,3 @@ +from .ghostfacenet import GhostFaceNetsV2 + +__all__ = ["GhostFaceNetsV2"] diff --git a/luxonis_train/nodes/backbones/ghostfacenet/blocks.py b/luxonis_train/nodes/backbones/ghostfacenet/blocks.py new file mode 100644 index 00000000..46a9ba27 --- /dev/null +++ b/luxonis_train/nodes/backbones/ghostfacenet/blocks.py @@ -0,0 +1,256 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from luxonis_train.nodes.backbones.micronet.blocks import _make_divisible +from luxonis_train.nodes.blocks import SqueezeExciteBlock + + +class ModifiedGDC(nn.Module): + def __init__(self, image_size, in_chs, num_classes, dropout, emb=512): + super().__init__() + + if image_size % 32 == 0: + self.conv_dw = nn.Conv2d( + in_chs, + in_chs, + kernel_size=(image_size // 32), + groups=in_chs, + bias=False, + ) + else: + self.conv_dw = nn.Conv2d( + in_chs, + in_chs, + kernel_size=(image_size // 32 + 1), + groups=in_chs, + bias=False, + ) + self.bn1 = nn.BatchNorm2d(in_chs) + self.dropout = nn.Dropout(dropout) + + self.conv = nn.Conv2d(in_chs, emb, kernel_size=1, bias=False) + self.bn2 = nn.BatchNorm1d(emb) + self.linear = ( + nn.Linear(emb, num_classes) if num_classes else nn.Identity() + ) + + def forward(self, inps): + x = inps + x = self.conv_dw(x) + x = self.bn1(x) + x = self.dropout(x) + x = self.conv(x) + x = x.view(x.size(0), -1) + x = self.bn2(x) + x = self.linear(x) + return x + + +class GhostModuleV2(nn.Module): + def __init__( + self, + inp, + oup, + kernel_size=1, + ratio=2, + dw_size=3, + stride=1, + prelu=True, + mode=None, + args=None, + ): + super(GhostModuleV2, self).__init__() + self.mode = mode + self.gate_fn = nn.Sigmoid() + + if self.mode in ["original"]: + self.oup = oup + init_channels = math.ceil(oup / ratio) + new_channels = init_channels * (ratio - 1) + self.primary_conv = nn.Sequential( + nn.Conv2d( + inp, + init_channels, + kernel_size, + stride, + kernel_size // 2, + bias=False, + ), + nn.BatchNorm2d(init_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + self.cheap_operation = nn.Sequential( + nn.Conv2d( + init_channels, + new_channels, + dw_size, + 1, + dw_size // 2, + groups=init_channels, + bias=False, + ), + nn.BatchNorm2d(new_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + elif self.mode in ["attn"]: # DFC + self.oup = oup + init_channels = math.ceil(oup / ratio) + new_channels = init_channels * (ratio - 1) + self.primary_conv = nn.Sequential( + nn.Conv2d( + inp, + init_channels, + kernel_size, + stride, + kernel_size // 2, + bias=False, + ), + nn.BatchNorm2d(init_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + self.cheap_operation = nn.Sequential( + nn.Conv2d( + init_channels, + new_channels, + dw_size, + 1, + dw_size // 2, + groups=init_channels, + bias=False, + ), + nn.BatchNorm2d(new_channels), + nn.PReLU() if prelu else nn.Sequential(), + ) + self.short_conv = nn.Sequential( + nn.Conv2d( + inp, oup, kernel_size, stride, kernel_size // 2, bias=False + ), + nn.BatchNorm2d(oup), + nn.Conv2d( + oup, + oup, + kernel_size=(1, 5), + stride=1, + padding=(0, 2), + groups=oup, + bias=False, + ), + nn.BatchNorm2d(oup), + nn.Conv2d( + oup, + oup, + kernel_size=(5, 1), + stride=1, + padding=(2, 0), + groups=oup, + bias=False, + ), + nn.BatchNorm2d(oup), + ) + + def forward(self, x): + if self.mode in ["original"]: + x1 = self.primary_conv(x) + x2 = self.cheap_operation(x1) + out = torch.cat([x1, x2], dim=1) + return out[:, : self.oup, :, :] + elif self.mode in ["attn"]: + res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2)) + x1 = self.primary_conv(x) + x2 = self.cheap_operation(x1) + out = torch.cat([x1, x2], dim=1) + return out[:, : self.oup, :, :] * F.interpolate( + self.gate_fn(res), + size=(out.shape[-2], out.shape[-1]), + mode="nearest", + ) + + +class GhostBottleneckV2(nn.Module): + def __init__( + self, + in_chs, + mid_chs, + out_chs, + dw_kernel_size=3, + stride=1, + act_layer=nn.PReLU, + se_ratio=0.0, + layer_id=None, + args=None, + ): + super(GhostBottleneckV2, self).__init__() + has_se = se_ratio is not None and se_ratio > 0.0 + self.stride = stride + + assert layer_id is not None, "Layer ID must be explicitly provided" + + # Point-wise expansion + if layer_id <= 1: + self.ghost1 = GhostModuleV2( + in_chs, mid_chs, prelu=True, mode="original", args=args + ) + else: + self.ghost1 = GhostModuleV2( + in_chs, mid_chs, prelu=True, mode="attn", args=args + ) + + # Depth-wise convolution + if self.stride > 1: + self.conv_dw = nn.Conv2d( + mid_chs, + mid_chs, + dw_kernel_size, + stride=stride, + padding=(dw_kernel_size - 1) // 2, + groups=mid_chs, + bias=False, + ) + self.bn_dw = nn.BatchNorm2d(mid_chs) + + # Squeeze-and-excitation + if has_se: + reduced_chs = _make_divisible(mid_chs * se_ratio, 4) + self.se = SqueezeExciteBlock( + mid_chs, reduced_chs, True, activation=nn.PReLU() + ) + else: + self.se = None + + self.ghost2 = GhostModuleV2( + mid_chs, out_chs, prelu=False, mode="original", args=args + ) + + # shortcut + if in_chs == out_chs and self.stride == 1: + self.shortcut = nn.Sequential() + else: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_chs, + in_chs, + dw_kernel_size, + stride=stride, + padding=(dw_kernel_size - 1) // 2, + groups=in_chs, + bias=False, + ), + nn.BatchNorm2d(in_chs), + nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(out_chs), + ) + + def forward(self, x): + residual = x + x = self.ghost1(x) + if self.stride > 1: + x = self.conv_dw(x) + x = self.bn_dw(x) + if self.se is not None: + x = self.se(x) + x = self.ghost2(x) + x += self.shortcut(residual) + return x diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py new file mode 100644 index 00000000..8bb61fee --- /dev/null +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -0,0 +1,159 @@ +# Original source: https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py +import math +from typing import Literal + +import torch +import torch.nn as nn + +from luxonis_train.nodes.backbones.ghostfacenet.blocks import ( + GhostBottleneckV2, + ModifiedGDC, +) +from luxonis_train.nodes.backbones.ghostfacenet.variants import get_variant +from luxonis_train.nodes.backbones.micronet.blocks import _make_divisible +from luxonis_train.nodes.base_node import BaseNode +from luxonis_train.nodes.blocks import ConvModule + + +class GhostFaceNetsV2(BaseNode[torch.Tensor, list[torch.Tensor]]): + in_channels: list[int] + in_width: list[int] + + def __init__( + self, + embedding_size=512, + num_classes=-1, + variant: Literal["V2"] = "V2", + *args, + **kwargs, + ): + """GhostFaceNetsV2 backbone. + + GhostFaceNetsV2 is a convolutional neural network architecture focused on face recognition, but it is + adaptable to generic embedding tasks. It is based on the GhostNet architecture and uses Ghost BottleneckV2 blocks. + + Source: U{https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py} + + @license: U{MIT License + } + + @see: U{GhostFaceNets: Lightweight Face Recognition Model From Cheap Operations + } + + @type embedding_size: int + @param embedding_size: Size of the embedding. Defaults to 512. + @type num_classes: int + @param num_classes: Number of classes. Defaults to -1, which leaves the default variant value in. Otherwise it can be used to + have the network return raw embeddings (=0) or add another linear layer to the network, which is useful for training using + ArcFace or similar classification-based losses that require the user to drop the last layer of the network. + @type variant: Literal["V2"] + @param variant: Variant of the GhostFaceNets embedding model. Defaults to "V2" (which is the only variant available). + """ + super().__init__(*args, **kwargs) + + image_size = self.in_width[0] + channels = self.in_channels[0] + var = get_variant(variant) + if num_classes >= 0: + var.num_classes = num_classes + self.cfgs = var.cfgs + + # Building first layer + output_channel = _make_divisible(int(16 * var.width), 4) + self.conv_stem = nn.Conv2d( + channels, output_channel, 3, 2, 1, bias=False + ) + self.bn1 = nn.BatchNorm2d(output_channel) + self.act1 = nn.PReLU() + input_channel = output_channel + + # Building Ghost BottleneckV2 blocks + stages = [] + layer_id = 0 + for cfg in self.cfgs: + layers = [] + for b_cfg in cfg: + output_channel = _make_divisible( + b_cfg.output_channels * var.width, 4 + ) + hidden_channel = _make_divisible( + b_cfg.expand_size * var.width, 4 + ) + if var.block == GhostBottleneckV2: + layers.append( + var.block( + input_channel, + hidden_channel, + output_channel, + b_cfg.kernel_size, + b_cfg.stride, + se_ratio=b_cfg.se_ratio, + layer_id=layer_id, + args=var.block_args, + ) + ) + input_channel = output_channel + layer_id += 1 + stages.append(nn.Sequential(*layers)) + + output_channel = _make_divisible(b_cfg.expand_size * var.width, 4) + stages.append( + nn.Sequential( + ConvModule( + input_channel, + output_channel, + kernel_size=1, + activation=nn.PReLU(), + ) + ) + ) + + self.blocks = nn.Sequential(*stages) + + # Building pointwise convolution + pointwise_conv = [] + if var.add_pointwise_conv: + pointwise_conv.append( + nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=True) + ) + pointwise_conv.append(nn.BatchNorm2d(output_channel)) + pointwise_conv.append(nn.PReLU()) + else: + pointwise_conv.append(nn.Sequential()) + + self.pointwise_conv = nn.Sequential(*pointwise_conv) + self.classifier = ModifiedGDC( + image_size, + output_channel, + var.num_classes, + var.dropout, + embedding_size, + ) + + # Initializing weights + for m in self.modules(): + if var.init_kaiming: + if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) + negative_slope = 0.25 + m.weight.data.normal_( + 0, math.sqrt(2.0 / (fan_in * (1 + negative_slope**2))) + ) + if isinstance(m, nn.BatchNorm2d): + m.momentum, m.eps = var.bn_momentum, var.bn_epsilon + + def unwrap(self, inputs): + return [inputs[0]["features"][0]] + + def wrap(self, outputs): + return {"features": [outputs]} + + def forward(self, inps): + x = inps[0] + x = self.conv_stem(x) + x = self.bn1(x) + x = self.act1(x) + x = self.blocks(x) + x = self.pointwise_conv(x) + x = self.classifier(x) + return x diff --git a/luxonis_train/nodes/backbones/ghostfacenet/variants.py b/luxonis_train/nodes/backbones/ghostfacenet/variants.py new file mode 100644 index 00000000..0e88ecfc --- /dev/null +++ b/luxonis_train/nodes/backbones/ghostfacenet/variants.py @@ -0,0 +1,214 @@ +from typing import List, Literal + +from pydantic import BaseModel +from torch import nn + +from luxonis_train.nodes.backbones.ghostfacenet.blocks import GhostBottleneckV2 + + +class BlockConfig(BaseModel): + kernel_size: int + expand_size: int + output_channels: int + se_ratio: float + stride: int + + +class GhostFaceNetsVariant(BaseModel): + """Variant of the GhostFaceNets embedding model. + + @type cfgs: List[List[BlockConfig]] + @param cfgs: List of Ghost BottleneckV2 configurations. + @type num_classes: int + @param num_classes: Number of classes. Defaults to 0, which makes + the network output the raw embeddings. Otherwise it can be used + to add another linear layer to the network, which is useful for + training using ArcFace or similar classification-based losses + that require the user to drop the last layer of the network. + @type width: int + @param width: Width multiplier. Increases complexity and number of + parameters. Defaults to 1.0. + @type dropout: float + @param dropout: Dropout rate. Defaults to 0.2. + @type block: nn.Module + @param block: Ghost BottleneckV2 block. Defaults to + GhostBottleneckV2. + @type add_pointwise_conv: bool + @param add_pointwise_conv: If True, adds a pointwise convolution + layer at the end of the network. Defaults to False. + @type bn_momentum: float + @param bn_momentum: Batch normalization momentum. Defaults to 0.9. + @type bn_epsilon: float + @param bn_epsilon: Batch normalization epsilon. Defaults to 1e-5. + @type init_kaiming: bool + @param init_kaiming: If True, initializes the weights using the + Kaiming initialization. Defaults to True. + @type block_args: dict + @param block_args: Arguments to pass to the block. Defaults to None. + """ + + num_classes: int + width: int + dropout: float + block: type[nn.Module] + add_pointwise_conv: bool + bn_momentum: float + bn_epsilon: float + init_kaiming: bool + block_args: dict | None + cfgs: List[List[BlockConfig]] + + +V2 = GhostFaceNetsVariant( + num_classes=0, + width=1, + dropout=0.2, + block=GhostBottleneckV2, + add_pointwise_conv=False, + bn_momentum=0.9, + bn_epsilon=1e-5, + init_kaiming=True, + block_args=None, + cfgs=[ + [ + BlockConfig( + kernel_size=3, + expand_size=16, + output_channels=16, + se_ratio=0.0, + stride=1, + ) + ], + [ + BlockConfig( + kernel_size=3, + expand_size=48, + output_channels=24, + se_ratio=0.0, + stride=2, + ) + ], + [ + BlockConfig( + kernel_size=3, + expand_size=72, + output_channels=24, + se_ratio=0.0, + stride=1, + ) + ], + [ + BlockConfig( + kernel_size=5, + expand_size=72, + output_channels=40, + se_ratio=0.25, + stride=2, + ) + ], + [ + BlockConfig( + kernel_size=5, + expand_size=120, + output_channels=40, + se_ratio=0.25, + stride=1, + ) + ], + [ + BlockConfig( + kernel_size=3, + expand_size=240, + output_channels=80, + se_ratio=0.0, + stride=2, + ) + ], + [ + BlockConfig( + kernel_size=3, + expand_size=200, + output_channels=80, + se_ratio=0.0, + stride=1, + ), + BlockConfig( + kernel_size=3, + expand_size=184, + output_channels=80, + se_ratio=0.0, + stride=1, + ), + BlockConfig( + kernel_size=3, + expand_size=184, + output_channels=80, + se_ratio=0.0, + stride=1, + ), + BlockConfig( + kernel_size=3, + expand_size=480, + output_channels=112, + se_ratio=0.25, + stride=1, + ), + BlockConfig( + kernel_size=3, + expand_size=672, + output_channels=112, + se_ratio=0.25, + stride=1, + ), + ], + [ + BlockConfig( + kernel_size=5, + expand_size=672, + output_channels=160, + se_ratio=0.25, + stride=2, + ) + ], + [ + BlockConfig( + kernel_size=5, + expand_size=960, + output_channels=160, + se_ratio=0.0, + stride=1, + ), + BlockConfig( + kernel_size=5, + expand_size=960, + output_channels=160, + se_ratio=0.25, + stride=1, + ), + BlockConfig( + kernel_size=5, + expand_size=960, + output_channels=160, + se_ratio=0.0, + stride=1, + ), + BlockConfig( + kernel_size=5, + expand_size=960, + output_channels=160, + se_ratio=0.25, + stride=1, + ), + ], + ], +) + + +def get_variant(variant: Literal["V2"]) -> GhostFaceNetsVariant: + variants = {"V2": V2} + if variant not in variants: # pragma: no cover + raise ValueError( + "GhostFaceNets model variant should be in " + f"{list(variants.keys())}, got {variant}." + ) + return variants[variant] diff --git a/luxonis_train/nodes/backbones/micronet/blocks.py b/luxonis_train/nodes/backbones/micronet/blocks.py index 3da5e15e..b29082cf 100644 --- a/luxonis_train/nodes/backbones/micronet/blocks.py +++ b/luxonis_train/nodes/backbones/micronet/blocks.py @@ -357,7 +357,7 @@ def __init__( self.avg_pool = nn.AdaptiveAvgPool2d(1) - squeeze_channels = self._make_divisible(in_channels // reduction, 4) + squeeze_channels = _make_divisible(in_channels // reduction, 4) self.fc = nn.Sequential( nn.Linear(in_channels, squeeze_channels), @@ -413,16 +413,17 @@ def forward(self, x: Tensor) -> Tensor: return out - def _make_divisible( - self, value: int, divisor: int, min_value: int | None = None - ) -> int: - if min_value is None: - min_value = divisor - new_v = max(min_value, int(value + divisor / 2) // divisor * divisor) - # Make sure that round down does not go down by more than 10%. - if new_v < 0.9 * value: - new_v += divisor - return new_v + +def _make_divisible( + value: int, divisor: int, min_value: int | None = None +) -> int: + if min_value is None: + min_value = divisor + new_v = max(min_value, int(value + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_v < 0.9 * value: + new_v += divisor + return new_v class SpatialSepConvSF(nn.Module): diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml index d9c0ec11..21ca2748 100644 --- a/tests/configs/reid.yaml +++ b/tests/configs/reid.yaml @@ -11,7 +11,7 @@ model: embedding_size: &embedding_size 512 losses: - - name: MetricLearningLoss + - name: EmbeddingLossWrapper params: loss_name: SupConLoss embedding_size: *embedding_size diff --git a/tests/integration/test_reid.py b/tests/integration/test_reid.py index 9ed4e867..0d006072 100644 --- a/tests/integration/test_reid.py +++ b/tests/integration/test_reid.py @@ -5,6 +5,10 @@ import pytest import torch +from luxonis_train.attached_modules.losses.pml_loss import ( + ALL_EMBEDDING_LOSSES, + CLASS_EMBEDDING_LOSSES, +) from luxonis_train.core import LuxonisModel from luxonis_train.enums import TaskType from luxonis_train.loaders import BaseLoaderTorch @@ -15,6 +19,8 @@ ONNX_PATH = Path("tests/integration/_model.onnx") STUDY_PATH = Path("study_local.db") +NUM_INDIVIDUALS = 100 + class CustomReIDLoader(BaseLoaderTorch): def __init__(self, *args, **kwargs): @@ -35,7 +41,7 @@ def __getitem__(self, _): # pragma: no cover } # Fake labels - id = torch.randint(0, 1000, (1,), dtype=torch.int64) + id = torch.randint(0, NUM_INDIVIDUALS, (1,), dtype=torch.int64) labels = { "id": (id, TaskType.LABEL), } @@ -76,8 +82,26 @@ def clear_files(): ONNX_PATH.unlink(missing_ok=True) -def test_reid(opts: dict[str, Any], infer_path: Path): +not_class_based_losses = ALL_EMBEDDING_LOSSES.copy() +for loss in CLASS_EMBEDDING_LOSSES: + not_class_based_losses.remove(loss) + + +@pytest.mark.parametrize("loss_name", not_class_based_losses) +def test_reid(opts: dict[str, Any], infer_path: Path, loss_name: str): config_file = "tests/configs/reid.yaml" + opts["model.losses.0.params.loss_name"] = loss_name + + # if loss_name in CLASS_EMBEDDING_LOSSES: + # opts["model.losses.0.params.num_classes"] = NUM_INDIVIDUALS + # opts["model.nodes.0.params.num_classes"] = NUM_INDIVIDUALS + # else: + # opts["model.losses.0.params.num_classes"] = 0 + # opts["model.nodes.0.params.num_classes"] = 0 + + if loss_name == "RankedListLoss": + opts["model.losses.0.params.loss_kwargs"] = {"margin": 1.0, "Tn": 0.5} + model = LuxonisModel(config_file, opts) model.train() model.test(view="val") From 06899357c0ba236114f778a3300c92d79b426a36 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Mon, 16 Dec 2024 11:54:55 +0100 Subject: [PATCH 06/57] refactor: update type hint for GhostFaceNetsV2 class to use Tensor from torch --- luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index 8bb61fee..2188645f 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -2,8 +2,8 @@ import math from typing import Literal -import torch import torch.nn as nn +from torch import Tensor from luxonis_train.nodes.backbones.ghostfacenet.blocks import ( GhostBottleneckV2, @@ -15,7 +15,7 @@ from luxonis_train.nodes.blocks import ConvModule -class GhostFaceNetsV2(BaseNode[torch.Tensor, list[torch.Tensor]]): +class GhostFaceNetsV2(BaseNode[Tensor, list[Tensor]]): in_channels: list[int] in_width: list[int] From 94639972c720172a27505fbe4439480bc640d824 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Mon, 16 Dec 2024 11:55:52 +0100 Subject: [PATCH 07/57] refactor: remove unused unwrap and wrap methods from GhostFaceNetsV2 class --- luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index 2188645f..cb065c43 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -142,12 +142,6 @@ def __init__( if isinstance(m, nn.BatchNorm2d): m.momentum, m.eps = var.bn_momentum, var.bn_epsilon - def unwrap(self, inputs): - return [inputs[0]["features"][0]] - - def wrap(self, outputs): - return {"features": [outputs]} - def forward(self, inps): x = inps[0] x = self.conv_stem(x) From 555fe2aa483d2493e1a07bc4e99df5c05954be00 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Mon, 16 Dec 2024 12:17:50 +0100 Subject: [PATCH 08/57] fix: correct formatting in __all__ list in metrics module --- luxonis_train/attached_modules/metrics/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/luxonis_train/attached_modules/metrics/__init__.py b/luxonis_train/attached_modules/metrics/__init__.py index 10f993ee..59e9cc57 100644 --- a/luxonis_train/attached_modules/metrics/__init__.py +++ b/luxonis_train/attached_modules/metrics/__init__.py @@ -17,6 +17,6 @@ "Precision", "Recall", "ClosestIsPositiveAccuracy", - "ConfusionMatrix", + "ConfusionMatrix", "MedianDistances", ] From 9fe0b798a121d480b067fe9cb5f4b9c939e1afe8 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Mon, 16 Dec 2024 14:02:28 +0100 Subject: [PATCH 09/57] Improved coverage, explicitly set mdformat github version --- .pre-commit-config.yaml | 2 +- tests/configs/reid.yaml | 2 +- tests/integration/test_reid.py | 29 ++++++++++++++++++++++++++++- 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c7779beb..226a18b8 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,4 +20,4 @@ repos: hooks: - id: mdformat additional_dependencies: - - mdformat-gfm + - mdformat-gfm==0.3.7 diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml index 21ca2748..c79e4f8e 100644 --- a/tests/configs/reid.yaml +++ b/tests/configs/reid.yaml @@ -15,7 +15,7 @@ model: params: loss_name: SupConLoss embedding_size: *embedding_size - cross_batch_memory_size: &memory_size 200 + cross_batch_memory_size: &memory_size 4 attached_to: GhostFaceNetsV2 metrics: diff --git a/tests/integration/test_reid.py b/tests/integration/test_reid.py index 0d006072..8094dd80 100644 --- a/tests/integration/test_reid.py +++ b/tests/integration/test_reid.py @@ -88,7 +88,9 @@ def clear_files(): @pytest.mark.parametrize("loss_name", not_class_based_losses) -def test_reid(opts: dict[str, Any], infer_path: Path, loss_name: str): +def test_available_losses( + opts: dict[str, Any], infer_path: Path, loss_name: str +): config_file = "tests/configs/reid.yaml" opts["model.losses.0.params.loss_name"] = loss_name @@ -113,3 +115,28 @@ def test_reid(opts: dict[str, Any], infer_path: Path, loss_name: str): assert len(list(infer_path.iterdir())) == 0 model.infer(view="val", save_dir=infer_path) assert infer_path.exists() + + +@pytest.mark.parametrize("loss_name", CLASS_EMBEDDING_LOSSES) +@pytest.mark.parametrize("num_classes", [-2, NUM_INDIVIDUALS]) +def test_unsupported_class_based_losses( + opts: dict[str, Any], loss_name: str, num_classes: int +): + config_file = "tests/configs/reid.yaml" + opts["model.losses.0.params.loss_name"] = loss_name + opts["model.losses.0.params.num_classes"] = num_classes + opts["model.nodes.0.params.num_classes"] = num_classes + + with pytest.raises(ValueError): + model = LuxonisModel(config_file, opts) + model.train() + + +@pytest.mark.parametrize("loss_name", ["NonExistentLoss"]) +def test_nonexistent_losses(opts: dict[str, Any], loss_name: str): + config_file = "tests/configs/reid.yaml" + opts["model.losses.0.params.loss_name"] = loss_name + + with pytest.raises(ValueError): + model = LuxonisModel(config_file, opts) + model.train() From b47e79ea82771758745f0ef6c84605e5e7e72e4f Mon Sep 17 00:00:00 2001 From: CaptainTrojan <49991681+CaptainTrojan@users.noreply.github.com> Date: Tue, 17 Dec 2024 18:20:51 +0100 Subject: [PATCH 10/57] Reduced mdformat-gfm version to 0.3.6 to support Python 3.8 --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 226a18b8..c9355abb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,4 +20,4 @@ repos: hooks: - id: mdformat additional_dependencies: - - mdformat-gfm==0.3.7 + - mdformat-gfm==0.3.6 From 8e376a0367e85081d77227a8890ef5fde3bfa11a Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Wed, 1 Jan 2025 22:58:59 +0100 Subject: [PATCH 11/57] Coverage fixes --- .../visualizers/embeddings_visualizer.py | 24 ++++------ .../backbones/ghostfacenet/ghostfacenet.py | 11 +---- tests/integration/test_reid.py | 48 +++++++++++++++++-- 3 files changed, 55 insertions(+), 28 deletions(-) diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py index d8e5c940..f3591c83 100644 --- a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -41,7 +41,7 @@ def forward( label_canvas: Tensor, prediction_canvas: Tensor, embeddings: Tensor, - ids: Tensor | None, + ids: Tensor, **kwargs, ) -> Tensor: """Creates a visualization of the embeddings. @@ -73,20 +73,14 @@ def forward( # Plot the embeddings fig, ax = plt.subplots(figsize=(10, 10)) - if ids is not None: - scatter = ax.scatter( - embeddings_2d[:, 0], - embeddings_2d[:, 1], - c=ids.detach().cpu().numpy(), - cmap="viridis", - s=5, - ) - else: - scatter = ax.scatter( - embeddings_2d[:, 0], - embeddings_2d[:, 1], - s=5, - ) + scatter = ax.scatter( + embeddings_2d[:, 0], + embeddings_2d[:, 1], + c=ids.detach().cpu().numpy(), + cmap="viridis", + s=5, + ) + fig.colorbar(scatter, ax=ax) ax.set_title("Embeddings Visualization") ax.set_xlabel("Dimension 1") diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index cb065c43..5a99ae28 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -111,16 +111,7 @@ def __init__( self.blocks = nn.Sequential(*stages) # Building pointwise convolution - pointwise_conv = [] - if var.add_pointwise_conv: - pointwise_conv.append( - nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=True) - ) - pointwise_conv.append(nn.BatchNorm2d(output_channel)) - pointwise_conv.append(nn.PReLU()) - else: - pointwise_conv.append(nn.Sequential()) - + pointwise_conv = [nn.Sequential()] self.pointwise_conv = nn.Sequential(*pointwise_conv) self.classifier = ModifiedGDC( image_size, diff --git a/tests/integration/test_reid.py b/tests/integration/test_reid.py index 8094dd80..53355025 100644 --- a/tests/integration/test_reid.py +++ b/tests/integration/test_reid.py @@ -35,7 +35,7 @@ def input_shapes(self): def __getitem__(self, _): # pragma: no cover # Fake data - image = torch.rand(3, 256, 256, dtype=torch.float32) + image = torch.rand(self.input_shapes["image"], dtype=torch.float32) inputs = { "image": image, } @@ -55,6 +55,24 @@ def get_classes(self) -> dict[TaskType, list[str]]: return {TaskType.LABEL: ["id"]} +class CustomReIDLoaderNoID(CustomReIDLoader): + def __getitem__(self, _): + inputs, labels = super().__getitem__(_) + labels["something_else"] = labels["id"] + del labels["id"] + + return inputs, labels + + +class CustomReIDLoaderImageSize2(CustomReIDLoader): + @property + def input_shapes(self): + return { + "image": torch.Size([3, 200, 200]), + "id": torch.Size([1]), + } + + @pytest.fixture def infer_path() -> Path: if INFER_PATH.exists(): @@ -128,8 +146,7 @@ def test_unsupported_class_based_losses( opts["model.nodes.0.params.num_classes"] = num_classes with pytest.raises(ValueError): - model = LuxonisModel(config_file, opts) - model.train() + LuxonisModel(config_file, opts) @pytest.mark.parametrize("loss_name", ["NonExistentLoss"]) @@ -137,6 +154,31 @@ def test_nonexistent_losses(opts: dict[str, Any], loss_name: str): config_file = "tests/configs/reid.yaml" opts["model.losses.0.params.loss_name"] = loss_name + with pytest.raises(ValueError): + LuxonisModel(config_file, opts) + + +def test_bad_loader(opts: dict[str, Any]): + config_file = "tests/configs/reid.yaml" + opts["loader.name"] = "CustomReIDLoaderNoID" + with pytest.raises(ValueError): model = LuxonisModel(config_file, opts) model.train() + + +def test_not_enough_samples_for_metrics(opts: dict[str, Any]): + config_file = "tests/configs/reid.yaml" + opts["model.metrics.1.params.cross_batch_memory_size"] = 100 + + model = LuxonisModel(config_file, opts) + model.train() + + +def test_image_size_not_divisible_by_32(opts: dict[str, Any]): + config_file = "tests/configs/reid.yaml" + opts["loader.name"] = "CustomReIDLoaderImageSize2" + + # with pytest.raises(ValueError): + model = LuxonisModel(config_file, opts) + model.train() From 23e75001a56b810a2047aeac69e4c2faacd79860 Mon Sep 17 00:00:00 2001 From: Michal Sejak Date: Thu, 2 Jan 2025 00:44:48 +0100 Subject: [PATCH 12/57] fix: return a model copy for the specified GhostFaceNets variant --- luxonis_train/nodes/backbones/ghostfacenet/variants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/luxonis_train/nodes/backbones/ghostfacenet/variants.py b/luxonis_train/nodes/backbones/ghostfacenet/variants.py index 0e88ecfc..aa78daf8 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/variants.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/variants.py @@ -211,4 +211,4 @@ def get_variant(variant: Literal["V2"]) -> GhostFaceNetsVariant: "GhostFaceNets model variant should be in " f"{list(variants.keys())}, got {variant}." ) - return variants[variant] + return variants[variant].model_copy() From df89eef7382c21afdfeaadcdc044f27b6683ec38 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Sat, 11 Jan 2025 12:04:32 +0100 Subject: [PATCH 13/57] initial labels refactor support --- .../attached_modules/base_attached_module.py | 28 +++-- luxonis_train/config/config.py | 17 +-- .../predefined_models/classification_model.py | 2 +- .../predefined_models/detection_fomo_model.py | 2 +- .../predefined_models/detection_model.py | 4 +- .../keypoint_detection_model.py | 2 +- .../predefined_models/segmentation_model.py | 7 +- luxonis_train/core/core.py | 36 ++---- luxonis_train/enums.py | 1 + luxonis_train/loaders/base_loader.py | 3 - luxonis_train/loaders/luxonis_loader_torch.py | 43 +++---- luxonis_train/loaders/utils.py | 20 ++-- luxonis_train/models/luxonis_lightning.py | 29 ++--- luxonis_train/nodes/base_node.py | 112 ++++++------------ .../nodes/heads/ddrnet_segmentation_head.py | 2 +- .../nodes/heads/efficient_bbox_head.py | 4 +- luxonis_train/utils/dataset_metadata.py | 11 ++ luxonis_train/utils/types.py | 7 +- 18 files changed, 132 insertions(+), 198 deletions(-) diff --git a/luxonis_train/attached_modules/base_attached_module.py b/luxonis_train/attached_modules/base_attached_module.py index c84d8575..99461c96 100644 --- a/luxonis_train/attached_modules/base_attached_module.py +++ b/luxonis_train/attached_modules/base_attached_module.py @@ -73,13 +73,13 @@ def __init__(self, *, node: BaseNode | None = None): for label in self.supported_tasks ] module_supported = f"[{', '.join(module_supported)}]" - if not self.node.tasks: + if not self.node.task_types: raise IncompatibleException( f"Module '{self.name}' requires one of the following " f"labels or combinations of labels: {module_supported}, " f"but is connected to node '{self.node.name}' which does not specify any tasks." ) - node_tasks = set(self.node.tasks) + node_tasks = set(self.node.task_types) for required_labels in self.supported_tasks: if isinstance(required_labels, TaskType): required_labels = [required_labels] @@ -89,7 +89,7 @@ def __init__(self, *, node: BaseNode | None = None): self.required_labels = required_labels break else: - node_supported = [task.value for task in self.node.tasks] + node_supported = [task.value for task in self.node.task_types] raise IncompatibleException( f"Module '{self.name}' requires one of the following labels or combinations of labels: {module_supported}, " f"but is connected to node '{self.node.name}' which does not support any of them. " @@ -159,18 +159,18 @@ def class_names(self) -> list[str]: return self.node.class_names @property - def node_tasks(self) -> dict[TaskType, str]: + def node_tasks(self) -> list[TaskType]: """Getter for the tasks of the attached node. @type: dict[TaskType, str] @raises RuntimeError: If the node does not have the C{tasks} attribute set. """ - if self.node._tasks is None: + if self.node.task_types is None: raise RuntimeError( "Node must have the `tasks` attribute specified." ) - return self.node._tasks + return self.node.task_types def get_label( self, labels: Labels, task_type: TaskType | None = None @@ -210,13 +210,14 @@ def _get_label( if len(self.required_labels) == 1: task_type = self.required_labels[0] - if task_type is not None: - task_name = self.node.get_task_name(task_type) - if task_name not in labels: + if task_type is not None and self.node.task_name is not None: + task_name = self.node.task_name + task = f"{task_name}/{task_type.value}" + if task not in labels: raise IncompatibleException.from_missing_task( task_type.value, list(labels.keys()), self.name ) - return labels[task_name] + return labels[task], task_type raise ValueError( f"{self.name} requires multiple labels. You must provide the " @@ -260,7 +261,7 @@ def get_input_tensors( f"Task {task_type.value} is not supported by the node " f"{self.node.name}." ) - return inputs[self.node_tasks[task_type]] + return inputs[f"{self.node.task_name}/{task_type.value}"] else: if task_type not in inputs: raise IncompatibleException( @@ -273,7 +274,8 @@ def get_input_tensors( f"{self.name} requires multiple labels, " "you must provide the `task_type` argument to extract the desired input." ) - return inputs[self.node_tasks[self.required_labels[0]]] + task_type = self.node_tasks[0].value + return inputs[f"{self.node.task_name}/{task_type}"] def prepare( self, inputs: Packet[Tensor], labels: Labels | None @@ -305,7 +307,7 @@ def prepare( @raises RuntimeError: If the C{tasks} attribute is not set on the node. @raises RuntimeError: If the C{supported_tasks} attribute is not set on the module. """ - if self.node._tasks is None: + if self.node.task_types is None: raise RuntimeError( f"{self.node.name} must have the `tasks` attribute specified " f"for {self.name} to make use of the default `prepare` method." diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index fcbbf24a..f2a49f85 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -1,7 +1,7 @@ import logging import sys import warnings -from typing import Annotated, Any, Literal, TypeAlias +from typing import Annotated, Any, Literal, NamedTuple, TypeAlias from luxonis_ml.enums import DatasetType from luxonis_ml.utils import ( @@ -19,13 +19,16 @@ ) from typing_extensions import Self -from luxonis_train.enums import TaskType - logger = logging.getLogger(__name__) Params: TypeAlias = dict[str, Any] +class ImageSize(NamedTuple): + height: int + width: int + + class AttachedModuleConfig(BaseModelExtraForbid): name: str attached_to: str @@ -62,7 +65,7 @@ class ModelNodeConfig(BaseModelExtraForbid): input_sources: list[str] = [] # From data loader freezing: FreezingConfig = FreezingConfig() remove_on_export: bool = False - task: str | dict[TaskType, str] | None = None + task_name: str | None = None params: Params = {} @@ -303,10 +306,10 @@ class AugmentationConfig(BaseModelExtraForbid): class PreprocessingConfig(BaseModelExtraForbid): train_image_size: Annotated[ - list[int], Field(default=[256, 256], min_length=2, max_length=2) - ] = [256, 256] + ImageSize, Field(default=[256, 256], min_length=2, max_length=2) + ] = ImageSize(256, 256) keep_aspect_ratio: bool = True - train_rgb: bool = True + color_format: Literal["RGB", "BGR"] = "RGB" normalize: NormalizeAugmentationConfig = NormalizeAugmentationConfig() augmentations: list[AugmentationConfig] = [] diff --git a/luxonis_train/config/predefined_models/classification_model.py b/luxonis_train/config/predefined_models/classification_model.py index e028bba5..0d749b5a 100644 --- a/luxonis_train/config/predefined_models/classification_model.py +++ b/luxonis_train/config/predefined_models/classification_model.py @@ -84,7 +84,7 @@ def nodes(self) -> list[ModelNodeConfig]: inputs=[f"{self.backbone}-{self.task_name}"], freezing=self.head_params.pop("freezing", {}), params=self.head_params, - task=self.task_name, + task_name=self.task_name, ), ] diff --git a/luxonis_train/config/predefined_models/detection_fomo_model.py b/luxonis_train/config/predefined_models/detection_fomo_model.py index d9702ece..1a21a5a3 100644 --- a/luxonis_train/config/predefined_models/detection_fomo_model.py +++ b/luxonis_train/config/predefined_models/detection_fomo_model.py @@ -80,7 +80,7 @@ def nodes(self) -> list[ModelNodeConfig]: alias=f"FOMOHead-{self.kpt_task_name}", inputs=[f"{self.backbone}-{self.kpt_task_name}"], params=self.head_params, - task={ + task_name={ TaskType.BOUNDINGBOX: self.bbox_task_name, TaskType.KEYPOINTS: self.kpt_task_name, }, diff --git a/luxonis_train/config/predefined_models/detection_model.py b/luxonis_train/config/predefined_models/detection_model.py index dbbc8886..d1498845 100644 --- a/luxonis_train/config/predefined_models/detection_model.py +++ b/luxonis_train/config/predefined_models/detection_model.py @@ -80,7 +80,7 @@ def __init__( self.head_params = head_params or var_config.head_params self.loss_params = loss_params or {"n_warmup_epochs": 0} self.visualizer_params = visualizer_params or {} - self.task_name = task_name or "boundingbox" + self.task_name = task_name @property def nodes(self) -> list[ModelNodeConfig]: @@ -114,7 +114,7 @@ def nodes(self) -> list[ModelNodeConfig]: if self.use_neck else [f"{self.backbone}-{self.task_name}"], params=self.head_params, - task=self.task_name, + task_name=self.task_name, ) ) return nodes diff --git a/luxonis_train/config/predefined_models/keypoint_detection_model.py b/luxonis_train/config/predefined_models/keypoint_detection_model.py index 51d790a7..8882f338 100644 --- a/luxonis_train/config/predefined_models/keypoint_detection_model.py +++ b/luxonis_train/config/predefined_models/keypoint_detection_model.py @@ -122,7 +122,7 @@ def nodes(self) -> list[ModelNodeConfig]: ), freezing=self.head_params.pop("freezing", {}), params=self.head_params, - task=task, + task_name=task, ) ) return nodes diff --git a/luxonis_train/config/predefined_models/segmentation_model.py b/luxonis_train/config/predefined_models/segmentation_model.py index eff4fd02..0260e843 100644 --- a/luxonis_train/config/predefined_models/segmentation_model.py +++ b/luxonis_train/config/predefined_models/segmentation_model.py @@ -71,7 +71,7 @@ def __init__( self.loss_params = loss_params or {} self.visualizer_params = visualizer_params or {} self.task = task - self.task_name = task_name or "segmentation" + self.task_name = task_name @property def nodes(self) -> list[ModelNodeConfig]: @@ -85,6 +85,7 @@ def nodes(self) -> list[ModelNodeConfig]: alias=f"{self.backbone}-{self.task_name}", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, + task_name=self.task_name, ), ModelNodeConfig( name="DDRNetSegmentationHead", @@ -92,7 +93,7 @@ def nodes(self) -> list[ModelNodeConfig]: inputs=[f"{self.backbone}-{self.task_name}"], freezing=self.head_params.pop("freezing", {}), params=self.head_params, - task=self.task_name, + task_name=self.task_name, ), ] if self.backbone_params.get("use_aux_heads", True): @@ -103,7 +104,7 @@ def nodes(self) -> list[ModelNodeConfig]: inputs=[f"{self.backbone}-{self.task_name}"], freezing=self.aux_head_params.pop("freezing", {}), params=self.aux_head_params, - task=self.task_name, + task_name=self.task_name, remove_on_export=self.aux_head_params.pop( "remove_on_export", True ), diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py index 03ff5189..fe959349 100644 --- a/luxonis_train/core/core.py +++ b/luxonis_train/core/core.py @@ -13,7 +13,6 @@ import torch.utils.data as torch_data import yaml from lightning.pytorch.utilities import rank_zero_only -from luxonis_ml.data import Augmentations from luxonis_ml.nn_archive import ArchiveGenerator from luxonis_ml.nn_archive.config import CONFIG_VERSION from luxonis_ml.utils import LuxonisFileSystem, reset_logging, setup_logging @@ -113,26 +112,6 @@ def __init__( precision=self.cfg.trainer.precision, ) - self.train_augmentations = Augmentations( - image_size=self.cfg.trainer.preprocessing.train_image_size, - augmentations=[ - i.model_dump() - for i in self.cfg.trainer.preprocessing.get_active_augmentations() - ], - train_rgb=self.cfg.trainer.preprocessing.train_rgb, - keep_aspect_ratio=self.cfg.trainer.preprocessing.keep_aspect_ratio, - ) - self.val_augmentations = Augmentations( - image_size=self.cfg.trainer.preprocessing.train_image_size, - augmentations=[ - i.model_dump() - for i in self.cfg.trainer.preprocessing.get_active_augmentations() - ], - train_rgb=self.cfg.trainer.preprocessing.train_rgb, - keep_aspect_ratio=self.cfg.trainer.preprocessing.keep_aspect_ratio, - only_normalize=True, - ) - self.loaders: dict[str, BaseLoaderTorch] = {} for view in ["train", "val", "test"]: loader_name = self.cfg.loader.name @@ -141,17 +120,20 @@ def __init__( self.cfg.loader.params["delete_existing"] = False self.loaders[view] = Loader( - augmentations=( - self.train_augmentations - if view == "train" - else self.val_augmentations - ), view={ "train": self.cfg.loader.train_view, "val": self.cfg.loader.val_view, "test": self.cfg.loader.test_view, }[view], image_source=self.cfg.loader.image_source, + height=self.cfg.trainer.preprocessing.train_image_size.height, + width=self.cfg.trainer.preprocessing.train_image_size.width, + augmentation_config=[ + i.model_dump() + for i in self.cfg.trainer.preprocessing.get_active_augmentations() + ], + out_image_format=self.cfg.trainer.preprocessing.color_format, + keep_aspect_ratio=self.cfg.trainer.preprocessing.keep_aspect_ratio, **self.cfg.loader.params, ) @@ -739,7 +721,7 @@ def _mult(lst: list[float | int]) -> list[float]: self.cfg.trainer.preprocessing.normalize.params["std"] ), "dai_type": "RGB888p" - if self.cfg.trainer.preprocessing.train_rgb + if self.cfg.trainer.preprocessing.out_image_format else "BGR888p", } diff --git a/luxonis_train/enums.py b/luxonis_train/enums.py index b024d6a9..ea719e1c 100644 --- a/luxonis_train/enums.py +++ b/luxonis_train/enums.py @@ -6,6 +6,7 @@ class TaskType(str, Enum): CLASSIFICATION = "classification" SEGMENTATION = "segmentation" + INSTANCE_SEGMENTATION = "instance_segmentation" BOUNDINGBOX = "boundingbox" KEYPOINTS = "keypoints" LABEL = "label" diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index 0c056d98..3e3589dd 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -1,6 +1,5 @@ from abc import ABC, abstractmethod -from luxonis_ml.data import Augmentations from luxonis_ml.utils.registry import AutoRegisterMeta from torch import Size from torch.utils.data import Dataset @@ -23,11 +22,9 @@ class BaseLoaderTorch( def __init__( self, view: str | list[str], - augmentations: Augmentations | None = None, image_source: str | None = None, ): self.view = view if isinstance(view, list) else [view] - self.augmentations = augmentations self._image_source = image_source @property diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index 230128b5..445fe641 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -1,9 +1,9 @@ import logging -from typing import Literal +from pathlib import Path +from typing import Any, Literal import numpy as np from luxonis_ml.data import ( - Augmentations, BucketStorage, BucketType, LuxonisDataset, @@ -14,8 +14,6 @@ from torch import Size, Tensor from typeguard import typechecked -from luxonis_train.enums import TaskType - from .base_loader import BaseLoaderTorch, LuxonisLoaderTorchOutput logger = logging.getLogger(__name__) @@ -31,10 +29,15 @@ def __init__( team_id: str | None = None, bucket_type: Literal["internal", "external"] = "internal", bucket_storage: Literal["local", "s3", "gcs", "azure"] = "local", - stream: bool = False, delete_existing: bool = True, view: str | list[str] = "train", - augmentations: Augmentations | None = None, + augmentation_engine: str + | Literal["albumentations"] = "albumentations", + augmentation_config: Path | str | list[dict[str, Any]] | None = None, + height: int | None = None, + width: int | None = None, + keep_aspect_ratio: bool = True, + out_image_format: Literal["RGB", "BGR"] = "RGB", **kwargs, ): """Torch-compatible loader for Luxonis datasets. @@ -61,8 +64,6 @@ def __init__( Defaults to 'internal'. @type bucket_storage: Literal["local", "s3", "gcs", "azure"] @param bucket_storage: Type of the bucket storage. Defaults to 'local'. - @type stream: bool - @param stream: Flag for data streaming. Defaults to C{False}. @type delete_existing: bool @param delete_existing: Only relevant when C{dataset_dir} is provided. By default, the dataset is parsed again every time the loader is created @@ -74,10 +75,8 @@ def __init__( view of the dataset. Each split is a string that represents a subset of the dataset. The available splits depend on the dataset, but usually include 'train', 'val', and 'test'. Defaults to 'train'. - @type augmentations: Augmentations | None - @param augmentations: Augmentations to apply to the data. Defaults to C{None}. """ - super().__init__(view=view, augmentations=augmentations, **kwargs) + super().__init__(view=view, **kwargs) if dataset_dir is not None: self.dataset = self._parse_dataset( dataset_dir, dataset_name, dataset_type, delete_existing @@ -93,15 +92,19 @@ def __init__( bucket_type=BucketType(bucket_type), bucket_storage=BucketStorage(bucket_storage), ) - self.base_loader = LuxonisLoader( + self.loader = LuxonisLoader( dataset=self.dataset, - view=self.view, - stream=stream, - augmentations=self.augmentations, + view=view, + augmentation_engine=augmentation_engine, + augmentation_config=augmentation_config, + height=height, + width=width, + keep_aspect_ratio=keep_aspect_ratio, + out_image_format=out_image_format, ) def __len__(self) -> int: - return len(self.base_loader) + return len(self.loader) @property def input_shapes(self) -> dict[str, Size]: @@ -109,13 +112,13 @@ def input_shapes(self) -> dict[str, Size]: return {self.image_source: img.shape} def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: - img, labels = self.base_loader[idx] + img, labels = self.loader[idx] img = np.transpose(img, (2, 0, 1)) # HWC to CHW tensor_img = Tensor(img) - tensor_labels: dict[str, tuple[Tensor, TaskType]] = {} - for task, (array, label_type) in labels.items(): - tensor_labels[task] = (Tensor(array), TaskType(label_type.value)) + tensor_labels: dict[str, Tensor] = {} + for task, array in labels.items(): + tensor_labels[task] = Tensor(array) return {self.image_source: tensor_img}, tensor_labels diff --git a/luxonis_train/loaders/utils.py b/luxonis_train/loaders/utils.py index b030e218..10b4d17a 100644 --- a/luxonis_train/loaders/utils.py +++ b/luxonis_train/loaders/utils.py @@ -1,7 +1,7 @@ import torch +from luxonis_ml.data.utils import get_task_type from torch import Tensor -from luxonis_train.enums import TaskType from luxonis_train.utils.types import Labels LuxonisLoaderTorchOutput = tuple[dict[str, Tensor], Labels] @@ -32,22 +32,18 @@ def collate_fn( out_labels: Labels = {} for task in labels[0].keys(): - task_type = labels[0][task][1] - annos = [label[task][0] for label in labels] - if task_type in [ - TaskType.CLASSIFICATION, - TaskType.SEGMENTATION, - TaskType.ARRAY, - ]: - out_labels[task] = torch.stack(annos, 0), task_type - - elif task_type in [TaskType.KEYPOINTS, TaskType.BOUNDINGBOX]: + task_type = get_task_type(task) + annos = [label[task] for label in labels] + + if task_type in {"keypoints", "boundingbox"}: label_box: list[Tensor] = [] for i, box in enumerate(annos): l_box = torch.zeros((box.shape[0], box.shape[1] + 1)) l_box[:, 0] = i # add target image index for build_targets() l_box[:, 1:] = box label_box.append(l_box) - out_labels[task] = torch.cat(label_box, 0), task_type + out_labels[task] = torch.cat(label_box, 0) + else: + out_labels[task] = torch.stack(annos, 0) return out_inputs, out_labels diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index dae683ce..cebce43f 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -181,32 +181,17 @@ def __init__( ) frozen_nodes.append((node_name, unfreeze_after)) - if node_cfg.task is not None: - if Node.tasks is None: - raise ValueError( - f"Cannot define tasks for node {node_name}." - "This node doesn't specify any tasks." - ) - if isinstance(node_cfg.task, str): - assert Node.tasks - if len(Node.tasks) > 1: - raise ValueError( - f"Node {node_name} specifies multiple tasks, " - "but only one task is specified in the config. " - "Specify the tasks as a dictionary instead." - ) + if node_cfg.task_name is not None and Node.task_types is None: + raise ValueError( + f"Cannot define tasks for node {node_name}." + "This node doesn't specify any tasks." + ) - node_cfg.task = {next(iter(Node.tasks)): node_cfg.task} - else: - node_cfg.task = { - **Node._process_tasks(Node.tasks), - **node_cfg.task, - } nodes[node_name] = ( Node, { **node_cfg.params, - "_tasks": node_cfg.task, + "task_name": node_cfg.task_name, "remove_on_export": node_cfg.remove_on_export, }, ) @@ -1000,7 +985,7 @@ def _init_attached_module( loader = self._core.loaders["train"] dataset = getattr(loader, "dataset", None) if isinstance(dataset, LuxonisDataset): - n_classes = len(dataset.get_classes()[1][node.task]) + n_classes = len(dataset.get_classes()[1][node.task_name]) if n_classes == 1: cfg.params["task"] = "binary" else: diff --git a/luxonis_train/nodes/base_node.py b/luxonis_train/nodes/base_node.py index 748742dd..eff2c2d1 100644 --- a/luxonis_train/nodes/base_node.py +++ b/luxonis_train/nodes/base_node.py @@ -109,7 +109,7 @@ def wrap(output: Tensor) -> Packet[Tensor]: """ attach_index: AttachIndexType - tasks: list[TaskType] | dict[TaskType, str] | None = None + task_types: list[TaskType] | None = None def __init__( self, @@ -123,7 +123,7 @@ def __init__( remove_on_export: bool = False, export_output_names: list[str] | None = None, attach_index: AttachIndexType | None = None, - _tasks: dict[TaskType, str] | None = None, + task_name: str | None = None, ): """Constructor for the C{BaseNode}. @@ -168,11 +168,16 @@ def __init__( "Make sure this is intended." ) self.attach_index = attach_index - self._tasks = None - if _tasks is not None: - self._tasks = _tasks - elif self.tasks is not None: - self._tasks = self._process_tasks(self.tasks) + + self.task_name = task_name + if task_name is None and dataset_metadata is not None: + if len(dataset_metadata.task_names) == 1: + self.task_name = next(iter(dataset_metadata.task_names)) + else: + raise ValueError( + f"Dataset contain multiple tasks, but the `task_name` " + f"argument for node '{self.name}' was not provided." + ) if getattr(self, "attach_index", None) is None: parameters = inspect.signature(self.forward).parameters @@ -200,15 +205,6 @@ def __init__( self._check_type_overrides() - @staticmethod - def _process_tasks( - tasks: dict[TaskType, str] | list[TaskType], - ) -> dict[TaskType, str]: - if isinstance(tasks, dict): - return tasks - else: - return {task: task.value for task in tasks} - def _check_type_overrides(self) -> None: properties = [] for name, value in inspect.getmembers(self.__class__): @@ -228,67 +224,28 @@ def _check_type_overrides(self) -> None: "not compatible with its predecessor." ) from e - def get_task_name(self, task: TaskType) -> str: - """Gets the name of a task for a particular C{TaskType}. - - @type task: TaskType - @param task: Task to get the name for. - @rtype: str - @return: Name of the task. - @raises RuntimeError: If the node does not define any tasks. - @raises ValueError: If the task is not supported by the node. - """ - if not self._tasks: - raise RuntimeError(f"Node '{self.name}' does not define any task.") - - if task not in self._tasks: - raise ValueError( - f"Node '{self.name}' does not support the '{task.value}' task." - ) - return self._tasks[task] - @property def name(self) -> str: return self.__class__.__name__ @property - def task(self) -> str: - """Getter for the task. + def task_type(self) -> str: + """Getter for the task type. @type: str @raises RuntimeError: If the node doesn't define any task. @raises ValueError: If the node defines more than one task. In that case, use the L{get_task_name} method instead. """ - if not self._tasks: + if not self.task_types: raise RuntimeError(f"{self.name} does not define any task.") - if len(self._tasks) > 1: + if len(self.task_types) > 1: raise ValueError( f"Node {self.name} has multiple tasks defined. " "Use the `get_task_name` method instead." ) - return next(iter(self._tasks.values())) - - def get_n_classes(self, task: TaskType) -> int: - """Gets the number of classes for a particular task. - - @type task: TaskType - @param task: Task to get the number of classes for. - @rtype: int - @return: Number of classes for the task. - """ - return self.dataset_metadata.n_classes(self.get_task_name(task)) - - def get_class_names(self, task: TaskType) -> list[str]: - """Gets the class names for a particular task. - - @type task: TaskType - @param task: Task to get the class names for. - @rtype: list[str] - @return: Class names for the task. - """ - return self.dataset_metadata.classes(self.get_task_name(task)) + return self.task_types[0].value @property def n_keypoints(self) -> int: @@ -301,12 +258,10 @@ def n_keypoints(self) -> int: if self._n_keypoints is not None: return self._n_keypoints - if self._tasks: - if TaskType.KEYPOINTS not in self._tasks: + if self.task_types: + if TaskType.KEYPOINTS not in self.task_types: raise ValueError(f"{self.name} does not support keypoints.") - return self.dataset_metadata.n_keypoints( - self.get_task_name(TaskType.KEYPOINTS) - ) + return self.dataset_metadata.n_keypoints(self.task_name) raise RuntimeError( f"{self.name} does not have any tasks defined, " @@ -329,7 +284,7 @@ def n_classes(self) -> int: if self._n_classes is not None: return self._n_classes - if not self._tasks: + if not self.task_types: raise RuntimeError( f"{self.name} does not have any tasks defined, " "`BaseNode.n_classes` property cannot be used. " @@ -337,12 +292,12 @@ def n_classes(self) -> int: "pass the `n_classes` attribute to the constructor or call " "the `BaseNode.dataset_metadata.n_classes` method manually." ) - elif len(self._tasks) == 1: - return self.dataset_metadata.n_classes(self.task) + elif len(self.task_types) == 1: + return self.dataset_metadata.n_classes(self.task_name) else: n_classes = [ - self.dataset_metadata.n_classes(self.get_task_name(task)) - for task in self._tasks + self.dataset_metadata.n_classes(self.task_name) + for task in self.task_types ] if len(set(n_classes)) == 1: return n_classes[0] @@ -362,7 +317,7 @@ def class_names(self) -> list[str]: different tasks. In that case, use the L{get_class_names} method. """ - if not self._tasks: + if not self.task_types: raise RuntimeError( f"{self.name} does not have any tasks defined, " "`BaseNode.class_names` property cannot be used. " @@ -370,12 +325,12 @@ def class_names(self) -> list[str]: "pass the `n_classes` attribute to the constructor or call " "the `BaseNode.dataset_metadata.class_names` method manually." ) - elif len(self._tasks) == 1: - return self.dataset_metadata.classes(self.task) + elif len(self.task_types) == 1: + return self.dataset_metadata.classes(self.task_name) else: class_names = [ - self.dataset_metadata.classes(self.get_task_name(task)) - for task in self._tasks + self.dataset_metadata.classes(self.task_name) + for task in self.task_types ] if all(set(names) == set(class_names[0]) for names in class_names): return class_names[0] @@ -633,7 +588,7 @@ def wrap(self, output: ForwardOutputT) -> Packet[Tensor]: "Default `wrap` expects a single tensor or a list of tensors." ) try: - task = self.task + task = f"{self.task_name}/{self.task_type}" except RuntimeError: task = "features" return {task: outputs} @@ -654,11 +609,12 @@ def run(self, inputs: list[Packet[Tensor]]) -> Packet[Tensor]: unwrapped = self.unwrap(inputs) outputs = self(unwrapped) wrapped = self.wrap(outputs) - str_tasks = [task.value for task in self._tasks] if self._tasks else [] + str_tasks = [task.value for task in self.task_types or []] for key in list(wrapped.keys()): if key in str_tasks: + assert self.task_name is not None value = wrapped.pop(key) - wrapped[self.get_task_name(TaskType(key))] = value + wrapped[f"{self.task_name}/{key}"] = value return wrapped T = TypeVar("T", Tensor, Size) diff --git a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py index 2b313ab6..35c5e69c 100644 --- a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py +++ b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py @@ -18,7 +18,7 @@ class DDRNetSegmentationHead(BaseHead[Tensor, Tensor]): in_width: int in_channels: int - tasks: list[TaskType] = [TaskType.SEGMENTATION] + task_types: list[TaskType] = [TaskType.SEGMENTATION] parser: str = "SegmentationParser" def __init__( diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py index 76eb2e5a..2eb57dfb 100644 --- a/luxonis_train/nodes/heads/efficient_bbox_head.py +++ b/luxonis_train/nodes/heads/efficient_bbox_head.py @@ -21,7 +21,7 @@ class EfficientBBoxHead( BaseHead[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]], ): in_channels: list[int] - tasks: list[TaskType] = [TaskType.BOUNDINGBOX] + task_types: list[TaskType] = [TaskType.BOUNDINGBOX] parser = "YOLO" def __init__( @@ -171,7 +171,7 @@ def wrap( conf, _ = out_cls.max(1, keepdim=True) out = torch.cat([out_reg, conf, out_cls], dim=1) outputs.append(out) - return {self.task: outputs} + return {self.task_type: outputs} cls_tensor = torch.cat( [cls_score_list[i].flatten(2) for i in range(len(cls_score_list))], diff --git a/luxonis_train/utils/dataset_metadata.py b/luxonis_train/utils/dataset_metadata.py index 3a9cecdf..fdbec775 100644 --- a/luxonis_train/utils/dataset_metadata.py +++ b/luxonis_train/utils/dataset_metadata.py @@ -1,3 +1,5 @@ +from typing import Set + from luxonis_train.loaders import BaseLoaderTorch @@ -28,6 +30,15 @@ def __init__( self._n_keypoints = n_keypoints or {} self._loader = loader + @property + def task_names(self) -> Set[str]: + """Gets the names of the tasks present in the dataset. + + @rtype: set[str] + @return: Names of the tasks present in the dataset. + """ + return set(self._classes.keys()) + def n_classes(self, task: str | None = None) -> int: """Gets the number of classes for the specified task. diff --git a/luxonis_train/utils/types.py b/luxonis_train/utils/types.py index 8666751b..f1d8e6b5 100644 --- a/luxonis_train/utils/types.py +++ b/luxonis_train/utils/types.py @@ -2,14 +2,11 @@ from torch import Size, Tensor -from luxonis_train.enums import TaskType - Kwargs = dict[str, Any] """Kwargs is a dictionary containing keyword arguments.""" -Labels = dict[str, tuple[Tensor, TaskType]] -"""Labels is a dictionary containing a tuple of tensors and their -corresponding task type.""" +Labels = dict[str, Tensor] +"""Labels is a dictionary mapping task names to tensors.""" AttachIndexType = Literal["all"] | int | tuple[int, int] | tuple[int, int, int] """AttachIndexType is used to specify to which output of the prevoius From d01816bc5b551f0408036a38890f7c0360e20f98 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 07:14:04 -0600 Subject: [PATCH 14/57] updated docs --- configs/README.md | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/configs/README.md b/configs/README.md index 69d77243..37d150f9 100644 --- a/configs/README.md +++ b/configs/README.md @@ -280,14 +280,14 @@ We use [`Albumentations`](https://albumentations.ai/docs/) library for `augmenta Additionally, we support `Mosaic4` and `MixUp` batch augmentations and letterbox resizing if `keep_aspect_ratio: true`. -| Key | Type | Default value | Description | -| ------------------- | ------------ | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `train_image_size` | `list[int]` | `[256, 256]` | Image size used for training as `[height, width]` | -| `keep_aspect_ratio` | `bool` | `True` | Whether to keep the aspect ratio while resizing | -| `train_rgb` | `bool` | `True` | Whether to train on RGB or BGR images | -| `normalize.active` | `bool` | `True` | Whether to use normalization | -| `normalize.params` | `dict` | `{}` | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | -| `augmentations` | `list[dict]` | `[]` | List of `Albumentations` augmentations | +| Key | Type | Default value | Description | +| ------------------- | ----------------------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `train_image_size` | `list[int]` | `[256, 256]` | Image size used for training as `[height, width]` | +| `keep_aspect_ratio` | `bool` | `True` | Whether to keep the aspect ratio while resizing | +| `color_format` | `Literal["RGB", "BGR"]` | `"RGB"` | Whether to train on RGB or BGR images | +| `normalize.active` | `bool` | `True` | Whether to use normalization | +| `normalize.params` | `dict` | `{}` | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | +| `augmentations` | `list[dict]` | `[]` | List of `Albumentations` augmentations | #### Augmentations @@ -306,7 +306,7 @@ trainer: # using YAML capture to reuse the image size train_image_size: [&height 384, &width 384] keep_aspect_ratio: true - train_rgb: true + color_format: "RGB" normalize: active: true augmentations: @@ -418,7 +418,7 @@ Each training strategy is a dictionary with the following fields: ```yaml training_strategy: name: "TripleLRSGDStrategy" - params: + params: warmup_epochs: 3 warmup_bias_lr: 0.1 warmup_momentum: 0.8 From e34e893b83b5e03b2fbead9b1eb4e45b8d564cd1 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 07:43:38 -0600 Subject: [PATCH 15/57] updated predefined models --- configs/complex_model.yaml | 2 +- luxonis_train/config/config.py | 39 ++++++++++++------- .../anomaly_detection_model.py | 19 ++++----- .../predefined_models/classification_model.py | 23 +++++------ .../predefined_models/detection_fomo_model.py | 29 +++++--------- .../predefined_models/detection_model.py | 23 +++++------ .../keypoint_detection_model.py | 38 +++++++----------- .../predefined_models/segmentation_model.py | 35 ++++++----------- 8 files changed, 85 insertions(+), 123 deletions(-) diff --git a/configs/complex_model.yaml b/configs/complex_model.yaml index 149530ad..fed25c23 100644 --- a/configs/complex_model.yaml +++ b/configs/complex_model.yaml @@ -100,7 +100,7 @@ trainer: preprocessing: train_image_size: [&height 384, &width 384] keep_aspect_ratio: true - train_rgb: true + color_format: RGB normalize: active: true augmentations: diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index f2a49f85..014caaf2 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -65,7 +65,7 @@ class ModelNodeConfig(BaseModelExtraForbid): input_sources: list[str] = [] # From data loader freezing: FreezingConfig = FreezingConfig() remove_on_export: bool = False - task_name: str | None = None + task_name: str = "" params: Params = {} @@ -105,7 +105,7 @@ def validate_nodes(cls, nodes: Any) -> Any: if "Head" in name and last_body_index is None: last_body_index = i - 1 names.append(name) - if i > 0 and "inputs" not in node: + if i > 0 and "inputs" not in node and "input_sources" not in node: if last_body_index is not None: prev_name = names[last_body_index] else: @@ -202,23 +202,32 @@ def check_graph(self) -> Self: @model_validator(mode="after") def check_unique_names(self) -> Self: - for section, objects in [ - ("nodes", self.nodes), - ("losses", self.losses), - ("metrics", self.metrics), - ("visualizers", self.visualizers), + for modules in [ + self.nodes, + self.losses, + self.metrics, + self.visualizers, ]: names: set[str] = set() - for obj in objects: - obj: AttachedModuleConfig - name = obj.alias or obj.name + node_index = 0 + for module in modules: + module: AttachedModuleConfig | ModelNodeConfig + name = module.alias or module.name if name in names: - if obj.alias is None: - obj.alias = f"{name}_{obj.attached_to}" - if obj.alias in names: - raise ValueError( - f"Duplicate name `{name}` in `{section}` section." + if module.alias is None: + if isinstance(module, ModelNodeConfig): + module.alias = module.name + else: + module.alias = f"{name}_{module.attached_to}" + + if module.alias in names: + new_alias = f"{module.alias}_{node_index}" + logger.warning( + f"Duplicate name: {module.alias}. Renaming to {new_alias}." ) + module.alias = new_alias + node_index += 1 + names.add(name) return self diff --git a/luxonis_train/config/predefined_models/anomaly_detection_model.py b/luxonis_train/config/predefined_models/anomaly_detection_model.py index 6dbb1a1d..5dff9b93 100644 --- a/luxonis_train/config/predefined_models/anomaly_detection_model.py +++ b/luxonis_train/config/predefined_models/anomaly_detection_model.py @@ -5,7 +5,7 @@ from luxonis_train.config import ( AttachedModuleConfig, LossModuleConfig, - MetricModuleConfig, # Metrics support added + MetricModuleConfig, ModelNodeConfig, Params, ) @@ -51,7 +51,7 @@ def __init__( disc_subnet_params: Params | None = None, loss_params: Params | None = None, visualizer_params: Params | None = None, - task_name: str | None = None, + task_name: str = "", ): var_config = get_variant(variant) @@ -73,13 +73,13 @@ def nodes(self) -> list[ModelNodeConfig]: return [ ModelNodeConfig( name=self.backbone, - alias=f"{self.backbone}-{self.task_name}", + alias=f"{self.task_name}/{self.backbone}", params=self.backbone_params, ), ModelNodeConfig( name="DiscSubNetHead", - alias=f"DiscSubNetHead-{self.task_name}", - inputs=[f"{self.backbone}-{self.task_name}"], + alias=f"{self.task_name}/DiscSubNetHead", + inputs=[f"{self.task_name}/{self.backbone}"], params=self.disc_subnet_params, ), ] @@ -90,8 +90,7 @@ def losses(self) -> list[LossModuleConfig]: return [ LossModuleConfig( name="ReconstructionSegmentationLoss", - alias=f"ReconstructionSegmentationLoss-{self.task_name}", - attached_to=f"DiscSubNetHead-{self.task_name}", + attached_to=f"{self.task_name}/DiscSubNetHead", params=self.loss_params, weight=1.0, ) @@ -103,8 +102,7 @@ def metrics(self) -> list[MetricModuleConfig]: return [ MetricModuleConfig( name="JaccardIndex", - alias=f"JaccardIndex-{self.task_name}", - attached_to=f"DiscSubNetHead-{self.task_name}", + attached_to=f"{self.task_name}/DiscSubNetHead", params={"num_classes": 2, "task": "multiclass"}, is_main_metric=True, ), @@ -117,8 +115,7 @@ def visualizers(self) -> list[AttachedModuleConfig]: return [ AttachedModuleConfig( name="SegmentationVisualizer", - alias=f"SegmentationVisualizer-{self.task_name}", - attached_to=f"DiscSubNetHead-{self.task_name}", + attached_to=f"{self.task_name}/DiscSubNetHead", params=self.visualizer_params, ) ] diff --git a/luxonis_train/config/predefined_models/classification_model.py b/luxonis_train/config/predefined_models/classification_model.py index 0d749b5a..86964e0a 100644 --- a/luxonis_train/config/predefined_models/classification_model.py +++ b/luxonis_train/config/predefined_models/classification_model.py @@ -52,7 +52,7 @@ def __init__( loss_params: Params | None = None, visualizer_params: Params | None = None, task: Literal["multiclass", "multilabel"] = "multiclass", - task_name: str | None = None, + task_name: str = "", ): var_config = get_variant(variant) @@ -74,14 +74,14 @@ def nodes(self) -> list[ModelNodeConfig]: return [ ModelNodeConfig( name=self.backbone, - alias=f"{self.backbone}-{self.task_name}", + alias=f"{self.task_name}/{self.backbone}", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, ), ModelNodeConfig( name="ClassificationHead", - alias=f"ClassificationHead-{self.task_name}", - inputs=[f"{self.backbone}-{self.task_name}"], + alias=f"{self.task_name}/ClassificationHead", + inputs=[f"{self.task_name}/{self.backbone}"], freezing=self.head_params.pop("freezing", {}), params=self.head_params, task_name=self.task_name, @@ -94,8 +94,7 @@ def losses(self) -> list[LossModuleConfig]: return [ LossModuleConfig( name="CrossEntropyLoss", - alias=f"CrossEntropyLoss-{self.task_name}", - attached_to=f"ClassificationHead-{self.task_name}", + attached_to=f"{self.task_name}/ClassificationHead", params=self.loss_params, weight=1.0, ) @@ -107,21 +106,18 @@ def metrics(self) -> list[MetricModuleConfig]: return [ MetricModuleConfig( name="F1Score", - alias=f"F1Score-{self.task_name}", is_main_metric=True, - attached_to=f"ClassificationHead-{self.task_name}", + attached_to=f"{self.task_name}/ClassificationHead", params={"task": self.task}, ), MetricModuleConfig( name="Accuracy", - alias=f"Accuracy-{self.task_name}", - attached_to=f"ClassificationHead-{self.task_name}", + attached_to=f"{self.task_name}/ClassificationHead", params={"task": self.task}, ), MetricModuleConfig( name="Recall", - alias=f"Recall-{self.task_name}", - attached_to=f"ClassificationHead-{self.task_name}", + attached_to=f"{self.task_name}/ClassificationHead", params={"task": self.task}, ), ] @@ -132,8 +128,7 @@ def visualizers(self) -> list[AttachedModuleConfig]: return [ AttachedModuleConfig( name="ClassificationVisualizer", - alias=f"ClassificationVisualizer-{self.task_name}", - attached_to=f"ClassificationHead-{self.task_name}", + attached_to=f"{self.task_name}/ClassificationHead", params=self.visualizer_params, ) ] diff --git a/luxonis_train/config/predefined_models/detection_fomo_model.py b/luxonis_train/config/predefined_models/detection_fomo_model.py index 1a21a5a3..309a4572 100644 --- a/luxonis_train/config/predefined_models/detection_fomo_model.py +++ b/luxonis_train/config/predefined_models/detection_fomo_model.py @@ -9,7 +9,6 @@ ModelNodeConfig, Params, ) -from luxonis_train.enums import TaskType from .base_predefined_model import BasePredefinedModel @@ -51,8 +50,7 @@ def __init__( head_params: Params | None = None, loss_params: Params | None = None, kpt_visualizer_params: Params | None = None, - bbox_task_name: str | None = None, - kpt_task_name: str | None = None, + task_name: str = "", ): var_config = get_variant(variant) @@ -61,29 +59,23 @@ def __init__( self.head_params = head_params or var_config.head_params self.loss_params = loss_params or {} self.kpt_visualizer_params = kpt_visualizer_params or {} - self.bbox_task_name = ( - bbox_task_name or "boundingbox" - ) # Needed for OKS calculation - self.kpt_task_name = kpt_task_name or "keypoints" + self.task_name = task_name @property def nodes(self) -> list[ModelNodeConfig]: nodes = [ ModelNodeConfig( name=self.backbone, - alias=f"{self.backbone}-{self.kpt_task_name}", + alias=f"{self.task_name}/{self.backbone}", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, ), ModelNodeConfig( name="FOMOHead", - alias=f"FOMOHead-{self.kpt_task_name}", - inputs=[f"{self.backbone}-{self.kpt_task_name}"], + alias=f"{self.task_name}/FOMOHead", + inputs=[f"{self.task_name}/{self.backbone}"], params=self.head_params, - task_name={ - TaskType.BOUNDINGBOX: self.bbox_task_name, - TaskType.KEYPOINTS: self.kpt_task_name, - }, + task_name=self.task_name, ), ] return nodes @@ -93,8 +85,7 @@ def losses(self) -> list[LossModuleConfig]: return [ LossModuleConfig( name="FOMOLocalizationLoss", - alias=f"FOMOLocalizationLoss-{self.kpt_task_name}", - attached_to=f"FOMOHead-{self.kpt_task_name}", + attached_to=f"{self.task_name}/FOMOHead", params=self.loss_params, weight=1.0, ) @@ -105,8 +96,7 @@ def metrics(self) -> list[MetricModuleConfig]: return [ MetricModuleConfig( name="ObjectKeypointSimilarity", - alias=f"ObjectKeypointSimilarity-{self.kpt_task_name}", - attached_to=f"FOMOHead-{self.kpt_task_name}", + attached_to=f"{self.task_name}/FOMOHead", is_main_metric=True, ), ] @@ -116,8 +106,7 @@ def visualizers(self) -> list[AttachedModuleConfig]: return [ AttachedModuleConfig( name="MultiVisualizer", - alias=f"MultiVisualizer-{self.kpt_task_name}", - attached_to=f"FOMOHead-{self.kpt_task_name}", + attached_to=f"{self.task_name}/FOMOHead", params={ "visualizers": [ { diff --git a/luxonis_train/config/predefined_models/detection_model.py b/luxonis_train/config/predefined_models/detection_model.py index d1498845..cda9c503 100644 --- a/luxonis_train/config/predefined_models/detection_model.py +++ b/luxonis_train/config/predefined_models/detection_model.py @@ -65,7 +65,7 @@ def __init__( head_params: Params | None = None, loss_params: Params | None = None, visualizer_params: Params | None = None, - task_name: str | None = None, + task_name: str = "", ): var_config = get_variant(variant) @@ -89,7 +89,7 @@ def nodes(self) -> list[ModelNodeConfig]: nodes = [ ModelNodeConfig( name=self.backbone, - alias=f"{self.backbone}-{self.task_name}", + alias=f"{self.task_name}/{self.backbone}", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, ), @@ -98,8 +98,8 @@ def nodes(self) -> list[ModelNodeConfig]: nodes.append( ModelNodeConfig( name="RepPANNeck", - alias=f"RepPANNeck-{self.task_name}", - inputs=[f"{self.backbone}-{self.task_name}"], + alias=f"{self.task_name}/RepPANNeck", + inputs=[f"{self.task_name}/{self.backbone}"], freezing=self.neck_params.pop("freezing", {}), params=self.neck_params, ) @@ -108,11 +108,11 @@ def nodes(self) -> list[ModelNodeConfig]: nodes.append( ModelNodeConfig( name="EfficientBBoxHead", - alias=f"EfficientBBoxHead-{self.task_name}", + alias=f"{self.task_name}/EfficientBBoxHead", freezing=self.head_params.pop("freezing", {}), - inputs=[f"RepPANNeck-{self.task_name}"] + inputs=[f"{self.task_name}/RepPANNeck"] if self.use_neck - else [f"{self.backbone}-{self.task_name}"], + else [f"{self.task_name}/{self.backbone}"], params=self.head_params, task_name=self.task_name, ) @@ -125,8 +125,7 @@ def losses(self) -> list[LossModuleConfig]: return [ LossModuleConfig( name="AdaptiveDetectionLoss", - alias=f"AdaptiveDetectionLoss-{self.task_name}", - attached_to=f"EfficientBBoxHead-{self.task_name}", + attached_to=f"{self.task_name}/EfficientBBoxHead", params=self.loss_params, weight=1.0, ) @@ -138,8 +137,7 @@ def metrics(self) -> list[MetricModuleConfig]: return [ MetricModuleConfig( name="MeanAveragePrecision", - alias=f"MeanAveragePrecision-{self.task_name}", - attached_to=f"EfficientBBoxHead-{self.task_name}", + attached_to=f"{self.task_name}/EfficientBBoxHead", is_main_metric=True, ), ] @@ -150,8 +148,7 @@ def visualizers(self) -> list[AttachedModuleConfig]: return [ AttachedModuleConfig( name="BBoxVisualizer", - alias=f"BBoxVisualizer-{self.task_name}", - attached_to=f"EfficientBBoxHead-{self.task_name}", + attached_to=f"{self.task_name}/EfficientBBoxHead", params=self.visualizer_params, ) ] diff --git a/luxonis_train/config/predefined_models/keypoint_detection_model.py b/luxonis_train/config/predefined_models/keypoint_detection_model.py index 8882f338..88c7aa63 100644 --- a/luxonis_train/config/predefined_models/keypoint_detection_model.py +++ b/luxonis_train/config/predefined_models/keypoint_detection_model.py @@ -62,8 +62,7 @@ def __init__( loss_params: Params | None = None, kpt_visualizer_params: Params | None = None, bbox_visualizer_params: Params | None = None, - bbox_task_name: str | None = None, - kpt_task_name: str | None = None, + task_name: str = "", ): var_config = get_variant(variant) @@ -79,8 +78,7 @@ def __init__( self.loss_params = loss_params or {"n_warmup_epochs": 0} self.kpt_visualizer_params = kpt_visualizer_params or {} self.bbox_visualizer_params = bbox_visualizer_params or {} - self.bbox_task_name = bbox_task_name or "boundingbox" - self.kpt_task_name = kpt_task_name or "keypoints" + self.task_name = task_name @property def nodes(self) -> list[ModelNodeConfig]: @@ -89,7 +87,7 @@ def nodes(self) -> list[ModelNodeConfig]: nodes = [ ModelNodeConfig( name=self.backbone, - alias=f"{self.backbone}-{self.kpt_task_name}", + alias=f"{self.task_name}/{self.backbone}", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, ), @@ -98,31 +96,25 @@ def nodes(self) -> list[ModelNodeConfig]: nodes.append( ModelNodeConfig( name="RepPANNeck", - alias=f"RepPANNeck-{self.kpt_task_name}", - inputs=[f"{self.backbone}-{self.kpt_task_name}"], + alias=f"{self.task_name}/RepPANNeck", + inputs=[f"{self.task_name}/{self.backbone}"], freezing=self.neck_params.pop("freezing", {}), params=self.neck_params, ) ) - task = {} - if self.bbox_task_name is not None: - task["boundingbox"] = self.bbox_task_name - if self.kpt_task_name is not None: - task["keypoints"] = self.kpt_task_name - nodes.append( ModelNodeConfig( name="EfficientKeypointBBoxHead", - alias=f"EfficientKeypointBBoxHead-{self.kpt_task_name}", + alias=f"{self.task_name}/EfficientKeypointBBoxHead", inputs=( - [f"RepPANNeck-{self.kpt_task_name}"] + [f"{self.task_name}/RepPANNeck"] if self.use_neck - else [f"{self.backbone}-{self.kpt_task_name}"] + else [f"{self.task_name}/{self.backbone}"] ), freezing=self.head_params.pop("freezing", {}), params=self.head_params, - task_name=task, + task_name=self.task_name, ) ) return nodes @@ -133,8 +125,7 @@ def losses(self) -> list[LossModuleConfig]: return [ LossModuleConfig( name="EfficientKeypointBBoxLoss", - alias=f"EfficientKeypointBBoxLoss-{self.kpt_task_name}", - attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}", + attached_to=f"{self.task_name}/EfficientKeypointBBoxHead", params=self.loss_params, weight=1.0, ) @@ -146,14 +137,12 @@ def metrics(self) -> list[MetricModuleConfig]: return [ MetricModuleConfig( name="ObjectKeypointSimilarity", - alias=f"ObjectKeypointSimilarity-{self.kpt_task_name}", - attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}", + attached_to=f"{self.task_name}/EfficientKeypointBBoxHead", is_main_metric=True, ), MetricModuleConfig( name="MeanAveragePrecisionKeypoints", - alias=f"MeanAveragePrecisionKeypoints-{self.kpt_task_name}", - attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}", + attached_to=f"{self.task_name}/EfficientKeypointBBoxHead", ), ] @@ -164,8 +153,7 @@ def visualizers(self) -> list[AttachedModuleConfig]: return [ AttachedModuleConfig( name="MultiVisualizer", - alias=f"MultiVisualizer-{self.kpt_task_name}", - attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}", + attached_to=f"{self.task_name}/EfficientKeypointBBoxHead", params={ "visualizers": [ { diff --git a/luxonis_train/config/predefined_models/segmentation_model.py b/luxonis_train/config/predefined_models/segmentation_model.py index 0260e843..fb04d0f1 100644 --- a/luxonis_train/config/predefined_models/segmentation_model.py +++ b/luxonis_train/config/predefined_models/segmentation_model.py @@ -56,7 +56,7 @@ def __init__( loss_params: Params | None = None, visualizer_params: Params | None = None, task: Literal["binary", "multiclass"] = "binary", - task_name: str | None = None, + task_name: str = "", ): var_config = get_variant(variant) @@ -82,15 +82,15 @@ def nodes(self) -> list[ModelNodeConfig]: node_list = [ ModelNodeConfig( name=self.backbone, - alias=f"{self.backbone}-{self.task_name}", + alias=f"{self.task_name}/{self.backbone}", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, task_name=self.task_name, ), ModelNodeConfig( name="DDRNetSegmentationHead", - alias=f"DDRNetSegmentationHead-{self.task_name}", - inputs=[f"{self.backbone}-{self.task_name}"], + alias=f"{self.task_name}/DDRNetSegmentationHead", + inputs=[f"{self.task_name}/{self.backbone}"], freezing=self.head_params.pop("freezing", {}), params=self.head_params, task_name=self.task_name, @@ -100,8 +100,8 @@ def nodes(self) -> list[ModelNodeConfig]: node_list.append( ModelNodeConfig( name="DDRNetSegmentationHead", - alias=f"DDRNetSegmentationHead_aux-{self.task_name}", - inputs=[f"{self.backbone}-{self.task_name}"], + alias=f"{self.task_name}/DDRNetSegmentationHead_aux", + inputs=[f"{self.task_name}/{self.backbone}"], freezing=self.aux_head_params.pop("freezing", {}), params=self.aux_head_params, task_name=self.task_name, @@ -122,12 +122,7 @@ def losses(self) -> list[LossModuleConfig]: if self.task == "binary" else "OHEMCrossEntropyLoss" ), - alias=( - f"OHEMBCEWithLogitsLoss-{self.task_name}" - if self.task == "binary" - else f"OHEMCrossEntropyLoss-{self.task_name}" - ), - attached_to=f"DDRNetSegmentationHead-{self.task_name}", + attached_to=f"{self.task_name}/DDRNetSegmentationHead", params=self.loss_params, weight=1.0, ), @@ -140,12 +135,7 @@ def losses(self) -> list[LossModuleConfig]: if self.task == "binary" else "OHEMCrossEntropyLoss" ), - alias=( - f"OHEMBCEWithLogitsLoss_aux-{self.task_name}" - if self.task == "binary" - else f"OHEMCrossEntropyLoss_aux-{self.task_name}" - ), - attached_to=f"DDRNetSegmentationHead_aux-{self.task_name}", + attached_to=f"{self.task_name}/DDRNetSegmentationHead_aux", params=self.loss_params, weight=0.4, ) @@ -158,15 +148,13 @@ def metrics(self) -> list[MetricModuleConfig]: return [ MetricModuleConfig( name="JaccardIndex", - alias=f"JaccardIndex-{self.task_name}", - attached_to=f"DDRNetSegmentationHead-{self.task_name}", + attached_to=f"{self.task_name}/DDRNetSegmentationHead", is_main_metric=True, params={"task": self.task}, ), MetricModuleConfig( name="F1Score", - alias=f"F1Score-{self.task_name}", - attached_to=f"DDRNetSegmentationHead-{self.task_name}", + attached_to=f"{self.task_name}/DDRNetSegmentationHead", params={"task": self.task}, ), ] @@ -177,8 +165,7 @@ def visualizers(self) -> list[AttachedModuleConfig]: return [ AttachedModuleConfig( name="SegmentationVisualizer", - alias=f"SegmentationVisualizer-{self.task_name}", - attached_to=f"DDRNetSegmentationHead-{self.task_name}", + attached_to=f"{self.task_name}/DDRNetSegmentationHead", params=self.visualizer_params, ) ] From 82abeae938fc771ac534f8ad9349b5e9d6b817af Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 07:43:51 -0600 Subject: [PATCH 16/57] updated attached modules --- .../attached_modules/base_attached_module.py | 34 ++++++++++++++----- .../losses/efficient_keypoint_bbox_loss.py | 4 ++- .../losses/fomo_localization_loss.py | 3 ++ .../mean_average_precision_keypoints.py | 22 ++++++------ .../metrics/object_keypoint_similarity.py | 2 ++ .../attached_modules/visualizers/utils.py | 2 +- 6 files changed, 46 insertions(+), 21 deletions(-) diff --git a/luxonis_train/attached_modules/base_attached_module.py b/luxonis_train/attached_modules/base_attached_module.py index 99461c96..a5f14761 100644 --- a/luxonis_train/attached_modules/base_attached_module.py +++ b/luxonis_train/attached_modules/base_attached_module.py @@ -73,13 +73,13 @@ def __init__(self, *, node: BaseNode | None = None): for label in self.supported_tasks ] module_supported = f"[{', '.join(module_supported)}]" - if not self.node.task_types: + if not self.node.tasks: raise IncompatibleException( f"Module '{self.name}' requires one of the following " f"labels or combinations of labels: {module_supported}, " f"but is connected to node '{self.node.name}' which does not specify any tasks." ) - node_tasks = set(self.node.task_types) + node_tasks = set(self.node.tasks) for required_labels in self.supported_tasks: if isinstance(required_labels, TaskType): required_labels = [required_labels] @@ -89,7 +89,7 @@ def __init__(self, *, node: BaseNode | None = None): self.required_labels = required_labels break else: - node_supported = [task.value for task in self.node.task_types] + node_supported = [task.value for task in self.node.tasks] raise IncompatibleException( f"Module '{self.name}' requires one of the following labels or combinations of labels: {module_supported}, " f"but is connected to node '{self.node.name}' which does not support any of them. " @@ -166,11 +166,11 @@ def node_tasks(self) -> list[TaskType]: @raises RuntimeError: If the node does not have the C{tasks} attribute set. """ - if self.node.task_types is None: + if self.node.tasks is None: raise RuntimeError( "Node must have the `tasks` attribute specified." ) - return self.node.task_types + return self.node.tasks def get_label( self, labels: Labels, task_type: TaskType | None = None @@ -214,8 +214,12 @@ def _get_label( task_name = self.node.task_name task = f"{task_name}/{task_type.value}" if task not in labels: - raise IncompatibleException.from_missing_task( - task_type.value, list(labels.keys()), self.name + raise IncompatibleException( + f"Module '{self.name}' requires label of type " + f"'{task_type.value}' assigned to task '{task_name}', " + "but the label is missing from the dataset. " + f"Available labels: {list(labels.keys())}. " + f"Missing label: '{task}'." ) return labels[task], task_type @@ -274,7 +278,19 @@ def get_input_tensors( f"{self.name} requires multiple labels, " "you must provide the `task_type` argument to extract the desired input." ) - task_type = self.node_tasks[0].value + if len(self.node_tasks) == 1: + task_type = self.node_tasks[0].value + else: + required_label = self.required_labels[0] + for task in self.node_tasks: + if task.value == required_label: + task_type = task.value + break + else: + raise IncompatibleException( + f"Task {required_label} is not supported by the node " + f"{self.node.name}." + ) return inputs[f"{self.node.task_name}/{task_type}"] def prepare( @@ -307,7 +323,7 @@ def prepare( @raises RuntimeError: If the C{tasks} attribute is not set on the node. @raises RuntimeError: If the C{supported_tasks} attribute is not set on the module. """ - if self.node.task_types is None: + if self.node.tasks is None: raise RuntimeError( f"{self.node.name} must have the `tasks` attribute specified " f"for {self.name} to make use of the default `prepare` method." diff --git a/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py b/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py index 701a3c72..d9a191e9 100644 --- a/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py +++ b/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py @@ -16,6 +16,7 @@ get_with_default, ) from luxonis_train.utils.boundingbox import IoUType +from luxonis_train.utils.keypoints import insert_class from .bce_with_logits import BCEWithLogitsLoss @@ -100,8 +101,9 @@ def prepare( pred_distri = self.get_input_tensors(inputs, "distributions")[0] pred_kpts = self.get_input_tensors(inputs, "keypoints_raw")[0] - target_kpts = self.get_label(labels, TaskType.KEYPOINTS) target_bbox = self.get_label(labels, TaskType.BOUNDINGBOX) + target_kpts = self.get_label(labels, TaskType.KEYPOINTS) + target_kpts = insert_class(target_kpts, target_bbox) batch_size = pred_scores.shape[0] n_kpts = (target_kpts.shape[1] - 2) // 3 diff --git a/luxonis_train/attached_modules/losses/fomo_localization_loss.py b/luxonis_train/attached_modules/losses/fomo_localization_loss.py index 0ad1ea60..1b181077 100644 --- a/luxonis_train/attached_modules/losses/fomo_localization_loss.py +++ b/luxonis_train/attached_modules/losses/fomo_localization_loss.py @@ -8,6 +8,7 @@ from luxonis_train.enums import TaskType from luxonis_train.nodes import FOMOHead from luxonis_train.utils import Labels, Packet +from luxonis_train.utils.keypoints import insert_class from .base_loss import BaseLoss @@ -37,6 +38,8 @@ def prepare( ) -> tuple[Tensor, Tensor]: heatmap = self.get_input_tensors(inputs, "features")[0] target_kpts = self.get_label(labels, TaskType.KEYPOINTS) + target_bbox = self.get_label(labels, TaskType.BOUNDINGBOX) + target_kpts = insert_class(target_kpts, target_bbox) batch_size, num_classes, height, width = heatmap.shape target_heatmap = torch.zeros( (batch_size, num_classes, height, width), device=heatmap.device diff --git a/luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py b/luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py index 6a6440df..6be5f2ec 100644 --- a/luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py +++ b/luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py @@ -15,6 +15,7 @@ get_sigmas, get_with_default, ) +from luxonis_train.utils.keypoints import insert_class from .base_metric import BaseMetric @@ -107,16 +108,17 @@ def prepare( self, inputs: Packet[Tensor], labels: Labels ) -> tuple[list[dict[str, Tensor]], list[dict[str, Tensor]]]: assert self.node.tasks is not None - kpts = self.get_label(labels, TaskType.KEYPOINTS) - boxes = self.get_label(labels, TaskType.BOUNDINGBOX) - - nkpts = (kpts.shape[1] - 2) // 3 - label = torch.zeros((len(boxes), nkpts * 3 + 6)) - label[:, :2] = boxes[:, :2] - label[:, 2:6] = box_convert(boxes[:, 2:], "xywh", "xyxy") - label[:, 6::3] = kpts[:, 2::3] # x - label[:, 7::3] = kpts[:, 3::3] # y - label[:, 8::3] = kpts[:, 4::3] # visiblity + kpts_labels = self.get_label(labels, TaskType.KEYPOINTS) + bbox_labels = self.get_label(labels, TaskType.BOUNDINGBOX) + kpts_labels = insert_class(kpts_labels, bbox_labels) + + n_kpts = (kpts_labels.shape[1] - 2) // 3 + label = torch.zeros((len(bbox_labels), n_kpts * 3 + 6)) + label[:, :2] = bbox_labels[:, :2] + label[:, 2:6] = box_convert(bbox_labels[:, 2:], "xywh", "xyxy") + label[:, 6::3] = kpts_labels[:, 2::3] # x + label[:, 7::3] = kpts_labels[:, 3::3] # y + label[:, 8::3] = kpts_labels[:, 4::3] # visiblity output_list_kpt_map: list[dict[str, Tensor]] = [] label_list_kpt_map: list[dict[str, Tensor]] = [] diff --git a/luxonis_train/attached_modules/metrics/object_keypoint_similarity.py b/luxonis_train/attached_modules/metrics/object_keypoint_similarity.py index f32051b3..ec7b930d 100644 --- a/luxonis_train/attached_modules/metrics/object_keypoint_similarity.py +++ b/luxonis_train/attached_modules/metrics/object_keypoint_similarity.py @@ -13,6 +13,7 @@ get_sigmas, get_with_default, ) +from luxonis_train.utils.keypoints import insert_class from .base_metric import BaseMetric @@ -77,6 +78,7 @@ def prepare( ) -> tuple[list[dict[str, Tensor]], list[dict[str, Tensor]]]: kpts_labels = self.get_label(labels, TaskType.KEYPOINTS) bbox_labels = self.get_label(labels, TaskType.BOUNDINGBOX) + kpts_labels = insert_class(kpts_labels, bbox_labels) n_keypoints = (kpts_labels.shape[1] - 2) // 3 label = torch.zeros((len(bbox_labels), n_keypoints * 3 + 6)) label[:, :2] = bbox_labels[:, :2] diff --git a/luxonis_train/attached_modules/visualizers/utils.py b/luxonis_train/attached_modules/visualizers/utils.py index 45ec454b..1a571eca 100644 --- a/luxonis_train/attached_modules/visualizers/utils.py +++ b/luxonis_train/attached_modules/visualizers/utils.py @@ -160,7 +160,7 @@ def draw_keypoint_labels(img: Tensor, label: Tensor, **kwargs) -> Tensor: @return: Image with keypoint labels drawn on. """ _, H, W = img.shape - keypoints_unflat = label[:, 1:].reshape(-1, 3) + keypoints_unflat = label.reshape(-1, 3) keypoints_points = keypoints_unflat[:, :2] keypoints_points[:, 0] *= W keypoints_points[:, 1] *= H From f48622bf5a92288181b8d60d4111d52b234a4446 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 07:44:18 -0600 Subject: [PATCH 17/57] small changes --- luxonis_train/core/core.py | 2 +- luxonis_train/enums.py | 1 - luxonis_train/models/luxonis_lightning.py | 21 ++-- luxonis_train/nodes/base_node.py | 113 +++--------------- .../nodes/heads/ddrnet_segmentation_head.py | 2 +- .../nodes/heads/efficient_bbox_head.py | 4 +- luxonis_train/utils/__init__.py | 3 +- luxonis_train/utils/exceptions.py | 9 +- luxonis_train/utils/keypoints.py | 21 ++++ 9 files changed, 59 insertions(+), 117 deletions(-) diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py index fe959349..2fb59457 100644 --- a/luxonis_train/core/core.py +++ b/luxonis_train/core/core.py @@ -721,7 +721,7 @@ def _mult(lst: list[float | int]) -> list[float]: self.cfg.trainer.preprocessing.normalize.params["std"] ), "dai_type": "RGB888p" - if self.cfg.trainer.preprocessing.out_image_format + if self.cfg.trainer.preprocessing.color_format == "RGB" else "BGR888p", } diff --git a/luxonis_train/enums.py b/luxonis_train/enums.py index ea719e1c..09d38fb2 100644 --- a/luxonis_train/enums.py +++ b/luxonis_train/enums.py @@ -9,5 +9,4 @@ class TaskType(str, Enum): INSTANCE_SEGMENTATION = "instance_segmentation" BOUNDINGBOX = "boundingbox" KEYPOINTS = "keypoints" - LABEL = "label" ARRAY = "array" diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index cebce43f..afd422ae 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -180,12 +180,16 @@ def __init__( node_cfg.freezing.unfreeze_after * epochs ) frozen_nodes.append((node_name, unfreeze_after)) - - if node_cfg.task_name is not None and Node.task_types is None: - raise ValueError( - f"Cannot define tasks for node {node_name}." - "This node doesn't specify any tasks." - ) + task_names = list(self.dataset_metadata.task_names) + if not node_cfg.task_name: + if len(task_names) == 1: + node_cfg.task_name = task_names[0] + elif issubclass(Node, BaseHead): + raise ValueError( + f"Dataset contains multiple tasks: {task_names}. " + f"Node {node_name} does not have the `task_name` parameter set. " + "Please specify the `task_name` parameter for each head node. " + ) nodes[node_name] = ( Node, @@ -1030,7 +1034,10 @@ def _print_results( ) if self.main_metric is not None: - main_metric_node, main_metric_name = self.main_metric.split("/") + print(self.main_metric) + *main_metric_node, main_metric_name = self.main_metric.split("/") + main_metric_node = "/".join(main_metric_node) + main_metric = metrics[main_metric_node][main_metric_name] logger.info( f"{stage} main metric ({self.main_metric}): {main_metric:.4f}" diff --git a/luxonis_train/nodes/base_node.py b/luxonis_train/nodes/base_node.py index eff2c2d1..19ec33a7 100644 --- a/luxonis_train/nodes/base_node.py +++ b/luxonis_train/nodes/base_node.py @@ -109,7 +109,7 @@ def wrap(output: Tensor) -> Packet[Tensor]: """ attach_index: AttachIndexType - task_types: list[TaskType] | None = None + tasks: list[TaskType] | None = None def __init__( self, @@ -154,10 +154,6 @@ def __init__( of outputs or C{"all"} to specify all outputs. Defaults to "all". Python indexing conventions apply. If provided as a constructor argument, overrides the class attribute. - @type _tasks: dict[TaskType, str] | None - @param _tasks: Dictionary of tasks that the node supports. - Overrides the class L{tasks} attribute. Shouldn't be - provided by the user in most cases. """ super().__init__() @@ -169,15 +165,15 @@ def __init__( ) self.attach_index = attach_index - self.task_name = task_name if task_name is None and dataset_metadata is not None: if len(dataset_metadata.task_names) == 1: - self.task_name = next(iter(dataset_metadata.task_names)) + task_name = next(iter(dataset_metadata.task_names)) else: raise ValueError( f"Dataset contain multiple tasks, but the `task_name` " f"argument for node '{self.name}' was not provided." ) + self.task_name = task_name or "" if getattr(self, "attach_index", None) is None: parameters = inspect.signature(self.forward).parameters @@ -228,117 +224,38 @@ def _check_type_overrides(self) -> None: def name(self) -> str: return self.__class__.__name__ - @property - def task_type(self) -> str: - """Getter for the task type. - - @type: str - @raises RuntimeError: If the node doesn't define any task. - @raises ValueError: If the node defines more than one task. In - that case, use the L{get_task_name} method instead. - """ - if not self.task_types: - raise RuntimeError(f"{self.name} does not define any task.") - - if len(self.task_types) > 1: - raise ValueError( - f"Node {self.name} has multiple tasks defined. " - "Use the `get_task_name` method instead." - ) - return self.task_types[0].value - @property def n_keypoints(self) -> int: """Getter for the number of keypoints. @type: int @raises ValueError: If the node does not support keypoints. - @raises RuntimeError: If the node doesn't define any task. """ if self._n_keypoints is not None: return self._n_keypoints - if self.task_types: - if TaskType.KEYPOINTS not in self.task_types: - raise ValueError(f"{self.name} does not support keypoints.") - return self.dataset_metadata.n_keypoints(self.task_name) - - raise RuntimeError( - f"{self.name} does not have any tasks defined, " - "`BaseNode.n_keypoints` property cannot be used. " - "Either override the `tasks` class attribute, " - "pass the `n_keypoints` attribute to the constructor or call " - "the `BaseNode.dataset_metadata.get_n_keypoints` method manually." - ) + if TaskType.KEYPOINTS not in (self.tasks or []): + raise ValueError(f"{self.name} does not support keypoints.") + return self.dataset_metadata.n_keypoints(self.task_name) @property def n_classes(self) -> int: """Getter for the number of classes. @type: int - @raises RuntimeError: If the node doesn't define any task. - @raises ValueError: If the number of classes is different for - different tasks. In that case, use the L{get_n_classes} - method. """ if self._n_classes is not None: return self._n_classes - if not self.task_types: - raise RuntimeError( - f"{self.name} does not have any tasks defined, " - "`BaseNode.n_classes` property cannot be used. " - "Either override the `tasks` class attribute, " - "pass the `n_classes` attribute to the constructor or call " - "the `BaseNode.dataset_metadata.n_classes` method manually." - ) - elif len(self.task_types) == 1: - return self.dataset_metadata.n_classes(self.task_name) - else: - n_classes = [ - self.dataset_metadata.n_classes(self.task_name) - for task in self.task_types - ] - if len(set(n_classes)) == 1: - return n_classes[0] - raise ValueError( - "Node defines multiple tasks but they have different number of classes. " - "This is likely an error, as the number of classes should be the same." - "If it is intended, use `BaseNode.get_n_classes` instead." - ) + return self.dataset_metadata.n_classes(self.task_name) @property def class_names(self) -> list[str]: """Getter for the class names. @type: list[str] - @raises RuntimeError: If the node doesn't define any task. - @raises ValueError: If the class names are different for - different tasks. In that case, use the L{get_class_names} - method. """ - if not self.task_types: - raise RuntimeError( - f"{self.name} does not have any tasks defined, " - "`BaseNode.class_names` property cannot be used. " - "Either override the `tasks` class attribute, " - "pass the `n_classes` attribute to the constructor or call " - "the `BaseNode.dataset_metadata.class_names` method manually." - ) - elif len(self.task_types) == 1: - return self.dataset_metadata.classes(self.task_name) - else: - class_names = [ - self.dataset_metadata.classes(self.task_name) - for task in self.task_types - ] - if all(set(names) == set(class_names[0]) for names in class_names): - return class_names[0] - raise ValueError( - "Node defines multiple tasks but they have different class names. " - "This is likely an error, as the class names should be the same. " - "If it is intended, use `BaseNode.get_class_names` instead." - ) + return self.dataset_metadata.classes(self.task_name) @property def input_shapes(self) -> list[Packet[Size]]: @@ -587,10 +504,14 @@ def wrap(self, output: ForwardOutputT) -> Packet[Tensor]: raise ValueError( "Default `wrap` expects a single tensor or a list of tensors." ) - try: - task = f"{self.task_name}/{self.task_type}" - except RuntimeError: - task = "features" + if not self.tasks: + return {"features": outputs} + if len(self.tasks) > 1: + raise RuntimeError( + f"Node {self.name} defines multiple tasks. " + "The `wrap` method should be overridden." + ) + task = f"{self.task_name or ''}/{self.tasks[0].value}" return {task: outputs} def run(self, inputs: list[Packet[Tensor]]) -> Packet[Tensor]: @@ -609,7 +530,7 @@ def run(self, inputs: list[Packet[Tensor]]) -> Packet[Tensor]: unwrapped = self.unwrap(inputs) outputs = self(unwrapped) wrapped = self.wrap(outputs) - str_tasks = [task.value for task in self.task_types or []] + str_tasks = [task.value for task in self.tasks or []] for key in list(wrapped.keys()): if key in str_tasks: assert self.task_name is not None diff --git a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py index 35c5e69c..2b313ab6 100644 --- a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py +++ b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py @@ -18,7 +18,7 @@ class DDRNetSegmentationHead(BaseHead[Tensor, Tensor]): in_width: int in_channels: int - task_types: list[TaskType] = [TaskType.SEGMENTATION] + tasks: list[TaskType] = [TaskType.SEGMENTATION] parser: str = "SegmentationParser" def __init__( diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py index 2eb57dfb..8d1a55f4 100644 --- a/luxonis_train/nodes/heads/efficient_bbox_head.py +++ b/luxonis_train/nodes/heads/efficient_bbox_head.py @@ -21,7 +21,7 @@ class EfficientBBoxHead( BaseHead[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]], ): in_channels: list[int] - task_types: list[TaskType] = [TaskType.BOUNDINGBOX] + tasks: list[TaskType] = [TaskType.BOUNDINGBOX] parser = "YOLO" def __init__( @@ -171,7 +171,7 @@ def wrap( conf, _ = out_cls.max(1, keepdim=True) out = torch.cat([out_reg, conf, out_cls], dim=1) outputs.append(out) - return {self.task_type: outputs} + return {"boundingbox": outputs} cls_tensor = torch.cat( [cls_score_list[i].flatten(2) for i in range(len(cls_score_list))], diff --git a/luxonis_train/utils/__init__.py b/luxonis_train/utils/__init__.py index 2944dfde..d7c0be9f 100644 --- a/luxonis_train/utils/__init__.py +++ b/luxonis_train/utils/__init__.py @@ -16,7 +16,7 @@ to_shape_packet, ) from .graph import traverse_graph -from .keypoints import get_sigmas +from .keypoints import get_sigmas, insert_class from .tracker import LuxonisTrackerPL from .types import AttachIndexType, Kwargs, Labels, Packet @@ -41,4 +41,5 @@ "compute_iou_loss", "get_sigmas", "traverse_graph", + "insert_class", ] diff --git a/luxonis_train/utils/exceptions.py b/luxonis_train/utils/exceptions.py index bab8c1aa..811b1a29 100644 --- a/luxonis_train/utils/exceptions.py +++ b/luxonis_train/utils/exceptions.py @@ -2,11 +2,4 @@ class IncompatibleException(Exception): """Raised when two parts of the model are incompatible with each other.""" - @classmethod - def from_missing_task( - cls, task: str, present_tasks: list[str], class_name: str - ): - return cls( - f"{class_name} requires '{task}' label, but it was not found in " - f"the label dictionary. Available labels: {present_tasks}." - ) + pass diff --git a/luxonis_train/utils/keypoints.py b/luxonis_train/utils/keypoints.py index 8073c399..7eaac550 100644 --- a/luxonis_train/utils/keypoints.py +++ b/luxonis_train/utils/keypoints.py @@ -65,3 +65,24 @@ def get_sigmas( msg = f"[{caller_name}] {msg}" logger.info(msg) return torch.tensor([0.04] * n_keypoints, dtype=torch.float32) + + +def insert_class(keypoints: Tensor, bboxes: Tensor) -> Tensor: + """Insert class index into keypoints tensor. + + @type keypoints: Tensor + @param keypoints: Tensor of keypoints. + @type bboxes: Tensor + @param bboxes: Tensor of bounding boxes with class index. + @rtype: Tensor + @return: Tensor of keypoints with class index. + """ + classes = bboxes[:, 1] + return torch.cat( + ( + keypoints[:, :1], + classes.unsqueeze(-1), + keypoints[:, 1:], + ), + dim=-1, + ) From 7c244af626ac3160bc6e65b43e50da0eff84a254 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 07:44:25 -0600 Subject: [PATCH 18/57] updated tests --- tests/configs/cli_commands.yaml | 2 +- tests/configs/multi_input.yaml | 78 +++----- tests/configs/parking_lot_config.yaml | 127 +++++-------- tests/integration/conftest.py | 164 +---------------- tests/integration/multi_input_modules.py | 20 +- tests/integration/parking_lot.json | 171 +++++++----------- tests/integration/test_detection.py | 7 +- .../test_fixed_validation_batch_limit.py | 6 +- tests/integration/test_fomo_detection.py | 22 +-- tests/integration/test_segmentation.py | 12 +- .../test_unsupervised_anomaly_detection.py | 20 +- tests/unittests/test_base_attached_module.py | 74 +++++--- tests/unittests/test_base_node.py | 61 +------ .../test_loaders/test_base_loader.py | 45 ++--- 14 files changed, 251 insertions(+), 558 deletions(-) diff --git a/tests/configs/cli_commands.yaml b/tests/configs/cli_commands.yaml index 56f77ef9..7123534c 100644 --- a/tests/configs/cli_commands.yaml +++ b/tests/configs/cli_commands.yaml @@ -51,7 +51,7 @@ trainer: preprocessing: train_image_size: [256, 320] keep_aspect_ratio: true - train_rgb: true + color_format: RGB normalize: active: true diff --git a/tests/configs/multi_input.yaml b/tests/configs/multi_input.yaml index 7db03d90..dcf92cdd 100644 --- a/tests/configs/multi_input.yaml +++ b/tests/configs/multi_input.yaml @@ -12,83 +12,57 @@ model: name: example_multi_input nodes: - name: FullBackbone - alias: full_backbone - name: RGBDBackbone - alias: rgbd_backbone input_sources: - left - right - disparity - name: PointcloudBackbone - alias: pointcloud_backbone input_sources: - pointcloud - name: FusionNeck - alias: fusion_neck inputs: - - rgbd_backbone - - pointcloud_backbone + - RGBDBackbone + - PointcloudBackbone input_sources: - disparity - name: FusionNeck2 - alias: fusion_neck_2 inputs: - - rgbd_backbone - - pointcloud_backbone - - full_backbone + - RGBDBackbone + - PointcloudBackbone + - FullBackbone - name: CustomSegHead1 - alias: head_1 inputs: - - fusion_neck + - FusionNeck + losses: + - name: BCEWithLogitsLoss + metrics: + - name: JaccardIndex + is_main_metric: true + params: + task: binary + visualizers: + - name: SegmentationVisualizer - name: CustomSegHead2 - alias: head_2 inputs: - - fusion_neck - - fusion_neck_2 + - FusionNeck + - FusionNeck2 input_sources: - disparity - - losses: - - name: BCEWithLogitsLoss - alias: loss_1 - attached_to: head_1 - - - name: CrossEntropyLoss - alias: loss_2 - attached_to: head_2 - - metrics: - - name: JaccardIndex - alias: jaccard_index_1 - attached_to: head_1 - is_main_metric: True - params: - task: binary - - - name: JaccardIndex - alias: jaccard_index_2 - attached_to: head_2 - params: - task: binary - - visualizers: - - name: SegmentationVisualizer - alias: seg_vis_1 - attached_to: head_1 - params: - colors: "#FF5055" - - - name: SegmentationVisualizer - alias: seg_vis_2 - attached_to: head_2 - params: - colors: "#55AAFF" + losses: + - name: CrossEntropyLoss + metrics: + - name: JaccardIndex + params: + task: binary + visualizers: + - name: SegmentationVisualizer tracker: project_name: multi_input_example @@ -111,4 +85,4 @@ trainer: exporter: onnx: - opset_version: 11 \ No newline at end of file + opset_version: 11 diff --git a/tests/configs/parking_lot_config.yaml b/tests/configs/parking_lot_config.yaml index 5cda65c1..2754197e 100644 --- a/tests/configs/parking_lot_config.yaml +++ b/tests/configs/parking_lot_config.yaml @@ -4,103 +4,58 @@ model: nodes: - name: EfficientRep - alias: backbone - name: RepPANNeck - alias: neck - inputs: - - backbone - name: EfficientBBoxHead - alias: bbox-head - inputs: - - neck + task_name: vehicle_type + losses: + - name: AdaptiveDetectionLoss + metrics: + - name: MeanAveragePrecision + is_main_metric: true + visualizers: + - name: BBoxVisualizer - name: EfficientKeypointBBoxHead - alias: motorbike-detection-head - task: - keypoints: motorbike-keypoints - boundingbox: motorbike-boundingbox - inputs: - - neck + task_name: motorbike + losses: + - name: EfficientKeypointBBoxLoss + metrics: + - name: MeanAveragePrecisionKeypoints + visualizers: + - name: MultiVisualizer + params: + visualizers: + - name: KeypointVisualizer + - name: BBoxVisualizer - name: SegmentationHead - alias: color-segmentation-head - task: color-segmentation - inputs: - - neck - - - name: SegmentationHead - alias: any-vehicle-segmentation-head - task: vehicle-segmentation - inputs: - - neck + task_name: color + losses: + - name: CrossEntropyLoss + metrics: + - name: JaccardIndex + visualizers: + - name: SegmentationVisualizer - name: BiSeNetHead - alias: brand-segmentation-head - task: brand-segmentation - inputs: - - neck + task_name: brand + losses: + - name: CrossEntropyLoss + metrics: + - name: Precision + visualizers: + - name: SegmentationVisualizer - name: BiSeNetHead - alias: vehicle-type-segmentation-head - task: vehicle_type-segmentation - inputs: - - neck - - losses: - - name: AdaptiveDetectionLoss - attached_to: bbox-head - - name: BCEWithLogitsLoss - attached_to: any-vehicle-segmentation-head - - name: CrossEntropyLoss - attached_to: vehicle-type-segmentation-head - - name: CrossEntropyLoss - attached_to: color-segmentation-head - - name: EfficientKeypointBBoxLoss - attached_to: motorbike-detection-head - - metrics: - - name: MeanAveragePrecisionKeypoints - attached_to: motorbike-detection-head - - name: MeanAveragePrecision - attached_to: bbox-head - is_main_metric: true - - name: F1Score - attached_to: any-vehicle-segmentation-head - - name: JaccardIndex - attached_to: color-segmentation-head - - name: Accuracy - attached_to: vehicle-type-segmentation-head - - name: Precision - attached_to: brand-segmentation-head - - visualizers: - - name: MultiVisualizer - alias: multi-visualizer-motorbike - attached_to: motorbike-detection-head - params: - visualizers: - - name: KeypointVisualizer - params: - nonvisible_color: blue - - name: BBoxVisualizer - - - name: SegmentationVisualizer - alias: color-segmentation-visualizer - attached_to: color-segmentation-head - - name: SegmentationVisualizer - alias: vehicle-type-segmentation-visualizer - attached_to: vehicle-type-segmentation-head - - name: SegmentationVisualizer - alias: vehicle-segmentation-visualizer - attached_to: any-vehicle-segmentation-head - - name: SegmentationVisualizer - alias: brand-segmentation-visualizer - attached_to: brand-segmentation-head - - name: BBoxVisualizer - alias: bbox-visualizer - attached_to: bbox-head + task_name: vehicle_type + losses: + - name: CrossEntropyLoss + metrics: + - name: Accuracy + visualizers: + - name: SegmentationVisualizer tracker: project_name: Parking_Lot @@ -132,7 +87,7 @@ trainer: preprocessing: train_image_size: [256, 320] keep_aspect_ratio: false - train_rgb: true + color_format: RGB normalize: active: true augmentations: diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 97189476..ab2fb1e8 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -1,22 +1,17 @@ -import json import multiprocessing as mp import os import shutil -from collections import defaultdict from pathlib import Path from typing import Any -import cv2 import gdown -import numpy as np import pytest import torchvision from luxonis_ml.data import LuxonisDataset from luxonis_ml.data.parsers import LuxonisParser -from luxonis_ml.data.utils.data_utils import rgb_to_bool_masks -from luxonis_ml.utils import LuxonisFileSystem, environ +from luxonis_ml.utils import environ -WORK_DIR = Path("tests", "data") +WORK_DIR = Path("tests", "data").absolute() @pytest.fixture(scope="session") @@ -40,152 +35,14 @@ def train_overfit() -> bool: @pytest.fixture(scope="session") def parking_lot_dataset() -> LuxonisDataset: - url = "gs://luxonis-test-bucket/luxonis-ml-test-data/D1_ParkingSlotTest" - base_path = WORK_DIR / "D1_ParkingSlotTest" - if not base_path.exists(): - base_path = LuxonisFileSystem.download(url, WORK_DIR) - - mask_brand_path = base_path / "mask_brand" - mask_color_path = base_path / "mask_color" - kpt_mask_path = base_path / "keypoints_mask_vehicle" - - def generator(): - filenames: dict[int, Path] = {} - for base_path in [kpt_mask_path, mask_brand_path, mask_color_path]: - for sequence_path in sorted(list(base_path.glob("sequence.*"))): - frame_data = sequence_path / "step0.frame_data.json" - with open(frame_data) as f: - data = json.load(f)["captures"][0] - frame_data = data["annotations"] - sequence_num = int(sequence_path.suffix[1:]) - filename = data["filename"] - if filename is not None: - filename = sequence_path / filename - filenames[sequence_num] = filename - else: - filename = filenames[sequence_num] - W, H = data["dimension"] - - annotations = { - anno["@type"].split(".")[-1]: anno for anno in frame_data - } - - bbox_classes = {} - bboxes = {} - - for bbox_annotation in annotations.get( - "BoundingBox2DAnnotation", defaultdict(list) - )["values"]: - class_ = ( - bbox_annotation["labelName"].split("-")[-1].lower() - ) - if class_ == "motorbiek": - class_ = "motorbike" - x, y = bbox_annotation["origin"] - w, h = bbox_annotation["dimension"] - instance_id = bbox_annotation["instanceId"] - bbox_classes[instance_id] = class_ - bboxes[instance_id] = [x / W, y / H, w / W, h / H] - yield { - "file": filename, - "annotation": { - "type": "boundingbox", - "class": class_, - "x": x / W, - "y": y / H, - "w": w / W, - "h": h / H, - "instance_id": instance_id, - }, - } - - for kpt_annotation in annotations.get( - "KeypointAnnotation", defaultdict(list) - )["values"]: - keypoints = kpt_annotation["keypoints"] - instance_id = kpt_annotation["instanceId"] - class_ = bbox_classes[instance_id] - bbox = bboxes[instance_id] - kpts = [] - - if class_ == "motorbike": - keypoints = keypoints[:3] - else: - keypoints = keypoints[3:] - - for kp in keypoints: - x, y = kp["location"] - kpts.append([x / W, y / H, kp["state"]]) - - yield { - "file": filename, - "annotation": { - "type": "detection", - "class": class_, - "task": class_, - "keypoints": kpts, - "instance_id": instance_id, - "boundingbox": { - "x": bbox[0], - "y": bbox[1], - "w": bbox[2], - "h": bbox[3], - }, - }, - } - - vehicle_type_segmentation = annotations[ - "SemanticSegmentationAnnotation" - ] - mask = cv2.cvtColor( - cv2.imread( - str( - sequence_path - / vehicle_type_segmentation["filename"] - ) - ), - cv2.COLOR_BGR2RGB, - ) - classes = { - inst["labelName"]: inst["pixelValue"][:3] - for inst in vehicle_type_segmentation["instances"] - } - if base_path == kpt_mask_path: - task = "vehicle_type-segmentation" - elif base_path == mask_brand_path: - task = "brand-segmentation" - else: - task = "color-segmentation" - for class_, mask_ in rgb_to_bool_masks( - mask, classes, add_background_class=True - ): - yield { - "file": filename, - "annotation": { - "type": "mask", - "class": class_, - "task": task, - "mask": mask_, - }, - } - if base_path == mask_color_path: - yield { - "file": filename, - "annotation": { - "type": "mask", - "class": "vehicle", - "task": "vehicle-segmentation", - "mask": mask.astype(bool)[..., 0] - | mask.astype(bool)[..., 1] - | mask.astype(bool)[..., 2], - }, - } - - dataset = LuxonisDataset("_ParkingLot", delete_existing=True) - dataset.add(generator()) - np.random.seed(42) - dataset.make_splits() - return dataset + url = "gs://luxonis-test-bucket/luxonis-ml-test-data/D1_ParkingLot_Native.zip" + parser = LuxonisParser( + url, + dataset_name="_D1_ParkingLot", + delete_existing=True, + save_dir=WORK_DIR, + ) + return parser.parse(random_split=True) @pytest.fixture(scope="session") @@ -236,7 +93,6 @@ def CIFAR10_subset_generator(): yield { "file": path, "annotation": { - "type": "classification", "class": classes[label], }, } diff --git a/tests/integration/multi_input_modules.py b/tests/integration/multi_input_modules.py index 31db4e2f..3db26f5d 100644 --- a/tests/integration/multi_input_modules.py +++ b/tests/integration/multi_input_modules.py @@ -1,5 +1,6 @@ import torch from torch import Tensor, nn +from typing_extensions import override from luxonis_train.enums import TaskType from luxonis_train.loaders import BaseLoaderTorch @@ -8,8 +9,10 @@ class CustomMultiInputLoader(BaseLoaderTorch): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__( + self, view: str | list[str], image_source: str | None = None, **_ + ): + super().__init__(view=view, image_source=image_source) @property def input_shapes(self): @@ -36,17 +39,16 @@ def __getitem__(self, _): # pragma: no cover # Fake labels segmap = torch.zeros(1, 224, 224, dtype=torch.float32) segmap[0, 100:150, 100:150] = 1 - labels = { - "segmentation": (segmap, TaskType.SEGMENTATION), - } + labels = {"/segmentation": segmap} return inputs, labels def __len__(self): return 10 - def get_classes(self) -> dict[TaskType, list[str]]: - return {TaskType.SEGMENTATION: ["square"]} + @override + def get_classes(self) -> dict[str, list[str]]: + return {"": ["square"]} class MultiInputTestBaseNode(BaseNode): @@ -77,7 +79,7 @@ class FusionNeck2(MultiInputTestBaseNode): ... class CustomSegHead1(MultiInputTestBaseNode): - tasks = {TaskType.SEGMENTATION: "segmentation"} + tasks = [TaskType.SEGMENTATION] def __init__(self, **kwargs): super().__init__(**kwargs) @@ -92,7 +94,7 @@ def forward(self, inputs: Tensor): class CustomSegHead2(MultiInputTestBaseNode): - tasks = {TaskType.SEGMENTATION: "segmentation"} + tasks = [TaskType.SEGMENTATION] def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/tests/integration/parking_lot.json b/tests/integration/parking_lot.json index 4a3f4f6e..66ce962d 100644 --- a/tests/integration/parking_lot.json +++ b/tests/integration/parking_lot.json @@ -37,62 +37,51 @@ ], "outputs": [ { - "name": "any-vehicle-segmentation-head/vehicle-segmentation/0", + "name": "BiSeNetHead/brand/segmentation/0", "dtype": "float32", "shape": [ 1, - 1, + 23, 256, 320 ], "layout": "NCHW" }, { - "name": "output1_yolov6r2", + "name": "EfficientKeypointBBoxHead/outputs/0", "dtype": "float32", "shape": [ 1, - 7, + 14, 32, 40 ], "layout": "NCHW" }, { - "name": "output2_yolov6r2", + "name": "EfficientKeypointBBoxHead/outputs/1", "dtype": "float32", "shape": [ 1, - 7, + 14, 16, 20 ], "layout": "NCHW" }, { - "name": "output3_yolov6r2", + "name": "EfficientKeypointBBoxHead/outputs/2", "dtype": "float32", "shape": [ 1, - 7, + 14, 8, 10 ], - "layout": "NCHW" - }, - { - "name": "brand-segmentation-head/brand-segmentation/0", - "dtype": "float32", - "shape": [ - 1, - 23, - 256, - 320 - ], - "layout": "NCHW" + "layout": "NCDE" }, { - "name": "color-segmentation-head/color-segmentation/0", + "name": "SegmentationHead/color/segmentation/0", "dtype": "float32", "shape": [ 1, @@ -103,91 +92,100 @@ "layout": "NCHW" }, { - "name": "motorbike-detection-head/outputs/0", + "name": "output1_yolov6r2", "dtype": "float32", "shape": [ 1, - 14, + 7, 32, 40 ], "layout": "NCHW" }, { - "name": "motorbike-detection-head/outputs/1", + "name": "output2_yolov6r2", "dtype": "float32", "shape": [ 1, - 14, + 7, 16, 20 ], "layout": "NCHW" }, { - "name": "motorbike-detection-head/outputs/2", + "name": "output3_yolov6r2", "dtype": "float32", "shape": [ 1, - 14, + 7, 8, 10 ], - "layout": "NCDE" - }, - { - "name": "vehicle-type-segmentation-head/vehicle_type-segmentation/0", - "dtype": "float32", - "shape": [ - 1, - 3, - 256, - 320 - ], "layout": "NCHW" } ], "heads": [ { - "name": "vehicle-type-segmentation-head", + "name": "BiSeNetHead", "parser": "SegmentationParser", "metadata": { "postprocessor_path": null, "classes": [ "background", - "car", - "motorbike" + "alfa-romeo", + "buick", + "ducati", + "harley", + "ferrari", + "infiniti", + "jeep", + "land-rover", + "roll-royce", + "yamaha", + "aprilia", + "bmw", + "dodge", + "honda", + "moto", + "piaggio", + "isuzu", + "Kawasaki", + "truimph", + "pontiac", + "saab", + "chrysler" ], - "n_classes": 3, + "n_classes": 23, "is_softmax": false }, "outputs": [ - "vehicle-type-segmentation-head/vehicle_type-segmentation/0" + "BiSeNetHead/brand/segmentation/0" ] }, { - "name": "any-vehicle-segmentation-head", + "name": "BiSeNetHead_0", "parser": "SegmentationParser", "metadata": { "postprocessor_path": null, "classes": [ - "vehicle" + "background", + "car", + "motorbike" ], - "n_classes": 1, + "n_classes": 3, "is_softmax": false }, - "outputs": [ - "any-vehicle-segmentation-head/vehicle-segmentation/0" - ] + "outputs": [] }, { - "name": "bbox-head", + "name": "EfficientBBoxHead", "parser": "YOLO", "metadata": { "postprocessor_path": null, "classes": [ - "car", - "motorbike" + "motorbike", + "car" ], "n_classes": 2, "iou_threshold": 0.45, @@ -203,44 +201,29 @@ ] }, { - "name": "brand-segmentation-head", - "parser": "SegmentationParser", + "name": "EfficientKeypointBBoxHead", + "parser": "YOLOExtendedParser", "metadata": { "postprocessor_path": null, "classes": [ - "Kawasaki", - "alfa-romeo", - "aprilia", - "background", - "bmw", - "buick", - "chrysler", - "dodge", - "ducati", - "ferrari", - "harley", - "honda", - "infiniti", - "isuzu", - "jeep", - "land-rover", - "moto", - "piaggio", - "pontiac", - "roll-royce", - "saab", - "truimph", - "yamaha" + "motorbike" ], - "n_classes": 23, - "is_softmax": false + "n_classes": 1, + "iou_threshold": 0.45, + "conf_threshold": 0.25, + "max_det": 300, + "anchors": null, + "subtype": "yolov6", + "n_keypoints": 3 }, "outputs": [ - "brand-segmentation-head/brand-segmentation/0" + "EfficientKeypointBBoxHead/outputs/0", + "EfficientKeypointBBoxHead/outputs/1", + "EfficientKeypointBBoxHead/outputs/2" ] }, { - "name": "color-segmentation-head", + "name": "SegmentationHead", "parser": "SegmentationParser", "metadata": { "postprocessor_path": null, @@ -254,31 +237,9 @@ "is_softmax": false }, "outputs": [ - "color-segmentation-head/color-segmentation/0" - ] - }, - { - "name": "motorbike-detection-head", - "parser": "YOLOExtendedParser", - "metadata": { - "postprocessor_path": null, - "classes": [ - "motorbike" - ], - "n_classes": 1, - "iou_threshold": 0.45, - "conf_threshold": 0.25, - "max_det": 300, - "anchors": null, - "subtype": "yolov6", - "n_keypoints": 3 - }, - "outputs": [ - "motorbike-detection-head/outputs/0", - "motorbike-detection-head/outputs/1", - "motorbike-detection-head/outputs/2" + "SegmentationHead/color/segmentation/0" ] } ] } -} \ No newline at end of file +} diff --git a/tests/integration/test_detection.py b/tests/integration/test_detection.py index 45e83f0a..6360ec79 100644 --- a/tests/integration/test_detection.py +++ b/tests/integration/test_detection.py @@ -16,14 +16,12 @@ def get_opts_backbone(backbone: str) -> dict[str, Any]: }, { "name": "EfficientBBoxHead", + "task_name": "vehicle_type", "inputs": [backbone], }, { "name": "EfficientKeypointBBoxHead", - "task": { - "keypoints": "car-keypoints", - "boundingbox": "car-boundingbox", - }, + "task_name": "car", "inputs": [backbone], }, ], @@ -70,6 +68,7 @@ def get_opts_variant(variant: str) -> dict[str, Any]: }, { "name": "EfficientBBoxHead", + "task_name": "motorbike", "inputs": ["neck"], }, ], diff --git a/tests/integration/test_fixed_validation_batch_limit.py b/tests/integration/test_fixed_validation_batch_limit.py index 25794cec..5a7a398c 100644 --- a/tests/integration/test_fixed_validation_batch_limit.py +++ b/tests/integration/test_fixed_validation_batch_limit.py @@ -19,7 +19,11 @@ def get_config() -> dict[str, Any]: "alias": "neck", "inputs": ["backbone"], }, - {"name": "EfficientBBoxHead", "inputs": ["neck"]}, + { + "name": "EfficientBBoxHead", + "task_name": "motorbike", + "inputs": ["neck"], + }, ], "losses": [ { diff --git a/tests/integration/test_fomo_detection.py b/tests/integration/test_fomo_detection.py index cd1d89fd..49a81da8 100644 --- a/tests/integration/test_fomo_detection.py +++ b/tests/integration/test_fomo_detection.py @@ -50,26 +50,18 @@ def dummy_generator(image_paths: List[Path]): ) for i, bbox in enumerate(bboxes): - # Generate bounding box annotation yield { "file": path, "annotation": { - "type": "boundingbox", - "instance_id": i, "class": "object", - "x": bbox["x"], - "y": bbox["y"], - "w": bbox["w"], - "h": bbox["h"], - }, - } - yield { - "file": path, - "annotation": { - "type": "keypoints", "instance_id": i, - "class": "object", - "keypoints": [keypoints[i]], + "boundingbox": { + "x": bbox["x"], + "y": bbox["y"], + "w": bbox["w"], + "h": bbox["h"], + }, + "keypoints": {"keypoints": [keypoints[i]]}, }, } diff --git a/tests/integration/test_segmentation.py b/tests/integration/test_segmentation.py index a8b4df91..1c79dc29 100644 --- a/tests/integration/test_segmentation.py +++ b/tests/integration/test_segmentation.py @@ -17,37 +17,37 @@ def get_opts(backbone: str) -> dict[str, Any]: { "name": "SegmentationHead", "alias": "seg-color-segmentation", - "task": "color-segmentation", + "task_name": "color", "inputs": [backbone], }, { "name": "BiSeNetHead", "alias": "bi-color-segmentation", - "task": "color-segmentation", + "task_name": "color", "inputs": [backbone], }, { "name": "SegmentationHead", "alias": "seg-vehicle-segmentation", - "task": "vehicle-segmentation", + "task_name": "vehicles", "inputs": [backbone], }, { "name": "BiSeNetHead", "alias": "bi-vehicle-segmentation", - "task": "vehicle-segmentation", + "task_name": "vehicles", "inputs": [backbone], }, { "name": "SegmentationHead", "alias": "seg-vehicle-segmentation-2", - "task": "vehicle-segmentation", + "task_name": "vehicles", "inputs": [backbone], }, { "name": "SegmentationHead", "alias": "seg-vehicle-segmentation-3", - "task": "vehicle-segmentation", + "task_name": "vehicles", "inputs": [backbone], }, ], diff --git a/tests/integration/test_unsupervised_anomaly_detection.py b/tests/integration/test_unsupervised_anomaly_detection.py index a98faa27..10c0a6f9 100644 --- a/tests/integration/test_unsupervised_anomaly_detection.py +++ b/tests/integration/test_unsupervised_anomaly_detection.py @@ -70,11 +70,12 @@ def dummy_generator( yield { "file": path, "annotation": { - "type": "rle", "class": "object", - "height": 256, - "width": 256, - "counts": "0" * (256 * 256), + "segmentation": { + "height": 256, + "width": 256, + "counts": "0" * (256 * 256), + }, }, } @@ -94,11 +95,14 @@ def dummy_generator( yield { "file": path, "annotation": { - "type": "polyline", "class": "object", - "points": [ - pt for segment in poly_normalized for pt in segment - ], + "segmentation": { + "height": img_h, + "width": img_w, + "points": [ + pt for segment in poly_normalized for pt in segment + ], + }, }, } diff --git a/tests/unittests/test_base_attached_module.py b/tests/unittests/test_base_attached_module.py index 77cab14b..c7cd1508 100644 --- a/tests/unittests/test_base_attached_module.py +++ b/tests/unittests/test_base_attached_module.py @@ -1,8 +1,17 @@ import pytest +import torch +from torch import Tensor from luxonis_train import BaseLoss, BaseNode from luxonis_train.enums import TaskType from luxonis_train.utils.exceptions import IncompatibleException +from luxonis_train.utils.types import Labels, Packet + +SEGMENTATION_ARRAY = torch.tensor([0]) +KEYPOINT_ARRAY = torch.tensor([1]) +BOUNDINGBOX_ARRAY = torch.tensor([2]) +CLASSIFICATION_ARRAY = torch.tensor([3]) +FEATURES_ARRAY = torch.tensor([4]) class DummyBackbone(BaseNode): @@ -41,20 +50,20 @@ def forward(self, _): ... @pytest.fixture -def labels(): +def labels() -> Labels: return { - "segmentation": ("segmentation", TaskType.SEGMENTATION), - "keypoints": ("keypoints", TaskType.KEYPOINTS), - "boundingbox": ("boundingbox", TaskType.BOUNDINGBOX), - "classification": ("classification", TaskType.CLASSIFICATION), + "/segmentation": SEGMENTATION_ARRAY, + "/keypoints": KEYPOINT_ARRAY, + "/boundingbox": BOUNDINGBOX_ARRAY, + "/classification": CLASSIFICATION_ARRAY, } @pytest.fixture -def inputs(): +def inputs() -> Packet[Tensor]: return { - "features": ["features"], - "segmentation": ["segmentation"], + "features": [FEATURES_ARRAY], + "/segmentation": [SEGMENTATION_ARRAY], } @@ -63,10 +72,10 @@ def test_valid_properties(): loss = DummyLoss(node=head) no_labels_loss = NoLabelLoss(node=head) assert loss.node == head - assert loss.node_tasks == {TaskType.SEGMENTATION: "segmentation"} + assert loss.node_tasks == [TaskType.SEGMENTATION] assert loss.required_labels == [TaskType.SEGMENTATION] assert no_labels_loss.node == head - assert no_labels_loss.node_tasks == {TaskType.SEGMENTATION: "segmentation"} + assert no_labels_loss.node_tasks == [TaskType.SEGMENTATION] assert no_labels_loss.required_labels == [] @@ -82,45 +91,49 @@ def test_invalid_properties(): _ = NoLabelLoss(node=backbone).node_tasks -def test_get_label(labels): +def test_get_label(labels: Labels): seg_head = DummySegmentationHead() det_head = DummyDetectionHead() seg_loss = DummyLoss(node=seg_head) - assert seg_loss.get_label(labels) == "segmentation" - assert seg_loss.get_label(labels, TaskType.SEGMENTATION) == "segmentation" + assert seg_loss.get_label(labels) == SEGMENTATION_ARRAY + assert ( + seg_loss.get_label(labels, TaskType.SEGMENTATION) == SEGMENTATION_ARRAY + ) - del labels["segmentation"] - labels["segmentation-task"] = ("segmentation", TaskType.SEGMENTATION) + del labels["/segmentation"] + labels["task/segmentation"] = SEGMENTATION_ARRAY with pytest.raises(IncompatibleException): seg_loss.get_label(labels) det_loss = DummyLoss(node=det_head) - assert det_loss.get_label(labels, TaskType.KEYPOINTS) == "keypoints" - assert det_loss.get_label(labels, TaskType.BOUNDINGBOX) == "boundingbox" + assert det_loss.get_label(labels, TaskType.KEYPOINTS) == KEYPOINT_ARRAY + assert ( + det_loss.get_label(labels, TaskType.BOUNDINGBOX) == BOUNDINGBOX_ARRAY + ) with pytest.raises(ValueError): det_loss.get_label(labels) - with pytest.raises(ValueError): + with pytest.raises(IncompatibleException): det_loss.get_label(labels, TaskType.SEGMENTATION) -def test_input_tensors(inputs): +def test_input_tensors(inputs: Packet[Tensor]): seg_head = DummySegmentationHead() seg_loss = DummyLoss(node=seg_head) - assert seg_loss.get_input_tensors(inputs) == ["segmentation"] - assert seg_loss.get_input_tensors(inputs, "segmentation") == [ - "segmentation" + assert seg_loss.get_input_tensors(inputs) == [SEGMENTATION_ARRAY] + assert seg_loss.get_input_tensors(inputs, "/segmentation") == [ + SEGMENTATION_ARRAY ] assert seg_loss.get_input_tensors(inputs, TaskType.SEGMENTATION) == [ - "segmentation" + SEGMENTATION_ARRAY ] with pytest.raises(IncompatibleException): seg_loss.get_input_tensors(inputs, TaskType.KEYPOINTS) with pytest.raises(IncompatibleException): - seg_loss.get_input_tensors(inputs, "keypoints") + seg_loss.get_input_tensors(inputs, "/keypoints") det_head = DummyDetectionHead() det_loss = DummyLoss(node=det_head) @@ -128,17 +141,20 @@ def test_input_tensors(inputs): det_loss.get_input_tensors(inputs) -def test_prepare(inputs, labels): +def test_prepare(inputs: Packet[Tensor], labels: Labels): backbone = DummyBackbone() seg_head = DummySegmentationHead() seg_loss = DummyLoss(node=seg_head) det_head = DummyDetectionHead() - assert seg_loss.prepare(inputs, labels) == ("segmentation", "segmentation") - inputs["segmentation"].append("segmentation2") assert seg_loss.prepare(inputs, labels) == ( - "segmentation2", - "segmentation", + SEGMENTATION_ARRAY, + SEGMENTATION_ARRAY, + ) + inputs["/segmentation"].append(FEATURES_ARRAY) + assert seg_loss.prepare(inputs, labels) == ( + FEATURES_ARRAY, + SEGMENTATION_ARRAY, ) with pytest.raises(RuntimeError): diff --git a/tests/unittests/test_base_node.py b/tests/unittests/test_base_node.py index 3ed284c3..e6cfd21d 100644 --- a/tests/unittests/test_base_node.py +++ b/tests/unittests/test_base_node.py @@ -2,9 +2,8 @@ import torch from torch import Size, Tensor -from luxonis_train.enums import TaskType from luxonis_train.nodes import AttachIndexType, BaseNode -from luxonis_train.utils import DatasetMetadata, Packet +from luxonis_train.utils import Packet from luxonis_train.utils.exceptions import IncompatibleException @@ -100,61 +99,3 @@ def forward(self, _): ... {"features": [Size((3, 224, 224)) for _ in range(3)]} ] ) - - -def test_tasks(): - class DummyHead(DummyNode): - tasks = [TaskType.CLASSIFICATION] - - class DummyMultiHead(DummyNode): - tasks = [TaskType.CLASSIFICATION, TaskType.SEGMENTATION] - - dummy_head = DummyHead() - dummy_node = DummyNode() - dummy_multi_head = DummyMultiHead(n_keypoints=4) - assert ( - dummy_head.get_task_name(TaskType.CLASSIFICATION) == "classification" - ) - assert dummy_head.task == "classification" - with pytest.raises(ValueError): - dummy_head.get_task_name(TaskType.SEGMENTATION) - - with pytest.raises(RuntimeError): - dummy_node.get_task_name(TaskType.SEGMENTATION) - - with pytest.raises(RuntimeError): - _ = dummy_node.task - - with pytest.raises(ValueError): - _ = dummy_multi_head.task - - metadata = DatasetMetadata( - classes={ - "segmentation": ["car", "person", "dog"], - "classification": ["car-class", "person-class"], - }, - n_keypoints={"color-segmentation": 0, "detection": 0}, - ) - - dummy_multi_head._dataset_metadata = metadata - assert dummy_multi_head.get_class_names(TaskType.SEGMENTATION) == [ - "car", - "person", - "dog", - ] - assert dummy_multi_head.get_class_names(TaskType.CLASSIFICATION) == [ - "car-class", - "person-class", - ] - assert dummy_multi_head.get_n_classes(TaskType.SEGMENTATION) == 3 - assert dummy_multi_head.get_n_classes(TaskType.CLASSIFICATION) == 2 - assert dummy_multi_head.n_keypoints == 4 - with pytest.raises(ValueError): - _ = dummy_head.n_keypoints - with pytest.raises(RuntimeError): - _ = dummy_node.n_keypoints - - dummy_head = DummyHead(n_classes=5) - assert dummy_head.n_classes == 5 - with pytest.raises(ValueError): - _ = dummy_multi_head.n_classes diff --git a/tests/unittests/test_loaders/test_base_loader.py b/tests/unittests/test_loaders/test_base_loader.py index 293b3c10..3cc9231f 100644 --- a/tests/unittests/test_loaders/test_base_loader.py +++ b/tests/unittests/test_loaders/test_base_loader.py @@ -2,7 +2,6 @@ import torch from torch import Size -from luxonis_train.enums import TaskType from luxonis_train.loaders import collate_fn @@ -40,22 +39,12 @@ def build_batch_element(): inputs[name] = torch.rand(shape, dtype=torch.float32) labels = { - "classification": ( - torch.randint(0, 2, (2,), dtype=torch.int64), - TaskType.CLASSIFICATION, - ), - "segmentation": ( - torch.randint(0, 2, (1, 224, 224), dtype=torch.int64), - TaskType.SEGMENTATION, - ), - "keypoints": ( - torch.rand(1, 52, dtype=torch.float32), - TaskType.KEYPOINTS, - ), - "boundingbox": ( - torch.rand(1, 5, dtype=torch.float32), - TaskType.BOUNDINGBOX, + "/classification": (torch.randint(0, 2, (2,), dtype=torch.int64)), + "/segmentation": ( + torch.randint(0, 2, (1, 224, 224), dtype=torch.int64) ), + "/keypoints": (torch.rand(1, 52, dtype=torch.float32)), + "/boundingbox": (torch.rand(1, 5, dtype=torch.float32)), } return inputs, labels @@ -69,26 +58,26 @@ def build_batch_element(): assert inputs["features"].dtype == torch.float32 with subtests.test("classification"): - assert "classification" in annotations - assert annotations["classification"][0].shape == (batch_size, 2) - assert annotations["classification"][0].dtype == torch.int64 + assert "/classification" in annotations + assert annotations["/classification"].shape == (batch_size, 2) + assert annotations["/classification"].dtype == torch.int64 with subtests.test("segmentation"): - assert "segmentation" in annotations - assert annotations["segmentation"][0].shape == ( + assert "/segmentation" in annotations + assert annotations["/segmentation"].shape == ( batch_size, 1, 224, 224, ) - assert annotations["segmentation"][0].dtype == torch.int64 + assert annotations["/segmentation"].dtype == torch.int64 with subtests.test("keypoints"): - assert "keypoints" in annotations - assert annotations["keypoints"][0].shape == (batch_size, 53) - assert annotations["keypoints"][0].dtype == torch.float32 + assert "/keypoints" in annotations + assert annotations["/keypoints"].shape == (batch_size, 53) + assert annotations["/keypoints"].dtype == torch.float32 with subtests.test("boundingbox"): - assert "boundingbox" in annotations - assert annotations["boundingbox"][0].shape == (batch_size, 6) - assert annotations["boundingbox"][0].dtype == torch.float32 + assert "/boundingbox" in annotations + assert annotations["/boundingbox"].shape == (batch_size, 6) + assert annotations["/boundingbox"].dtype == torch.float32 From 1de6f7413fc617c10df15696fffaf75dd16bc105 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 07:47:55 -0600 Subject: [PATCH 19/57] fixed predefined classification --- luxonis_train/config/predefined_models/classification_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/luxonis_train/config/predefined_models/classification_model.py b/luxonis_train/config/predefined_models/classification_model.py index 86964e0a..dc67a96e 100644 --- a/luxonis_train/config/predefined_models/classification_model.py +++ b/luxonis_train/config/predefined_models/classification_model.py @@ -66,7 +66,7 @@ def __init__( self.loss_params = loss_params or {} self.visualizer_params = visualizer_params or {} self.task = task - self.task_name = task_name or "classification" + self.task_name = task_name or "" @property def nodes(self) -> list[ModelNodeConfig]: From 8c320146c86cf755aa031a16a281431e2c04c9b8 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 07:55:18 -0600 Subject: [PATCH 20/57] docs --- luxonis_train/nodes/README.md | 2 +- luxonis_train/nodes/base_node.py | 14 ++++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/luxonis_train/nodes/README.md b/luxonis_train/nodes/README.md index 31f1f6c2..92f1969e 100644 --- a/luxonis_train/nodes/README.md +++ b/luxonis_train/nodes/README.md @@ -40,7 +40,7 @@ In addition, the following class attributes can be overridden: | Key | Type | Default value | Description | | -------------- | ----------------------------------------------------------------- | ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `attach_index` | `int \| "all" \| tuple[int, int] \| tuple[int, int, int] \| None` | `None` | Index of previous output that the head attaches to. Each node has a sensible default. Usually should not be manually set in most cases. Can be either a single index, a slice (negative indexing is also supported), or `"all"` | -| `tasks` | `list[TaskType] \| Dict[TaskType, str] \| None` | `None` | Tasks supported by the node. Should be overridden for head nodes. Either a list of tasks or a dictionary mapping tasks to their default names | +| `tasks` | `list[TaskType] \| None` | `None` | List of tasks types supported by the node. Should be overridden for head nodes. | Additional parameters for specific nodes are listed below. diff --git a/luxonis_train/nodes/base_node.py b/luxonis_train/nodes/base_node.py index 19ec33a7..7dcdbcf3 100644 --- a/luxonis_train/nodes/base_node.py +++ b/luxonis_train/nodes/base_node.py @@ -93,18 +93,16 @@ def wrap(output: Tensor) -> Packet[Tensor]: # by the attached modules. return {"classification": [output]} - @type attach_index: AttachIndexType @ivar attach_index: Index of previous output that this node attaches to. Can be a single integer to specify a single output, a tuple of two or three integers to specify a range of outputs or C{"all"} to specify all outputs. Defaults to "all". Python indexing conventions apply. - @type tasks: list[TaskType] | dict[TaskType, str] | None - @ivar tasks: Dictionary of tasks that the node supports. Should be defined - by the user as a class attribute. The key is the task type and the value - is the name of the task. For example: - C{{TaskType.CLASSIFICATION: "classification"}}. + @type tasks: list[TaskType] | None + @ivar tasks: List of task types that the node supports. + Should be defined as a class attribute by the user. + For example C{[TaskType.CLASSIFICATION]}. Only needs to be defined for head nodes. """ @@ -154,6 +152,10 @@ def __init__( of outputs or C{"all"} to specify all outputs. Defaults to "all". Python indexing conventions apply. If provided as a constructor argument, overrides the class attribute. + @type task_name: str | None + @param task_name: Specifies which task group from the dataset to use + in case the dataset contains multiple tasks. Otherwise, the + task group is inferred from the dataset metadata. """ super().__init__() From 8d7685bb45cc7773e2a958f3894ba8225267cab4 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 14 Jan 2025 08:59:53 -0600 Subject: [PATCH 21/57] fix inspect --- luxonis_train/__main__.py | 55 ++++++++------------------------------- 1 file changed, 11 insertions(+), 44 deletions(-) diff --git a/luxonis_train/__main__.py b/luxonis_train/__main__.py index c0aae2dc..5fd11b1e 100644 --- a/luxonis_train/__main__.py +++ b/luxonis_train/__main__.py @@ -3,11 +3,10 @@ from pathlib import Path from typing import Annotated +import numpy as np import typer from luxonis_ml.utils import setup_logging -from luxonis_train.config import Config - setup_logging(use_rich=True) @@ -175,47 +174,16 @@ def inspect( To close the window press 'q' or 'Esc'. """ import cv2 - from luxonis_ml.data import Augmentations, LabelType from luxonis_ml.data.utils.visualizations import visualize - from luxonis_train.utils.registry import LOADERS - - cfg = Config.get_config(config, opts) - train_augmentations = Augmentations( - image_size=cfg.trainer.preprocessing.train_image_size, - augmentations=[ - i.model_dump() - for i in cfg.trainer.preprocessing.get_active_augmentations() - if i.name != "Normalize" - ], - train_rgb=cfg.trainer.preprocessing.train_rgb, - keep_aspect_ratio=cfg.trainer.preprocessing.keep_aspect_ratio, - ) - val_augmentations = Augmentations( - image_size=cfg.trainer.preprocessing.train_image_size, - augmentations=[ - i.model_dump() - for i in cfg.trainer.preprocessing.get_active_augmentations() - ], - train_rgb=cfg.trainer.preprocessing.train_rgb, - keep_aspect_ratio=cfg.trainer.preprocessing.keep_aspect_ratio, - only_normalize=True, - ) + from luxonis_train.core import LuxonisModel - Loader = LOADERS.get(cfg.loader.name) - loader = Loader( - augmentations=( - train_augmentations if view == "train" else val_augmentations - ), - view={ - "train": cfg.loader.train_view, - "val": cfg.loader.val_view, - "test": cfg.loader.test_view, - }[view], - image_source=cfg.loader.image_source, - **cfg.loader.params, - ) + opts = opts or [] + opts.extend(["trainer.preprocessing.normalize.active", "False"]) + + model = LuxonisModel(config, opts) + loader = model.loaders[view.value] for images, labels in loader: for img in images.values(): if len(img.shape) != 3: @@ -226,11 +194,10 @@ def inspect( k: v.numpy().transpose(1, 2, 0) for k, v in images.items() } main_image = np_images[loader.image_source] - main_image = cv2.cvtColor(main_image, cv2.COLOR_RGB2BGR) - np_labels = { - task: (label.numpy(), LabelType(task_type)) - for task, (label, task_type) in labels.items() - } + main_image = cv2.cvtColor(main_image, cv2.COLOR_RGB2BGR).astype( + np.uint8 + ) + np_labels = {task: label.numpy() for task, label in labels.items()} h, w, _ = main_image.shape new_h, new_w = int(h * size_multiplier), int(w * size_multiplier) From fbbbc26cd44fdbbc433587ae45f3a23b93b51aa1 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 15 Jan 2025 18:20:31 -0500 Subject: [PATCH 22/57] fixed tests --- .../reconstruction_segmentation_loss.py | 64 ++++--------- .../predefined_models/classification_model.py | 4 +- .../predefined_models/detection_model.py | 4 +- .../keypoint_detection_model.py | 4 +- .../predefined_models/segmentation_model.py | 4 +- luxonis_train/core/utils/infer_utils.py | 12 +-- luxonis_train/loaders/base_loader.py | 9 +- luxonis_train/loaders/luxonis_loader_torch.py | 11 ++- .../loaders/luxonis_perlin_loader_torch.py | 96 ++++++++----------- luxonis_train/loaders/perlin.py | 16 ++-- luxonis_train/models/luxonis_lightning.py | 3 +- tests/configs/smart_cfg_populate_config.yaml | 1 + tests/integration/parking_lot.json | 17 ++-- .../test_unsupervised_anomaly_detection.py | 23 +++-- .../test_metrics/test_confusion_matrix.py | 3 +- 15 files changed, 123 insertions(+), 148 deletions(-) diff --git a/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py b/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py index 1937ad0b..a5b50f2b 100644 --- a/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py +++ b/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py @@ -1,18 +1,14 @@ import logging from math import exp -from typing import Any, Literal, Union +from typing import Literal import torch -import torch.nn as nn import torch.nn.functional as F -from torch import Tensor +from torch import Tensor, nn from luxonis_train.enums import TaskType from luxonis_train.nodes import DiscSubNetHead -from luxonis_train.utils import ( - Labels, - Packet, -) +from luxonis_train.utils import Labels, Packet from .base_loss import BaseLoss from .softmax_focal_loss import SoftmaxFocalLoss @@ -30,7 +26,7 @@ def __init__( gamma: float = 2.0, reduction: Literal["none", "mean", "sum"] = "mean", smooth: float = 1e-5, - **kwargs: Any, + **kwargs, ): """ReconstructionSegmentationLoss implements a combined loss function for reconstruction and segmentation tasks. @@ -55,28 +51,17 @@ def __init__( self.loss_ssim = SSIM() def prepare( - self, - inputs: Packet[Tensor], - labels: Labels, + self, inputs: Packet[Tensor], labels: Labels ) -> tuple[Tensor, Tensor, Tensor, Tensor]: recon = self.get_input_tensors(inputs, "reconstructed")[0] - seg_out = self.get_input_tensors(inputs, "segmentation")[0] - an_mask = labels["segmentation"][0] - orig = labels["original"][0] - - return ( - orig, - recon, - seg_out, - an_mask, - ) + seg_out = self.get_input_tensors(inputs)[0] + an_mask = self.get_label(labels) + orig = labels[f"{self.node.task_name}/original/segmentation"] + + return orig, recon, seg_out, an_mask def forward( - self, - orig: Tensor, - recon: Tensor, - seg_out: Tensor, - an_mask: Tensor, + self, orig: Tensor, recon: Tensor, seg_out: Tensor, an_mask: Tensor ): l2 = self.loss_l2(recon, orig) ssim = self.loss_ssim(recon, orig) @@ -93,14 +78,14 @@ def forward( return total_loss, sub_losses -class SSIM(torch.nn.Module): +class SSIM(nn.Module): def __init__( self, window_size: int = 11, size_average: bool = True, val_range: float | None = None, ): - super(SSIM, self).__init__() + super().__init__() self.window_size = window_size self.size_average = size_average self.val_range = val_range @@ -123,7 +108,7 @@ def forward(self, img1: Tensor, img2: Tensor) -> Tensor: self.window = window self.channel = channel - s_score, ssim_map = ssim( + s_score = ssim( img1, img2, window=window, @@ -134,11 +119,9 @@ def forward(self, img1: Tensor, img2: Tensor) -> Tensor: def create_window(window_size: int, channel: int = 1) -> Tensor: - _1D_window = gaussian(window_size, 1.5).unsqueeze(1) - _2D_window = ( - _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) - ) - window = _2D_window.expand( + window_1d = gaussian(window_size, 1.5).unsqueeze(1) + widnow_2d = window_1d.mm(window_1d.t()).float().unsqueeze(0).unsqueeze(0) + window = widnow_2d.expand( channel, 1, window_size, window_size ).contiguous() return window @@ -160,9 +143,8 @@ def ssim( window_size: int = 11, window: Tensor | None = None, size_average=True, - full=False, val_range=None, -) -> Union[Tensor, tuple[Tensor, Tensor]]: +) -> Tensor: if val_range is None: if torch.max(img1) > 128: max_val = 255 @@ -205,15 +187,9 @@ def ssim( v1 = 2.0 * sigma12 + c2 v2 = sigma1_sq + sigma2_sq + c2 - cs = torch.mean(v1 / v2) # contrast sensitivity ssim_map = ((2 * mu1_mu2 + c1) * v1) / ((mu1_sq + mu2_sq + c1) * v2) if size_average: - ret = ssim_map.mean() - else: - ret = ssim_map.mean(1).mean(1).mean(1) - - if full: - return ret, cs - return ret, ssim_map + return ssim_map.mean() + return ssim_map.mean(1).mean(1).mean(1) diff --git a/luxonis_train/config/predefined_models/classification_model.py b/luxonis_train/config/predefined_models/classification_model.py index a7405890..7508a278 100644 --- a/luxonis_train/config/predefined_models/classification_model.py +++ b/luxonis_train/config/predefined_models/classification_model.py @@ -129,8 +129,8 @@ def metrics(self) -> list[MetricModuleConfig]: metrics.append( MetricModuleConfig( name="ConfusionMatrix", - alias=f"ConfusionMatrix-{self.task_name}", - attached_to=f"ClassificationHead-{self.task_name}", + alias=f"{self.task_name}/ConfusionMatrix", + attached_to=f"{self.task_name}/ClassificationHead", params={**self.confusion_matrix_params}, ) ) diff --git a/luxonis_train/config/predefined_models/detection_model.py b/luxonis_train/config/predefined_models/detection_model.py index 3c1cabb9..0a827d36 100644 --- a/luxonis_train/config/predefined_models/detection_model.py +++ b/luxonis_train/config/predefined_models/detection_model.py @@ -149,8 +149,8 @@ def metrics(self) -> list[MetricModuleConfig]: metrics.append( MetricModuleConfig( name="ConfusionMatrix", - alias=f"ConfusionMatrix-{self.task_name}", - attached_to=f"EfficientBBoxHead-{self.task_name}", + alias=f"{self.task_name}/ConfusionMatrix", + attached_to=f"{self.task_name}/EfficientBBoxHead", params={**self.confusion_matrix_params}, ) ) diff --git a/luxonis_train/config/predefined_models/keypoint_detection_model.py b/luxonis_train/config/predefined_models/keypoint_detection_model.py index d5706d40..8144de46 100644 --- a/luxonis_train/config/predefined_models/keypoint_detection_model.py +++ b/luxonis_train/config/predefined_models/keypoint_detection_model.py @@ -153,8 +153,8 @@ def metrics(self) -> list[MetricModuleConfig]: metrics.append( MetricModuleConfig( name="ConfusionMatrix", - alias=f"ConfusionMatrix-{self.kpt_task_name}", - attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}", + alias=f"{self.task_name}/ConfusionMatrix", + attached_to=f"{self.task_name}/EfficientKeypointBBoxHead", params={**self.confusion_matrix_params}, ) ) diff --git a/luxonis_train/config/predefined_models/segmentation_model.py b/luxonis_train/config/predefined_models/segmentation_model.py index 9fab84c1..54beeb1f 100644 --- a/luxonis_train/config/predefined_models/segmentation_model.py +++ b/luxonis_train/config/predefined_models/segmentation_model.py @@ -166,8 +166,8 @@ def metrics(self) -> list[MetricModuleConfig]: metrics.append( MetricModuleConfig( name="ConfusionMatrix", - alias=f"ConfusionMatrix-{self.task_name}", - attached_to=f"DDRNetSegmentationHead-{self.task_name}", + alias=f"{self.task_name}/ConfusionMatrix", + attached_to=f"{self.task_name}/DDRNetSegmentationHead", params={**self.confusion_matrix_params}, ) ) diff --git a/luxonis_train/core/utils/infer_utils.py b/luxonis_train/core/utils/infer_utils.py index b4996751..7f009d8f 100644 --- a/luxonis_train/core/utils/infer_utils.py +++ b/luxonis_train/core/utils/infer_utils.py @@ -47,12 +47,10 @@ def process_visualizations( def prepare_and_infer_image( - model: "luxonis_train.core.LuxonisModel", - img: np.ndarray, + model: "luxonis_train.core.LuxonisModel", img: Tensor ): """Prepares the image for inference and runs the model.""" - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - img, _ = model.val_augmentations([(img, {})]) + img = model.loaders["val"].augment_test_image(img) # type: ignore inputs = { "image": torch.tensor(img).unsqueeze(0).permute(0, 3, 1, 2).float() @@ -94,9 +92,11 @@ def infer_from_video( ret, frame = cap.read() if not ret: # pragma: no cover break + if model.cfg.trainer.preprocessing.color_format == "RGB": + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # TODO: batched inference - outputs = prepare_and_infer_image(model, frame) + outputs = prepare_and_infer_image(model, torch.tensor(frame)) renders = process_visualizations(outputs.visualizations, batch_size=1) for (node_name, viz_name), [viz] in renders.items(): @@ -213,8 +213,8 @@ def generator(): dataset_name=dataset_name, image_source="image", view="test", - augmentations=model.val_augmentations, ) + loader.loader.augmentations = model.loaders["val"].loader.augmentations # type: ignore loader = torch_data.DataLoader( loader, batch_size=model.cfg.trainer.batch_size ) diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index 3e3589dd..a9998e4b 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from luxonis_ml.utils.registry import AutoRegisterMeta -from torch import Size +from torch import Size, Tensor from torch.utils.data import Dataset from luxonis_train.utils.registry import LOADERS @@ -83,6 +83,13 @@ def input_shape(self) -> Size: """ return self.input_shapes[self.image_source] + def augment_test_image(self, img: Tensor) -> Tensor: + raise NotImplementedError( + f"{self.__class__.__name__} does not expose interface " + "for test-time augmentation. Implement " + "`augment_test_image` method to expose this functionality." + ) + @abstractmethod def __len__(self) -> int: """Returns length of the dataset.""" diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index 445fe641..d6efe97b 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -3,6 +3,7 @@ from typing import Any, Literal import numpy as np +import torch from luxonis_ml.data import ( BucketStorage, BucketType, @@ -123,13 +124,19 @@ def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: return {self.image_source: tensor_img}, tensor_labels def get_classes(self) -> dict[str, list[str]]: - _, classes = self.dataset.get_classes() - return {task: classes[task] for task in classes} + return self.dataset.get_classes() def get_n_keypoints(self) -> dict[str, int]: skeletons = self.dataset.get_skeletons() return {task: len(skeletons[task][0]) for task in skeletons} + def augment_test_image(self, img: Tensor) -> Tensor: + if self.loader.augmentations is None: + return img + return torch.tensor( + self.loader.augmentations.apply([(img.numpy(), {})])[0] + ) + def _parse_dataset( self, dataset_dir: str, diff --git a/luxonis_train/loaders/luxonis_perlin_loader_torch.py b/luxonis_train/loaders/luxonis_perlin_loader_torch.py index c1660ffb..05e2c449 100644 --- a/luxonis_train/loaders/luxonis_perlin_loader_torch.py +++ b/luxonis_train/loaders/luxonis_perlin_loader_torch.py @@ -1,16 +1,13 @@ -import glob -import os import random -from typing import Callable, List +from typing import cast import numpy as np import torch import torch.nn.functional as F +from luxonis_ml.data import AlbumentationsEngine from luxonis_ml.utils import LuxonisFileSystem from torch import Tensor -from luxonis_train.enums import TaskType - from .base_loader import LuxonisLoaderTorchOutput from .luxonis_loader_torch import LuxonisLoaderTorch from .perlin import apply_anomaly_to_img @@ -32,53 +29,45 @@ def __init__( @param noise_prob: The probability with which to apply Perlin noise (only used during training). """ - if not anomaly_source_path: - raise ValueError("anomaly_source_path must be a valid string.") - super().__init__(*args, **kwargs) - lux_fs = LuxonisFileSystem(path=anomaly_source_path) - if lux_fs.protocol in ["s3", "gcs"]: - anomaly_source_path = str( - lux_fs.get_dir( - remote_paths=[anomaly_source_path], local_dir="./data" - ) + + try: + self.anomaly_source_path = LuxonisFileSystem.download( + anomaly_source_path, dest="./data" ) - else: - anomaly_source_path = str(lux_fs.path) + except Exception as e: + raise FileNotFoundError( + "The anomaly source path is invalid." + ) from e + + from luxonis_train.core.utils.infer_utils import IMAGE_FORMATS - if anomaly_source_path and os.path.exists(anomaly_source_path): - self.anomaly_source_paths = sorted( - glob.glob(os.path.join(anomaly_source_path, "*/*.jpg")) + self.anomaly_files = [ + f + for f in self.anomaly_source_path.rglob("*") + if f.suffix.lower() in IMAGE_FORMATS + ] + if not self.anomaly_files: + raise FileNotFoundError( + "No image files found at the specified path." ) - if not self.anomaly_source_paths: - raise FileNotFoundError( - "No .jpg files found at the specified path." - ) - else: - raise ValueError("Invalid or unspecified anomaly source path.") - self.anomaly_source_path = anomaly_source_path self.noise_prob = noise_prob - self.base_loader.add_background = True # type: ignore - self.base_loader.class_mappings["segmentation"]["background"] = 0 - self.base_loader.class_mappings["segmentation"] = { - k: (v + 1 if k != "background" else v) - for k, v in self.base_loader.class_mappings["segmentation"].items() - } - - if ( - self.augmentations is None - or self.augmentations.pixel_transform is None - ): - self.pixel_augs: List[Callable] = [] + if len(self.loader.dataset.get_tasks()) > 1: + # TODO: Can be extended to multiple tasks + raise ValueError( + "This loader only supports datasets with a single task." + ) + self.task_name = next(iter(self.loader.dataset.get_tasks())) + + augmentations = cast(AlbumentationsEngine, self.loader.augmentations) + if augmentations is None or augmentations.pixel_transform is None: + self.pixel_augs = None else: - self.pixel_augs: List[Callable] = [ - transform - for transform in self.augmentations.pixel_transform.transforms - ] + self.pixel_augs = augmentations.pixel_transform def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: - img, labels = self.base_loader[idx] + img, labels = self.loader[idx] img = np.transpose(img, (2, 0, 1)) tensor_img = Tensor(img) @@ -86,7 +75,7 @@ def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: if self.view[0] == "train" and random.random() < self.noise_prob: aug_tensor_img, an_mask = apply_anomaly_to_img( tensor_img, - anomaly_source_paths=self.anomaly_source_paths, + anomaly_source_paths=self.anomaly_files, pixel_augs=self.pixel_augs, ) else: @@ -94,17 +83,12 @@ def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: h, w = aug_tensor_img.shape[-2:] an_mask = torch.zeros((h, w)) - tensor_labels = {"original": (tensor_img, TaskType.ARRAY)} - if self.view[0] == "train": - tensor_labels["segmentation"] = ( - F.one_hot(an_mask.long(), 2).permute(2, 0, 1).float(), - TaskType.SEGMENTATION, - ) - else: - for task, (array, label_type) in labels.items(): - tensor_labels[task] = ( - Tensor(array), - TaskType(label_type.value), - ) + tensor_labels = {f"{self.task_name}/original/segmentation": tensor_img} + for task, array in labels.items(): + tensor_labels[task] = Tensor(array) + + tensor_labels[f"{self.task_name}/segmentation"] = ( + F.one_hot(an_mask.long(), 2).permute(2, 0, 1).float() + ) return {self.image_source: aug_tensor_img}, tensor_labels diff --git a/luxonis_train/loaders/perlin.py b/luxonis_train/loaders/perlin.py index a54e3d7c..687fd015 100644 --- a/luxonis_train/loaders/perlin.py +++ b/luxonis_train/loaders/perlin.py @@ -1,4 +1,5 @@ import random +from pathlib import Path from typing import Callable, List, Tuple import cv2 @@ -135,17 +136,17 @@ def generate_perlin_noise( return perlin_mask -def load_image_as_numpy(img_path: str) -> np.ndarray: - image = cv2.imread(img_path, cv2.IMREAD_COLOR) +def load_image_as_numpy(img_path: Path) -> np.ndarray: + image = cv2.imread(str(img_path), cv2.IMREAD_COLOR) image = image.astype(np.float32) / 255.0 return image def apply_anomaly_to_img( img: torch.Tensor, - anomaly_source_paths: List[str], + anomaly_source_paths: List[Path], beta: float | None = None, - pixel_augs: List[Callable] | None = None, + pixel_augs: Callable | None = None, # type: ignore ) -> Tuple[torch.Tensor, torch.Tensor]: """Applies Perlin noise-based anomalies to a single image (C, H, W). @@ -165,7 +166,9 @@ def apply_anomaly_to_img( """ if pixel_augs is None: - pixel_augs = [] + + def pixel_augs(image): + return {"image": image} sampled_anomaly_image_path = random.choice(anomaly_source_paths) @@ -177,8 +180,7 @@ def apply_anomaly_to_img( interpolation=cv2.INTER_LINEAR, ) - for aug in pixel_augs: - anomaly_image = aug(image=anomaly_image)["image"] + anomaly_image = pixel_augs(image=anomaly_image)["image"] anomaly_image = torch.tensor(anomaly_image).permute(2, 0, 1) diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index afd422ae..c478c126 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -989,7 +989,7 @@ def _init_attached_module( loader = self._core.loaders["train"] dataset = getattr(loader, "dataset", None) if isinstance(dataset, LuxonisDataset): - n_classes = len(dataset.get_classes()[1][node.task_name]) + n_classes = len(dataset.get_classes()[node.task_name]) if n_classes == 1: cfg.params["task"] = "binary" else: @@ -1034,7 +1034,6 @@ def _print_results( ) if self.main_metric is not None: - print(self.main_metric) *main_metric_node, main_metric_name = self.main_metric.split("/") main_metric_node = "/".join(main_metric_node) diff --git a/tests/configs/smart_cfg_populate_config.yaml b/tests/configs/smart_cfg_populate_config.yaml index 0c81752b..87f8a96d 100644 --- a/tests/configs/smart_cfg_populate_config.yaml +++ b/tests/configs/smart_cfg_populate_config.yaml @@ -9,6 +9,7 @@ model: - EfficientRep - name: EfficientBBoxHead + task_name: vehicle_type inputs: - RepPANNeck diff --git a/tests/integration/parking_lot.json b/tests/integration/parking_lot.json index 66ce962d..bf3e3835 100644 --- a/tests/integration/parking_lot.json +++ b/tests/integration/parking_lot.json @@ -96,7 +96,7 @@ "dtype": "float32", "shape": [ 1, - 7, + 8, 32, 40 ], @@ -107,7 +107,7 @@ "dtype": "float32", "shape": [ 1, - 7, + 8, 16, 20 ], @@ -118,11 +118,11 @@ "dtype": "float32", "shape": [ 1, - 7, + 8, 8, 10 ], - "layout": "NCHW" + "layout": "NCDE" } ], "heads": [ @@ -169,9 +169,9 @@ "metadata": { "postprocessor_path": null, "classes": [ - "background", + "motorbike", "car", - "motorbike" + "background" ], "n_classes": 3, "is_softmax": false @@ -185,9 +185,10 @@ "postprocessor_path": null, "classes": [ "motorbike", - "car" + "car", + "background" ], - "n_classes": 2, + "n_classes": 3, "iou_threshold": 0.45, "conf_threshold": 0.25, "max_det": 300, diff --git a/tests/integration/test_unsupervised_anomaly_detection.py b/tests/integration/test_unsupervised_anomaly_detection.py index 10c0a6f9..bf9347b5 100644 --- a/tests/integration/test_unsupervised_anomaly_detection.py +++ b/tests/integration/test_unsupervised_anomaly_detection.py @@ -11,7 +11,7 @@ PathType = Union[str, Path] -def get_opts() -> dict[str, Any]: +def get_config() -> dict[str, Any]: return { "model": { "name": "DREAM", @@ -38,12 +38,15 @@ def get_opts() -> dict[str, Any]: "keep_aspect_ratio": False, "normalize": {"active": True}, }, - "batch_size": 1, - "epochs": 1, + "batch_size": 4, + "epochs": 800, "num_workers": 0, "validation_interval": 10, "num_sanity_val_steps": 0, }, + "tracker": { + "save_directory": "tests/integration/save-directory", + }, } @@ -67,22 +70,19 @@ def dummy_generator( train_paths: List[PathType], test_paths: List[PathType] ): for path in train_paths: + img = cv2.imread(str(path)) + img_h, img_w, _ = img.shape + mask = np.zeros((img_h, img_w), dtype=np.uint8) yield { "file": path, "annotation": { "class": "object", - "segmentation": { - "height": 256, - "width": 256, - "counts": "0" * (256 * 256), - }, + "segmentation": {"mask": mask}, }, } for path in test_paths: img = cv2.imread(str(path)) - if img is None: - continue img_h, img_w, _ = img.shape mask = random_square_mask((img_h, img_w)) poly = cv2.findContours( @@ -130,7 +130,6 @@ def test_anomaly_detection(): create_dummy_anomaly_detection_dataset( Path("tests/data/COCO_people_subset/person_val2017_subset/*") ) - config = get_opts() - model = LuxonisModel(config) + model = LuxonisModel(get_config()) model.train() model.test() diff --git a/tests/unittests/test_metrics/test_confusion_matrix.py b/tests/unittests/test_metrics/test_confusion_matrix.py index 07429d8b..40a59dbb 100644 --- a/tests/unittests/test_metrics/test_confusion_matrix.py +++ b/tests/unittests/test_metrics/test_confusion_matrix.py @@ -11,8 +11,7 @@ def test_compute_detection_confusion_matrix_specific_case(): class DummyNodeDetection(BaseNode): tasks = [TaskType.BOUNDINGBOX] - def forward(self, _): - pass + def forward(self, _): ... metric = ConfusionMatrix( node=DummyNodeDetection(n_classes=3), iou_threshold=0.5 From bb5e88295a25a08e0725f53de6ef40033bce8112 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 15 Jan 2025 19:11:46 -0500 Subject: [PATCH 23/57] fix debug config --- tests/integration/test_unsupervised_anomaly_detection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_unsupervised_anomaly_detection.py b/tests/integration/test_unsupervised_anomaly_detection.py index bf9347b5..be00cfb0 100644 --- a/tests/integration/test_unsupervised_anomaly_detection.py +++ b/tests/integration/test_unsupervised_anomaly_detection.py @@ -38,8 +38,8 @@ def get_config() -> dict[str, Any]: "keep_aspect_ratio": False, "normalize": {"active": True}, }, - "batch_size": 4, - "epochs": 800, + "batch_size": 1, + "epochs": 1, "num_workers": 0, "validation_interval": 10, "num_sanity_val_steps": 0, From 785f2f819d203492e8fbc2264563fd40e69b72a3 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 15 Jan 2025 19:14:02 -0500 Subject: [PATCH 24/57] updated perlin --- luxonis_train/loaders/perlin.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/luxonis_train/loaders/perlin.py b/luxonis_train/loaders/perlin.py index 687fd015..cd201d35 100644 --- a/luxonis_train/loaders/perlin.py +++ b/luxonis_train/loaders/perlin.py @@ -14,7 +14,7 @@ def compute_gradients(res: tuple[int, int]) -> torch.Tensor: @torch.jit.script -def lerp_torch( +def lerp_torch( # pragma: no cover x: torch.Tensor, y: torch.Tensor, w: torch.Tensor ) -> torch.Tensor: return (y - x) * w + x @@ -92,7 +92,7 @@ def rand_perlin_2d( @torch.jit.script -def rotate_noise(noise: torch.Tensor) -> torch.Tensor: +def rotate_noise(noise: torch.Tensor) -> torch.Tensor: # pragma: no cover angle = torch.rand(1) * 2 * torch.pi h, w = noise.shape center_y, center_x = h // 2, w // 2 @@ -165,11 +165,6 @@ def apply_anomaly_to_img( - perlin_mask (torch.Tensor): The Perlin noise mask applied to the image. """ - if pixel_augs is None: - - def pixel_augs(image): - return {"image": image} - sampled_anomaly_image_path = random.choice(anomaly_source_paths) anomaly_image = load_image_as_numpy(sampled_anomaly_image_path) @@ -180,7 +175,8 @@ def pixel_augs(image): interpolation=cv2.INTER_LINEAR, ) - anomaly_image = pixel_augs(image=anomaly_image)["image"] + if pixel_augs is not None: + anomaly_image = pixel_augs(image=anomaly_image)["image"] anomaly_image = torch.tensor(anomaly_image).permute(2, 0, 1) From c09336363f499a3f03c3f13aeab890c5a3fd193e Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 15 Jan 2025 19:18:07 -0500 Subject: [PATCH 25/57] missing doc --- luxonis_train/loaders/luxonis_loader_torch.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index d6efe97b..d44c92ea 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -76,6 +76,39 @@ def __init__( view of the dataset. Each split is a string that represents a subset of the dataset. The available splits depend on the dataset, but usually include 'train', 'val', and 'test'. Defaults to 'train'. + @type augmentation_engine: Union[Literal["albumentations"], str] + @param augmentation_engine: The augmentation engine to use. + Defaults to C{"albumentations"}. + @type augmentation_config: Optional[Union[List[Dict[str, Any]], + PathType]] + @param augmentation_config: The configuration for the + augmentations. This can be either a list of C{Dict[str, Any]} or + a path to a configuration file. + The config member is a dictionary with two keys: C{name} and + C{params}. C{name} is the name of the augmentation to + instantiate and C{params} is an optional dictionary + of parameters to pass to the augmentation. + + Example:: + + [ + {"name": "HorizontalFlip", "params": {"p": 0.5}}, + {"name": "RandomBrightnessContrast", "params": {"p": 0.1}}, + {"name": "Defocus"} + ] + + @type height: Optional[int] + @param height: The height of the output images. Defaults to + C{None}. + @type width: Optional[int] + @param width: The width of the output images. Defaults to + C{None}. + @type keep_aspect_ratio: bool + @param keep_aspect_ratio: Whether to keep the aspect ratio of the + images. Defaults to C{True}. + @type out_image_format: Literal["RGB", "BGR"] + @param out_image_format: The format of the output images. Defaults + to C{"RGB"}. """ super().__init__(view=view, **kwargs) if dataset_dir is not None: From f2cdfa3ba3c884b2a7c10d1f718573f7c9eb815e Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 15 Jan 2025 19:21:17 -0500 Subject: [PATCH 26/57] reverted bacj to train_rgb --- configs/README.md | 18 +++++++++--------- configs/complex_model.yaml | 1 - luxonis_train/config/config.py | 2 +- luxonis_train/core/core.py | 6 ++++-- luxonis_train/core/utils/infer_utils.py | 2 +- tests/configs/cli_commands.yaml | 1 - tests/configs/parking_lot_config.yaml | 1 - 7 files changed, 15 insertions(+), 16 deletions(-) diff --git a/configs/README.md b/configs/README.md index 37d150f9..e5ae8c81 100644 --- a/configs/README.md +++ b/configs/README.md @@ -280,14 +280,14 @@ We use [`Albumentations`](https://albumentations.ai/docs/) library for `augmenta Additionally, we support `Mosaic4` and `MixUp` batch augmentations and letterbox resizing if `keep_aspect_ratio: true`. -| Key | Type | Default value | Description | -| ------------------- | ----------------------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `train_image_size` | `list[int]` | `[256, 256]` | Image size used for training as `[height, width]` | -| `keep_aspect_ratio` | `bool` | `True` | Whether to keep the aspect ratio while resizing | -| `color_format` | `Literal["RGB", "BGR"]` | `"RGB"` | Whether to train on RGB or BGR images | -| `normalize.active` | `bool` | `True` | Whether to use normalization | -| `normalize.params` | `dict` | `{}` | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | -| `augmentations` | `list[dict]` | `[]` | List of `Albumentations` augmentations | +| Key | Type | Default value | Description | +| ------------------- | ------------ | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `train_image_size` | `list[int]` | `[256, 256]` | Image size used for training as `[height, width]` | +| `keep_aspect_ratio` | `bool` | `True` | Whether to keep the aspect ratio while resizing | +| `train_rgb` | `bool` | `"RGB"` | Whether to train on RGB or BGR images | +| `normalize.active` | `bool` | `True` | Whether to use normalization | +| `normalize.params` | `dict` | `{}` | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | +| `augmentations` | `list[dict]` | `[]` | List of `Albumentations` augmentations | #### Augmentations @@ -306,7 +306,7 @@ trainer: # using YAML capture to reuse the image size train_image_size: [&height 384, &width 384] keep_aspect_ratio: true - color_format: "RGB" + train_rgb: true normalize: active: true augmentations: diff --git a/configs/complex_model.yaml b/configs/complex_model.yaml index fed25c23..ee8fa037 100644 --- a/configs/complex_model.yaml +++ b/configs/complex_model.yaml @@ -100,7 +100,6 @@ trainer: preprocessing: train_image_size: [&height 384, &width 384] keep_aspect_ratio: true - color_format: RGB normalize: active: true augmentations: diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index 014caaf2..3d3690e0 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -318,7 +318,7 @@ class PreprocessingConfig(BaseModelExtraForbid): ImageSize, Field(default=[256, 256], min_length=2, max_length=2) ] = ImageSize(256, 256) keep_aspect_ratio: bool = True - color_format: Literal["RGB", "BGR"] = "RGB" + train_rgb: bool = True normalize: NormalizeAugmentationConfig = NormalizeAugmentationConfig() augmentations: list[AugmentationConfig] = [] diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py index 2fb59457..a0529030 100644 --- a/luxonis_train/core/core.py +++ b/luxonis_train/core/core.py @@ -132,7 +132,9 @@ def __init__( i.model_dump() for i in self.cfg.trainer.preprocessing.get_active_augmentations() ], - out_image_format=self.cfg.trainer.preprocessing.color_format, + out_image_format="RGB" + if self.cfg.trainer.preprocessing.train_rgb + else "BGR", keep_aspect_ratio=self.cfg.trainer.preprocessing.keep_aspect_ratio, **self.cfg.loader.params, ) @@ -721,7 +723,7 @@ def _mult(lst: list[float | int]) -> list[float]: self.cfg.trainer.preprocessing.normalize.params["std"] ), "dai_type": "RGB888p" - if self.cfg.trainer.preprocessing.color_format == "RGB" + if self.cfg.trainer.preprocessing.train_rgb else "BGR888p", } diff --git a/luxonis_train/core/utils/infer_utils.py b/luxonis_train/core/utils/infer_utils.py index 7f009d8f..26e994e5 100644 --- a/luxonis_train/core/utils/infer_utils.py +++ b/luxonis_train/core/utils/infer_utils.py @@ -92,7 +92,7 @@ def infer_from_video( ret, frame = cap.read() if not ret: # pragma: no cover break - if model.cfg.trainer.preprocessing.color_format == "RGB": + if model.cfg.trainer.preprocessing.train_rgb: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # TODO: batched inference diff --git a/tests/configs/cli_commands.yaml b/tests/configs/cli_commands.yaml index 7123534c..b8e604ae 100644 --- a/tests/configs/cli_commands.yaml +++ b/tests/configs/cli_commands.yaml @@ -51,7 +51,6 @@ trainer: preprocessing: train_image_size: [256, 320] keep_aspect_ratio: true - color_format: RGB normalize: active: true diff --git a/tests/configs/parking_lot_config.yaml b/tests/configs/parking_lot_config.yaml index 2754197e..a5ee50cb 100644 --- a/tests/configs/parking_lot_config.yaml +++ b/tests/configs/parking_lot_config.yaml @@ -87,7 +87,6 @@ trainer: preprocessing: train_image_size: [256, 320] keep_aspect_ratio: false - color_format: RGB normalize: active: true augmentations: From e32f6ea7b93fd5c8f9660512ba1fe8f47a632c26 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 16 Jan 2025 17:55:32 -0500 Subject: [PATCH 27/57] fix type issues --- README.md | 2 +- configs/README.md | 18 +-- .../losses/adaptive_detection_loss.py | 6 +- luxonis_train/callbacks/__init__.py | 24 +-- .../callbacks/archive_on_train_end.py | 2 +- luxonis_train/callbacks/ema.py | 7 +- luxonis_train/callbacks/gpu_stats_monitor.py | 10 +- luxonis_train/callbacks/gradcam_visializer.py | 4 +- .../callbacks/luxonis_progress_bar.py | 4 +- luxonis_train/callbacks/metadata_logger.py | 2 +- luxonis_train/callbacks/test_on_train_end.py | 2 +- luxonis_train/callbacks/training_manager.py | 2 +- luxonis_train/config/config.py | 21 ++- luxonis_train/core/core.py | 33 ++-- luxonis_train/core/utils/infer_utils.py | 2 +- luxonis_train/loaders/base_loader.py | 148 ++++++++++++++++-- luxonis_train/loaders/luxonis_loader_torch.py | 4 +- luxonis_train/optimizers/optimizers.py | 2 +- luxonis_train/schedulers/schedulers.py | 2 +- luxonis_train/strategies/base_strategy.py | 2 +- luxonis_train/strategies/triple_lr_sgd.py | 2 +- tests/unittests/test_callbacks/test_ema.py | 2 +- 22 files changed, 216 insertions(+), 85 deletions(-) diff --git a/README.md b/README.md index 6cb85a00..10a69acd 100644 --- a/README.md +++ b/README.md @@ -595,7 +595,7 @@ from luxonis_train import LuxonisLightningModule from luxonis_train.utils.registry import CALLBACKS -@CALLBACKS.register_module() +@CALLBACKS.register() class CustomCallback(pl.Callback): def __init__(self, message: str, **kwargs): super().__init__(**kwargs) diff --git a/configs/README.md b/configs/README.md index e5ae8c81..b336ef98 100644 --- a/configs/README.md +++ b/configs/README.md @@ -280,14 +280,14 @@ We use [`Albumentations`](https://albumentations.ai/docs/) library for `augmenta Additionally, we support `Mosaic4` and `MixUp` batch augmentations and letterbox resizing if `keep_aspect_ratio: true`. -| Key | Type | Default value | Description | -| ------------------- | ------------ | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `train_image_size` | `list[int]` | `[256, 256]` | Image size used for training as `[height, width]` | -| `keep_aspect_ratio` | `bool` | `True` | Whether to keep the aspect ratio while resizing | -| `train_rgb` | `bool` | `"RGB"` | Whether to train on RGB or BGR images | -| `normalize.active` | `bool` | `True` | Whether to use normalization | -| `normalize.params` | `dict` | `{}` | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | -| `augmentations` | `list[dict]` | `[]` | List of `Albumentations` augmentations | +| Key | Type | Default value | Description | +| ------------------- | ----------------------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `train_image_size` | `list[int]` | `[256, 256]` | Image size used for training as `[height, width]` | +| `keep_aspect_ratio` | `bool` | `True` | Whether to keep the aspect ratio while resizing | +| `color_space` | `Literal["RGB", "BGR"]` | `"RGB"` | Whether to train on RGB or BGR images | +| `normalize.active` | `bool` | `True` | Whether to use normalization | +| `normalize.params` | `dict` | `{}` | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | +| `augmentations` | `list[dict]` | `[]` | List of `Albumentations` augmentations | #### Augmentations @@ -306,7 +306,7 @@ trainer: # using YAML capture to reuse the image size train_image_size: [&height 384, &width 384] keep_aspect_ratio: true - train_rgb: true + color_space: "RGB" normalize: active: true augmentations: diff --git a/luxonis_train/attached_modules/losses/adaptive_detection_loss.py b/luxonis_train/attached_modules/losses/adaptive_detection_loss.py index a81d5a45..952d11c3 100644 --- a/luxonis_train/attached_modules/losses/adaptive_detection_loss.py +++ b/luxonis_train/attached_modules/losses/adaptive_detection_loss.py @@ -3,7 +3,7 @@ import torch import torch.nn.functional as F -from torch import Tensor, nn +from torch import Tensor, amp, nn from torchvision.ops import box_convert from luxonis_train.assigners import ATSSAssigner, TaskAlignedAssigner @@ -270,9 +270,7 @@ def forward( self.alpha * pred_score.pow(self.gamma) * (1 - label) + target_score * label ) - with torch.amp.autocast( - device_type=pred_score.device.type, enabled=False - ): + with amp.autocast(device_type=pred_score.device.type, enabled=False): ce_loss = F.binary_cross_entropy( pred_score.float(), target_score.float(), reduction="none" ) diff --git a/luxonis_train/callbacks/__init__.py b/luxonis_train/callbacks/__init__.py index 7bea71a9..d374916c 100644 --- a/luxonis_train/callbacks/__init__.py +++ b/luxonis_train/callbacks/__init__.py @@ -28,18 +28,18 @@ from .training_manager import TrainingManager from .upload_checkpoint import UploadCheckpoint -CALLBACKS.register_module(module=EarlyStopping) -CALLBACKS.register_module(module=LearningRateMonitor) -CALLBACKS.register_module(module=ModelCheckpoint) -CALLBACKS.register_module(module=RichModelSummary) -CALLBACKS.register_module(module=DeviceStatsMonitor) -CALLBACKS.register_module(module=GradientAccumulationScheduler) -CALLBACKS.register_module(module=StochasticWeightAveraging) -CALLBACKS.register_module(module=Timer) -CALLBACKS.register_module(module=ModelPruning) -CALLBACKS.register_module(module=GradCamCallback) -CALLBACKS.register_module(module=EMACallback) -CALLBACKS.register_module(module=TrainingManager) +CALLBACKS.register(module=EarlyStopping) +CALLBACKS.register(module=LearningRateMonitor) +CALLBACKS.register(module=ModelCheckpoint) +CALLBACKS.register(module=RichModelSummary) +CALLBACKS.register(module=DeviceStatsMonitor) +CALLBACKS.register(module=GradientAccumulationScheduler) +CALLBACKS.register(module=StochasticWeightAveraging) +CALLBACKS.register(module=Timer) +CALLBACKS.register(module=ModelPruning) +CALLBACKS.register(module=GradCamCallback) +CALLBACKS.register(module=EMACallback) +CALLBACKS.register(module=TrainingManager) __all__ = [ diff --git a/luxonis_train/callbacks/archive_on_train_end.py b/luxonis_train/callbacks/archive_on_train_end.py index 0ed69bb5..67e27ab7 100644 --- a/luxonis_train/callbacks/archive_on_train_end.py +++ b/luxonis_train/callbacks/archive_on_train_end.py @@ -10,7 +10,7 @@ logger = logging.getLogger(__name__) -@CALLBACKS.register_module() +@CALLBACKS.register() class ArchiveOnTrainEnd(NeedsCheckpoint): def on_train_end( self, diff --git a/luxonis_train/callbacks/ema.py b/luxonis_train/callbacks/ema.py index 20c01c04..e76e8c1a 100644 --- a/luxonis_train/callbacks/ema.py +++ b/luxonis_train/callbacks/ema.py @@ -3,10 +3,9 @@ from copy import deepcopy from typing import Any -import pytorch_lightning as pl +import lightning.pytorch as pl import torch -from pytorch_lightning.callbacks import Callback -from pytorch_lightning.utilities.types import STEP_OUTPUT +from lightning.pytorch.utilities.types import STEP_OUTPUT from torch import nn logger = logging.getLogger(__name__) @@ -101,7 +100,7 @@ def update(self, model: pl.LightningModule) -> None: ) -class EMACallback(Callback): +class EMACallback(pl.Callback): """Callback that updates the stored parameters using a moving average.""" diff --git a/luxonis_train/callbacks/gpu_stats_monitor.py b/luxonis_train/callbacks/gpu_stats_monitor.py index a189ed3f..c8774aa6 100644 --- a/luxonis_train/callbacks/gpu_stats_monitor.py +++ b/luxonis_train/callbacks/gpu_stats_monitor.py @@ -25,20 +25,20 @@ import time from typing import Any, Dict, List, Optional, Tuple -import pytorch_lightning as pl +import lightning.pytorch as pl import torch from lightning.pytorch.accelerators.cuda import CUDAAccelerator +from lightning.pytorch.utilities import rank_zero_only +from lightning.pytorch.utilities.parsing import AttributeDict +from lightning.pytorch.utilities.types import STEP_OUTPUT from lightning_fabric.utilities.exceptions import ( MisconfigurationException, # noqa: F401 ) -from pytorch_lightning.utilities import rank_zero_only -from pytorch_lightning.utilities.parsing import AttributeDict -from pytorch_lightning.utilities.types import STEP_OUTPUT from luxonis_train.utils.registry import CALLBACKS -@CALLBACKS.register_module() +@CALLBACKS.register() class GPUStatsMonitor(pl.Callback): def __init__( self, diff --git a/luxonis_train/callbacks/gradcam_visializer.py b/luxonis_train/callbacks/gradcam_visializer.py index 28863502..1d9616f2 100644 --- a/luxonis_train/callbacks/gradcam_visializer.py +++ b/luxonis_train/callbacks/gradcam_visializer.py @@ -1,16 +1,16 @@ import logging from typing import Any, Union +import lightning.pytorch as pl import numpy as np -import pytorch_lightning as pl import torch +from lightning.pytorch.utilities.types import STEP_OUTPUT from pytorch_grad_cam import HiResCAM from pytorch_grad_cam.utils.image import show_cam_on_image from pytorch_grad_cam.utils.model_targets import ( ClassifierOutputTarget, SemanticSegmentationTarget, ) -from pytorch_lightning.utilities.types import STEP_OUTPUT from luxonis_train.attached_modules.visualizers import ( get_denormalized_images, diff --git a/luxonis_train/callbacks/luxonis_progress_bar.py b/luxonis_train/callbacks/luxonis_progress_bar.py index b8bf6512..20665ced 100644 --- a/luxonis_train/callbacks/luxonis_progress_bar.py +++ b/luxonis_train/callbacks/luxonis_progress_bar.py @@ -46,7 +46,7 @@ def print_results( ... -@CALLBACKS.register_module() +@CALLBACKS.register() class LuxonisTQDMProgressBar(TQDMProgressBar, BaseLuxonisProgressBar): """Custom text progress bar based on TQDMProgressBar from Pytorch Lightning.""" @@ -104,7 +104,7 @@ def print_results( self._rule() -@CALLBACKS.register_module() +@CALLBACKS.register() class LuxonisRichProgressBar(RichProgressBar, BaseLuxonisProgressBar): """Custom rich text progress bar based on RichProgressBar from Pytorch Lightning.""" diff --git a/luxonis_train/callbacks/metadata_logger.py b/luxonis_train/callbacks/metadata_logger.py index 997ccbcd..0d7e3905 100644 --- a/luxonis_train/callbacks/metadata_logger.py +++ b/luxonis_train/callbacks/metadata_logger.py @@ -10,7 +10,7 @@ from luxonis_train.utils.registry import CALLBACKS -@CALLBACKS.register_module() +@CALLBACKS.register() class MetadataLogger(pl.Callback): def __init__(self, hyperparams: list[str]): """Callback that logs training metadata. diff --git a/luxonis_train/callbacks/test_on_train_end.py b/luxonis_train/callbacks/test_on_train_end.py index a60a16dd..9a437ff4 100644 --- a/luxonis_train/callbacks/test_on_train_end.py +++ b/luxonis_train/callbacks/test_on_train_end.py @@ -5,7 +5,7 @@ from luxonis_train.utils.registry import CALLBACKS -@CALLBACKS.register_module() +@CALLBACKS.register() class TestOnTrainEnd(pl.Callback): """Callback to perform a test run at the end of the training.""" diff --git a/luxonis_train/callbacks/training_manager.py b/luxonis_train/callbacks/training_manager.py index 9131fa84..d9cc7002 100644 --- a/luxonis_train/callbacks/training_manager.py +++ b/luxonis_train/callbacks/training_manager.py @@ -1,4 +1,4 @@ -import pytorch_lightning as pl +import lightning.pytorch as pl from luxonis_train.strategies.base_strategy import BaseTrainingStrategy diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index 3d3690e0..a73737cb 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -4,6 +4,7 @@ from typing import Annotated, Any, Literal, NamedTuple, TypeAlias from luxonis_ml.enums import DatasetType +from luxonis_ml.typing import ConfigItem from luxonis_ml.utils import ( BaseModelExtraForbid, Environ, @@ -318,10 +319,20 @@ class PreprocessingConfig(BaseModelExtraForbid): ImageSize, Field(default=[256, 256], min_length=2, max_length=2) ] = ImageSize(256, 256) keep_aspect_ratio: bool = True - train_rgb: bool = True + color_space: Literal["RGB", "BGR"] = "RGB" normalize: NormalizeAugmentationConfig = NormalizeAugmentationConfig() augmentations: list[AugmentationConfig] = [] + @model_validator(mode="before") + @classmethod + def validate_train_rgb(cls, data: dict[str, Any]) -> dict[str, Any]: + if "train_rgb" in data: + warnings.warn( + "Field `train_rgb` is deprecated. Use `color_space` instead." + ) + data["color_space"] = "RGB" if data.pop("train_rgb") else "BGR" + return data + @model_validator(mode="after") def check_normalize(self) -> Self: if self.normalize.active: @@ -332,13 +343,17 @@ def check_normalize(self) -> Self: ) return self - def get_active_augmentations(self) -> list[AugmentationConfig]: + def get_active_augmentations(self) -> list[ConfigItem]: """Returns list of augmentations that are active. @rtype: list[AugmentationConfig] @return: Filtered list of active augmentation configs """ - return [aug for aug in self.augmentations if aug.active] + return [ + ConfigItem(name=aug.name, params=aug.params) + for aug in self.augmentations + if aug.active + ] class CallbackConfig(BaseModelExtraForbid): diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py index a0529030..2ae78162 100644 --- a/luxonis_train/core/core.py +++ b/luxonis_train/core/core.py @@ -75,6 +75,8 @@ def __init__( else: self.cfg = Config.get_config(cfg, opts) + self.cfg_preprocessing = self.cfg.trainer.preprocessing + rich.traceback.install(suppress=[pl, torch], show_locals=False) self.tracker = LuxonisTrackerPL( @@ -126,16 +128,11 @@ def __init__( "test": self.cfg.loader.test_view, }[view], image_source=self.cfg.loader.image_source, - height=self.cfg.trainer.preprocessing.train_image_size.height, - width=self.cfg.trainer.preprocessing.train_image_size.width, - augmentation_config=[ - i.model_dump() - for i in self.cfg.trainer.preprocessing.get_active_augmentations() - ], - out_image_format="RGB" - if self.cfg.trainer.preprocessing.train_rgb - else "BGR", - keep_aspect_ratio=self.cfg.trainer.preprocessing.keep_aspect_ratio, + height=self.cfg_preprocessing.train_image_size.height, + width=self.cfg_preprocessing.train_image_size.width, + augmentation_config=self.cfg_preprocessing.get_active_augmentations(), + color_space=self.cfg_preprocessing.color_space, + keep_aspect_ratio=self.cfg_preprocessing.keep_aspect_ratio, **self.cfg.loader.params, ) @@ -598,9 +595,7 @@ def _objective(trial: optuna.trial.Trial) -> float: "You have to specify the `tuner` section in config." ) - all_augs = [ - a.name for a in self.cfg.trainer.preprocessing.augmentations - ] + all_augs = [a.name for a in self.cfg_preprocessing.augmentations] rank = rank_zero_only.rank cfg_tracker = self.cfg.tracker tracker_params = cfg_tracker.model_dump() @@ -716,15 +711,9 @@ def _mult(lst: list[float | int]) -> list[float]: return [round(x * 255.0, 5) for x in lst] preprocessing = { # TODO: keep preprocessing same for each input? - "mean": _mult( - self.cfg.trainer.preprocessing.normalize.params["mean"] - ), - "scale": _mult( - self.cfg.trainer.preprocessing.normalize.params["std"] - ), - "dai_type": "RGB888p" - if self.cfg.trainer.preprocessing.train_rgb - else "BGR888p", + "mean": _mult(self.cfg_preprocessing.normalize.params["mean"]), + "scale": _mult(self.cfg_preprocessing.normalize.params["std"]), + "dai_type": f"{self.cfg_preprocessing.color_space}888p", } inputs_dict = get_inputs(path) diff --git a/luxonis_train/core/utils/infer_utils.py b/luxonis_train/core/utils/infer_utils.py index 26e994e5..22cd7521 100644 --- a/luxonis_train/core/utils/infer_utils.py +++ b/luxonis_train/core/utils/infer_utils.py @@ -92,7 +92,7 @@ def infer_from_video( ret, frame = cap.read() if not ret: # pragma: no cover break - if model.cfg.trainer.preprocessing.train_rgb: + if model.cfg.trainer.preprocessing.color_space == "RGB": frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # TODO: batched inference diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index a9998e4b..19fc4dc6 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -1,5 +1,7 @@ from abc import ABC, abstractmethod +from typing import Any, Literal +from luxonis_ml.typing import ConfigItem from luxonis_ml.utils.registry import AutoRegisterMeta from torch import Size, Tensor from torch.utils.data import Dataset @@ -16,28 +18,136 @@ class BaseLoaderTorch( register=False, registry=LOADERS, ): - """Base abstract loader class that enforces LuxonisLoaderTorchOutput - output label structure.""" - def __init__( self, - view: str | list[str], - image_source: str | None = None, + view: list[str], + height: int, + width: int, + augmentation_engine: str = "albumentations", + augmentation_config: list[ConfigItem] | None = None, + image_source: str = "default", + keep_aspect_ratio: bool = True, + color_space: Literal["RGB", "BGR"] = "RGB", ): - self.view = view if isinstance(view, list) else [view] + """Base abstract loader class that enforces + LuxonisLoaderTorchOutput output label structure. + + @type view: list[str] + @param view: List of view names. Usually contains only one element, + e.g. C{["train"]} or C{["test"]}. However, more complex datasets + can make use of multiple views, e.g. C{["train_synthetic", + "train_real"]} + + @type height: int + @param height: Height of the output image. + + @type width: int + @param width: Width of the output image. + + @type augmentation_engine: str + @param augmentation_engine: Name of the augmentation engine. Can + be used to enable swapping between different augmentation engines or making use of pre-defined engines, e.g. C{AlbumentationsEngine}. + + @type augmentation_config: list[ConfigItem] | None + @param augmentation_config: List of augmentation configurations. + Individual configurations are in the form of:: + + class ConfigItem: + name: str + params: dict[str, JsonValue] + + Where C{name} is the name of the augmentation and C{params} is a + dictionary of its parameters. + + Example:: + + ConfigItem( + name="HorizontalFlip", + params={"p": 0.5}, + ) + + @type image_source: str + @param image_source: Name of the input image group. This can be used for datasets with multiple image sources, e.g. left and right cameras or RGB and depth images. Irrelevant for datasets with only one image source. + + @type keep_aspect_ratio: bool + @param keep_aspect_ratio: Whether to keep the aspect ratio of the output image after resizing. + + @type color_space: Literal["RGB", "BGR"] + @param color_space: Color space of the output image. + """ + self._view = view self._image_source = image_source + self._augmentation_engine = augmentation_engine + self._augmentation_config = augmentation_config + self._height = height + self._width = width + self._keep_aspect_ratio = keep_aspect_ratio + self._color_space = color_space @property def image_source(self) -> str: """Name of the input image group. - Example: C{"image"} + @type: str + """ + return self._getter_check_none("image_source") + + @property + def view(self) -> list[str]: + """List of view names. + + @type: list[str] + """ + return self._view + + @property + def augmentation_engine(self) -> str: + """Name of the augmentation engine. @type: str """ - if self._image_source is None: - raise ValueError("image_source is not set") - return self._image_source + return self._getter_check_none("augmentation_engine") + + @property + def augmentation_config(self) -> list[ConfigItem]: + """List of augmentation configurations. + + @type: list[ConfigItem] + """ + return self._getter_check_none("augmentation_config") + + @property + def height(self) -> int: + """Height of the output image. + + @type: int + """ + return self._getter_check_none("height") + + @property + def width(self) -> int: + """Width of the output image. + + @type: int + """ + return self._getter_check_none("width") + + @property + def keep_aspect_ratio(self) -> bool: + """Whether to keep the aspect ratio of the output image after + resizing. + + @type: bool + """ + return self._getter_check_none("keep_aspect_ratio") + + @property + def color_space(self) -> str: + """Color space of the output image. + + @type: str + """ + return self._getter_check_none("color_space") @property @abstractmethod @@ -124,3 +234,21 @@ def get_n_keypoints(self) -> dict[str, int] | None: definitions. """ return None + + def _getter_check_none( + self, + attribute: Literal[ + "view", + "image_source", + "augmentation_engine", + "augmentation_config", + "height", + "width", + "keep_aspect_ratio", + "color_space", + ], + ) -> Any: + value = getattr(self, f"_{attribute}") + if value is None: + raise ValueError(f"{attribute} is not set") + return value diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index d44c92ea..7514ff91 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -110,7 +110,9 @@ def __init__( @param out_image_format: The format of the output images. Defaults to C{"RGB"}. """ - super().__init__(view=view, **kwargs) + super().__init__( + view=view if isinstance(view, list) else [view], **kwargs + ) if dataset_dir is not None: self.dataset = self._parse_dataset( dataset_dir, dataset_name, dataset_type, delete_existing diff --git a/luxonis_train/optimizers/optimizers.py b/luxonis_train/optimizers/optimizers.py index c2a4bf12..42f63f8a 100644 --- a/luxonis_train/optimizers/optimizers.py +++ b/luxonis_train/optimizers/optimizers.py @@ -16,4 +16,4 @@ optim.RMSprop, optim.SGD, ]: - OPTIMIZERS.register_module(module=optimizer) + OPTIMIZERS.register(module=optimizer) diff --git a/luxonis_train/schedulers/schedulers.py b/luxonis_train/schedulers/schedulers.py index 488a7498..0497a3c3 100644 --- a/luxonis_train/schedulers/schedulers.py +++ b/luxonis_train/schedulers/schedulers.py @@ -19,4 +19,4 @@ lr_scheduler.OneCycleLR, lr_scheduler.CosineAnnealingWarmRestarts, ]: - SCHEDULERS.register_module(module=scheduler) + SCHEDULERS.register(module=scheduler) diff --git a/luxonis_train/strategies/base_strategy.py b/luxonis_train/strategies/base_strategy.py index 8de6386d..5b812ebf 100644 --- a/luxonis_train/strategies/base_strategy.py +++ b/luxonis_train/strategies/base_strategy.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod -import pytorch_lightning as pl +import lightning.pytorch as pl from luxonis_ml.utils.registry import AutoRegisterMeta from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler diff --git a/luxonis_train/strategies/triple_lr_sgd.py b/luxonis_train/strategies/triple_lr_sgd.py index 33f7dfe3..b76b5ee8 100644 --- a/luxonis_train/strategies/triple_lr_sgd.py +++ b/luxonis_train/strategies/triple_lr_sgd.py @@ -1,8 +1,8 @@ # strategies/triple_lr_sgd.py import math +import lightning.pytorch as pl import numpy as np -import pytorch_lightning as pl import torch from torch.optim import SGD from torch.optim.lr_scheduler import LambdaLR diff --git a/tests/unittests/test_callbacks/test_ema.py b/tests/unittests/test_callbacks/test_ema.py index 0780e783..d117eb88 100644 --- a/tests/unittests/test_callbacks/test_ema.py +++ b/tests/unittests/test_callbacks/test_ema.py @@ -2,7 +2,7 @@ import pytest import torch -from pytorch_lightning import LightningModule, Trainer +from lightning.pytorch import LightningModule, Trainer from luxonis_train.callbacks.ema import EMACallback, ModelEma From eef219a7132f0d927c9af759d903f9bf3dfaa2ae Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 16 Jan 2025 18:26:21 -0500 Subject: [PATCH 28/57] replaced deprecated `register_module` --- luxonis_train/callbacks/export_on_train_end.py | 2 +- luxonis_train/callbacks/upload_checkpoint.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/luxonis_train/callbacks/export_on_train_end.py b/luxonis_train/callbacks/export_on_train_end.py index 80d2a648..195524c7 100644 --- a/luxonis_train/callbacks/export_on_train_end.py +++ b/luxonis_train/callbacks/export_on_train_end.py @@ -10,7 +10,7 @@ logger = logging.getLogger(__name__) -@CALLBACKS.register_module() +@CALLBACKS.register() class ExportOnTrainEnd(NeedsCheckpoint): def on_train_end( self, diff --git a/luxonis_train/callbacks/upload_checkpoint.py b/luxonis_train/callbacks/upload_checkpoint.py index b9753e94..0954737a 100644 --- a/luxonis_train/callbacks/upload_checkpoint.py +++ b/luxonis_train/callbacks/upload_checkpoint.py @@ -10,7 +10,7 @@ from luxonis_train.utils.registry import CALLBACKS -@CALLBACKS.register_module() +@CALLBACKS.register() class UploadCheckpoint(pl.Callback): """Callback that uploads best checkpoint based on the validation loss.""" From 0379b2a21842556eb6aaabb204dabf9a75b48f35 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 16 Jan 2025 18:26:42 -0500 Subject: [PATCH 29/57] removed init arguments --- luxonis_train/loaders/luxonis_loader_torch.py | 110 ++++++------------ 1 file changed, 34 insertions(+), 76 deletions(-) diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index 7514ff91..be31e28e 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -1,6 +1,5 @@ import logging -from pathlib import Path -from typing import Any, Literal +from typing import Literal import numpy as np import torch @@ -31,88 +30,45 @@ def __init__( bucket_type: Literal["internal", "external"] = "internal", bucket_storage: Literal["local", "s3", "gcs", "azure"] = "local", delete_existing: bool = True, - view: str | list[str] = "train", - augmentation_engine: str - | Literal["albumentations"] = "albumentations", - augmentation_config: Path | str | list[dict[str, Any]] | None = None, - height: int | None = None, - width: int | None = None, - keep_aspect_ratio: bool = True, - out_image_format: Literal["RGB", "BGR"] = "RGB", **kwargs, ): """Torch-compatible loader for Luxonis datasets. - Can either use an already existing dataset or parse a new one from a directory. + Can either use an already existing dataset or parse a new one + from a directory. @type dataset_name: str | None - @param dataset_name: Name of the dataset to load. If not provided, the - C{dataset_dir} argument must be provided instead. If both C{dataset_dir} and - C{dataset_name} are provided, the dataset will be parsed from the directory - and saved with the provided name. + @param dataset_name: Name of the dataset to load. If not + provided, the C{dataset_dir} argument must be provided + instead. If both C{dataset_dir} and C{dataset_name} are + provided, the dataset will be parsed from the directory and + saved with the provided name. @type dataset_dir: str | None - @param dataset_dir: Path to the dataset directory. It can be either a local path - or a URL. The data can be in a zip file. If not provided, C{dataset_name} of - an existing dataset must be provided. + @param dataset_dir: Path to the dataset directory. It can be + either a local path or a URL. The data can be in a zip file. + If not provided, C{dataset_name} of an existing dataset must + be provided. @type dataset_type: str | None - @param dataset_type: Type of the dataset. Only relevant when C{dataset_dir} is - provided. If not provided, the type will be inferred from the directory - structure. + @param dataset_type: Type of the dataset. Only relevant when + C{dataset_dir} is provided. If not provided, the type will + be inferred from the directory structure. @type team_id: str | None @param team_id: Optional unique team identifier for the cloud. @type bucket_type: Literal["internal", "external"] - @param bucket_type: Type of the bucket. Only relevant for remote datasets. - Defaults to 'internal'. + @param bucket_type: Type of the bucket. Only relevant for remote + datasets. Defaults to 'internal'. @type bucket_storage: Literal["local", "s3", "gcs", "azure"] - @param bucket_storage: Type of the bucket storage. Defaults to 'local'. + @param bucket_storage: Type of the bucket storage. Defaults to + 'local'. @type delete_existing: bool - @param delete_existing: Only relevant when C{dataset_dir} is provided. By - default, the dataset is parsed again every time the loader is created - because the underlying data might have changed. If C{delete_existing} is set - to C{False} and a dataset of the same name already exists, the existing + @param delete_existing: Only relevant when C{dataset_dir} is + provided. By default, the dataset is parsed again every time + the loader is created because the underlying data might have + changed. If C{delete_existing} is set to C{False} and a + dataset of the same name already exists, the existing dataset will be used instead of re-parsing the data. - @type view: str | list[str] - @param view: A single split or a list of splits that will be used to create a - view of the dataset. Each split is a string that represents a subset of the - dataset. The available splits depend on the dataset, but usually include - 'train', 'val', and 'test'. Defaults to 'train'. - @type augmentation_engine: Union[Literal["albumentations"], str] - @param augmentation_engine: The augmentation engine to use. - Defaults to C{"albumentations"}. - @type augmentation_config: Optional[Union[List[Dict[str, Any]], - PathType]] - @param augmentation_config: The configuration for the - augmentations. This can be either a list of C{Dict[str, Any]} or - a path to a configuration file. - The config member is a dictionary with two keys: C{name} and - C{params}. C{name} is the name of the augmentation to - instantiate and C{params} is an optional dictionary - of parameters to pass to the augmentation. - - Example:: - - [ - {"name": "HorizontalFlip", "params": {"p": 0.5}}, - {"name": "RandomBrightnessContrast", "params": {"p": 0.1}}, - {"name": "Defocus"} - ] - - @type height: Optional[int] - @param height: The height of the output images. Defaults to - C{None}. - @type width: Optional[int] - @param width: The width of the output images. Defaults to - C{None}. - @type keep_aspect_ratio: bool - @param keep_aspect_ratio: Whether to keep the aspect ratio of the - images. Defaults to C{True}. - @type out_image_format: Literal["RGB", "BGR"] - @param out_image_format: The format of the output images. Defaults - to C{"RGB"}. """ - super().__init__( - view=view if isinstance(view, list) else [view], **kwargs - ) + super().__init__(**kwargs) if dataset_dir is not None: self.dataset = self._parse_dataset( dataset_dir, dataset_name, dataset_type, delete_existing @@ -130,13 +86,15 @@ def __init__( ) self.loader = LuxonisLoader( dataset=self.dataset, - view=view, - augmentation_engine=augmentation_engine, - augmentation_config=augmentation_config, - height=height, - width=width, - keep_aspect_ratio=keep_aspect_ratio, - out_image_format=out_image_format, + view=self.view, + augmentation_engine=self.augmentation_engine, + augmentation_config=[ + aug.model_dump() for aug in self.augmentation_config + ], + height=self.height, + width=self.width, + keep_aspect_ratio=self.keep_aspect_ratio, + out_image_format=self.color_space, ) def __len__(self) -> int: From d6344efa3cad6cb5cd19a2c8c18b106692ceb52d Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 16 Jan 2025 19:18:36 -0500 Subject: [PATCH 30/57] added missing types --- luxonis_train/assigners/tal_assigner.py | 8 ++++++-- .../losses/adaptive_detection_loss.py | 6 +++--- .../losses/efficient_keypoint_bbox_loss.py | 4 ++-- .../losses/reconstruction_segmentation_loss.py | 6 +++--- .../metrics/mean_average_precision.py | 2 +- luxonis_train/callbacks/training_manager.py | 2 +- luxonis_train/config/config.py | 6 ++++-- luxonis_train/core/core.py | 4 ++-- luxonis_train/core/utils/export_utils.py | 3 ++- luxonis_train/core/utils/infer_utils.py | 7 ++++--- luxonis_train/models/luxonis_lightning.py | 17 ++++++++--------- luxonis_train/nodes/backbones/ddrnet/ddrnet.py | 2 +- .../backbones/efficientrep/efficientrep.py | 2 +- .../nodes/backbones/mobileone/blocks.py | 2 +- .../nodes/backbones/mobileone/mobileone.py | 6 ++++-- .../nodes/backbones/recsubnet/blocks.py | 2 +- luxonis_train/nodes/base_node.py | 2 +- .../nodes/heads/efficient_bbox_head.py | 4 ++-- .../nodes/necks/reppan_neck/reppan_neck.py | 2 +- luxonis_train/strategies/base_strategy.py | 6 ++---- luxonis_train/strategies/triple_lr_sgd.py | 6 +++--- 21 files changed, 53 insertions(+), 46 deletions(-) diff --git a/luxonis_train/assigners/tal_assigner.py b/luxonis_train/assigners/tal_assigner.py index c9435afa..b289fbd6 100644 --- a/luxonis_train/assigners/tal_assigner.py +++ b/luxonis_train/assigners/tal_assigner.py @@ -143,7 +143,7 @@ def _get_alignment_metric( pred_bboxes: Tensor, gt_labels: Tensor, gt_bboxes: Tensor, - ): + ) -> tuple[Tensor, Tensor]: """Calculates anchor alignment metric and IoU between GTs and predicted bboxes. @@ -155,7 +155,11 @@ def _get_alignment_metric( @param gt_labels: Initial GT labels [bs, n_max_boxes, 1] @type gt_bboxes: Tensor @param gt_bboxes: Initial GT bboxes [bs, n_max_boxes, 4] + @rtype: tuple[Tensor, Tensor] + @return: Anchor alignment metric and IoU between GTs and + predicted bboxes. """ + pred_scores = pred_scores.permute(0, 2, 1) gt_labels = gt_labels.to(torch.long) ind = torch.zeros([2, self.bs, self.n_max_boxes], dtype=torch.long) @@ -175,7 +179,7 @@ def _select_topk_candidates( metrics: Tensor, largest: bool = True, topk_mask: Tensor | None = None, - ): + ) -> Tensor: """Selects k anchors based on provided metrics tensor. @type metrics: Tensor diff --git a/luxonis_train/attached_modules/losses/adaptive_detection_loss.py b/luxonis_train/attached_modules/losses/adaptive_detection_loss.py index 952d11c3..6a7f57f2 100644 --- a/luxonis_train/attached_modules/losses/adaptive_detection_loss.py +++ b/luxonis_train/attached_modules/losses/adaptive_detection_loss.py @@ -132,7 +132,7 @@ def forward( assigned_labels: Tensor, assigned_scores: Tensor, mask_positive: Tensor, - ): + ) -> tuple[Tensor, dict[str, Tensor]]: one_hot_label = F.one_hot(assigned_labels.long(), self.n_classes + 1)[ ..., :-1 ] @@ -161,7 +161,7 @@ def forward( return loss, sub_losses - def _init_parameters(self, features: list[Tensor]): + def _init_parameters(self, features: list[Tensor]) -> None: if not hasattr(self, "gt_bboxes_scale"): self.gt_bboxes_scale = torch.tensor( [ @@ -235,7 +235,7 @@ def _preprocess_bbox_target( out_target[..., 1:] = box_convert(scaled_target, "xywh", "xyxy") return out_target - def _log_assigner_change(self): + def _log_assigner_change(self) -> None: if self._logged_assigner_change: return diff --git a/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py b/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py index d9a191e9..98630742 100644 --- a/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py +++ b/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py @@ -187,7 +187,7 @@ def forward( gt_kpts: Tensor, pred_kpts: Tensor, area: Tensor, - ): + ) -> tuple[Tensor, dict[str, Tensor]]: device = pred_bboxes.device sigmas = self.sigmas.to(device) d = (gt_kpts[..., 0] - pred_kpts[..., 0]).pow(2) + ( @@ -272,7 +272,7 @@ def dist2kpts_noscale(self, anchor_points: Tensor, kpts: Tensor) -> Tensor: adj_kpts[..., 1] += y_adj return adj_kpts - def _init_parameters(self, features: list[Tensor]): + def _init_parameters(self, features: list[Tensor]) -> None: device = features[0].device super()._init_parameters(features) self.gt_kpts_scale = torch.tensor( diff --git a/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py b/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py index a5b50f2b..6ba109d0 100644 --- a/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py +++ b/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py @@ -62,7 +62,7 @@ def prepare( def forward( self, orig: Tensor, recon: Tensor, seg_out: Tensor, an_mask: Tensor - ): + ) -> tuple[Tensor, dict[str, Tensor]]: l2 = self.loss_l2(recon, orig) ssim = self.loss_ssim(recon, orig) focal = self.loss_focal(seg_out, an_mask) @@ -142,8 +142,8 @@ def ssim( img2: Tensor, window_size: int = 11, window: Tensor | None = None, - size_average=True, - val_range=None, + size_average: bool = True, + val_range: float | None = None, ) -> Tensor: if val_range is None: if torch.max(img1) > 128: diff --git a/luxonis_train/attached_modules/metrics/mean_average_precision.py b/luxonis_train/attached_modules/metrics/mean_average_precision.py index d4731988..c082ee39 100644 --- a/luxonis_train/attached_modules/metrics/mean_average_precision.py +++ b/luxonis_train/attached_modules/metrics/mean_average_precision.py @@ -31,7 +31,7 @@ def update( self, outputs: list[dict[str, Tensor]], labels: list[dict[str, Tensor]], - ): + ) -> None: self.metric.update(outputs, labels) def prepare( diff --git a/luxonis_train/callbacks/training_manager.py b/luxonis_train/callbacks/training_manager.py index d9cc7002..390f49b6 100644 --- a/luxonis_train/callbacks/training_manager.py +++ b/luxonis_train/callbacks/training_manager.py @@ -15,7 +15,7 @@ def __init__(self, strategy: BaseTrainingStrategy | None = None): def on_after_backward( self, trainer: pl.Trainer, pl_module: pl.LightningModule - ): + ) -> None: """PyTorch Lightning hook that is called after the backward pass. diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index a73737cb..d44ac480 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -498,7 +498,9 @@ class ExportConfig(ArchiveConfig): @model_validator(mode="after") def check_values(self) -> Self: - def pad_values(values: float | list[float] | None): + def pad_values( + values: float | list[float] | None, + ) -> list[float] | None: if values is None: return None if isinstance(values, float): @@ -644,7 +646,7 @@ def is_acyclic(graph: dict[str, list[str]]) -> bool: """ graph = graph.copy() - def dfs(node: str, visited: set[str], recursion_stack: set[str]): + def dfs(node: str, visited: set[str], recursion_stack: set[str]) -> bool: visited.add(node) recursion_stack.add(node) diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py index 2ae78162..5f469590 100644 --- a/luxonis_train/core/core.py +++ b/luxonis_train/core/core.py @@ -208,7 +208,7 @@ def __init__( self._exported_models: dict[str, Path] = {} - def _train(self, resume: str | None, *args, **kwargs): + def _train(self, resume: str | None, *args, **kwargs) -> None: status = "success" try: self.pl_trainer.fit(*args, ckpt_path=resume, **kwargs) @@ -245,7 +245,7 @@ def train( LuxonisFileSystem.download(resume_weights, self.run_save_dir) ) - def graceful_exit(signum: int, _): # pragma: no cover + def graceful_exit(signum: int, _: Any) -> None: # pragma: no cover logger.info( f"{signal.Signals(signum).name} received, stopping training..." ) diff --git a/luxonis_train/core/utils/export_utils.py b/luxonis_train/core/utils/export_utils.py index 25e1a3ff..7190c889 100644 --- a/luxonis_train/core/utils/export_utils.py +++ b/luxonis_train/core/utils/export_utils.py @@ -1,4 +1,5 @@ import logging +from collections.abc import Generator from contextlib import contextmanager from pathlib import Path @@ -12,7 +13,7 @@ def replace_weights( module: "luxonis_train.models.LuxonisLightningModule", weights: str | Path | None = None, -): +) -> Generator[None, None, None]: old_weights = None if weights is not None: old_weights = module.state_dict() diff --git a/luxonis_train/core/utils/infer_utils.py b/luxonis_train/core/utils/infer_utils.py index 22cd7521..c4f9085f 100644 --- a/luxonis_train/core/utils/infer_utils.py +++ b/luxonis_train/core/utils/infer_utils.py @@ -7,12 +7,13 @@ import numpy as np import torch import torch.utils.data as torch_data -from luxonis_ml.data import LuxonisDataset +from luxonis_ml.data import DatasetIterator, LuxonisDataset from torch import Tensor import luxonis_train from luxonis_train.attached_modules.visualizers import get_denormalized_images from luxonis_train.loaders import LuxonisLoaderTorch +from luxonis_train.models.luxonis_output import LuxonisOutput IMAGE_FORMATS = { ".bmp", @@ -48,7 +49,7 @@ def process_visualizations( def prepare_and_infer_image( model: "luxonis_train.core.LuxonisModel", img: Tensor -): +) -> LuxonisOutput: """Prepares the image for inference and runs the model.""" img = model.loaders["val"].augment_test_image(img) # type: ignore @@ -196,7 +197,7 @@ def infer_from_directory( """ img_paths = list(img_paths) - def generator(): + def generator() -> DatasetIterator: for img_path in img_paths: yield { "file": img_path, diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index c478c126..9643a6f2 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -9,6 +9,7 @@ from lightning.pytorch.callbacks import ModelCheckpoint, RichModelSummary from lightning.pytorch.utilities import rank_zero_only # type: ignore from luxonis_ml.data import LuxonisDataset +from luxonis_ml.typing import ConfigItem from torch import Size, Tensor, nn import luxonis_train @@ -601,7 +602,7 @@ def export_onnx(self, save_path: str, **kwargs) -> list[str]: old_forward = self.forward - def export_forward(inputs) -> tuple[Tensor, ...]: + def export_forward(inputs: dict[str, Tensor]) -> tuple[Tensor, ...]: outputs = old_forward( inputs, None, @@ -904,14 +905,12 @@ def configure_optimizers( } optimizer = OPTIMIZERS.get(cfg_optimizer.name)(**optim_params) - def get_scheduler(scheduler_cfg, optimizer): - scheduler_class = SCHEDULERS.get( - scheduler_cfg["name"] - ) # For dictionary access - scheduler_params = scheduler_cfg["params"] | { - "optimizer": optimizer - } # Dictionary access for params - return scheduler_class(**scheduler_params) + def get_scheduler( + scheduler_cfg: ConfigItem, optimizer: torch.optim.Optimizer + ) -> torch.optim.lr_scheduler._LRScheduler: + scheduler_class = SCHEDULERS.get(scheduler_cfg.name) + scheduler_params = scheduler_cfg.params | {"optimizer": optimizer} + return scheduler_class(**scheduler_params) # type: ignore if cfg_scheduler.name == "SequentialLR": schedulers_list = [ diff --git a/luxonis_train/nodes/backbones/ddrnet/ddrnet.py b/luxonis_train/nodes/backbones/ddrnet/ddrnet.py index b029dfff..2698c26d 100644 --- a/luxonis_train/nodes/backbones/ddrnet/ddrnet.py +++ b/luxonis_train/nodes/backbones/ddrnet/ddrnet.py @@ -297,7 +297,7 @@ def forward(self, inputs: Tensor) -> list[Tensor]: else: return [x] - def init_params(self): + def init_params(self) -> None: for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_( diff --git a/luxonis_train/nodes/backbones/efficientrep/efficientrep.py b/luxonis_train/nodes/backbones/efficientrep/efficientrep.py index 76b0cf18..529cd816 100644 --- a/luxonis_train/nodes/backbones/efficientrep/efficientrep.py +++ b/luxonis_train/nodes/backbones/efficientrep/efficientrep.py @@ -140,7 +140,7 @@ def __init__( f"No checkpoint available for {self.name}, skipping." ) - def initialize_weights(self): + def initialize_weights(self) -> None: for m in self.modules(): if isinstance(m, nn.Conv2d): pass diff --git a/luxonis_train/nodes/backbones/mobileone/blocks.py b/luxonis_train/nodes/backbones/mobileone/blocks.py index a9006c7e..4b926038 100644 --- a/luxonis_train/nodes/backbones/mobileone/blocks.py +++ b/luxonis_train/nodes/backbones/mobileone/blocks.py @@ -133,7 +133,7 @@ def forward(self, inputs: Tensor) -> Tensor: return self.activation(self.se(out)) - def reparameterize(self): + def reparameterize(self) -> None: """Following works like U{RepVGG: Making VGG-style ConvNets Great Again } architecture used at training time to obtain a plain CNN-like structure diff --git a/luxonis_train/nodes/backbones/mobileone/mobileone.py b/luxonis_train/nodes/backbones/mobileone/mobileone.py index 1ed476cf..c2fe93c0 100644 --- a/luxonis_train/nodes/backbones/mobileone/mobileone.py +++ b/luxonis_train/nodes/backbones/mobileone/mobileone.py @@ -142,7 +142,9 @@ def set_export_mode(self, mode: bool = True) -> None: if hasattr(module, "reparameterize"): module.reparameterize() - def _make_stage(self, planes: int, n_blocks: int, n_se_blocks: int): + def _make_stage( + self, planes: int, n_blocks: int, n_se_blocks: int + ) -> nn.Sequential: """Build a stage of MobileOne model. @type planes: int @@ -161,7 +163,7 @@ def _make_stage(self, planes: int, n_blocks: int, n_se_blocks: int): use_se = False if n_se_blocks > n_blocks: raise ValueError( - "Number of SE blocks cannot " "exceed number of layers." + "Number of SE blocks cannot exceed number of layers." ) if ix >= (n_blocks - n_se_blocks): use_se = True diff --git a/luxonis_train/nodes/backbones/recsubnet/blocks.py b/luxonis_train/nodes/backbones/recsubnet/blocks.py index 0090a3ca..1d557aff 100644 --- a/luxonis_train/nodes/backbones/recsubnet/blocks.py +++ b/luxonis_train/nodes/backbones/recsubnet/blocks.py @@ -118,7 +118,7 @@ def __init__(self, input_channels: int, width: int) -> None: self.encoder_block2 = ConvBlock(width, int(width * 1.1)) self.pool2 = nn.MaxPool2d(2) - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: enc1 = self.encoder_block1(x) enc1_pool = self.pool1(enc1) enc2 = self.encoder_block2(enc1_pool) diff --git a/luxonis_train/nodes/base_node.py b/luxonis_train/nodes/base_node.py index 7dcdbcf3..0a9a208a 100644 --- a/luxonis_train/nodes/base_node.py +++ b/luxonis_train/nodes/base_node.py @@ -378,7 +378,7 @@ def in_width(self) -> int | list[int]: """ return self._get_nth_size(-1) - def load_checkpoint(self, path: str, strict: bool = True): + def load_checkpoint(self, path: str, strict: bool = True) -> None: """Loads checkpoint for the module. If path is url then it downloads it locally and stores it in cache. diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py index 8d1a55f4..95ebe1be 100644 --- a/luxonis_train/nodes/heads/efficient_bbox_head.py +++ b/luxonis_train/nodes/heads/efficient_bbox_head.py @@ -111,7 +111,7 @@ def __init__( f"No checkpoint available for {self.name}, skipping." ) - def initialize_weights(self): + def initialize_weights(self) -> None: for m in self.modules(): if isinstance(m, nn.Conv2d): pass @@ -201,7 +201,7 @@ def wrap( "distributions": [reg_tensor], } - def _fit_stride_to_n_heads(self): + def _fit_stride_to_n_heads(self) -> Tensor: """Returns correct stride for number of heads and attach index.""" stride = torch.tensor( diff --git a/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py b/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py index 383160e3..f7e0552e 100644 --- a/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py +++ b/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py @@ -180,7 +180,7 @@ def __init__( f"No checkpoint available for {self.name}, skipping." ) - def initialize_weights(self): + def initialize_weights(self) -> None: for m in self.modules(): if isinstance(m, nn.Conv2d): pass diff --git a/luxonis_train/strategies/base_strategy.py b/luxonis_train/strategies/base_strategy.py index 5b812ebf..09bc5392 100644 --- a/luxonis_train/strategies/base_strategy.py +++ b/luxonis_train/strategies/base_strategy.py @@ -20,9 +20,7 @@ def __init__(self, pl_module: pl.LightningModule): @abstractmethod def configure_optimizers( self, - ) -> tuple[list[Optimizer], list[LRScheduler]]: - pass + ) -> tuple[list[Optimizer], list[LRScheduler]]: ... @abstractmethod - def update_parameters(self, *args, **kwargs): - pass + def update_parameters(self, *args, **kwargs) -> None: ... diff --git a/luxonis_train/strategies/triple_lr_sgd.py b/luxonis_train/strategies/triple_lr_sgd.py index b76b5ee8..570c7bc7 100644 --- a/luxonis_train/strategies/triple_lr_sgd.py +++ b/luxonis_train/strategies/triple_lr_sgd.py @@ -51,7 +51,7 @@ def __init__( + 1 ) - def create_scheduler(self): + def create_scheduler(self) -> LambdaLR: scheduler = LambdaLR(self.optimizer, lr_lambda=self.lf) return scheduler @@ -103,7 +103,7 @@ def __init__(self, model: torch.nn.Module, params: dict) -> None: if params: self.params.update(params) - def create_optimizer(self): + def create_optimizer(self) -> torch.optim.Optimizer: batch_norm_weights, regular_weights, biases = [], [], [] for module in self.model.modules(): @@ -166,6 +166,6 @@ def __init__(self, pl_module: pl.LightningModule, params: dict): def configure_optimizers(self) -> tuple[list[Optimizer], list[LambdaLR]]: return [self.optimizer], [self.scheduler.create_scheduler()] - def update_parameters(self, *args, **kwargs): + def update_parameters(self, *args, **kwargs) -> None: current_epoch = self.model.current_epoch self.scheduler.update_learning_rate(current_epoch) From 44adfcb8aa6559f57aa1ce39879e752dc345a86e Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 16 Jan 2025 21:41:14 -0500 Subject: [PATCH 31/57] fixed anomaly detection --- .../reconstruction_segmentation_loss.py | 2 +- luxonis_train/core/utils/export_utils.py | 2 +- luxonis_train/loaders/base_loader.py | 50 +++++++++++-- luxonis_train/loaders/luxonis_loader_torch.py | 22 +++--- .../loaders/luxonis_perlin_loader_torch.py | 64 +++++++++++------ luxonis_train/loaders/perlin.py | 71 ++++++------------- luxonis_train/utils/general.py | 4 +- 7 files changed, 128 insertions(+), 87 deletions(-) diff --git a/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py b/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py index 6ba109d0..1e3ff449 100644 --- a/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py +++ b/luxonis_train/attached_modules/losses/reconstruction_segmentation_loss.py @@ -55,7 +55,7 @@ def prepare( ) -> tuple[Tensor, Tensor, Tensor, Tensor]: recon = self.get_input_tensors(inputs, "reconstructed")[0] seg_out = self.get_input_tensors(inputs)[0] - an_mask = self.get_label(labels) + an_mask = labels[f"{self.node.task_name}/segmentation"] orig = labels[f"{self.node.task_name}/original/segmentation"] return orig, recon, seg_out, an_mask diff --git a/luxonis_train/core/utils/export_utils.py b/luxonis_train/core/utils/export_utils.py index 7190c889..fb1af27c 100644 --- a/luxonis_train/core/utils/export_utils.py +++ b/luxonis_train/core/utils/export_utils.py @@ -13,7 +13,7 @@ def replace_weights( module: "luxonis_train.models.LuxonisLightningModule", weights: str | Path | None = None, -) -> Generator[None, None, None]: +) -> Generator: old_weights = None if weights is not None: old_weights = module.state_dict() diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index 19fc4dc6..bade09b4 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -1,12 +1,17 @@ from abc import ABC, abstractmethod from typing import Any, Literal +import cv2 +import numpy as np +import numpy.typing as npt +import torch from luxonis_ml.typing import ConfigItem from luxonis_ml.utils.registry import AutoRegisterMeta from torch import Size, Tensor from torch.utils.data import Dataset from luxonis_train.utils.registry import LOADERS +from luxonis_train.utils.types import Labels from .utils import LuxonisLoaderTorchOutput @@ -67,7 +72,9 @@ class ConfigItem: ) @type image_source: str - @param image_source: Name of the input image group. This can be used for datasets with multiple image sources, e.g. left and right cameras or RGB and depth images. Irrelevant for datasets with only one image source. + @param image_source: Name of the image source. Only relevant for + datasets with multiple image sources, e.g. C{"left"} and C{"right"}. This parameter defines which of these sources is used for + visualizations. @type keep_aspect_ratio: bool @param keep_aspect_ratio: Whether to keep the aspect ratio of the output image after resizing. @@ -142,10 +149,10 @@ def keep_aspect_ratio(self) -> bool: return self._getter_check_none("keep_aspect_ratio") @property - def color_space(self) -> str: + def color_space(self) -> Literal["RGB", "BGR"]: """Color space of the output image. - @type: str + @type: Literal["RGB", "BGR"] """ return self._getter_check_none("color_space") @@ -200,13 +207,19 @@ def augment_test_image(self, img: Tensor) -> Tensor: "`augment_test_image` method to expose this functionality." ) + def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: + img, labels = self.get(idx) + if isinstance(img, Tensor): + img = {self.image_source: img} + return img, labels + @abstractmethod def __len__(self) -> int: """Returns length of the dataset.""" ... @abstractmethod - def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: + def get(self, idx: int) -> tuple[Tensor | dict[str, Tensor], Labels]: """Loads sample from dataset. @type idx: int @@ -235,6 +248,35 @@ def get_n_keypoints(self) -> dict[str, int] | None: """ return None + def dict_numpy_to_torch( + self, numpy_dictionary: dict[str, np.ndarray] + ) -> dict[str, Tensor]: + """Converts a dictionary of numpy arrays to a dictionary of + torch tensors. + + @type numpy_dictionary: dict[str, np.ndarray] + @param numpy_dictionary: Dictionary of numpy arrays. + @rtype: dict[str, torch.Tensor] + @return: Dictionary of torch tensors. + """ + return { + task: torch.tensor(array) + for task, array in numpy_dictionary.items() + } + + def read_image(self, path: str) -> npt.NDArray[np.float32]: + """Reads an image from a file. + + @type path: str + @param path: Path to the image file. + @rtype: np.ndarray[np.float32] + @return: Image as a numpy array. + """ + img = cv2.imread(path, cv2.IMREAD_COLOR) + if self.color_space == "RGB": + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return img.astype(np.float32) / 255.0 + def _getter_check_none( self, attribute: Literal[ diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index be31e28e..c21d5230 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -12,15 +12,17 @@ from luxonis_ml.data.parsers import LuxonisParser from luxonis_ml.enums import DatasetType from torch import Size, Tensor -from typeguard import typechecked +from typing_extensions import override -from .base_loader import BaseLoaderTorch, LuxonisLoaderTorchOutput +from luxonis_train.utils.types import Labels + +from .base_loader import BaseLoaderTorch logger = logging.getLogger(__name__) class LuxonisLoaderTorch(BaseLoaderTorch): - @typechecked + @override def __init__( self, dataset_name: str | None = None, @@ -97,28 +99,30 @@ def __init__( out_image_format=self.color_space, ) + @override def __len__(self) -> int: return len(self.loader) @property + @override def input_shapes(self) -> dict[str, Size]: img = self[0][0][self.image_source] return {self.image_source: img.shape} - def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: + @override + def get(self, idx: int) -> tuple[Tensor, Labels]: img, labels = self.loader[idx] img = np.transpose(img, (2, 0, 1)) # HWC to CHW - tensor_img = Tensor(img) - tensor_labels: dict[str, Tensor] = {} - for task, array in labels.items(): - tensor_labels[task] = Tensor(array) + tensor_img = torch.tensor(img) - return {self.image_source: tensor_img}, tensor_labels + return tensor_img, self.dict_numpy_to_torch(labels) + @override def get_classes(self) -> dict[str, list[str]]: return self.dataset.get_classes() + @override def get_n_keypoints(self) -> dict[str, int]: skeletons = self.dataset.get_skeletons() return {task: len(skeletons[task][0]) for task in skeletons} diff --git a/luxonis_train/loaders/luxonis_perlin_loader_torch.py b/luxonis_train/loaders/luxonis_perlin_loader_torch.py index 05e2c449..68b5512b 100644 --- a/luxonis_train/loaders/luxonis_perlin_loader_torch.py +++ b/luxonis_train/loaders/luxonis_perlin_loader_torch.py @@ -1,24 +1,28 @@ import random -from typing import cast +from collections.abc import Generator +from contextlib import contextmanager import numpy as np import torch import torch.nn.functional as F -from luxonis_ml.data import AlbumentationsEngine from luxonis_ml.utils import LuxonisFileSystem from torch import Tensor +from typing_extensions import override + +from luxonis_train.utils.types import Labels -from .base_loader import LuxonisLoaderTorchOutput from .luxonis_loader_torch import LuxonisLoaderTorch from .perlin import apply_anomaly_to_img class LuxonisLoaderPerlinNoise(LuxonisLoaderTorch): + @override def __init__( self, *args, anomaly_source_path: str, noise_prob: float = 0.5, + beta: float | None = None, **kwargs, ): """Custom loader for Luxonis datasets that adds Perlin noise @@ -58,37 +62,51 @@ def __init__( raise ValueError( "This loader only supports datasets with a single task." ) + self.beta = beta self.task_name = next(iter(self.loader.dataset.get_tasks())) + self.augmentations = self.loader.augmentations - augmentations = cast(AlbumentationsEngine, self.loader.augmentations) - if augmentations is None or augmentations.pixel_transform is None: - self.pixel_augs = None - else: - self.pixel_augs = augmentations.pixel_transform + @override + def get(self, idx: int) -> tuple[Tensor, Labels]: + with _freeze_seed(): + img, labels = self.loader[idx] - def __getitem__(self, idx: int) -> LuxonisLoaderTorchOutput: - img, labels = self.loader[idx] + an_mask = torch.tensor(labels.pop(f"{self.task_name}/segmentation"))[ + 0, ... + ] img = np.transpose(img, (2, 0, 1)) - tensor_img = Tensor(img) + tensor_img = torch.tensor(img) + tensor_labels = self.dict_numpy_to_torch(labels) if self.view[0] == "train" and random.random() < self.noise_prob: + anomaly_path = random.choice(self.anomaly_files) + anomaly_img = self.read_image(str(anomaly_path)) + + if self.augmentations is not None: + anomaly_img = self.augmentations.apply([(anomaly_img, {})])[0] + + anomaly_img = torch.tensor(anomaly_img).permute(2, 0, 1) aug_tensor_img, an_mask = apply_anomaly_to_img( - tensor_img, - anomaly_source_paths=self.anomaly_files, - pixel_augs=self.pixel_augs, + tensor_img, anomaly_img, self.beta ) else: aug_tensor_img = tensor_img - h, w = aug_tensor_img.shape[-2:] - an_mask = torch.zeros((h, w)) - tensor_labels = {f"{self.task_name}/original/segmentation": tensor_img} - for task, array in labels.items(): - tensor_labels[task] = Tensor(array) + an_mask = F.one_hot(an_mask.long(), 2).permute(2, 0, 1).float() + + tensor_labels = { + f"{self.task_name}/original/segmentation": tensor_img, + f"{self.task_name}/segmentation": an_mask, + } + + return aug_tensor_img, tensor_labels - tensor_labels[f"{self.task_name}/segmentation"] = ( - F.one_hot(an_mask.long(), 2).permute(2, 0, 1).float() - ) - return {self.image_source: aug_tensor_img}, tensor_labels +@contextmanager +def _freeze_seed() -> Generator: + python_seed = random.getstate() + numpy_seed = np.random.get_state() + yield + random.setstate(python_seed) + np.random.set_state(numpy_seed) diff --git a/luxonis_train/loaders/perlin.py b/luxonis_train/loaders/perlin.py index cd201d35..6a973d85 100644 --- a/luxonis_train/loaders/perlin.py +++ b/luxonis_train/loaders/perlin.py @@ -1,13 +1,10 @@ -import random -from pathlib import Path -from typing import Callable, List, Tuple +from typing import Callable, Tuple -import cv2 -import numpy as np import torch +from torch import Tensor -def compute_gradients(res: tuple[int, int]) -> torch.Tensor: +def compute_gradients(res: tuple[int, int]) -> Tensor: angles = 2 * torch.pi * torch.rand(res[0] + 1, res[1] + 1) gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1) return gradients @@ -15,21 +12,21 @@ def compute_gradients(res: tuple[int, int]) -> torch.Tensor: @torch.jit.script def lerp_torch( # pragma: no cover - x: torch.Tensor, y: torch.Tensor, w: torch.Tensor -) -> torch.Tensor: + x: Tensor, y: Tensor, w: Tensor +) -> Tensor: return (y - x) * w + x -def fade_function(t: torch.Tensor) -> torch.Tensor: +def fade_function(t: Tensor) -> Tensor: return 6 * t**5 - 15 * t**4 + 10 * t**3 def tile_grads( slice1: Tuple[int, int | None], slice2: Tuple[int, int | None], - gradients: torch.Tensor, + gradients: Tensor, d: Tuple[int, int], -) -> torch.Tensor: +) -> Tensor: return ( gradients[slice1[0] : slice1[1], slice2[0] : slice2[1]] .repeat_interleave(d[0], 0) @@ -38,11 +35,11 @@ def tile_grads( def dot( - grad: torch.Tensor, + grad: Tensor, shift: Tuple[int, int], - grid: torch.Tensor, + grid: Tensor, shape: Tuple[int, int], -) -> torch.Tensor: +) -> Tensor: return ( torch.stack( ( @@ -58,8 +55,8 @@ def dot( def rand_perlin_2d( shape: Tuple[int, int], res: Tuple[int, int], - fade: Callable[[torch.Tensor], torch.Tensor] = fade_function, -) -> torch.Tensor: + fade: Callable[[Tensor], Tensor] = fade_function, +) -> Tensor: delta = (res[0] / shape[0], res[1] / shape[1]) d = (shape[0] // res[0], shape[1] // res[1]) grid_x, grid_y = torch.meshgrid( @@ -92,7 +89,7 @@ def rand_perlin_2d( @torch.jit.script -def rotate_noise(noise: torch.Tensor) -> torch.Tensor: # pragma: no cover +def rotate_noise(noise: Tensor) -> Tensor: # pragma: no cover angle = torch.rand(1) * 2 * torch.pi h, w = noise.shape center_y, center_x = h // 2, w // 2 @@ -117,7 +114,7 @@ def generate_perlin_noise( min_perlin_scale: int = 0, perlin_scale: int = 6, threshold: float = 0.5, -) -> torch.Tensor: +) -> Tensor: perlin_scalex = 2 ** int( torch.randint(min_perlin_scale, perlin_scale, (1,)).item() ) @@ -136,21 +133,14 @@ def generate_perlin_noise( return perlin_mask -def load_image_as_numpy(img_path: Path) -> np.ndarray: - image = cv2.imread(str(img_path), cv2.IMREAD_COLOR) - image = image.astype(np.float32) / 255.0 - return image - - def apply_anomaly_to_img( - img: torch.Tensor, - anomaly_source_paths: List[Path], + img: Tensor, + anomaly_img: Tensor, beta: float | None = None, - pixel_augs: Callable | None = None, # type: ignore -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[Tensor, Tensor]: """Applies Perlin noise-based anomalies to a single image (C, H, W). - @type img: torch.Tensor + @type img: Tensor @param img: The input image tensor of shape (C, H, W). @type anomaly_source_paths: List[str] @param anomaly_source_paths: List of file paths to the anomaly images. @@ -159,27 +149,12 @@ def apply_anomaly_to_img( @type beta: float | None @param beta: A blending factor for anomaly and noise. If None, a random value in the range [0, 0.8] is used. Defaults to C{None}. - @rtype: Tuple[torch.Tensor, torch.Tensor] + @rtype: Tuple[Tensor, Tensor] @return: A tuple containing: - - augmented_img (torch.Tensor): The augmented image with applied anomaly and Perlin noise. - - perlin_mask (torch.Tensor): The Perlin noise mask applied to the image. + - augmented_img (Tensor): The augmented image with applied anomaly and Perlin noise. + - perlin_mask (Tensor): The Perlin noise mask applied to the image. """ - sampled_anomaly_image_path = random.choice(anomaly_source_paths) - - anomaly_image = load_image_as_numpy(sampled_anomaly_image_path) - - anomaly_image = cv2.resize( - anomaly_image, - (img.shape[2], img.shape[1]), - interpolation=cv2.INTER_LINEAR, - ) - - if pixel_augs is not None: - anomaly_image = pixel_augs(image=anomaly_image)["image"] - - anomaly_image = torch.tensor(anomaly_image).permute(2, 0, 1) - perlin_mask = generate_perlin_noise( shape=(img.shape[1], img.shape[2]), ) @@ -189,7 +164,7 @@ def apply_anomaly_to_img( augmented_img = ( (1 - perlin_mask).unsqueeze(0) * img - + (1 - beta) * perlin_mask.unsqueeze(0) * anomaly_image + + (1 - beta) * perlin_mask.unsqueeze(0) * anomaly_img + beta * perlin_mask.unsqueeze(0) * img ) diff --git a/luxonis_train/utils/general.py b/luxonis_train/utils/general.py index 390ddfd2..746df792 100644 --- a/luxonis_train/utils/general.py +++ b/luxonis_train/utils/general.py @@ -187,7 +187,9 @@ def safe_download( if i == retry: logger.warning("Download failed, retry limit reached.") return None - logger.warning(f"Download failed, retrying {i+1}/{retry} ...") + logger.warning( + f"Download failed, retrying {i + 1}/{retry} ..." + ) def clean_url(url: str) -> str: From c76135c5cfacddcc77de9c9182c7ec52a7de368f Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 17 Jan 2025 00:24:02 -0500 Subject: [PATCH 32/57] converting to float --- luxonis_train/loaders/base_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index bade09b4..bb3b9c13 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -260,7 +260,7 @@ def dict_numpy_to_torch( @return: Dictionary of torch tensors. """ return { - task: torch.tensor(array) + task: torch.tensor(array).float() for task, array in numpy_dictionary.items() } From 058f449770cab742d1623ad375d3f4d3dff802df Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 17 Jan 2025 04:24:43 -0500 Subject: [PATCH 33/57] helper function --- luxonis_train/core/core.py | 2 +- luxonis_train/loaders/base_loader.py | 6 ++-- luxonis_train/utils/__init__.py | 2 ++ luxonis_train/utils/general.py | 41 +++++++++++++++++++++++++++- 4 files changed, 45 insertions(+), 6 deletions(-) diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py index 5f469590..1ea49bab 100644 --- a/luxonis_train/core/core.py +++ b/luxonis_train/core/core.py @@ -133,7 +133,7 @@ def __init__( augmentation_config=self.cfg_preprocessing.get_active_augmentations(), color_space=self.cfg_preprocessing.color_space, keep_aspect_ratio=self.cfg_preprocessing.keep_aspect_ratio, - **self.cfg.loader.params, + **self.cfg.loader.params, # type: ignore ) for name, loader in self.loaders.items(): diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index bb3b9c13..92e54794 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -10,6 +10,7 @@ from torch import Size, Tensor from torch.utils.data import Dataset +from luxonis_train.utils.general import get_attribute_check_none from luxonis_train.utils.registry import LOADERS from luxonis_train.utils.types import Labels @@ -290,7 +291,4 @@ def _getter_check_none( "color_space", ], ) -> Any: - value = getattr(self, f"_{attribute}") - if value is None: - raise ValueError(f"{attribute} is not set") - return value + return get_attribute_check_none(self, attribute) diff --git a/luxonis_train/utils/__init__.py b/luxonis_train/utils/__init__.py index d7c0be9f..2f2b550a 100644 --- a/luxonis_train/utils/__init__.py +++ b/luxonis_train/utils/__init__.py @@ -9,6 +9,7 @@ from .dataset_metadata import DatasetMetadata from .exceptions import IncompatibleException from .general import ( + get_attribute_check_none, get_with_default, infer_upscale_factor, make_divisible, @@ -42,4 +43,5 @@ "get_sigmas", "traverse_graph", "insert_class", + "get_attribute_check_none", ] diff --git a/luxonis_train/utils/general.py b/luxonis_train/utils/general.py index 746df792..94cd1613 100644 --- a/luxonis_train/utils/general.py +++ b/luxonis_train/utils/general.py @@ -3,7 +3,7 @@ import os import urllib.parse from pathlib import Path, PurePosixPath -from typing import TypeVar +from typing import Any, TypeVar import torch from torch import Size, Tensor @@ -205,3 +205,42 @@ def clean_url(url: str) -> str: def url2file(url: str) -> str: """Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt.""" return Path(clean_url(url)).name + + +def get_attribute_check_none(obj: object, attribute: str) -> Any: + """Get private attribute from object and check if it is not None. + + Example: + + >>> class Person: + ... def __init__(self, age: int | None = None): + ... self._age = age + ... + ... @property + ... def age(self): + ... return get_attribute_check_none(self, "age") + + >>> mike = Person(20) + >>> print(mike.age) + 20 + + >>> amanda = Person() + >>> print(amanda.age) + Traceback (most recent call last): + ValueError: attribute 'age' was not set + + @type obj: object + @param obj: Object to get attribute from. + + @type attribute: str + @param attribute: Name of the attribute to get. + + @rtype: Any + @return: Value of the attribute. + + @raise ValueError: If the attribute is None. + """ + value = getattr(obj, f"_{attribute}") + if value is None: + raise ValueError(f"attribute '{attribute}' was not set") + return value From 732ad1f5c07888e1b68d317953d3b26e88af95a8 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 17 Jan 2025 07:20:11 -0500 Subject: [PATCH 34/57] changes for latest luxonis-ml --- luxonis_train/loaders/luxonis_loader_torch.py | 2 +- luxonis_train/loaders/utils.py | 13 ++++++++----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index c21d5230..4267cced 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -96,7 +96,7 @@ def __init__( height=self.height, width=self.width, keep_aspect_ratio=self.keep_aspect_ratio, - out_image_format=self.color_space, + color_space=self.color_space, ) @override diff --git a/luxonis_train/loaders/utils.py b/luxonis_train/loaders/utils.py index 10b4d17a..9c9e1d45 100644 --- a/luxonis_train/loaders/utils.py +++ b/luxonis_train/loaders/utils.py @@ -37,12 +37,15 @@ def collate_fn( if task_type in {"keypoints", "boundingbox"}: label_box: list[Tensor] = [] - for i, box in enumerate(annos): - l_box = torch.zeros((box.shape[0], box.shape[1] + 1)) - l_box[:, 0] = i # add target image index for build_targets() - l_box[:, 1:] = box - label_box.append(l_box) + for i, ann in enumerate(annos): + new_ann = torch.zeros((ann.shape[0], ann.shape[1] + 1)) + # add target image index for build_targets() + new_ann[:, 0] = i + new_ann[:, 1:] = ann + label_box.append(new_ann) out_labels[task] = torch.cat(label_box, 0) + elif task_type == "instance_segmentation": + out_labels[task] = torch.cat(annos, 0) else: out_labels[task] = torch.stack(annos, 0) From 09b1e58259258d8f76a925536a41caba41d644b8 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 17 Jan 2025 17:20:43 -0500 Subject: [PATCH 35/57] fixed tests --- luxonis_train/core/utils/infer_utils.py | 7 +++-- luxonis_train/loaders/base_loader.py | 6 ++--- luxonis_train/models/luxonis_lightning.py | 4 +-- pyproject.toml | 2 +- tests/integration/multi_input_modules.py | 10 +++---- tests/integration/test_scheduler.py | 27 ++++++++++--------- .../test_utils/test_dataset_metadata.py | 8 +++--- 7 files changed, 35 insertions(+), 29 deletions(-) diff --git a/luxonis_train/core/utils/infer_utils.py b/luxonis_train/core/utils/infer_utils.py index c4f9085f..a4aefdc9 100644 --- a/luxonis_train/core/utils/infer_utils.py +++ b/luxonis_train/core/utils/infer_utils.py @@ -212,10 +212,13 @@ def generator() -> DatasetIterator: loader = LuxonisLoaderTorch( dataset_name=dataset_name, - image_source="image", view="test", + height=model.cfg_preprocessing.train_image_size.height, + width=model.cfg_preprocessing.train_image_size.width, + augmentation_config=model.cfg_preprocessing.get_active_augmentations(), + color_space=model.cfg_preprocessing.color_space, + keep_aspect_ratio=model.cfg_preprocessing.keep_aspect_ratio, ) - loader.loader.augmentations = model.loaders["val"].loader.augmentations # type: ignore loader = torch_data.DataLoader( loader, batch_size=model.cfg.trainer.batch_size ) diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index 92e54794..752a15d2 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -27,11 +27,11 @@ class BaseLoaderTorch( def __init__( self, view: list[str], - height: int, - width: int, + height: int | None = None, + width: int | None = None, augmentation_engine: str = "albumentations", augmentation_config: list[ConfigItem] | None = None, - image_source: str = "default", + image_source: str = "image", keep_aspect_ratio: bool = True, color_space: Literal["RGB", "BGR"] = "RGB", ): diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index 9643a6f2..d2d3fe87 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -907,14 +907,14 @@ def configure_optimizers( def get_scheduler( scheduler_cfg: ConfigItem, optimizer: torch.optim.Optimizer - ) -> torch.optim.lr_scheduler._LRScheduler: + ) -> torch.optim.lr_scheduler.LRScheduler: scheduler_class = SCHEDULERS.get(scheduler_cfg.name) scheduler_params = scheduler_cfg.params | {"optimizer": optimizer} return scheduler_class(**scheduler_params) # type: ignore if cfg_scheduler.name == "SequentialLR": schedulers_list = [ - get_scheduler(scheduler_cfg, optimizer) + get_scheduler(ConfigItem(**scheduler_cfg), optimizer) for scheduler_cfg in cfg_scheduler.params["schedulers"] ] diff --git a/pyproject.toml b/pyproject.toml index 42457ed6..18b98796 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,7 +64,7 @@ reportUnnecessaryIsInstance = "none" [tool.pytest.ini_options] testpaths = ["tests"] -addopts = "--disable-warnings" +addopts = "--disable-warnings -x" markers = [ "unit: mark a test as a unit test", "integration: mark a test as an integration test", diff --git a/tests/integration/multi_input_modules.py b/tests/integration/multi_input_modules.py index 3db26f5d..362c163f 100644 --- a/tests/integration/multi_input_modules.py +++ b/tests/integration/multi_input_modules.py @@ -9,10 +9,10 @@ class CustomMultiInputLoader(BaseLoaderTorch): - def __init__( - self, view: str | list[str], image_source: str | None = None, **_ - ): - super().__init__(view=view, image_source=image_source) + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._height = 224 + self._width = 224 @property def input_shapes(self): @@ -23,7 +23,7 @@ def input_shapes(self): "pointcloud": torch.Size([1000, 3]), } - def __getitem__(self, _): # pragma: no cover + def get(self, _): # pragma: no cover # Fake data left = torch.rand(3, 224, 224, dtype=torch.float32) right = torch.rand(3, 224, 224, dtype=torch.float32) diff --git a/tests/integration/test_scheduler.py b/tests/integration/test_scheduler.py index ba3f4364..cd46e199 100644 --- a/tests/integration/test_scheduler.py +++ b/tests/integration/test_scheduler.py @@ -3,6 +3,7 @@ import pytest from luxonis_ml.data import LuxonisDataset +from luxonis_train.config.config import SchedulerConfig from luxonis_train.core import LuxonisModel @@ -32,10 +33,10 @@ def create_model_config(): } -def sequential_scheduler(): - return { - "name": "SequentialLR", - "params": { +def sequential_scheduler() -> SchedulerConfig: + return SchedulerConfig( + name="SequentialLR", + params={ "schedulers": [ { "name": "LinearLR", @@ -48,14 +49,14 @@ def sequential_scheduler(): ], "milestones": [1], }, - } + ) -def cosine_annealing_scheduler(): - return { - "name": "CosineAnnealingLR", - "params": {"T_max": 2, "eta_min": 0.001}, - } +def cosine_annealing_scheduler() -> SchedulerConfig: + return SchedulerConfig( + name="CosineAnnealingLR", + params={"T_max": 2, "eta_min": 0.001}, + ) @pytest.mark.parametrize( @@ -63,7 +64,9 @@ def cosine_annealing_scheduler(): ) def test_scheduler(coco_dataset: LuxonisDataset, scheduler_config): config = create_model_config() - config["trainer"]["scheduler"] = scheduler_config - opts = {"loader.params.dataset_name": coco_dataset.dataset_name} + opts = { + "loader.params.dataset_name": coco_dataset.dataset_name, + "trainer.scheduler": scheduler_config, + } model = LuxonisModel(config, opts) model.train() diff --git a/tests/unittests/test_utils/test_dataset_metadata.py b/tests/unittests/test_utils/test_dataset_metadata.py index daf01725..c2e582dd 100644 --- a/tests/unittests/test_utils/test_dataset_metadata.py +++ b/tests/unittests/test_utils/test_dataset_metadata.py @@ -4,7 +4,7 @@ @pytest.fixture -def metadata(): +def metadata() -> DatasetMetadata: return DatasetMetadata( classes={ "color-segmentation": ["car", "person"], @@ -14,7 +14,7 @@ def metadata(): ) -def test_n_classes(metadata): +def test_n_classes(metadata: DatasetMetadata): assert metadata.n_classes("color-segmentation") == 2 assert metadata.n_classes("detection") == 2 assert metadata.n_classes() == 2 @@ -25,7 +25,7 @@ def test_n_classes(metadata): metadata.n_classes() -def test_n_keypoints(metadata): +def test_n_keypoints(metadata: DatasetMetadata): assert metadata.n_keypoints("color-segmentation") == 0 assert metadata.n_keypoints("detection") == 0 assert metadata.n_keypoints() == 0 @@ -36,7 +36,7 @@ def test_n_keypoints(metadata): metadata.n_keypoints() -def test_class_names(metadata): +def test_class_names(metadata: DatasetMetadata): assert metadata.classes("color-segmentation") == ["car", "person"] assert metadata.classes("detection") == ["car", "person"] assert metadata.classes() == ["car", "person"] From 5a10d61ed5d1dbed5ee7707612d8c9aec9e736fe Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 22 Jan 2025 21:10:16 -0500 Subject: [PATCH 36/57] reid fixes --- .../attached_modules/base_attached_module.py | 59 ++-- .../attached_modules/losses/pml_loss.py | 113 ++----- .../attached_modules/metrics/pml_metrics.py | 74 +---- .../visualizers/embeddings_visualizer.py | 29 +- luxonis_train/config/config.py | 2 +- luxonis_train/enums.py | 30 +- luxonis_train/loaders/utils.py | 4 +- .../nodes/backbones/ghostfacenet/blocks.py | 292 ++++++++---------- .../backbones/ghostfacenet/ghostfacenet.py | 36 +-- .../nodes/backbones/micronet/blocks.py | 7 +- luxonis_train/nodes/base_node.py | 4 +- luxonis_train/nodes/blocks/blocks.py | 8 +- luxonis_train/utils/dataset_metadata.py | 4 +- tests/configs/reid.yaml | 63 ++-- tests/integration/test_reid.py | 184 ----------- 15 files changed, 287 insertions(+), 622 deletions(-) delete mode 100644 tests/integration/test_reid.py diff --git a/luxonis_train/attached_modules/base_attached_module.py b/luxonis_train/attached_modules/base_attached_module.py index a5f14761..866898fe 100644 --- a/luxonis_train/attached_modules/base_attached_module.py +++ b/luxonis_train/attached_modules/base_attached_module.py @@ -1,13 +1,13 @@ import logging from abc import ABC from contextlib import suppress -from typing import Generic +from typing import Generic, get_args from luxonis_ml.utils.registry import AutoRegisterMeta from torch import Size, Tensor, nn from typing_extensions import TypeVarTuple, Unpack -from luxonis_train.enums import TaskType +from luxonis_train.enums import Task, TaskType from luxonis_train.nodes import BaseNode from luxonis_train.utils import IncompatibleException, Labels, Packet @@ -57,19 +57,19 @@ class BaseAttachedModule( labels I{or} segmentation labels. """ - supported_tasks: list[TaskType | tuple[TaskType, ...]] | None = None + supported_tasks: list[Task | tuple[Task, ...]] | None = None def __init__(self, *, node: BaseNode | None = None): super().__init__() self._node = node self._epoch = 0 - self.required_labels: list[TaskType] = [] + self.required_labels: list[Task] = [] if self._node and self.supported_tasks: module_supported = [ label.value - if isinstance(label, TaskType) - else f"({' + '.join(label)})" + if isinstance(label, Task) + else f"({' + '.join(map(str, label))})" for label in self.supported_tasks ] module_supported = f"[{', '.join(module_supported)}]" @@ -81,7 +81,7 @@ def __init__(self, *, node: BaseNode | None = None): ) node_tasks = set(self.node.tasks) for required_labels in self.supported_tasks: - if isinstance(required_labels, TaskType): + if isinstance(required_labels, Task): required_labels = [required_labels] else: required_labels = list(required_labels) @@ -159,7 +159,7 @@ def class_names(self) -> list[str]: return self.node.class_names @property - def node_tasks(self) -> list[TaskType]: + def node_tasks(self) -> list[Task]: """Getter for the tasks of the attached node. @type: dict[TaskType, str] @@ -201,11 +201,11 @@ def get_label( @raises ValueError: If the module requires multiple labels and the C{task_type} is not provided. @raises IncompatibleException: If the label is not found in the labels dictionary. """ - return self._get_label(labels, task_type)[0] + return self._get_label(labels, task_type) def _get_label( - self, labels: Labels, task_type: TaskType | None = None - ) -> tuple[Tensor, TaskType]: + self, labels: Labels, task_type: Task | None = None + ) -> Tensor: if task_type is None: if len(self.required_labels) == 1: task_type = self.required_labels[0] @@ -221,7 +221,7 @@ def _get_label( f"Available labels: {list(labels.keys())}. " f"Missing label: '{task}'." ) - return labels[task], task_type + return labels[task] raise ValueError( f"{self.name} requires multiple labels. You must provide the " @@ -229,7 +229,7 @@ def _get_label( ) def get_input_tensors( - self, inputs: Packet[Tensor], task_type: TaskType | str | None = None + self, inputs: Packet[Tensor], task_type: Task | str | None = None ) -> list[Tensor]: """Extracts the input tensors from the packet. @@ -259,7 +259,7 @@ def get_input_tensors( For such cases, the C{prepare} method should be overridden. """ if task_type is not None: - if isinstance(task_type, TaskType): + if isinstance(task_type, Task): if task_type not in self.node_tasks: raise IncompatibleException( f"Task {task_type.value} is not supported by the node " @@ -345,24 +345,45 @@ def prepare( set(self.supported_tasks) & set(self.node_tasks) ) x = self.get_input_tensors(inputs) - if labels is None or len(labels) == 0: + if labels is None or not labels: return x, None # type: ignore - label, task_type = self._get_label(labels) - if task_type in [TaskType.CLASSIFICATION, TaskType.SEGMENTATION]: + + label = self._get_label(labels) + generics = self._get_generic_params() + if generics is None: + return x, label # type: ignore + + if len(generics) != 2: + raise RuntimeError( + f"The type signature of '{self.name}' implies a complicated " + f"custom module ({self.name}[{', '.join(g.__name__ for g in generics)}]). " + "Please implement your own `prepare` method. The default " + "`prepare` works only when the generic type of the module " + "is `[Tensor | list[Tensor], Tensor]`." + ) + + if generics[0] is Tensor: if len(x) == 1: x = x[0] else: logger.warning( - f"Module {self.name} expects a single tensor as input, " + f"Module '{self.name}' expects a single tensor as input, " f"but got {len(x)} tensors. Using the last tensor. " f"If this is not the desired behavior, please override the " "`prepare` method of the attached module or the `wrap` " - f"method of {self.node.name}." + f"method of '{self.node.name}'." ) x = x[-1] return x, label # type: ignore + def _get_generic_params(self) -> tuple[type, ...] | None: + cls = type(self) + try: + return get_args(cls.__orig_bases__[0]) # type: ignore + except Exception: + return None + def _check_node_type_override(self) -> None: if "node" not in self.__annotations__: return diff --git a/luxonis_train/attached_modules/losses/pml_loss.py b/luxonis_train/attached_modules/losses/pml_loss.py index 959a5f68..087963c0 100644 --- a/luxonis_train/attached_modules/losses/pml_loss.py +++ b/luxonis_train/attached_modules/losses/pml_loss.py @@ -4,116 +4,63 @@ from pytorch_metric_learning.losses import CrossBatchMemory from torch import Tensor +from luxonis_train.enums import Metadata +from luxonis_train.nodes.backbones.ghostfacenet import GhostFaceNetsV2 +from luxonis_train.nodes.base_node import BaseNode + from .base_loss import BaseLoss logger = logging.getLogger(__name__) -ALL_EMBEDDING_LOSSES = [ +EMBEDDING_LOSSES = [ "AngularLoss", - "ArcFaceLoss", "CircleLoss", "ContrastiveLoss", - "CosFaceLoss", "DynamicSoftMarginLoss", "FastAPLoss", "HistogramLoss", "InstanceLoss", "IntraPairVarianceLoss", - "LargeMarginSoftmaxLoss", "GeneralizedLiftedStructureLoss", "LiftedStructureLoss", "MarginLoss", "MultiSimilarityLoss", "NPairsLoss", "NCALoss", - "NormalizedSoftmaxLoss", "NTXentLoss", "PNPLoss", - "ProxyAnchorLoss", - "ProxyNCALoss", "RankedListLoss", "SignalToNoiseRatioContrastiveLoss", - "SoftTripleLoss", - "SphereFaceLoss", - "SubCenterArcFaceLoss", "SupConLoss", "ThresholdConsistentMarginLoss", "TripletMarginLoss", "TupletMarginLoss", ] -CLASS_EMBEDDING_LOSSES = [ - "ArcFaceLoss", - "CosFaceLoss", - "LargeMarginSoftmaxLoss", - "NormalizedSoftmaxLoss", - "ProxyAnchorLoss", - "ProxyNCALoss", - "SoftTripleLoss", - "SphereFaceLoss", - "SubCenterArcFaceLoss", -] - +for loss_name in EMBEDDING_LOSSES: -class EmbeddingLossWrapper(BaseLoss): - def __init__( - self, - loss_name: str, - embedding_size: int = 512, - cross_batch_memory_size=0, - num_classes: int = 0, - loss_kwargs: dict | None = None, - *args, - **kwargs, + class EmbeddingLossWrapper( + BaseLoss[Tensor, Tensor], register_name=loss_name ): - super().__init__(*args, **kwargs) - if loss_kwargs is None: - loss_kwargs = {} - - try: - loss_cls = getattr(pml_losses, loss_name) - except AttributeError as e: - raise ValueError( - f"Loss {loss_name} not found in pytorch_metric_learning" - ) from e - - if loss_name in CLASS_EMBEDDING_LOSSES: - if num_classes < 0: - raise ValueError( - f"Loss {loss_name} requires num_classes to be set to a positive value" - ) - loss_kwargs["num_classes"] = num_classes - loss_kwargs["embedding_size"] = embedding_size - - # If we wanted to support these losses, we would need to add a separate optimizer for them. - # They may be useful in some scenarios, so leaving this here for future reference. - raise ValueError( - f"Loss {loss_name} requires its own optimizer, and that is not currently supported." - ) - - self.loss_func = loss_cls(**loss_kwargs) - - if cross_batch_memory_size > 0: - if loss_name in CrossBatchMemory.supported_losses(): - self.loss_func = CrossBatchMemory( - self.loss_func, embedding_size=embedding_size - ) - else: - logger.warning( - f"Cross batch memory is not supported for {loss_name}. Ignoring cross_batch_memory_size." - ) - - def prepare( - self, inputs: dict[str, list[Tensor]], labels: dict[str, list[Tensor]] - ) -> tuple[Tensor, Tensor]: - embeddings = self.get_input_tensors(inputs, "features")[0] - - if labels is None or "id" not in labels: - raise ValueError("Labels must contain 'id' key") - - ids = labels["id"][0][:, 0] - return embeddings, ids - - def forward(self, inputs: Tensor, target: Tensor) -> Tensor: - loss = self.loss_func(inputs, target) - return loss + node: GhostFaceNetsV2 + supported_tasks = [Metadata("id")] + + def __init__(self, *, node: BaseNode | None = None, **kwargs): + super().__init__(node=node) + + Loss = getattr(pml_losses, loss_name) # noqa: B023 + self.loss_func = Loss(**kwargs) + + if self.node.embedding_size is not None: + if loss_name in CrossBatchMemory.supported_losses(): # noqa: B023 + self.loss_func = CrossBatchMemory( + self.loss_func, embedding_size=self.node.embedding_size + ) + else: + logger.warning( + f"CrossBatchMemory is not supported for {loss_name}. " # noqa: B023 + "Ignoring cross_batch_memory_size." + ) + + def forward(self, inputs: Tensor, target: Tensor) -> Tensor: + return self.loss_func(inputs, target) diff --git a/luxonis_train/attached_modules/metrics/pml_metrics.py b/luxonis_train/attached_modules/metrics/pml_metrics.py index ad8b0d88..9e4ab50b 100644 --- a/luxonis_train/attached_modules/metrics/pml_metrics.py +++ b/luxonis_train/attached_modules/metrics/pml_metrics.py @@ -1,13 +1,17 @@ import torch from torch import Tensor +from luxonis_train.enums import Metadata + from .base_metric import BaseMetric # Converted from https://omoindrot.github.io/triplet-loss#offline-and-online-triplet-mining # to PyTorch from TensorFlow -class ClosestIsPositiveAccuracy(BaseMetric): +class ClosestIsPositiveAccuracy(BaseMetric[Tensor, Tensor]): + supported_tasks = [Metadata("id")] + def __init__(self, cross_batch_memory_size=0, **kwargs): super().__init__(**kwargs) self.cross_batch_memory_size = cross_batch_memory_size @@ -21,70 +25,45 @@ def __init__(self, cross_batch_memory_size=0, **kwargs): "total_predictions", default=torch.tensor(0), dist_reduce_fx="sum" ) - def prepare(self, inputs, labels): - embeddings = inputs["features"][0] - - assert ( - labels is not None and "id" in labels - ), "ID labels are required for metric learning losses" - ids = labels["id"][0][:, 0] - return embeddings, ids - def update(self, inputs: Tensor, target: Tensor): embeddings, labels = inputs, target if self.cross_batch_memory_size > 0: - # Append embedding and labels to the memory self.cross_batch_memory.extend(list(zip(embeddings, labels))) - # If the memory is full, remove the oldest elements if len(self.cross_batch_memory) > self.cross_batch_memory_size: self.cross_batch_memory = self.cross_batch_memory[ -self.cross_batch_memory_size : ] - # If the memory is not full, return if len(self.cross_batch_memory) < self.cross_batch_memory_size: return - # Get the embeddings and labels from the memory embeddings, labels = zip(*self.cross_batch_memory) embeddings = torch.stack(embeddings) labels = torch.stack(labels) - # print(f"Calculating accuracy for {len(embeddings)} embeddings") - - # Get the pairwise distances between all embeddings pairwise_distances = _pairwise_distances(embeddings) - - # Set diagonal to infinity so that the closest embedding is not the same embedding pairwise_distances.fill_diagonal_(float("inf")) - # Find the closest embedding for each query embedding closest_indices = torch.argmin(pairwise_distances, dim=1) - - # Get the labels of the closest embeddings closest_labels = labels[closest_indices] - # Filter out embeddings that don't have both positive and negative examples positive_mask = _get_anchor_positive_triplet_mask(labels) num_positives = positive_mask.sum(dim=1) has_at_least_one_positive_and_negative = (num_positives > 0) & ( num_positives < len(labels) ) - # Filter embeddings, labels, and closest indices based on valid indices filtered_labels = labels[has_at_least_one_positive_and_negative] filtered_closest_labels = closest_labels[ has_at_least_one_positive_and_negative ] - # Calculate the number of correct predictions where the closest is positive correct_predictions = ( filtered_labels == filtered_closest_labels ).sum() - # Update the metric state self.correct_predictions += correct_predictions self.total_predictions += len(filtered_labels) @@ -92,7 +71,9 @@ def compute(self): return self.correct_predictions / self.total_predictions -class MedianDistances(BaseMetric): +class MedianDistances(BaseMetric[Tensor, Tensor]): + supported_tasks = [Metadata("id")] + def __init__(self, cross_batch_memory_size=0, **kwargs): super().__init__(**kwargs) self.cross_batch_memory_size = cross_batch_memory_size @@ -104,40 +85,25 @@ def __init__(self, cross_batch_memory_size=0, **kwargs): "closest_vs_positive_distances", default=[], dist_reduce_fx="cat" ) - def prepare(self, inputs, labels): - embeddings = inputs["features"][0] - - assert ( - labels is not None and "id" in labels - ), "ID labels are required for metric learning losses" - ids = labels["id"][0][:, 0] - return embeddings, ids - def update(self, inputs: Tensor, target: Tensor): embeddings, labels = inputs, target if self.cross_batch_memory_size > 0: - # Append embedding and labels to the memory self.cross_batch_memory.extend(list(zip(embeddings, labels))) - # If the memory is full, remove the oldest elements if len(self.cross_batch_memory) > self.cross_batch_memory_size: self.cross_batch_memory = self.cross_batch_memory[ -self.cross_batch_memory_size : ] - # If the memory is not full, return if len(self.cross_batch_memory) < self.cross_batch_memory_size: return - # Get the embeddings and labels from the memory embeddings, labels = zip(*self.cross_batch_memory) embeddings = torch.stack(embeddings) labels = torch.stack(labels) - # Get the pairwise distances between all embeddings pairwise_distances = _pairwise_distances(embeddings) - # Append only upper triangular part of the matrix self.all_distances.append( pairwise_distances[ torch.triu(torch.ones_like(pairwise_distances), diagonal=1) @@ -145,33 +111,24 @@ def update(self, inputs: Tensor, target: Tensor): ].flatten() ) - # Set diagonal to infinity so that the closest embedding is not the same embedding pairwise_distances.fill_diagonal_(float("inf")) - # Get the closest distance for each query embedding closest_distances, _ = torch.min(pairwise_distances, dim=1) self.closest_distances.append(closest_distances) - # Get the positive mask and convert it to boolean positive_mask = _get_anchor_positive_triplet_mask(labels).bool() - # Filter out distances to negative elements w.r.t. each query embedding only_positive_distances = pairwise_distances.clone() only_positive_distances[~positive_mask] = float("inf") - # From the positive distances, get the closest positive distance for each query embedding closest_positive_distances, _ = torch.min( only_positive_distances, dim=1 ) - # Calculate the difference between the closest distance (any) and closest positive distances - # - this tells us how much closer should the closest positive be in order for the embedding - # to be considered correct non_inf_mask = closest_positive_distances != float("inf") difference = closest_positive_distances - closest_distances difference = difference[non_inf_mask] - # Update the metric state self.closest_vs_positive_distances.append(difference) self.positive_distances.append( closest_positive_distances[non_inf_mask] @@ -179,7 +136,6 @@ def update(self, inputs: Tensor, target: Tensor): def compute(self): if len(self.all_distances) == 0: - # Return NaN tensor if no distances were calculated return { "MedianDistance": torch.tensor(float("nan")), "MedianClosestDistance": torch.tensor(float("nan")), @@ -196,7 +152,6 @@ def compute(self): self.closest_vs_positive_distances ) - # Return medians return { "MedianDistance": torch.median(all_distances), "MedianClosestDistance": torch.median(closest_distances), @@ -220,34 +175,21 @@ def _pairwise_distances(embeddings, squared=False): batch_size) @rtype: torch.Tensor """ - # Get the dot product between all embeddings - # shape (batch_size, batch_size) dot_product = torch.matmul(embeddings, embeddings.t()) - # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`. - # This also provides more numerical stability (the diagonal of the result will be exactly 0). - # shape (batch_size,) square_norm = torch.diag(dot_product) - # Compute the pairwise distance matrix as we have: - # ||a - b||^2 = ||a||^2 - 2 + ||b||^2 - # shape (batch_size, batch_size) distances = ( square_norm.unsqueeze(0) - 2.0 * dot_product + square_norm.unsqueeze(1) ) - - # Because of computation errors, some distances might be negative so we put everything >= 0.0 distances = torch.max(distances, torch.tensor(0.0)) if not squared: - # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal) - # we need to add a small epsilon where distances == 0.0 mask = (distances == 0.0).float() distances = distances + mask * 1e-16 distances = torch.sqrt(distances) - # Correct the epsilon added: set the distances on the mask to be exactly 0.0 distances = distances * (1.0 - mask) return distances diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py index f3591c83..b368b2c5 100644 --- a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -4,19 +4,17 @@ from sklearn.manifold import TSNE from torch import Tensor -from luxonis_train.utils import Labels, Packet +from luxonis_train.enums import Metadata from .base_visualizer import BaseVisualizer -from .utils import ( - figure_to_torch, -) +from .utils import figure_to_torch logger = logging.getLogger(__name__) log_disable = False class EmbeddingsVisualizer(BaseVisualizer[Tensor, Tensor]): - # supported_tasks: list[TaskType] = [TaskType.LABEL] + supported_tasks = [Metadata("id")] def __init__( self, @@ -25,17 +23,6 @@ def __init__( """Visualizer for embedding tasks like reID.""" super().__init__(**kwargs) - def prepare( - self, inputs: Packet[Tensor], labels: Labels | None - ) -> tuple[Tensor, Tensor]: - embeddings = inputs["features"][0] - - assert ( - labels is not None and "id" in labels - ), "ID labels are required for metric learning losses" - ids = labels["id"][0] - return embeddings, ids - def forward( self, label_canvas: Tensor, @@ -58,20 +45,13 @@ def forward( @return: An embedding space projection. """ - # Embeddings: [B, D], D = e.g. 512 - # ids: [B, 1], corresponding to the embeddings - - # Convert embeddings to numpy array embeddings_np = embeddings.detach().cpu().numpy() - # Perplexity must be less than the number of samples perplexity = min(30, embeddings_np.shape[0] - 1) - # Reduce dimensionality to 2D using t-SNE tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity) embeddings_2d = tsne.fit_transform(embeddings_np) - # Plot the embeddings fig, ax = plt.subplots(figsize=(10, 10)) scatter = ax.scatter( embeddings_2d[:, 0], @@ -86,15 +66,12 @@ def forward( ax.set_xlabel("Dimension 1") ax.set_ylabel("Dimension 2") - # Convert figure to tensor image_tensor = figure_to_torch( fig, width=label_canvas.shape[3], height=label_canvas.shape[2] ) - # Close the figure to free memory plt.close(fig) - # Add fake batch dimension image_tensor = image_tensor.unsqueeze(0) return image_tensor diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index d44ac480..9ee9cd1d 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -243,7 +243,7 @@ def check_attached_modules(cls, data: Params) -> Params: else: warnings.warn( f"Field `model.{section}` is deprecated. " - f"Please specify `{section}`under " + f"Please specify `{section}` under " "the node they are attached to." ) for node in data["nodes"]: diff --git a/luxonis_train/enums.py b/luxonis_train/enums.py index 09d38fb2..04d4a241 100644 --- a/luxonis_train/enums.py +++ b/luxonis_train/enums.py @@ -1,12 +1,38 @@ +from dataclasses import dataclass from enum import Enum +from typing import TypeAlias class TaskType(str, Enum): - """Tasks supported by nodes in LuxonisTrain.""" - CLASSIFICATION = "classification" SEGMENTATION = "segmentation" INSTANCE_SEGMENTATION = "instance_segmentation" BOUNDINGBOX = "boundingbox" KEYPOINTS = "keypoints" ARRAY = "array" + + +@dataclass +class Metadata: + name: str + + @property + def value(self): + return f"metadata/{self.name}" + + def __str__(self) -> str: + return self.value + + def __hash__(self) -> int: + return hash(self.name) + + +Task: TypeAlias = TaskType | Metadata +# class TaskType: +# +# CLASSIFICATION = SimpleTask.CLASSIFICATION +# SEGMENTATION = SimpleTask.SEGMENTATION +# INSTANCE_SEGMENTATION = SimpleTask.INSTANCE_SEGMENTATION +# BOUNDINGBOX = SimpleTask.BOUNDINGBOX +# KEYPOINTS = SimpleTask.KEYPOINTS +# ARRAY = SimpleTask.ARRAY diff --git a/luxonis_train/loaders/utils.py b/luxonis_train/loaders/utils.py index 9c9e1d45..e96b15bf 100644 --- a/luxonis_train/loaders/utils.py +++ b/luxonis_train/loaders/utils.py @@ -1,5 +1,5 @@ import torch -from luxonis_ml.data.utils import get_task_type +from luxonis_ml.data.utils import get_task_type, task_is_metadata from torch import Tensor from luxonis_train.utils.types import Labels @@ -44,7 +44,7 @@ def collate_fn( new_ann[:, 1:] = ann label_box.append(new_ann) out_labels[task] = torch.cat(label_box, 0) - elif task_type == "instance_segmentation": + elif task_type == "instance_segmentation" or task_is_metadata(task): out_labels[task] = torch.cat(annos, 0) else: out_labels[task] = torch.stack(annos, 0) diff --git a/luxonis_train/nodes/backbones/ghostfacenet/blocks.py b/luxonis_train/nodes/backbones/ghostfacenet/blocks.py index 46a9ba27..2cbbfae4 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/blocks.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/blocks.py @@ -1,246 +1,204 @@ import math +from typing import Literal import torch -import torch.nn as nn import torch.nn.functional as F +from torch import Tensor, nn from luxonis_train.nodes.backbones.micronet.blocks import _make_divisible from luxonis_train.nodes.blocks import SqueezeExciteBlock +from luxonis_train.nodes.blocks.blocks import ConvModule -class ModifiedGDC(nn.Module): - def __init__(self, image_size, in_chs, num_classes, dropout, emb=512): - super().__init__() - - if image_size % 32 == 0: - self.conv_dw = nn.Conv2d( - in_chs, - in_chs, - kernel_size=(image_size // 32), - groups=in_chs, - bias=False, - ) - else: - self.conv_dw = nn.Conv2d( - in_chs, - in_chs, - kernel_size=(image_size // 32 + 1), - groups=in_chs, - bias=False, - ) - self.bn1 = nn.BatchNorm2d(in_chs) - self.dropout = nn.Dropout(dropout) - - self.conv = nn.Conv2d(in_chs, emb, kernel_size=1, bias=False) - self.bn2 = nn.BatchNorm1d(emb) - self.linear = ( - nn.Linear(emb, num_classes) if num_classes else nn.Identity() - ) - - def forward(self, inps): - x = inps - x = self.conv_dw(x) - x = self.bn1(x) - x = self.dropout(x) - x = self.conv(x) - x = x.view(x.size(0), -1) - x = self.bn2(x) - x = self.linear(x) - return x +class ModifiedGDC(nn.Sequential): + def __init__( + self, + image_size: int, + in_channels: int, + dropout: float, + embedding_size: int = 512, + ): + modules = [ + ConvModule( + in_channels, + in_channels, + kernel_size=(image_size // 32) + if image_size % 32 == 0 + else (image_size // 32 + 1), + groups=in_channels, + activation=nn.Identity(), + ), + nn.Dropout(dropout), + nn.Conv2d(in_channels, embedding_size, kernel_size=1, bias=False), + nn.Flatten(), + nn.BatchNorm1d(embedding_size), + ] + super().__init__(*modules) class GhostModuleV2(nn.Module): def __init__( self, - inp, - oup, - kernel_size=1, - ratio=2, - dw_size=3, - stride=1, - prelu=True, - mode=None, - args=None, + in_channels: int, + out_channels: int, + mode: Literal["original", "attn"], + kernel_size: int = 1, + ratio: int = 2, + dw_size: int = 3, + stride: int = 1, + use_prelu: bool = True, ): - super(GhostModuleV2, self).__init__() + super().__init__() self.mode = mode - self.gate_fn = nn.Sigmoid() + self.out_channels = out_channels + intermediate_channels = math.ceil(out_channels / ratio) + new_channels = intermediate_channels * (ratio - 1) + self.primary_conv = ConvModule( + in_channels, + intermediate_channels, + kernel_size, + stride, + kernel_size // 2, + activation=nn.PReLU() if use_prelu else nn.Identity(), + ) + self.cheap_operation = ConvModule( + intermediate_channels, + new_channels, + dw_size, + 1, + dw_size // 2, + groups=intermediate_channels, + activation=nn.PReLU() if use_prelu else nn.Identity(), + ) - if self.mode in ["original"]: - self.oup = oup - init_channels = math.ceil(oup / ratio) - new_channels = init_channels * (ratio - 1) - self.primary_conv = nn.Sequential( - nn.Conv2d( - inp, - init_channels, - kernel_size, - stride, - kernel_size // 2, - bias=False, - ), - nn.BatchNorm2d(init_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - self.cheap_operation = nn.Sequential( - nn.Conv2d( - init_channels, - new_channels, - dw_size, - 1, - dw_size // 2, - groups=init_channels, - bias=False, - ), - nn.BatchNorm2d(new_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - elif self.mode in ["attn"]: # DFC - self.oup = oup - init_channels = math.ceil(oup / ratio) - new_channels = init_channels * (ratio - 1) - self.primary_conv = nn.Sequential( - nn.Conv2d( - inp, - init_channels, + if self.mode == "attn": + self.short_conv = nn.Sequential( + ConvModule( + in_channels, + out_channels, kernel_size, stride, kernel_size // 2, - bias=False, - ), - nn.BatchNorm2d(init_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - self.cheap_operation = nn.Sequential( - nn.Conv2d( - init_channels, - new_channels, - dw_size, - 1, - dw_size // 2, - groups=init_channels, - bias=False, - ), - nn.BatchNorm2d(new_channels), - nn.PReLU() if prelu else nn.Sequential(), - ) - self.short_conv = nn.Sequential( - nn.Conv2d( - inp, oup, kernel_size, stride, kernel_size // 2, bias=False + activation=nn.Identity(), ), - nn.BatchNorm2d(oup), - nn.Conv2d( - oup, - oup, + ConvModule( + out_channels, + out_channels, kernel_size=(1, 5), stride=1, padding=(0, 2), - groups=oup, - bias=False, + groups=out_channels, + activation=nn.Identity(), ), - nn.BatchNorm2d(oup), - nn.Conv2d( - oup, - oup, + ConvModule( + out_channels, + out_channels, kernel_size=(5, 1), stride=1, padding=(2, 0), - groups=oup, - bias=False, + groups=out_channels, + activation=nn.Identity(), ), - nn.BatchNorm2d(oup), + nn.AvgPool2d(kernel_size=2, stride=2), + nn.Sigmoid(), ) - def forward(self, x): - if self.mode in ["original"]: - x1 = self.primary_conv(x) - x2 = self.cheap_operation(x1) - out = torch.cat([x1, x2], dim=1) - return out[:, : self.oup, :, :] - elif self.mode in ["attn"]: - res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2)) - x1 = self.primary_conv(x) - x2 = self.cheap_operation(x1) - out = torch.cat([x1, x2], dim=1) - return out[:, : self.oup, :, :] * F.interpolate( - self.gate_fn(res), - size=(out.shape[-2], out.shape[-1]), - mode="nearest", - ) + def forward(self, x: Tensor) -> Tensor: + x1 = self.primary_conv(x) + x2 = self.cheap_operation(x1) + out = torch.cat([x1, x2], dim=1) + if self.mode == "original": + return out[:, : self.out_channels, ...] + + return out[:, : self.out_channels, ...] * F.interpolate( + self.short_conv(x), + size=(out.shape[-2], out.shape[-1]), + mode="nearest", + ) class GhostBottleneckV2(nn.Module): def __init__( self, - in_chs, - mid_chs, - out_chs, - dw_kernel_size=3, - stride=1, - act_layer=nn.PReLU, - se_ratio=0.0, - layer_id=None, - args=None, + in_channels: int, + intermediate_channels: int, + out_channels: int, + dw_kernel_size: int = 3, + stride: int = 1, + se_ratio: float = 0.0, + *, + layer_id: int, ): - super(GhostBottleneckV2, self).__init__() + super().__init__() has_se = se_ratio is not None and se_ratio > 0.0 self.stride = stride - assert layer_id is not None, "Layer ID must be explicitly provided" - # Point-wise expansion if layer_id <= 1: self.ghost1 = GhostModuleV2( - in_chs, mid_chs, prelu=True, mode="original", args=args + in_channels, + intermediate_channels, + use_prelu=True, + mode="original", ) else: self.ghost1 = GhostModuleV2( - in_chs, mid_chs, prelu=True, mode="attn", args=args + in_channels, intermediate_channels, use_prelu=True, mode="attn" ) # Depth-wise convolution if self.stride > 1: self.conv_dw = nn.Conv2d( - mid_chs, - mid_chs, + intermediate_channels, + intermediate_channels, dw_kernel_size, stride=stride, padding=(dw_kernel_size - 1) // 2, - groups=mid_chs, + groups=intermediate_channels, bias=False, ) - self.bn_dw = nn.BatchNorm2d(mid_chs) + self.bn_dw = nn.BatchNorm2d(intermediate_channels) # Squeeze-and-excitation if has_se: - reduced_chs = _make_divisible(mid_chs * se_ratio, 4) + reduced_chs = _make_divisible(intermediate_channels * se_ratio, 4) self.se = SqueezeExciteBlock( - mid_chs, reduced_chs, True, activation=nn.PReLU() + intermediate_channels, reduced_chs, True, activation=nn.PReLU() ) else: self.se = None self.ghost2 = GhostModuleV2( - mid_chs, out_chs, prelu=False, mode="original", args=args + intermediate_channels, + out_channels, + use_prelu=False, + mode="original", ) # shortcut - if in_chs == out_chs and self.stride == 1: - self.shortcut = nn.Sequential() + if in_channels == out_channels and self.stride == 1: + self.shortcut = nn.Identity() else: self.shortcut = nn.Sequential( nn.Conv2d( - in_chs, - in_chs, + in_channels, + in_channels, dw_kernel_size, stride=stride, padding=(dw_kernel_size - 1) // 2, - groups=in_chs, + groups=in_channels, + bias=False, + ), + nn.BatchNorm2d(in_channels), + nn.Conv2d( + in_channels, + out_channels, + 1, + stride=1, + padding=0, bias=False, ), - nn.BatchNorm2d(in_chs), - nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False), - nn.BatchNorm2d(out_chs), + nn.BatchNorm2d(out_channels), ) def forward(self, x): diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index 5a99ae28..3b9dc7b2 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -5,6 +5,7 @@ import torch.nn as nn from torch import Tensor +from luxonis_train.enums import Metadata from luxonis_train.nodes.backbones.ghostfacenet.blocks import ( GhostBottleneckV2, ModifiedGDC, @@ -16,15 +17,14 @@ class GhostFaceNetsV2(BaseNode[Tensor, list[Tensor]]): - in_channels: list[int] - in_width: list[int] + in_channels: int + in_width: int + tasks = [Metadata("id")] def __init__( self, - embedding_size=512, - num_classes=-1, + embedding_size: int = 512, variant: Literal["V2"] = "V2", - *args, **kwargs, ): """GhostFaceNetsV2 backbone. @@ -42,20 +42,15 @@ def __init__( @type embedding_size: int @param embedding_size: Size of the embedding. Defaults to 512. - @type num_classes: int - @param num_classes: Number of classes. Defaults to -1, which leaves the default variant value in. Otherwise it can be used to - have the network return raw embeddings (=0) or add another linear layer to the network, which is useful for training using - ArcFace or similar classification-based losses that require the user to drop the last layer of the network. @type variant: Literal["V2"] @param variant: Variant of the GhostFaceNets embedding model. Defaults to "V2" (which is the only variant available). """ - super().__init__(*args, **kwargs) + super().__init__(**kwargs) + self.embedding_size = embedding_size - image_size = self.in_width[0] - channels = self.in_channels[0] + image_size = self.in_width + channels = self.in_channels var = get_variant(variant) - if num_classes >= 0: - var.num_classes = num_classes self.cfgs = var.cfgs # Building first layer @@ -89,7 +84,6 @@ def __init__( b_cfg.stride, se_ratio=b_cfg.se_ratio, layer_id=layer_id, - args=var.block_args, ) ) input_channel = output_channel @@ -110,13 +104,9 @@ def __init__( self.blocks = nn.Sequential(*stages) - # Building pointwise convolution - pointwise_conv = [nn.Sequential()] - self.pointwise_conv = nn.Sequential(*pointwise_conv) - self.classifier = ModifiedGDC( + self.head = ModifiedGDC( image_size, output_channel, - var.num_classes, var.dropout, embedding_size, ) @@ -133,12 +123,10 @@ def __init__( if isinstance(m, nn.BatchNorm2d): m.momentum, m.eps = var.bn_momentum, var.bn_epsilon - def forward(self, inps): - x = inps[0] + def forward(self, x: Tensor) -> Tensor: x = self.conv_stem(x) x = self.bn1(x) x = self.act1(x) x = self.blocks(x) - x = self.pointwise_conv(x) - x = self.classifier(x) + x = self.head(x) return x diff --git a/luxonis_train/nodes/backbones/micronet/blocks.py b/luxonis_train/nodes/backbones/micronet/blocks.py index b29082cf..72beb66e 100644 --- a/luxonis_train/nodes/backbones/micronet/blocks.py +++ b/luxonis_train/nodes/backbones/micronet/blocks.py @@ -414,11 +414,8 @@ def forward(self, x: Tensor) -> Tensor: return out -def _make_divisible( - value: int, divisor: int, min_value: int | None = None -) -> int: - if min_value is None: - min_value = divisor +def _make_divisible(value: int, divisor: int) -> int: + min_value = divisor new_v = max(min_value, int(value + divisor / 2) // divisor * divisor) # Make sure that round down does not go down by more than 10%. if new_v < 0.9 * value: diff --git a/luxonis_train/nodes/base_node.py b/luxonis_train/nodes/base_node.py index 0a9a208a..a35c5edd 100644 --- a/luxonis_train/nodes/base_node.py +++ b/luxonis_train/nodes/base_node.py @@ -9,7 +9,7 @@ from torch import Size, Tensor, nn from typeguard import TypeCheckError, check_type -from luxonis_train.enums import TaskType +from luxonis_train.enums import Task, TaskType from luxonis_train.utils import ( AttachIndexType, DatasetMetadata, @@ -107,7 +107,7 @@ def wrap(output: Tensor) -> Packet[Tensor]: """ attach_index: AttachIndexType - tasks: list[TaskType] | None = None + tasks: list[Task] | None = None def __init__( self, diff --git a/luxonis_train/nodes/blocks/blocks.py b/luxonis_train/nodes/blocks/blocks.py index 25bea7c5..e873dccd 100644 --- a/luxonis_train/nodes/blocks/blocks.py +++ b/luxonis_train/nodes/blocks/blocks.py @@ -86,10 +86,10 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: int, - stride: int = 1, - padding: int = 0, - dilation: int = 1, + kernel_size: int | tuple[int, int], + stride: int | tuple[int, int] = 1, + padding: int | tuple[int, int] = 0, + dilation: int | tuple[int, int] = 1, groups: int = 1, bias: bool = False, activation: nn.Module | None = None, diff --git a/luxonis_train/utils/dataset_metadata.py b/luxonis_train/utils/dataset_metadata.py index fdbec775..2c10905c 100644 --- a/luxonis_train/utils/dataset_metadata.py +++ b/luxonis_train/utils/dataset_metadata.py @@ -62,7 +62,7 @@ def n_classes(self, task: str | None = None) -> int: for classes in self._classes.values(): if len(classes) != n_classes: raise RuntimeError( - "The dataset contains different number of classes for different tasks." + "The dataset contains different number of classes for different tasks. " "Please specify the 'task' argument to get the number of classes." ) return n_classes @@ -90,7 +90,7 @@ def n_keypoints(self, task: str | None = None) -> int: for n in self._n_keypoints.values(): if n != n_keypoints: raise RuntimeError( - "The dataset contains different number of keypoints for different tasks." + "The dataset contains different number of keypoints for different tasks. " "Please specify the 'task' argument to get the number of keypoints." ) return n_keypoints diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml index c79e4f8e..e3f5dfad 100644 --- a/tests/configs/reid.yaml +++ b/tests/configs/reid.yaml @@ -5,43 +5,44 @@ model: name: reid_test nodes: - name: GhostFaceNetsV2 - input_sources: - - image params: embedding_size: &embedding_size 512 - - losses: - - name: EmbeddingLossWrapper - params: - loss_name: SupConLoss - embedding_size: *embedding_size - cross_batch_memory_size: &memory_size 4 - attached_to: GhostFaceNetsV2 - - metrics: - - name: ClosestIsPositiveAccuracy - params: - cross_batch_memory_size: *memory_size - attached_to: GhostFaceNetsV2 - is_main_metric: True - - name: MedianDistances - params: - cross_batch_memory_size: *memory_size - attached_to: GhostFaceNetsV2 - is_main_metric: False - visualizers: - - name: EmbeddingsVisualizer - attached_to: GhostFaceNetsV2 + losses: + - name: SupConLoss + # params: + # embedding_size: *embedding_size + # cross_batch_memory_size: &memory_size 4 + + metrics: + - name: ClosestIsPositiveAccuracy + # params: + # cross_batch_memory_size: *memory_size + is_main_metric: True + + - name: MedianDistances + # params: + # cross_batch_memory_size: *memory_size + + visualizers: + - name: EmbeddingsVisualizer + +loader: + train_view: train + val_view: [test_query, test_gallery] + test_view: [test_query, test_gallery] + + params: + dataset_name: reid_dataset trainer: preprocessing: train_image_size: [256, 256] batch_size: 16 - epochs: 10 + epochs: 1 n_workers: 0 - validation_interval: 10 + validation_interval: 1 callbacks: - name: ExportOnTrainEnd @@ -50,11 +51,3 @@ trainer: name: Adam params: lr: 0.01 - -tracker: - project_name: reid_example - is_tensorboard: True - -exporter: - onnx: - opset_version: 11 \ No newline at end of file diff --git a/tests/integration/test_reid.py b/tests/integration/test_reid.py deleted file mode 100644 index 53355025..00000000 --- a/tests/integration/test_reid.py +++ /dev/null @@ -1,184 +0,0 @@ -import shutil -from pathlib import Path -from typing import Any - -import pytest -import torch - -from luxonis_train.attached_modules.losses.pml_loss import ( - ALL_EMBEDDING_LOSSES, - CLASS_EMBEDDING_LOSSES, -) -from luxonis_train.core import LuxonisModel -from luxonis_train.enums import TaskType -from luxonis_train.loaders import BaseLoaderTorch - -from .multi_input_modules import * - -INFER_PATH = Path("tests/integration/infer-save-directory") -ONNX_PATH = Path("tests/integration/_model.onnx") -STUDY_PATH = Path("study_local.db") - -NUM_INDIVIDUALS = 100 - - -class CustomReIDLoader(BaseLoaderTorch): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - @property - def input_shapes(self): - return { - "image": torch.Size([3, 256, 256]), - "id": torch.Size([1]), - } - - def __getitem__(self, _): # pragma: no cover - # Fake data - image = torch.rand(self.input_shapes["image"], dtype=torch.float32) - inputs = { - "image": image, - } - - # Fake labels - id = torch.randint(0, NUM_INDIVIDUALS, (1,), dtype=torch.int64) - labels = { - "id": (id, TaskType.LABEL), - } - - return inputs, labels - - def __len__(self): - return 10 - - def get_classes(self) -> dict[TaskType, list[str]]: - return {TaskType.LABEL: ["id"]} - - -class CustomReIDLoaderNoID(CustomReIDLoader): - def __getitem__(self, _): - inputs, labels = super().__getitem__(_) - labels["something_else"] = labels["id"] - del labels["id"] - - return inputs, labels - - -class CustomReIDLoaderImageSize2(CustomReIDLoader): - @property - def input_shapes(self): - return { - "image": torch.Size([3, 200, 200]), - "id": torch.Size([1]), - } - - -@pytest.fixture -def infer_path() -> Path: - if INFER_PATH.exists(): - shutil.rmtree(INFER_PATH) - INFER_PATH.mkdir() - return INFER_PATH - - -@pytest.fixture -def opts(test_output_dir: Path) -> dict[str, Any]: - return { - "trainer.epochs": 1, - "trainer.batch_size": 2, - "trainer.validation_interval": 1, - "trainer.callbacks": "[]", - "tracker.save_directory": str(test_output_dir), - "tuner.n_trials": 4, - } - - -@pytest.fixture(scope="function", autouse=True) -def clear_files(): - yield - STUDY_PATH.unlink(missing_ok=True) - ONNX_PATH.unlink(missing_ok=True) - - -not_class_based_losses = ALL_EMBEDDING_LOSSES.copy() -for loss in CLASS_EMBEDDING_LOSSES: - not_class_based_losses.remove(loss) - - -@pytest.mark.parametrize("loss_name", not_class_based_losses) -def test_available_losses( - opts: dict[str, Any], infer_path: Path, loss_name: str -): - config_file = "tests/configs/reid.yaml" - opts["model.losses.0.params.loss_name"] = loss_name - - # if loss_name in CLASS_EMBEDDING_LOSSES: - # opts["model.losses.0.params.num_classes"] = NUM_INDIVIDUALS - # opts["model.nodes.0.params.num_classes"] = NUM_INDIVIDUALS - # else: - # opts["model.losses.0.params.num_classes"] = 0 - # opts["model.nodes.0.params.num_classes"] = 0 - - if loss_name == "RankedListLoss": - opts["model.losses.0.params.loss_kwargs"] = {"margin": 1.0, "Tn": 0.5} - - model = LuxonisModel(config_file, opts) - model.train() - model.test(view="val") - - assert not ONNX_PATH.exists() - model.export(str(ONNX_PATH)) - assert ONNX_PATH.exists() - - assert len(list(infer_path.iterdir())) == 0 - model.infer(view="val", save_dir=infer_path) - assert infer_path.exists() - - -@pytest.mark.parametrize("loss_name", CLASS_EMBEDDING_LOSSES) -@pytest.mark.parametrize("num_classes", [-2, NUM_INDIVIDUALS]) -def test_unsupported_class_based_losses( - opts: dict[str, Any], loss_name: str, num_classes: int -): - config_file = "tests/configs/reid.yaml" - opts["model.losses.0.params.loss_name"] = loss_name - opts["model.losses.0.params.num_classes"] = num_classes - opts["model.nodes.0.params.num_classes"] = num_classes - - with pytest.raises(ValueError): - LuxonisModel(config_file, opts) - - -@pytest.mark.parametrize("loss_name", ["NonExistentLoss"]) -def test_nonexistent_losses(opts: dict[str, Any], loss_name: str): - config_file = "tests/configs/reid.yaml" - opts["model.losses.0.params.loss_name"] = loss_name - - with pytest.raises(ValueError): - LuxonisModel(config_file, opts) - - -def test_bad_loader(opts: dict[str, Any]): - config_file = "tests/configs/reid.yaml" - opts["loader.name"] = "CustomReIDLoaderNoID" - - with pytest.raises(ValueError): - model = LuxonisModel(config_file, opts) - model.train() - - -def test_not_enough_samples_for_metrics(opts: dict[str, Any]): - config_file = "tests/configs/reid.yaml" - opts["model.metrics.1.params.cross_batch_memory_size"] = 100 - - model = LuxonisModel(config_file, opts) - model.train() - - -def test_image_size_not_divisible_by_32(opts: dict[str, Any]): - config_file = "tests/configs/reid.yaml" - opts["loader.name"] = "CustomReIDLoaderImageSize2" - - # with pytest.raises(ValueError): - model = LuxonisModel(config_file, opts) - model.train() From 3dfb8b2aec9ecb83b0343294c67e338629dbfd96 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 22 Jan 2025 22:25:22 -0500 Subject: [PATCH 37/57] renamed --- luxonis_train/attached_modules/losses/pml_loss.py | 4 ++-- luxonis_train/nodes/backbones/__init__.py | 4 ++-- luxonis_train/nodes/backbones/ghostfacenet/__init__.py | 4 ++-- luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py | 2 +- tests/configs/reid.yaml | 2 +- tests/integration/test_detection.py | 2 +- tests/integration/test_segmentation.py | 2 +- 7 files changed, 10 insertions(+), 10 deletions(-) diff --git a/luxonis_train/attached_modules/losses/pml_loss.py b/luxonis_train/attached_modules/losses/pml_loss.py index 087963c0..e0efdc55 100644 --- a/luxonis_train/attached_modules/losses/pml_loss.py +++ b/luxonis_train/attached_modules/losses/pml_loss.py @@ -5,7 +5,7 @@ from torch import Tensor from luxonis_train.enums import Metadata -from luxonis_train.nodes.backbones.ghostfacenet import GhostFaceNetsV2 +from luxonis_train.nodes.backbones.ghostfacenet import GhostFaceNetV2 from luxonis_train.nodes.base_node import BaseNode from .base_loss import BaseLoss @@ -42,7 +42,7 @@ class EmbeddingLossWrapper( BaseLoss[Tensor, Tensor], register_name=loss_name ): - node: GhostFaceNetsV2 + node: GhostFaceNetV2 supported_tasks = [Metadata("id")] def __init__(self, *, node: BaseNode | None = None, **kwargs): diff --git a/luxonis_train/nodes/backbones/__init__.py b/luxonis_train/nodes/backbones/__init__.py index da063a5e..a15d2591 100644 --- a/luxonis_train/nodes/backbones/__init__.py +++ b/luxonis_train/nodes/backbones/__init__.py @@ -2,7 +2,7 @@ from .ddrnet import DDRNet from .efficientnet import EfficientNet from .efficientrep import EfficientRep -from .ghostfacenet.ghostfacenet import GhostFaceNetsV2 +from .ghostfacenet import GhostFaceNetV2 from .micronet import MicroNet from .mobilenetv2 import MobileNetV2 from .mobileone import MobileOne @@ -23,5 +23,5 @@ "ResNet", "DDRNet", "RecSubNet", - "GhostFaceNetsV2", + "GhostFaceNetV2", ] diff --git a/luxonis_train/nodes/backbones/ghostfacenet/__init__.py b/luxonis_train/nodes/backbones/ghostfacenet/__init__.py index 85ed4447..c24d9afb 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/__init__.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/__init__.py @@ -1,3 +1,3 @@ -from .ghostfacenet import GhostFaceNetsV2 +from .ghostfacenet import GhostFaceNetV2 -__all__ = ["GhostFaceNetsV2"] +__all__ = ["GhostFaceNetV2"] diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index 3b9dc7b2..82872ae3 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -16,7 +16,7 @@ from luxonis_train.nodes.blocks import ConvModule -class GhostFaceNetsV2(BaseNode[Tensor, list[Tensor]]): +class GhostFaceNetV2(BaseNode[Tensor, Tensor]): in_channels: int in_width: int tasks = [Metadata("id")] diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml index e3f5dfad..0eea37bd 100644 --- a/tests/configs/reid.yaml +++ b/tests/configs/reid.yaml @@ -4,7 +4,7 @@ loader: model: name: reid_test nodes: - - name: GhostFaceNetsV2 + - name: GhostFaceNetV2 params: embedding_size: &embedding_size 512 diff --git a/tests/integration/test_detection.py b/tests/integration/test_detection.py index 15edf777..db7fe004 100644 --- a/tests/integration/test_detection.py +++ b/tests/integration/test_detection.py @@ -103,7 +103,7 @@ def train_and_test( @pytest.mark.parametrize( - "backbone", [b for b in BACKBONES if b != "GhostFaceNetsV2"] + "backbone", [b for b in BACKBONES if b != "GhostFaceNetV2"] ) def test_backbones( backbone: str, diff --git a/tests/integration/test_segmentation.py b/tests/integration/test_segmentation.py index f66c2602..275b35e3 100644 --- a/tests/integration/test_segmentation.py +++ b/tests/integration/test_segmentation.py @@ -124,7 +124,7 @@ def train_and_test( @pytest.mark.parametrize( - "backbone", [b for b in BACKBONES if b != "GhostFaceNetsV2"] + "backbone", [b for b in BACKBONES if b != "GhostFaceNetV2"] ) def test_backbones( backbone: str, From f16aad401266231481346c7d48d590c8f737b073 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 22 Jan 2025 22:34:05 -0500 Subject: [PATCH 38/57] simplified --- .../backbones/ghostfacenet/ghostfacenet.py | 52 +++++++++---------- .../nodes/backbones/ghostfacenet/variants.py | 12 +---- 2 files changed, 26 insertions(+), 38 deletions(-) diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index 82872ae3..87e12c21 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -6,15 +6,13 @@ from torch import Tensor from luxonis_train.enums import Metadata -from luxonis_train.nodes.backbones.ghostfacenet.blocks import ( - GhostBottleneckV2, - ModifiedGDC, -) from luxonis_train.nodes.backbones.ghostfacenet.variants import get_variant from luxonis_train.nodes.backbones.micronet.blocks import _make_divisible from luxonis_train.nodes.base_node import BaseNode from luxonis_train.nodes.blocks import ConvModule +from .blocks import GhostBottleneckV2, ModifiedGDC + class GhostFaceNetV2(BaseNode[Tensor, Tensor]): in_channels: int @@ -49,23 +47,22 @@ def __init__( self.embedding_size = embedding_size image_size = self.in_width - channels = self.in_channels var = get_variant(variant) - self.cfgs = var.cfgs - - # Building first layer output_channel = _make_divisible(int(16 * var.width), 4) - self.conv_stem = nn.Conv2d( - channels, output_channel, 3, 2, 1, bias=False - ) - self.bn1 = nn.BatchNorm2d(output_channel) - self.act1 = nn.PReLU() input_channel = output_channel - # Building Ghost BottleneckV2 blocks - stages = [] + stages: list[nn.Module] = [ + ConvModule( + self.in_channels, + output_channel, + kernel_size=3, + stride=2, + padding=1, + activation=nn.PReLU(), + ) + ] layer_id = 0 - for cfg in self.cfgs: + for cfg in var.block_configs: layers = [] for b_cfg in cfg: output_channel = _make_divisible( @@ -111,22 +108,21 @@ def __init__( embedding_size, ) - # Initializing weights + self._init_weights() + + def _init_weights(self): for m in self.modules(): - if var.init_kaiming: - if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): - fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) - negative_slope = 0.25 - m.weight.data.normal_( - 0, math.sqrt(2.0 / (fan_in * (1 + negative_slope**2))) - ) + if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) + negative_slope = 0.25 + m.weight.data.normal_( + 0, math.sqrt(2.0 / (fan_in * (1 + negative_slope**2))) + ) if isinstance(m, nn.BatchNorm2d): - m.momentum, m.eps = var.bn_momentum, var.bn_epsilon + m.momentum = 0.9 + m.eps = 1e-5 def forward(self, x: Tensor) -> Tensor: - x = self.conv_stem(x) - x = self.bn1(x) - x = self.act1(x) x = self.blocks(x) x = self.head(x) return x diff --git a/luxonis_train/nodes/backbones/ghostfacenet/variants.py b/luxonis_train/nodes/backbones/ghostfacenet/variants.py index aa78daf8..4e032675 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/variants.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/variants.py @@ -52,11 +52,7 @@ class GhostFaceNetsVariant(BaseModel): dropout: float block: type[nn.Module] add_pointwise_conv: bool - bn_momentum: float - bn_epsilon: float - init_kaiming: bool - block_args: dict | None - cfgs: List[List[BlockConfig]] + block_configs: List[List[BlockConfig]] V2 = GhostFaceNetsVariant( @@ -65,11 +61,7 @@ class GhostFaceNetsVariant(BaseModel): dropout=0.2, block=GhostBottleneckV2, add_pointwise_conv=False, - bn_momentum=0.9, - bn_epsilon=1e-5, - init_kaiming=True, - block_args=None, - cfgs=[ + block_configs=[ [ BlockConfig( kernel_size=3, From 2fe723e0e177fcbcb3d47674a3cff00b372e1884 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 22 Jan 2025 22:42:59 -0500 Subject: [PATCH 39/57] separated head --- .../attached_modules/losses/pml_loss.py | 8 ++- .../backbones/ghostfacenet/ghostfacenet.py | 22 +------ luxonis_train/nodes/heads/__init__.py | 2 + .../nodes/heads/ghostfacenet_head.py | 63 +++++++++++++++++++ tests/configs/reid.yaml | 2 + 5 files changed, 76 insertions(+), 21 deletions(-) create mode 100644 luxonis_train/nodes/heads/ghostfacenet_head.py diff --git a/luxonis_train/attached_modules/losses/pml_loss.py b/luxonis_train/attached_modules/losses/pml_loss.py index e0efdc55..70a87329 100644 --- a/luxonis_train/attached_modules/losses/pml_loss.py +++ b/luxonis_train/attached_modules/losses/pml_loss.py @@ -5,8 +5,8 @@ from torch import Tensor from luxonis_train.enums import Metadata -from luxonis_train.nodes.backbones.ghostfacenet import GhostFaceNetV2 from luxonis_train.nodes.base_node import BaseNode +from luxonis_train.nodes.heads.ghostfacenet_head import GhostFaceNetHead from .base_loss import BaseLoss @@ -42,7 +42,7 @@ class EmbeddingLossWrapper( BaseLoss[Tensor, Tensor], register_name=loss_name ): - node: GhostFaceNetV2 + node: GhostFaceNetHead supported_tasks = [Metadata("id")] def __init__(self, *, node: BaseNode | None = None, **kwargs): @@ -64,3 +64,7 @@ def __init__(self, *, node: BaseNode | None = None, **kwargs): def forward(self, inputs: Tensor, target: Tensor) -> Tensor: return self.loss_func(inputs, target) + + @property + def name(self) -> str: + return loss_name # noqa: B023 diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index 87e12c21..e223608b 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -5,23 +5,20 @@ import torch.nn as nn from torch import Tensor -from luxonis_train.enums import Metadata from luxonis_train.nodes.backbones.ghostfacenet.variants import get_variant from luxonis_train.nodes.backbones.micronet.blocks import _make_divisible from luxonis_train.nodes.base_node import BaseNode from luxonis_train.nodes.blocks import ConvModule -from .blocks import GhostBottleneckV2, ModifiedGDC +from .blocks import GhostBottleneckV2 class GhostFaceNetV2(BaseNode[Tensor, Tensor]): in_channels: int in_width: int - tasks = [Metadata("id")] def __init__( self, - embedding_size: int = 512, variant: Literal["V2"] = "V2", **kwargs, ): @@ -38,15 +35,11 @@ def __init__( @see: U{GhostFaceNets: Lightweight Face Recognition Model From Cheap Operations } - @type embedding_size: int - @param embedding_size: Size of the embedding. Defaults to 512. @type variant: Literal["V2"] @param variant: Variant of the GhostFaceNets embedding model. Defaults to "V2" (which is the only variant available). """ super().__init__(**kwargs) - self.embedding_size = embedding_size - image_size = self.in_width var = get_variant(variant) output_channel = _make_divisible(int(16 * var.width), 4) input_channel = output_channel @@ -101,16 +94,9 @@ def __init__( self.blocks = nn.Sequential(*stages) - self.head = ModifiedGDC( - image_size, - output_channel, - var.dropout, - embedding_size, - ) - self._init_weights() - def _init_weights(self): + def _init_weights(self) -> None: for m in self.modules(): if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) @@ -123,6 +109,4 @@ def _init_weights(self): m.eps = 1e-5 def forward(self, x: Tensor) -> Tensor: - x = self.blocks(x) - x = self.head(x) - return x + return self.blocks(x) diff --git a/luxonis_train/nodes/heads/__init__.py b/luxonis_train/nodes/heads/__init__.py index e5abd973..842efdc4 100644 --- a/luxonis_train/nodes/heads/__init__.py +++ b/luxonis_train/nodes/heads/__init__.py @@ -6,6 +6,7 @@ from .efficient_bbox_head import EfficientBBoxHead from .efficient_keypoint_bbox_head import EfficientKeypointBBoxHead from .fomo_head import FOMOHead +from .ghostfacenet_head import GhostFaceNetHead from .segmentation_head import SegmentationHead __all__ = [ @@ -17,5 +18,6 @@ "SegmentationHead", "DDRNetSegmentationHead", "DiscSubNetHead", + "GhostFaceNetHead", "FOMOHead", ] diff --git a/luxonis_train/nodes/heads/ghostfacenet_head.py b/luxonis_train/nodes/heads/ghostfacenet_head.py new file mode 100644 index 00000000..e4753d03 --- /dev/null +++ b/luxonis_train/nodes/heads/ghostfacenet_head.py @@ -0,0 +1,63 @@ +# Original source: https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py +import math + +import torch.nn as nn +from torch import Tensor + +from luxonis_train.enums import Metadata +from luxonis_train.nodes.backbones.ghostfacenet.blocks import ( + ModifiedGDC, +) +from luxonis_train.nodes.base_node import BaseNode + + +class GhostFaceNetHead(BaseNode[Tensor, list[Tensor]]): + in_channels: int + in_width: int + tasks = [Metadata("id")] + + def __init__( + self, embedding_size: int = 512, dropout: float = 0.2, **kwargs + ): + """GhostFaceNetV2 backbone. + + GhostFaceNetV2 is a convolutional neural network architecture focused on face recognition, but it is + adaptable to generic embedding tasks. It is based on the GhostNet architecture and uses Ghost BottleneckV2 blocks. + + Source: U{https://github.com/Hazqeel09/ellzaf_ml/blob/main/ellzaf_ml/models/ghostfacenetsv2.py} + + @license: U{MIT License + } + + @see: U{GhostFaceNets: Lightweight Face Recognition Model From Cheap Operations + } + + @type embedding_size: int + @param embedding_size: Size of the embedding. Defaults to 512. + """ + super().__init__(**kwargs) + self.embedding_size = embedding_size + + self.head = ModifiedGDC( + self.original_in_shape[1], + self.in_channels, + dropout, + self.embedding_size, + ) + + self._init_weights() + + def _init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) + negative_slope = 0.25 + m.weight.data.normal_( + 0, math.sqrt(2.0 / (fan_in * (1 + negative_slope**2))) + ) + if isinstance(m, nn.BatchNorm2d): + m.momentum = 0.9 + m.eps = 1e-5 + + def forward(self, x: Tensor) -> Tensor: + return self.head(x) diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml index 0eea37bd..88ffad66 100644 --- a/tests/configs/reid.yaml +++ b/tests/configs/reid.yaml @@ -5,6 +5,8 @@ model: name: reid_test nodes: - name: GhostFaceNetV2 + + - name: GhostFaceNetHead params: embedding_size: &embedding_size 512 From 45ade94a188cde9bad639e8bdbafbb4bab1919cf Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Wed, 22 Jan 2025 23:48:30 -0500 Subject: [PATCH 40/57] simplified --- .../nodes/backbones/ddrnet/ddrnet.py | 8 ++--- .../nodes/backbones/ghostfacenet/blocks.py | 36 +++---------------- .../backbones/ghostfacenet/ghostfacenet.py | 35 ++++++++---------- .../nodes/backbones/ghostfacenet/variants.py | 13 ++----- .../nodes/backbones/micronet/blocks.py | 8 ++--- .../nodes/backbones/mobileone/blocks.py | 4 +-- luxonis_train/nodes/backbones/rexnetv1.py | 2 +- luxonis_train/nodes/blocks/blocks.py | 23 ++++++------ .../nodes/heads/ghostfacenet_head.py | 27 +++++++++----- 9 files changed, 64 insertions(+), 92 deletions(-) diff --git a/luxonis_train/nodes/backbones/ddrnet/ddrnet.py b/luxonis_train/nodes/backbones/ddrnet/ddrnet.py index 2698c26d..6bee5dfc 100644 --- a/luxonis_train/nodes/backbones/ddrnet/ddrnet.py +++ b/luxonis_train/nodes/backbones/ddrnet/ddrnet.py @@ -148,7 +148,7 @@ def __init__( out_channels=highres_channels, kernel_size=1, bias=False, - activation=nn.Identity(), + activation=False, ) ) self.down3.append( @@ -159,7 +159,7 @@ def __init__( stride=2, padding=1, bias=False, - activation=nn.Identity(), + activation=False, ) ) self.layer3_skip.append( @@ -180,7 +180,7 @@ def __init__( out_channels=highres_channels, kernel_size=1, bias=False, - activation=nn.Identity(), + activation=False, ) self.down4 = nn.Sequential( @@ -200,7 +200,7 @@ def __init__( stride=2, padding=1, bias=False, - activation=nn.Identity(), + activation=False, ), ) diff --git a/luxonis_train/nodes/backbones/ghostfacenet/blocks.py b/luxonis_train/nodes/backbones/ghostfacenet/blocks.py index 2cbbfae4..118d61ac 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/blocks.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/blocks.py @@ -10,32 +10,6 @@ from luxonis_train.nodes.blocks.blocks import ConvModule -class ModifiedGDC(nn.Sequential): - def __init__( - self, - image_size: int, - in_channels: int, - dropout: float, - embedding_size: int = 512, - ): - modules = [ - ConvModule( - in_channels, - in_channels, - kernel_size=(image_size // 32) - if image_size % 32 == 0 - else (image_size // 32 + 1), - groups=in_channels, - activation=nn.Identity(), - ), - nn.Dropout(dropout), - nn.Conv2d(in_channels, embedding_size, kernel_size=1, bias=False), - nn.Flatten(), - nn.BatchNorm1d(embedding_size), - ] - super().__init__(*modules) - - class GhostModuleV2(nn.Module): def __init__( self, @@ -59,7 +33,7 @@ def __init__( kernel_size, stride, kernel_size // 2, - activation=nn.PReLU() if use_prelu else nn.Identity(), + activation=nn.PReLU() if use_prelu else False, ) self.cheap_operation = ConvModule( intermediate_channels, @@ -68,7 +42,7 @@ def __init__( 1, dw_size // 2, groups=intermediate_channels, - activation=nn.PReLU() if use_prelu else nn.Identity(), + activation=nn.PReLU() if use_prelu else False, ) if self.mode == "attn": @@ -79,7 +53,7 @@ def __init__( kernel_size, stride, kernel_size // 2, - activation=nn.Identity(), + activation=False, ), ConvModule( out_channels, @@ -88,7 +62,7 @@ def __init__( stride=1, padding=(0, 2), groups=out_channels, - activation=nn.Identity(), + activation=False, ), ConvModule( out_channels, @@ -97,7 +71,7 @@ def __init__( stride=1, padding=(2, 0), groups=out_channels, - activation=nn.Identity(), + activation=False, ), nn.AvgPool2d(kernel_size=2, stride=2), nn.Sigmoid(), diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index e223608b..66b71679 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -10,8 +10,6 @@ from luxonis_train.nodes.base_node import BaseNode from luxonis_train.nodes.blocks import ConvModule -from .blocks import GhostBottleneckV2 - class GhostFaceNetV2(BaseNode[Tensor, Tensor]): in_channels: int @@ -64,31 +62,28 @@ def __init__( hidden_channel = _make_divisible( b_cfg.expand_size * var.width, 4 ) - if var.block == GhostBottleneckV2: - layers.append( - var.block( - input_channel, - hidden_channel, - output_channel, - b_cfg.kernel_size, - b_cfg.stride, - se_ratio=b_cfg.se_ratio, - layer_id=layer_id, - ) + layers.append( + var.block( + input_channel, + hidden_channel, + output_channel, + b_cfg.kernel_size, + b_cfg.stride, + se_ratio=b_cfg.se_ratio, + layer_id=layer_id, ) + ) input_channel = output_channel layer_id += 1 stages.append(nn.Sequential(*layers)) output_channel = _make_divisible(b_cfg.expand_size * var.width, 4) stages.append( - nn.Sequential( - ConvModule( - input_channel, - output_channel, - kernel_size=1, - activation=nn.PReLU(), - ) + ConvModule( + input_channel, + output_channel, + kernel_size=1, + activation=nn.PReLU(), ) ) diff --git a/luxonis_train/nodes/backbones/ghostfacenet/variants.py b/luxonis_train/nodes/backbones/ghostfacenet/variants.py index 4e032675..9e09befc 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/variants.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/variants.py @@ -3,15 +3,15 @@ from pydantic import BaseModel from torch import nn -from luxonis_train.nodes.backbones.ghostfacenet.blocks import GhostBottleneckV2 +from .blocks import GhostBottleneckV2 class BlockConfig(BaseModel): kernel_size: int expand_size: int output_channels: int - se_ratio: float stride: int + se_ratio: float class GhostFaceNetsVariant(BaseModel): @@ -33,9 +33,6 @@ class GhostFaceNetsVariant(BaseModel): @type block: nn.Module @param block: Ghost BottleneckV2 block. Defaults to GhostBottleneckV2. - @type add_pointwise_conv: bool - @param add_pointwise_conv: If True, adds a pointwise convolution - layer at the end of the network. Defaults to False. @type bn_momentum: float @param bn_momentum: Batch normalization momentum. Defaults to 0.9. @type bn_epsilon: float @@ -47,20 +44,14 @@ class GhostFaceNetsVariant(BaseModel): @param block_args: Arguments to pass to the block. Defaults to None. """ - num_classes: int width: int - dropout: float block: type[nn.Module] - add_pointwise_conv: bool block_configs: List[List[BlockConfig]] V2 = GhostFaceNetsVariant( - num_classes=0, width=1, - dropout=0.2, block=GhostBottleneckV2, - add_pointwise_conv=False, block_configs=[ [ BlockConfig( diff --git a/luxonis_train/nodes/backbones/micronet/blocks.py b/luxonis_train/nodes/backbones/micronet/blocks.py index 72beb66e..a1fd8f13 100644 --- a/luxonis_train/nodes/backbones/micronet/blocks.py +++ b/luxonis_train/nodes/backbones/micronet/blocks.py @@ -145,7 +145,7 @@ def _create_lite_block( out_channels=out_channels, kernel_size=1, groups=group2, - activation=nn.Identity(), + activation=False, ), DYShiftMax( out_channels, @@ -179,7 +179,7 @@ def _create_transition_block( out_channels=intermediate_channels, kernel_size=1, groups=group1, - activation=nn.Identity(), + activation=False, ), DYShiftMax( intermediate_channels, @@ -217,7 +217,7 @@ def _create_full_block( out_channels=intermediate_channels, kernel_size=1, groups=groups_1[0], - activation=nn.Identity(), + activation=False, ), DYShiftMax( intermediate_channels, @@ -256,7 +256,7 @@ def _create_full_block( out_channels=out_channels, kernel_size=1, groups=group1, - activation=nn.Identity(), + activation=False, ), DYShiftMax( out_channels, diff --git a/luxonis_train/nodes/backbones/mobileone/blocks.py b/luxonis_train/nodes/backbones/mobileone/blocks.py index 4b926038..54017aa0 100644 --- a/luxonis_train/nodes/backbones/mobileone/blocks.py +++ b/luxonis_train/nodes/backbones/mobileone/blocks.py @@ -91,7 +91,7 @@ def __init__( stride=self.stride, padding=padding, groups=self.groups, - activation=nn.Identity(), + activation=False, ) ) self.rbr_conv: list[nn.Sequential] = nn.ModuleList(rbr_conv) # type: ignore @@ -106,7 +106,7 @@ def __init__( stride=self.stride, padding=0, groups=self.groups, - activation=nn.Identity(), + activation=False, ) def forward(self, inputs: Tensor) -> Tensor: diff --git a/luxonis_train/nodes/backbones/rexnetv1.py b/luxonis_train/nodes/backbones/rexnetv1.py index 6567586a..c34dbafe 100644 --- a/luxonis_train/nodes/backbones/rexnetv1.py +++ b/luxonis_train/nodes/backbones/rexnetv1.py @@ -202,7 +202,7 @@ def __init__( in_channels=dw_channels, out_channels=channels, kernel_size=1, - activation=nn.Identity(), + activation=False, ) ) diff --git a/luxonis_train/nodes/blocks/blocks.py b/luxonis_train/nodes/blocks/blocks.py index e873dccd..e8a8e251 100644 --- a/luxonis_train/nodes/blocks/blocks.py +++ b/luxonis_train/nodes/blocks/blocks.py @@ -92,7 +92,7 @@ def __init__( dilation: int | tuple[int, int] = 1, groups: int = 1, bias: bool = False, - activation: nn.Module | None = None, + activation: nn.Module | None | Literal[False] = None, ): """Conv2d + BN + Activation. @@ -112,10 +112,11 @@ def __init__( @param groups: Groups. Defaults to 1. @type bias: bool @param bias: Whether to use bias. Defaults to False. - @type activation: L{nn.Module} | None - @param activation: Activation function. If None then nn.ReLU. + @type activation: L{nn.Module} | None | Literal[False] + @param activation: Activation function. If None then nn.ReLU. If + False then no activation. Defaults to None. """ - super().__init__( + blocks = [ nn.Conv2d( in_channels, out_channels, @@ -127,8 +128,10 @@ def __init__( bias, ), nn.BatchNorm2d(out_channels), - activation or nn.ReLU(), - ) + ] + if activation is not False: + blocks.append(activation or nn.ReLU()) + super().__init__(*blocks) class UpBlock(nn.Sequential): @@ -316,7 +319,7 @@ def __init__( stride=stride, padding=padding, groups=groups, - activation=nn.Identity(), + activation=False, ) self.rbr_1x1 = ConvModule( in_channels=in_channels, @@ -325,7 +328,7 @@ def __init__( stride=stride, padding=padding_11, groups=groups, - activation=nn.Identity(), + activation=False, ) def forward(self, x: Tensor) -> Tensor: @@ -601,7 +604,7 @@ def __init__(self, in_channels: int, out_channels: int): in_channels=out_channels, out_channels=out_channels, kernel_size=1, - activation=nn.Identity(), + activation=False, ), nn.Sigmoid(), ) @@ -641,7 +644,7 @@ def __init__( in_channels=out_channels, out_channels=out_channels // reduction, kernel_size=1, - activation=nn.Identity(), + activation=False, ), nn.Sigmoid(), ) diff --git a/luxonis_train/nodes/heads/ghostfacenet_head.py b/luxonis_train/nodes/heads/ghostfacenet_head.py index e4753d03..8d9b66f4 100644 --- a/luxonis_train/nodes/heads/ghostfacenet_head.py +++ b/luxonis_train/nodes/heads/ghostfacenet_head.py @@ -5,10 +5,8 @@ from torch import Tensor from luxonis_train.enums import Metadata -from luxonis_train.nodes.backbones.ghostfacenet.blocks import ( - ModifiedGDC, -) from luxonis_train.nodes.base_node import BaseNode +from luxonis_train.nodes.blocks.blocks import ConvModule class GhostFaceNetHead(BaseNode[Tensor, list[Tensor]]): @@ -37,14 +35,25 @@ def __init__( """ super().__init__(**kwargs) self.embedding_size = embedding_size + image_size = self.original_in_shape[1] - self.head = ModifiedGDC( - self.original_in_shape[1], - self.in_channels, - dropout, - self.embedding_size, + self.head = nn.Sequential( + ConvModule( + self.in_channels, + self.in_channels, + kernel_size=(image_size // 32) + if image_size % 32 == 0 + else (image_size // 32 + 1), + groups=self.in_channels, + activation=False, + ), + nn.Dropout(dropout), + nn.Conv2d( + self.in_channels, embedding_size, kernel_size=1, bias=False + ), + nn.Flatten(), + nn.BatchNorm1d(embedding_size), ) - self._init_weights() def _init_weights(self): From 368188b0b2beb1a7b6e97b2c8aca96799ce45f70 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 23 Jan 2025 00:00:51 -0500 Subject: [PATCH 41/57] updated config --- tests/configs/reid.yaml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml index 88ffad66..cdf5058a 100644 --- a/tests/configs/reid.yaml +++ b/tests/configs/reid.yaml @@ -30,10 +30,6 @@ model: - name: EmbeddingsVisualizer loader: - train_view: train - val_view: [test_query, test_gallery] - test_view: [test_query, test_gallery] - params: dataset_name: reid_dataset From 7b96ab8574a0c338c6206b7a9aa27f77ce48214d Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 23 Jan 2025 00:31:00 -0500 Subject: [PATCH 42/57] small changes --- luxonis_train/attached_modules/metrics/torchmetrics.py | 2 +- luxonis_train/callbacks/ema.py | 2 +- luxonis_train/loaders/utils.py | 2 +- tests/integration/test_detection.py | 4 ---- tests/integration/test_segmentation.py | 5 ----- tests/unittests/test_callbacks/test_ema.py | 2 +- 6 files changed, 4 insertions(+), 13 deletions(-) diff --git a/luxonis_train/attached_modules/metrics/torchmetrics.py b/luxonis_train/attached_modules/metrics/torchmetrics.py index c222cb78..553ce31c 100644 --- a/luxonis_train/attached_modules/metrics/torchmetrics.py +++ b/luxonis_train/attached_modules/metrics/torchmetrics.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) -class TorchMetricWrapper(BaseMetric[Tensor]): +class TorchMetricWrapper(BaseMetric[Tensor, Tensor]): Metric: type[torchmetrics.Metric] def __init__(self, **kwargs: Any): diff --git a/luxonis_train/callbacks/ema.py b/luxonis_train/callbacks/ema.py index e76e8c1a..d280f4f4 100644 --- a/luxonis_train/callbacks/ema.py +++ b/luxonis_train/callbacks/ema.py @@ -39,7 +39,7 @@ def __init__( @type device: str | None @param device: Device to perform EMA on. """ - super(ModelEma, self).__init__() + super().__init__() model.eval() self.state_dict_ema = deepcopy(model.state_dict()) model.train() diff --git a/luxonis_train/loaders/utils.py b/luxonis_train/loaders/utils.py index e96b15bf..aa6b9fb4 100644 --- a/luxonis_train/loaders/utils.py +++ b/luxonis_train/loaders/utils.py @@ -39,7 +39,7 @@ def collate_fn( label_box: list[Tensor] = [] for i, ann in enumerate(annos): new_ann = torch.zeros((ann.shape[0], ann.shape[1] + 1)) - # add target image index for build_targets() + # add batch index to separate boxes from different images new_ann[:, 0] = i new_ann[:, 1:] = ann label_box.append(new_ann) diff --git a/tests/integration/test_detection.py b/tests/integration/test_detection.py index db7fe004..f59448b4 100644 --- a/tests/integration/test_detection.py +++ b/tests/integration/test_detection.py @@ -4,7 +4,6 @@ from luxonis_ml.data import LuxonisDataset from luxonis_train.core import LuxonisModel -from luxonis_train.nodes.backbones import __all__ as BACKBONES def get_opts_backbone(backbone: str) -> dict[str, Any]: @@ -102,9 +101,6 @@ def train_and_test( assert value > 0.8, f"{name} = {value} (expected > 0.8)" -@pytest.mark.parametrize( - "backbone", [b for b in BACKBONES if b != "GhostFaceNetV2"] -) def test_backbones( backbone: str, config: dict[str, Any], diff --git a/tests/integration/test_segmentation.py b/tests/integration/test_segmentation.py index 275b35e3..dcbc2db7 100644 --- a/tests/integration/test_segmentation.py +++ b/tests/integration/test_segmentation.py @@ -1,10 +1,8 @@ from typing import Any -import pytest from luxonis_ml.data import LuxonisDataset from luxonis_train.core import LuxonisModel -from luxonis_train.nodes.backbones import __all__ as BACKBONES def get_opts(backbone: str) -> dict[str, Any]: @@ -123,9 +121,6 @@ def train_and_test( assert value > 0.8, f"{name} = {value} (expected > 0.8)" -@pytest.mark.parametrize( - "backbone", [b for b in BACKBONES if b != "GhostFaceNetV2"] -) def test_backbones( backbone: str, config: dict[str, Any], diff --git a/tests/unittests/test_callbacks/test_ema.py b/tests/unittests/test_callbacks/test_ema.py index d117eb88..8ec6a37d 100644 --- a/tests/unittests/test_callbacks/test_ema.py +++ b/tests/unittests/test_callbacks/test_ema.py @@ -9,7 +9,7 @@ class SimpleModel(LightningModule): def __init__(self): - super(SimpleModel, self).__init__() + super().__init__() self.layer = torch.nn.Linear(2, 2) def forward(self, x): From b14a76c695f0c30f51cf897830ae1f599c77d478 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 23 Jan 2025 00:37:36 -0500 Subject: [PATCH 43/57] fix for rectangular images --- luxonis_train/nodes/heads/ghostfacenet_head.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/luxonis_train/nodes/heads/ghostfacenet_head.py b/luxonis_train/nodes/heads/ghostfacenet_head.py index 8d9b66f4..56afb889 100644 --- a/luxonis_train/nodes/heads/ghostfacenet_head.py +++ b/luxonis_train/nodes/heads/ghostfacenet_head.py @@ -35,15 +35,16 @@ def __init__( """ super().__init__(**kwargs) self.embedding_size = embedding_size - image_size = self.original_in_shape[1] + _, H, W = self.original_in_shape self.head = nn.Sequential( ConvModule( self.in_channels, self.in_channels, - kernel_size=(image_size // 32) - if image_size % 32 == 0 - else (image_size // 32 + 1), + kernel_size=( + H // 32 if H % 32 == 0 else H // 32 + 1, + W // 32 if W % 32 == 0 else W // 32 + 1, + ), groups=self.in_channels, activation=False, ), From 1bce803eb836e51a7190cba0bd0e656f6053728a Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 23 Jan 2025 00:51:23 -0500 Subject: [PATCH 44/57] renamed --- luxonis_train/attached_modules/losses/__init__.py | 2 +- .../losses/{pml_loss.py => embedding_losses.py} | 5 +++++ luxonis_train/attached_modules/metrics/__init__.py | 2 +- .../metrics/{pml_metrics.py => embedding_metrics.py} | 0 4 files changed, 7 insertions(+), 2 deletions(-) rename luxonis_train/attached_modules/losses/{pml_loss.py => embedding_losses.py} (90%) rename luxonis_train/attached_modules/metrics/{pml_metrics.py => embedding_metrics.py} (100%) diff --git a/luxonis_train/attached_modules/losses/__init__.py b/luxonis_train/attached_modules/losses/__init__.py index 2d0c77e1..520a42c6 100644 --- a/luxonis_train/attached_modules/losses/__init__.py +++ b/luxonis_train/attached_modules/losses/__init__.py @@ -3,11 +3,11 @@ from .bce_with_logits import BCEWithLogitsLoss from .cross_entropy import CrossEntropyLoss from .efficient_keypoint_bbox_loss import EfficientKeypointBBoxLoss +from .embedding_losses import EmbeddingLossWrapper from .fomo_localization_loss import FOMOLocalizationLoss from .ohem_bce_with_logits import OHEMBCEWithLogitsLoss from .ohem_cross_entropy import OHEMCrossEntropyLoss from .ohem_loss import OHEMLoss -from .pml_loss import EmbeddingLossWrapper from .reconstruction_segmentation_loss import ReconstructionSegmentationLoss from .sigmoid_focal_loss import SigmoidFocalLoss from .smooth_bce_with_logits import SmoothBCEWithLogitsLoss diff --git a/luxonis_train/attached_modules/losses/pml_loss.py b/luxonis_train/attached_modules/losses/embedding_losses.py similarity index 90% rename from luxonis_train/attached_modules/losses/pml_loss.py rename to luxonis_train/attached_modules/losses/embedding_losses.py index 70a87329..efba8394 100644 --- a/luxonis_train/attached_modules/losses/pml_loss.py +++ b/luxonis_train/attached_modules/losses/embedding_losses.py @@ -48,6 +48,11 @@ class EmbeddingLossWrapper( def __init__(self, *, node: BaseNode | None = None, **kwargs): super().__init__(node=node) + if not hasattr(pml_losses, loss_name): # noqa: B023 + raise ValueError( + f"Loss {loss_name} not found in " # noqa: B023 + "pytorch-metric-learning" + ) Loss = getattr(pml_losses, loss_name) # noqa: B023 self.loss_func = Loss(**kwargs) diff --git a/luxonis_train/attached_modules/metrics/__init__.py b/luxonis_train/attached_modules/metrics/__init__.py index 59e9cc57..df72a785 100644 --- a/luxonis_train/attached_modules/metrics/__init__.py +++ b/luxonis_train/attached_modules/metrics/__init__.py @@ -1,9 +1,9 @@ from .base_metric import BaseMetric from .confusion_matrix import ConfusionMatrix +from .embedding_metrics import ClosestIsPositiveAccuracy, MedianDistances from .mean_average_precision import MeanAveragePrecision from .mean_average_precision_keypoints import MeanAveragePrecisionKeypoints from .object_keypoint_similarity import ObjectKeypointSimilarity -from .pml_metrics import ClosestIsPositiveAccuracy, MedianDistances from .torchmetrics import Accuracy, F1Score, JaccardIndex, Precision, Recall __all__ = [ diff --git a/luxonis_train/attached_modules/metrics/pml_metrics.py b/luxonis_train/attached_modules/metrics/embedding_metrics.py similarity index 100% rename from luxonis_train/attached_modules/metrics/pml_metrics.py rename to luxonis_train/attached_modules/metrics/embedding_metrics.py From 3c0423e30e46688c2e1a33b5b6521c373dd97f50 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 23 Jan 2025 00:53:40 -0500 Subject: [PATCH 45/57] type simplification --- .../losses/embedding_losses.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/luxonis_train/attached_modules/losses/embedding_losses.py b/luxonis_train/attached_modules/losses/embedding_losses.py index efba8394..4d8c0a0d 100644 --- a/luxonis_train/attached_modules/losses/embedding_losses.py +++ b/luxonis_train/attached_modules/losses/embedding_losses.py @@ -37,33 +37,33 @@ "TupletMarginLoss", ] -for loss_name in EMBEDDING_LOSSES: +for _loss_name in EMBEDDING_LOSSES: class EmbeddingLossWrapper( - BaseLoss[Tensor, Tensor], register_name=loss_name + BaseLoss[Tensor, Tensor], register_name=_loss_name ): node: GhostFaceNetHead supported_tasks = [Metadata("id")] def __init__(self, *, node: BaseNode | None = None, **kwargs): super().__init__(node=node) + loss_name = _loss_name # noqa: B023 - if not hasattr(pml_losses, loss_name): # noqa: B023 + if not hasattr(pml_losses, loss_name): raise ValueError( - f"Loss {loss_name} not found in " # noqa: B023 - "pytorch-metric-learning" + f"Loss {loss_name} not found in pytorch-metric-learning" ) - Loss = getattr(pml_losses, loss_name) # noqa: B023 + Loss = getattr(pml_losses, loss_name) self.loss_func = Loss(**kwargs) if self.node.embedding_size is not None: - if loss_name in CrossBatchMemory.supported_losses(): # noqa: B023 + if loss_name in CrossBatchMemory.supported_losses(): self.loss_func = CrossBatchMemory( self.loss_func, embedding_size=self.node.embedding_size ) else: logger.warning( - f"CrossBatchMemory is not supported for {loss_name}. " # noqa: B023 + f"'CrossBatchMemory' is not supported for {loss_name}. " "Ignoring cross_batch_memory_size." ) @@ -72,4 +72,4 @@ def forward(self, inputs: Tensor, target: Tensor) -> Tensor: @property def name(self) -> str: - return loss_name # noqa: B023 + return _loss_name # noqa: B023 From cb970fa0840b36a80a752e34ce70baf8749d9ea1 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 23 Jan 2025 01:01:19 -0500 Subject: [PATCH 46/57] added cross batch memory --- .../attached_modules/losses/embedding_losses.py | 2 +- .../attached_modules/metrics/embedding_metrics.py | 15 +++++++++------ luxonis_train/nodes/heads/ghostfacenet_head.py | 11 ++++++++++- tests/configs/reid.yaml | 14 +++++--------- 4 files changed, 25 insertions(+), 17 deletions(-) diff --git a/luxonis_train/attached_modules/losses/embedding_losses.py b/luxonis_train/attached_modules/losses/embedding_losses.py index 4d8c0a0d..8a5cd301 100644 --- a/luxonis_train/attached_modules/losses/embedding_losses.py +++ b/luxonis_train/attached_modules/losses/embedding_losses.py @@ -56,7 +56,7 @@ def __init__(self, *, node: BaseNode | None = None, **kwargs): Loss = getattr(pml_losses, loss_name) self.loss_func = Loss(**kwargs) - if self.node.embedding_size is not None: + if self.node.cross_batch_memory_size is not None: if loss_name in CrossBatchMemory.supported_losses(): self.loss_func = CrossBatchMemory( self.loss_func, embedding_size=self.node.embedding_size diff --git a/luxonis_train/attached_modules/metrics/embedding_metrics.py b/luxonis_train/attached_modules/metrics/embedding_metrics.py index 9e4ab50b..b09d42f6 100644 --- a/luxonis_train/attached_modules/metrics/embedding_metrics.py +++ b/luxonis_train/attached_modules/metrics/embedding_metrics.py @@ -2,6 +2,7 @@ from torch import Tensor from luxonis_train.enums import Metadata +from luxonis_train.nodes.heads.ghostfacenet_head import GhostFaceNetHead from .base_metric import BaseMetric @@ -11,10 +12,11 @@ class ClosestIsPositiveAccuracy(BaseMetric[Tensor, Tensor]): supported_tasks = [Metadata("id")] + node: GhostFaceNetHead - def __init__(self, cross_batch_memory_size=0, **kwargs): + def __init__(self, **kwargs): super().__init__(**kwargs) - self.cross_batch_memory_size = cross_batch_memory_size + self.cross_batch_memory_size = self.node.cross_batch_memory_size self.add_state("cross_batch_memory", default=[], dist_reduce_fx="cat") self.add_state( "correct_predictions", @@ -28,7 +30,7 @@ def __init__(self, cross_batch_memory_size=0, **kwargs): def update(self, inputs: Tensor, target: Tensor): embeddings, labels = inputs, target - if self.cross_batch_memory_size > 0: + if self.cross_batch_memory_size is not None: self.cross_batch_memory.extend(list(zip(embeddings, labels))) if len(self.cross_batch_memory) > self.cross_batch_memory_size: @@ -73,10 +75,11 @@ def compute(self): class MedianDistances(BaseMetric[Tensor, Tensor]): supported_tasks = [Metadata("id")] + node: GhostFaceNetHead - def __init__(self, cross_batch_memory_size=0, **kwargs): + def __init__(self, **kwargs): super().__init__(**kwargs) - self.cross_batch_memory_size = cross_batch_memory_size + self.cross_batch_memory_size = self.node.cross_batch_memory_size self.add_state("cross_batch_memory", default=[], dist_reduce_fx="cat") self.add_state("all_distances", default=[], dist_reduce_fx="cat") self.add_state("closest_distances", default=[], dist_reduce_fx="cat") @@ -88,7 +91,7 @@ def __init__(self, cross_batch_memory_size=0, **kwargs): def update(self, inputs: Tensor, target: Tensor): embeddings, labels = inputs, target - if self.cross_batch_memory_size > 0: + if self.cross_batch_memory_size is not None: self.cross_batch_memory.extend(list(zip(embeddings, labels))) if len(self.cross_batch_memory) > self.cross_batch_memory_size: diff --git a/luxonis_train/nodes/heads/ghostfacenet_head.py b/luxonis_train/nodes/heads/ghostfacenet_head.py index 56afb889..4b7dcb06 100644 --- a/luxonis_train/nodes/heads/ghostfacenet_head.py +++ b/luxonis_train/nodes/heads/ghostfacenet_head.py @@ -15,7 +15,11 @@ class GhostFaceNetHead(BaseNode[Tensor, list[Tensor]]): tasks = [Metadata("id")] def __init__( - self, embedding_size: int = 512, dropout: float = 0.2, **kwargs + self, + embedding_size: int = 512, + cross_batch_memory_size: int | None = None, + dropout: float = 0.2, + **kwargs, ): """GhostFaceNetV2 backbone. @@ -32,9 +36,14 @@ def __init__( @type embedding_size: int @param embedding_size: Size of the embedding. Defaults to 512. + @type cross_batch_memory_size: int | None + @param cross_batch_memory_size: Size of the cross-batch memory. Defaults to None. + @type dropout: float + @param dropout: Dropout rate. Defaults to 0.2. """ super().__init__(**kwargs) self.embedding_size = embedding_size + self.cross_batch_memory_size = cross_batch_memory_size _, H, W = self.original_in_shape self.head = nn.Sequential( diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml index cdf5058a..91ac5dfb 100644 --- a/tests/configs/reid.yaml +++ b/tests/configs/reid.yaml @@ -8,34 +8,30 @@ model: - name: GhostFaceNetHead params: - embedding_size: &embedding_size 512 + embedding_size: 512 + cross_batch_memory_size: 4 losses: - name: SupConLoss - # params: - # embedding_size: *embedding_size - # cross_batch_memory_size: &memory_size 4 metrics: - name: ClosestIsPositiveAccuracy - # params: - # cross_batch_memory_size: *memory_size is_main_metric: True - name: MedianDistances - # params: - # cross_batch_memory_size: *memory_size visualizers: - name: EmbeddingsVisualizer loader: + train_view: val + params: dataset_name: reid_dataset trainer: preprocessing: - train_image_size: [256, 256] + train_image_size: [256, 320] batch_size: 16 epochs: 1 From c368723baeaa5bf7ca28a82e0aeddcc4638c3ab6 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 03:35:59 -0500 Subject: [PATCH 47/57] attached modules improvememnt --- .../losses/embedding_losses.py | 69 ++++++++- .../visualizers/embeddings_visualizer.py | 140 +++++++++++++----- .../attached_modules/visualizers/utils.py | 7 +- luxonis_train/config/config.py | 1 + luxonis_train/enums.py | 9 +- luxonis_train/loaders/base_loader.py | 12 +- luxonis_train/models/luxonis_lightning.py | 43 +++--- requirements.txt | 3 +- tests/configs/reid.yaml | 47 ------ 9 files changed, 214 insertions(+), 117 deletions(-) delete mode 100644 tests/configs/reid.yaml diff --git a/luxonis_train/attached_modules/losses/embedding_losses.py b/luxonis_train/attached_modules/losses/embedding_losses.py index 8a5cd301..b7b3518e 100644 --- a/luxonis_train/attached_modules/losses/embedding_losses.py +++ b/luxonis_train/attached_modules/losses/embedding_losses.py @@ -1,12 +1,17 @@ import logging +import pytorch_metric_learning.distances as pml_distances import pytorch_metric_learning.losses as pml_losses +import pytorch_metric_learning.miners as pml_miners +import pytorch_metric_learning.reducers as pml_reducers +import pytorch_metric_learning.regularizers as pml_regularizers from pytorch_metric_learning.losses import CrossBatchMemory from torch import Tensor from luxonis_train.enums import Metadata from luxonis_train.nodes.base_node import BaseNode from luxonis_train.nodes.heads.ghostfacenet_head import GhostFaceNetHead +from luxonis_train.utils.types import Kwargs from .base_loss import BaseLoss @@ -44,8 +49,22 @@ class EmbeddingLossWrapper( ): node: GhostFaceNetHead supported_tasks = [Metadata("id")] + miner: pml_miners.BaseMiner | None - def __init__(self, *, node: BaseNode | None = None, **kwargs): + def __init__( + self, + *, + miner: str | None = None, + miner_params: Kwargs | None = None, + distance: str | None = None, + distance_params: Kwargs | None = None, + reducer: str | None = None, + reducer_params: Kwargs | None = None, + regularizer: str | None = None, + regularizer_params: Kwargs | None = None, + node: BaseNode | None = None, + **kwargs, + ): super().__init__(node=node) loss_name = _loss_name # noqa: B023 @@ -54,12 +73,49 @@ def __init__(self, *, node: BaseNode | None = None, **kwargs): f"Loss {loss_name} not found in pytorch-metric-learning" ) Loss = getattr(pml_losses, loss_name) - self.loss_func = Loss(**kwargs) + + if reducer is not None: + if not hasattr(pml_reducers, reducer): + raise ValueError( + f"Reducer {reducer} not found in pytorch-metric-learning" + ) + Reducer = getattr(pml_reducers, reducer) + kwargs["reducer"] = Reducer(**(reducer_params or {})) + if regularizer is not None: + if not hasattr(pml_regularizers, regularizer): + raise ValueError( + f"Regularizer {regularizer} not found in pytorch-metric-learning" + ) + Regularizer = getattr(pml_regularizers, regularizer) + kwargs["embedding_regularizer"] = Regularizer( + **(regularizer_params or {}) + ) + if distance is not None: + if not hasattr(pml_distances, distance): + raise ValueError( + f"Distance {distance} not found in pytorch-metric-learning" + ) + Distance = getattr(pml_distances, distance) + kwargs["distance"] = Distance(**(distance_params or {})) + + if miner is not None: + if not hasattr(pml_miners, miner): + raise ValueError( + f"Miner {miner} not found in pytorch-metric-learning" + ) + Miner = getattr(pml_miners, miner) + self.miner = Miner(**(miner_params or {})) + else: + self.miner = None + + self.loss = Loss(**kwargs) if self.node.cross_batch_memory_size is not None: if loss_name in CrossBatchMemory.supported_losses(): - self.loss_func = CrossBatchMemory( - self.loss_func, embedding_size=self.node.embedding_size + self.loss = CrossBatchMemory( + self.loss, + embedding_size=self.node.embedding_size, + miner=self.miner, ) else: logger.warning( @@ -68,7 +124,10 @@ def __init__(self, *, node: BaseNode | None = None, **kwargs): ) def forward(self, inputs: Tensor, target: Tensor) -> Tensor: - return self.loss_func(inputs, target) + if self.miner is not None: + hard_pairs = self.miner(inputs, target) + return self.loss(inputs, target, hard_pairs) + return self.loss(inputs, target) @property def name(self) -> str: diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py index b368b2c5..aef19aac 100644 --- a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -1,7 +1,12 @@ +import colorsys import logging +from collections.abc import Callable +import numpy as np +import seaborn as sns from matplotlib import pyplot as plt -from sklearn.manifold import TSNE +from scipy.stats import zscore +from sklearn.decomposition import PCA from torch import Tensor from luxonis_train.enums import Metadata @@ -18,10 +23,20 @@ class EmbeddingsVisualizer(BaseVisualizer[Tensor, Tensor]): def __init__( self, + accumulate_n_batches: int = 2, **kwargs, ): - """Visualizer for embedding tasks like reID.""" + """Visualizer for embedding tasks like reID. + + @type accumulate_n_batches: int + @param accumulate_n_batches: Number of batches to accumulate + before visualizing. + """ super().__init__(**kwargs) + # self.memory = [] + # self.memory_size = accumulate_n_batches + self.color_dict = {} + self.gen = self._distinct_color_generator() def forward( self, @@ -29,8 +44,7 @@ def forward( prediction_canvas: Tensor, embeddings: Tensor, ids: Tensor, - **kwargs, - ) -> Tensor: + ) -> tuple[Tensor, Tensor]: """Creates a visualization of the embeddings. @type label_canvas: Tensor @@ -46,32 +60,92 @@ def forward( """ embeddings_np = embeddings.detach().cpu().numpy() - - perplexity = min(30, embeddings_np.shape[0] - 1) - - tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity) - embeddings_2d = tsne.fit_transform(embeddings_np) - - fig, ax = plt.subplots(figsize=(10, 10)) - scatter = ax.scatter( - embeddings_2d[:, 0], - embeddings_2d[:, 1], - c=ids.detach().cpu().numpy(), - cmap="viridis", - s=5, - ) - - fig.colorbar(scatter, ax=ax) - ax.set_title("Embeddings Visualization") - ax.set_xlabel("Dimension 1") - ax.set_ylabel("Dimension 2") - - image_tensor = figure_to_torch( - fig, width=label_canvas.shape[3], height=label_canvas.shape[2] - ) - - plt.close(fig) - - image_tensor = image_tensor.unsqueeze(0) - - return image_tensor + ids_np = ids.detach().cpu().numpy().astype(int) + # if len(self.memory) < self.memory_size: + # self.memory.append((embeddings_np, ids_np)) + # return None + # + # else: + # embeddings_np = np.concatenate( + # [mem[0] for mem in self.memory], axis=0 + # ) + # ids_np = np.concatenate([mem[1] for mem in self.memory], axis=0) + # self.memory = [] + + pca = PCA(n_components=2) + embeddings_2d = pca.fit_transform(embeddings_np) + + z = np.abs(zscore(embeddings_2d)) + mask = (z < 3).all(axis=1) + embeddings_2d = embeddings_2d[mask] + ids_np = ids_np[mask] + + def plot_to_tensor( + embeddings_2d: np.ndarray, + ids_np: np.ndarray, + plot_func: Callable[[plt.Axes, np.ndarray, np.ndarray], None], + ) -> Tensor: + fig, ax = plt.subplots(figsize=(10, 10)) + ax.set_xlim(embeddings_2d[:, 0].min(), embeddings_2d[:, 0].max()) + ax.set_ylim(embeddings_2d[:, 1].min(), embeddings_2d[:, 1].max()) + + plot_func(ax, embeddings_2d, ids_np) + ax.axis("off") + + tensor_image = figure_to_torch( + fig, width=512, height=512 + ).unsqueeze(0) + plt.close(fig) + return tensor_image + + def kde_plot( + ax: plt.Axes, emb: np.ndarray, labels: np.ndarray + ) -> None: + for label in np.unique(labels): + subset = emb[labels == label] + color = self._get_color(label) + sns.kdeplot( + x=subset[:, 0], + y=subset[:, 1], + color=color, + alpha=0.9, + fill=True, + warn_singular=False, + ax=ax, + ) + + def scatter_plot( + ax: plt.Axes, emb: np.ndarray, labels: np.ndarray + ) -> None: + unique_labels = np.unique(labels) + palette = {lbl: self._get_color(lbl) for lbl in unique_labels} + sns.scatterplot( + x=emb[:, 0], + y=emb[:, 1], + hue=labels, + palette=palette, + alpha=0.9, + s=300, + legend=False, + ax=ax, + ) + + kdeplot = plot_to_tensor(embeddings_2d, ids_np, kde_plot) + scatterplot = plot_to_tensor(embeddings_2d, ids_np, scatter_plot) + + return kdeplot, scatterplot + + def _get_color(self, label: int) -> tuple[float, float, float]: + if label not in self.color_dict: + self.color_dict[label] = next(self.gen) + return self.color_dict[label] + + @staticmethod + def _distinct_color_generator(): + golden_ratio = 0.618033988749895 + hue = 0.0 + while True: + hue = (hue + golden_ratio) % 1 + saturation = 0.8 + value = 0.95 + yield colorsys.hsv_to_rgb(hue, saturation, value) diff --git a/luxonis_train/attached_modules/visualizers/utils.py b/luxonis_train/attached_modules/visualizers/utils.py index 1a571eca..28e2ebf0 100644 --- a/luxonis_train/attached_modules/visualizers/utils.py +++ b/luxonis_train/attached_modules/visualizers/utils.py @@ -306,10 +306,15 @@ def combine_visualizations( visualization: Tensor | tuple[Tensor, Tensor] | tuple[Tensor, list[Tensor]], + # | None, ) -> Tensor: + # ) -> Tensor | None: """Default way of combining multiple visualizations into one final image.""" + # if visualization is None: + # return None + def resize_to_match( fst: Tensor, snd: Tensor, @@ -427,6 +432,6 @@ def resize_to_match( case _: raise ValueError( "Visualization should be either a single tensor or a tuple of " - "two tensors or a tuple of a tensor and a list of tensors." + "two tensors or a tuple of a tensor and a list of tensors. " f"Got: `{type(visualization)}`." ) diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index 9ee9cd1d..88f4a654 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -67,6 +67,7 @@ class ModelNodeConfig(BaseModelExtraForbid): freezing: FreezingConfig = FreezingConfig() remove_on_export: bool = False task_name: str = "" + metadata_task_override: str | dict[str, str] | None = None params: Params = {} diff --git a/luxonis_train/enums.py b/luxonis_train/enums.py index 04d4a241..ec68267f 100644 --- a/luxonis_train/enums.py +++ b/luxonis_train/enums.py @@ -14,6 +14,7 @@ class TaskType(str, Enum): @dataclass class Metadata: + # typ: type[float] | type[int] | type[str] | type[Category] name: str @property @@ -28,11 +29,3 @@ def __hash__(self) -> int: Task: TypeAlias = TaskType | Metadata -# class TaskType: -# -# CLASSIFICATION = SimpleTask.CLASSIFICATION -# SEGMENTATION = SimpleTask.SEGMENTATION -# INSTANCE_SEGMENTATION = SimpleTask.INSTANCE_SEGMENTATION -# BOUNDINGBOX = SimpleTask.BOUNDINGBOX -# KEYPOINTS = SimpleTask.KEYPOINTS -# ARRAY = SimpleTask.ARRAY diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index 752a15d2..8a375e36 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -260,10 +260,14 @@ def dict_numpy_to_torch( @rtype: dict[str, torch.Tensor] @return: Dictionary of torch tensors. """ - return { - task: torch.tensor(array).float() - for task, array in numpy_dictionary.items() - } + torch_dictionary = {} + + for task, array in numpy_dictionary.items(): + if array.dtype.kind in "U": + array = np.array([ord(c) for c in array[0]], dtype=np.int32) + torch_dictionary[task] = torch.tensor(array, dtype=torch.float32) + + return torch_dictionary def read_image(self, path: str) -> npt.NDArray[np.float32]: """Reads an image from a file. diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index d2d3fe87..49ebb44c 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -161,7 +161,7 @@ def __init__( dict ) - self._logged_images = 0 + self._logged_images = defaultdict(int) frozen_nodes: list[tuple[str, int]] = [] nodes: dict[str, tuple[type[BaseNode], Kwargs]] = {} @@ -191,6 +191,12 @@ def __init__( f"Node {node_name} does not have the `task_name` parameter set. " "Please specify the `task_name` parameter for each head node. " ) + # if node_cfg.metadata_task_name is not None: + # if Node.tasks is None: + # raise ValueError( + # f"`metadata_task_name` is set for node {node_name}, " + # "but the node does not define any tasks." + # ) nodes[node_name] = ( Node, @@ -312,15 +318,10 @@ def _initiate_nodes( for source_name, shape in shapes.items() } - for ( - node_name, - ( - Node, - node_kwargs, - ), - node_input_names, - _, - ) in traverse_graph(self.graph, nodes): + for node_name, ( + Node, + node_kwargs, + ), node_input_names, _ in traverse_graph(self.graph, nodes): node_dummy_inputs: list[Packet[Tensor]] = [] """List of dummy input packets for the node. @@ -766,8 +767,13 @@ def _evaluation_step( ) -> dict[str, Tensor]: inputs, labels = batch images = None - if self._logged_images < self.cfg.trainer.n_log_images: + if not self._logged_images: images = get_denormalized_images(self.cfg, inputs) + for value in self._logged_images.values(): + if value < self.cfg.trainer.n_log_images: + images = get_denormalized_images(self.cfg, inputs) + break + outputs = self.forward( inputs, labels, @@ -782,17 +788,18 @@ def _evaluation_step( logged_images = self._logged_images for node_name, visualizations in outputs.visualizations.items(): for viz_name, viz_batch in visualizations.items(): - logged_images = self._logged_images + # if viz_batch is None: + # continue for viz in viz_batch: - if logged_images >= self.cfg.trainer.n_log_images: - break + name = f"{mode}/visualizations/{node_name}/{viz_name}" + if logged_images[name] >= self.cfg.trainer.n_log_images: + continue self.logger.log_image( - f"{mode}/visualizations/{node_name}/{viz_name}/{logged_images}", + f"{name}/{logged_images[name]}", viz.detach().cpu().numpy().transpose(1, 2, 0), step=self.current_epoch, ) - logged_images += 1 - self._logged_images = logged_images + logged_images[name] += 1 return step_output @@ -832,7 +839,7 @@ def _evaluation_epoch_end(self, mode: Literal["test", "val"]) -> None: ) self.validation_step_outputs.clear() - self._logged_images = 0 + self._logged_images.clear() def configure_callbacks(self) -> list[pl.Callback]: """Configures Pytorch Lightning callbacks.""" diff --git a/requirements.txt b/requirements.txt index b49d9b3e..bd4e663a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,4 +19,5 @@ psutil>=5.0.0 tabulate>=0.9.0 grad-cam>=1.5.4 pytorch_metric_learning>=2.7.0 -scikit-learn>=1.5.0 \ No newline at end of file +scikit-learn>=1.5.0 +seaborn>=1.16.0 diff --git a/tests/configs/reid.yaml b/tests/configs/reid.yaml deleted file mode 100644 index 91ac5dfb..00000000 --- a/tests/configs/reid.yaml +++ /dev/null @@ -1,47 +0,0 @@ -loader: - name: CustomReIDLoader - -model: - name: reid_test - nodes: - - name: GhostFaceNetV2 - - - name: GhostFaceNetHead - params: - embedding_size: 512 - cross_batch_memory_size: 4 - - losses: - - name: SupConLoss - - metrics: - - name: ClosestIsPositiveAccuracy - is_main_metric: True - - - name: MedianDistances - - visualizers: - - name: EmbeddingsVisualizer - -loader: - train_view: val - - params: - dataset_name: reid_dataset - -trainer: - preprocessing: - train_image_size: [256, 320] - - batch_size: 16 - epochs: 1 - n_workers: 0 - validation_interval: 1 - - callbacks: - - name: ExportOnTrainEnd - - optimizer: - name: Adam - params: - lr: 0.01 From 2383283184e3dec98687574980fc557355fa5c06 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 20:35:18 -0500 Subject: [PATCH 48/57] metadata task override --- .../attached_modules/base_attached_module.py | 12 +++++++++++- luxonis_train/enums.py | 6 +----- luxonis_train/models/luxonis_lightning.py | 8 +------- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/luxonis_train/attached_modules/base_attached_module.py b/luxonis_train/attached_modules/base_attached_module.py index 866898fe..bd2e29e2 100644 --- a/luxonis_train/attached_modules/base_attached_module.py +++ b/luxonis_train/attached_modules/base_attached_module.py @@ -65,7 +65,17 @@ def __init__(self, *, node: BaseNode | None = None): self._epoch = 0 self.required_labels: list[Task] = [] - if self._node and self.supported_tasks: + if self._node is not None and self.supported_tasks: + for tasks in self.supported_tasks: + if not isinstance(tasks, tuple): + tasks = (tasks,) + for task in tasks: + if isinstance(task, TaskType): + continue + task.name = self.node.metadata_task_override.get( + task.name, task.name + ) + module_supported = [ label.value if isinstance(label, Task) diff --git a/luxonis_train/enums.py b/luxonis_train/enums.py index ec68267f..d82e8378 100644 --- a/luxonis_train/enums.py +++ b/luxonis_train/enums.py @@ -12,9 +12,8 @@ class TaskType(str, Enum): ARRAY = "array" -@dataclass +@dataclass(unsafe_hash=True) class Metadata: - # typ: type[float] | type[int] | type[str] | type[Category] name: str @property @@ -24,8 +23,5 @@ def value(self): def __str__(self) -> str: return self.value - def __hash__(self) -> int: - return hash(self.name) - Task: TypeAlias = TaskType | Metadata diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index 49ebb44c..8a480f38 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -191,19 +191,13 @@ def __init__( f"Node {node_name} does not have the `task_name` parameter set. " "Please specify the `task_name` parameter for each head node. " ) - # if node_cfg.metadata_task_name is not None: - # if Node.tasks is None: - # raise ValueError( - # f"`metadata_task_name` is set for node {node_name}, " - # "but the node does not define any tasks." - # ) - nodes[node_name] = ( Node, { **node_cfg.params, "task_name": node_cfg.task_name, "remove_on_export": node_cfg.remove_on_export, + "metadata_task_override": node_cfg.metadata_task_override, }, ) From dec365b597896e823257e7c0ee64dbebad1d425a Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 20:35:36 -0500 Subject: [PATCH 49/57] fix automatic inputs --- luxonis_train/config/config.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index 88f4a654..ea06d3ca 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -99,13 +99,14 @@ def validate_nodes(cls, nodes: Any) -> Any: names = [] last_body_index: int | None = None for i, node in enumerate(nodes): - name = node.get("alias", node.get("name")) + name = node.get("name") if name is None: raise ValueError( f"Node {i} does not specify the `name` field." ) if "Head" in name and last_body_index is None: last_body_index = i - 1 + name = node.get("alias") or name names.append(name) if i > 0 and "inputs" not in node and "input_sources" not in node: if last_body_index is not None: From 1f05da0d4a4aaa90a0530a85525751b798ed932e Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 20:36:00 -0500 Subject: [PATCH 50/57] cleaned --- .../visualizers/embeddings_visualizer.py | 126 ++++++++++-------- 1 file changed, 68 insertions(+), 58 deletions(-) diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py index aef19aac..ba4d7a98 100644 --- a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -5,7 +5,6 @@ import numpy as np import seaborn as sns from matplotlib import pyplot as plt -from scipy.stats import zscore from sklearn.decomposition import PCA from torch import Tensor @@ -24,6 +23,7 @@ class EmbeddingsVisualizer(BaseVisualizer[Tensor, Tensor]): def __init__( self, accumulate_n_batches: int = 2, + z_score_threshold: float = 3, **kwargs, ): """Visualizer for embedding tasks like reID. @@ -37,6 +37,7 @@ def __init__( # self.memory_size = accumulate_n_batches self.color_dict = {} self.gen = self._distinct_color_generator() + self.z_score_threshold = z_score_threshold def forward( self, @@ -74,64 +75,12 @@ def forward( pca = PCA(n_components=2) embeddings_2d = pca.fit_transform(embeddings_np) + embeddings_2d, ids_np = self._filter_outliers(embeddings_2d, ids_np) - z = np.abs(zscore(embeddings_2d)) - mask = (z < 3).all(axis=1) - embeddings_2d = embeddings_2d[mask] - ids_np = ids_np[mask] - - def plot_to_tensor( - embeddings_2d: np.ndarray, - ids_np: np.ndarray, - plot_func: Callable[[plt.Axes, np.ndarray, np.ndarray], None], - ) -> Tensor: - fig, ax = plt.subplots(figsize=(10, 10)) - ax.set_xlim(embeddings_2d[:, 0].min(), embeddings_2d[:, 0].max()) - ax.set_ylim(embeddings_2d[:, 1].min(), embeddings_2d[:, 1].max()) - - plot_func(ax, embeddings_2d, ids_np) - ax.axis("off") - - tensor_image = figure_to_torch( - fig, width=512, height=512 - ).unsqueeze(0) - plt.close(fig) - return tensor_image - - def kde_plot( - ax: plt.Axes, emb: np.ndarray, labels: np.ndarray - ) -> None: - for label in np.unique(labels): - subset = emb[labels == label] - color = self._get_color(label) - sns.kdeplot( - x=subset[:, 0], - y=subset[:, 1], - color=color, - alpha=0.9, - fill=True, - warn_singular=False, - ax=ax, - ) - - def scatter_plot( - ax: plt.Axes, emb: np.ndarray, labels: np.ndarray - ) -> None: - unique_labels = np.unique(labels) - palette = {lbl: self._get_color(lbl) for lbl in unique_labels} - sns.scatterplot( - x=emb[:, 0], - y=emb[:, 1], - hue=labels, - palette=palette, - alpha=0.9, - s=300, - legend=False, - ax=ax, - ) - - kdeplot = plot_to_tensor(embeddings_2d, ids_np, kde_plot) - scatterplot = plot_to_tensor(embeddings_2d, ids_np, scatter_plot) + kdeplot = self.plot_to_tensor(embeddings_2d, ids_np, self.kde_plot) + scatterplot = self.plot_to_tensor( + embeddings_2d, ids_np, self.scatter_plot + ) return kdeplot, scatterplot @@ -149,3 +98,64 @@ def _distinct_color_generator(): saturation = 0.8 value = 0.95 yield colorsys.hsv_to_rgb(hue, saturation, value) + + def _filter_outliers( + self, points: np.ndarray, ids: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + mean = np.mean(points, axis=0) + std_dev = np.std(points, axis=0) + z_scores = (points - mean) / std_dev + + mask = (np.abs(z_scores) < self.z_score_threshold).all(axis=1) + logger.info(f"Filtered out {len(points) - mask.sum()} outliers") + return points[mask], ids[mask] + + @staticmethod + def plot_to_tensor( + embeddings_2d: np.ndarray, + ids_np: np.ndarray, + plot_func: Callable[[plt.Axes, np.ndarray, np.ndarray], None], + ) -> Tensor: + fig, ax = plt.subplots(figsize=(10, 10)) + ax.set_xlim(embeddings_2d[:, 0].min(), embeddings_2d[:, 0].max()) + ax.set_ylim(embeddings_2d[:, 1].min(), embeddings_2d[:, 1].max()) + + plot_func(ax, embeddings_2d, ids_np) + ax.axis("off") + + tensor_image = figure_to_torch(fig, width=512, height=512).unsqueeze(0) + plt.close(fig) + return tensor_image + + def kde_plot( + self, ax: plt.Axes, emb: np.ndarray, labels: np.ndarray + ) -> None: + for label in np.unique(labels): + subset = emb[labels == label] + color = self._get_color(label) + sns.kdeplot( + x=subset[:, 0], + y=subset[:, 1], + color=color, + alpha=0.9, + bw_adjust=1.5, + fill=True, + warn_singular=False, + ax=ax, + ) + + def scatter_plot( + self, ax: plt.Axes, emb: np.ndarray, labels: np.ndarray + ) -> None: + unique_labels = np.unique(labels) + palette = {lbl: self._get_color(lbl) for lbl in unique_labels} + sns.scatterplot( + x=emb[:, 0], + y=emb[:, 1], + hue=labels, + palette=palette, + alpha=0.9, + s=300, + legend=False, + ax=ax, + ) From 6537153c94651c5a8c6112cffa8c9a430080e0bb Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 20:36:19 -0500 Subject: [PATCH 51/57] metadata overriding --- luxonis_train/nodes/base_node.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/luxonis_train/nodes/base_node.py b/luxonis_train/nodes/base_node.py index a35c5edd..409a9fe3 100644 --- a/luxonis_train/nodes/base_node.py +++ b/luxonis_train/nodes/base_node.py @@ -9,7 +9,7 @@ from torch import Size, Tensor, nn from typeguard import TypeCheckError, check_type -from luxonis_train.enums import Task, TaskType +from luxonis_train.enums import Metadata, Task, TaskType from luxonis_train.utils import ( AttachIndexType, DatasetMetadata, @@ -122,6 +122,7 @@ def __init__( export_output_names: list[str] | None = None, attach_index: AttachIndexType | None = None, task_name: str | None = None, + metadata_task_override: str | dict[str, str] | None = None, ): """Constructor for the C{BaseNode}. @@ -176,6 +177,35 @@ def __init__( f"argument for node '{self.name}' was not provided." ) self.task_name = task_name or "" + self.metadata_task_override = {} + if metadata_task_override is not None: + if self.tasks is None: + raise ValueError( + "Metadata task override can only be used with nodes that define tasks." + ) + n_metadata_tasks = sum( + 1 for task in self.tasks if isinstance(task, Metadata) + ) + if n_metadata_tasks > 1 and isinstance( + metadata_task_override, str + ): + raise ValueError( + f"Node '{self.name}' defines multiple metadata tasks, " + "but only a single task name was provided for " + "`metadata_task_override`. Provide a dictionary " + "mapping default names to new names instead ." + ) + for task in self.tasks: + if not isinstance(task, Metadata): + continue + + if isinstance(metadata_task_override, dict): + new_name = metadata_task_override.get(task.name, task.name) + else: + new_name = metadata_task_override + + self.metadata_task_override[task.name] = new_name + task.name = new_name if getattr(self, "attach_index", None) is None: parameters = inspect.signature(self.forward).parameters From fb15dfffaa5774716cf07a62f110e3251a3c1ca4 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 20:36:29 -0500 Subject: [PATCH 52/57] type checking --- luxonis_train/loaders/base_loader.py | 5 +++++ luxonis_train/loaders/luxonis_loader_torch.py | 9 +++++++++ 2 files changed, 14 insertions(+) diff --git a/luxonis_train/loaders/base_loader.py b/luxonis_train/loaders/base_loader.py index 8a375e36..bea27386 100644 --- a/luxonis_train/loaders/base_loader.py +++ b/luxonis_train/loaders/base_loader.py @@ -249,6 +249,11 @@ def get_n_keypoints(self) -> dict[str, int] | None: """ return None + def get_metadata_types( + self, + ) -> dict[str, dict[str, type[int] | type[float] | type[str]]]: + return {} + def dict_numpy_to_torch( self, numpy_dictionary: dict[str, np.ndarray] ) -> dict[str, Tensor]: diff --git a/luxonis_train/loaders/luxonis_loader_torch.py b/luxonis_train/loaders/luxonis_loader_torch.py index 4267cced..531b4b1f 100644 --- a/luxonis_train/loaders/luxonis_loader_torch.py +++ b/luxonis_train/loaders/luxonis_loader_torch.py @@ -127,6 +127,15 @@ def get_n_keypoints(self) -> dict[str, int]: skeletons = self.dataset.get_skeletons() return {task: len(skeletons[task][0]) for task in skeletons} + @override + def get_metadata_types( + self, + ) -> dict[str, dict[str, type[int] | type[float] | type[str]]]: + return { + k: {"float": float, "int": int, "str": str, "Category": int}[v] + for k, v in self.dataset.get_metadata_types().items() + } + def augment_test_image(self, img: Tensor) -> Tensor: if self.loader.augmentations is None: return img From d680406c67e459d0ec47b1b030f83d449f40e496 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 21:14:11 -0500 Subject: [PATCH 53/57] embedding tests --- configs/embeddings_model.yaml | 47 ++++++++++++++++++++++++++++ tests/integration/conftest.py | 33 +++++++++++++++++-- tests/integration/test_embeddings.py | 15 +++++++++ 3 files changed, 93 insertions(+), 2 deletions(-) create mode 100644 configs/embeddings_model.yaml create mode 100644 tests/integration/test_embeddings.py diff --git a/configs/embeddings_model.yaml b/configs/embeddings_model.yaml new file mode 100644 index 00000000..0e5308d9 --- /dev/null +++ b/configs/embeddings_model.yaml @@ -0,0 +1,47 @@ +loader: + name: CustomReIDLoader + +model: + name: reid_test + nodes: + - name: GhostFaceNetV2 + + - name: GhostFaceNetHead + alias: color-embeddings + metadata_task_override: color + params: + embedding_size: 16 + + losses: + - name: SupConLoss + params: + miner: MultiSimilarityMiner + distance: CosineSimilarity + reducer: ThresholdReducer + reducer_params: + high: 0.3 + regularizer: LpRegularizer + + metrics: + - name: ClosestIsPositiveAccuracy + + - name: MedianDistances + + visualizers: + - name: EmbeddingsVisualizer + +loader: + params: + dataset_name: ParkingLot + +trainer: + preprocessing: + train_image_size: [256, 256] + + batch_size: 16 + epochs: 100 + validation_interval: 10 + n_log_images: 8 + + callbacks: + - name: ExportOnTrainEnd diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index ab2fb1e8..92fc8720 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -4,10 +4,12 @@ from pathlib import Path from typing import Any +import cv2 import gdown +import numpy as np import pytest import torchvision -from luxonis_ml.data import LuxonisDataset +from luxonis_ml.data import Category, LuxonisDataset from luxonis_ml.data.parsers import LuxonisParser from luxonis_ml.utils import environ @@ -38,13 +40,40 @@ def parking_lot_dataset() -> LuxonisDataset: url = "gs://luxonis-test-bucket/luxonis-ml-test-data/D1_ParkingLot_Native.zip" parser = LuxonisParser( url, - dataset_name="_D1_ParkingLot", + dataset_name="D1_ParkingLot", delete_existing=True, save_dir=WORK_DIR, ) return parser.parse(random_split=True) +@pytest.fixture(scope="session") +def embedding_dataset() -> LuxonisDataset: + img_dir = WORK_DIR / "embedding_images" + img_dir.mkdir(exist_ok=True) + + def generator(): + for i in range(100): + color = [(255, 0, 0), (0, 255, 0), (0, 0, 255)][i % 3] + img = np.full((100, 100, 3), color, dtype=np.uint8) + img[i, i] = 255 + cv2.imwrite(str(img_dir / f"image_{i}.png"), img) + + yield { + "file": img_dir / f"image_{i}.png", + "annotation": { + "metadata": { + "color": Category(["red", "green", "blue"][i % 3]), + }, + }, + } + + dataset = LuxonisDataset("embedding_test", delete_existing=True) + dataset.add(generator()) + dataset.make_splits() + return dataset + + @pytest.fixture(scope="session") def coco_dataset() -> LuxonisDataset: dataset_name = "coco_test" diff --git a/tests/integration/test_embeddings.py b/tests/integration/test_embeddings.py new file mode 100644 index 00000000..ea1a4868 --- /dev/null +++ b/tests/integration/test_embeddings.py @@ -0,0 +1,15 @@ +from luxonis_ml.data import LuxonisDataset + +from luxonis_train.core import LuxonisModel + + +def test_embeddings_model(embedding_dataset: LuxonisDataset): + model = LuxonisModel( + cfg="configs/embeddings_model.yaml", + opts={ + "loader.params.dataset_name": embedding_dataset.dataset_name, + "trainer.epochs": 1, + "trainer.validation_interval": 1, + }, + ) + model.train() From f844d199bc9694d8ace17047f029ae1ee2705e5a Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 21:32:15 -0500 Subject: [PATCH 54/57] fix --- luxonis_train/attached_modules/base_attached_module.py | 2 +- tests/unittests/test_base_attached_module.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/luxonis_train/attached_modules/base_attached_module.py b/luxonis_train/attached_modules/base_attached_module.py index bd2e29e2..7e177735 100644 --- a/luxonis_train/attached_modules/base_attached_module.py +++ b/luxonis_train/attached_modules/base_attached_module.py @@ -360,7 +360,7 @@ def prepare( label = self._get_label(labels) generics = self._get_generic_params() - if generics is None: + if generics is None or generics[0].__name__ == "Unpack": return x, label # type: ignore if len(generics) != 2: diff --git a/tests/unittests/test_base_attached_module.py b/tests/unittests/test_base_attached_module.py index c7cd1508..450242b9 100644 --- a/tests/unittests/test_base_attached_module.py +++ b/tests/unittests/test_base_attached_module.py @@ -148,12 +148,12 @@ def test_prepare(inputs: Packet[Tensor], labels: Labels): det_head = DummyDetectionHead() assert seg_loss.prepare(inputs, labels) == ( - SEGMENTATION_ARRAY, + [SEGMENTATION_ARRAY], SEGMENTATION_ARRAY, ) inputs["/segmentation"].append(FEATURES_ARRAY) assert seg_loss.prepare(inputs, labels) == ( - FEATURES_ARRAY, + [SEGMENTATION_ARRAY, FEATURES_ARRAY], SEGMENTATION_ARRAY, ) From 1670566817fe66085a3e9ebabd161fa38c3877ff Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 21:44:00 -0500 Subject: [PATCH 55/57] parametrized tests --- tests/integration/test_detection.py | 2 ++ tests/integration/test_segmentation.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/tests/integration/test_detection.py b/tests/integration/test_detection.py index f59448b4..6360ec79 100644 --- a/tests/integration/test_detection.py +++ b/tests/integration/test_detection.py @@ -4,6 +4,7 @@ from luxonis_ml.data import LuxonisDataset from luxonis_train.core import LuxonisModel +from luxonis_train.nodes.backbones import __all__ as BACKBONES def get_opts_backbone(backbone: str) -> dict[str, Any]: @@ -101,6 +102,7 @@ def train_and_test( assert value > 0.8, f"{name} = {value} (expected > 0.8)" +@pytest.mark.parametrize("backbone", BACKBONES) def test_backbones( backbone: str, config: dict[str, Any], diff --git a/tests/integration/test_segmentation.py b/tests/integration/test_segmentation.py index dcbc2db7..1c79dc29 100644 --- a/tests/integration/test_segmentation.py +++ b/tests/integration/test_segmentation.py @@ -1,8 +1,10 @@ from typing import Any +import pytest from luxonis_ml.data import LuxonisDataset from luxonis_train.core import LuxonisModel +from luxonis_train.nodes.backbones import __all__ as BACKBONES def get_opts(backbone: str) -> dict[str, Any]: @@ -121,6 +123,7 @@ def train_and_test( assert value > 0.8, f"{name} = {value} (expected > 0.8)" +@pytest.mark.parametrize("backbone", BACKBONES) def test_backbones( backbone: str, config: dict[str, Any], From 01de24b51495da049d4b0cc22673c43a243dfc41 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Fri, 24 Jan 2025 22:36:36 -0500 Subject: [PATCH 56/57] docs --- .../attached_modules/losses/README.md | 43 +++++++++++++++++++ .../attached_modules/metrics/README.md | 12 ++++++ .../attached_modules/visualizers/README.md | 10 +++++ luxonis_train/nodes/README.md | 18 ++++++++ .../backbones/ghostfacenet/ghostfacenet.py | 6 +-- 5 files changed, 84 insertions(+), 5 deletions(-) diff --git a/luxonis_train/attached_modules/losses/README.md b/luxonis_train/attached_modules/losses/README.md index aa1b9ca6..b65e48ce 100644 --- a/luxonis_train/attached_modules/losses/README.md +++ b/luxonis_train/attached_modules/losses/README.md @@ -12,6 +12,7 @@ List of all the available loss functions. - [`AdaptiveDetectionLoss`](#adaptivedetectionloss) - [`EfficientKeypointBBoxLoss`](#efficientkeypointbboxloss) - [`FOMOLocalizationLoss`](#fomolocalizationLoss) +- [Embedding Losses](#embedding-losses) ## `CrossEntropyLoss` @@ -121,3 +122,45 @@ 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. Training with a larger `object_weight` in the loss parameters may result in more false positives (FP), but it will improve accuracy. | + +## Embedding Losses + +We support the following losses taken from [pytorch-metric-learning](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/): + +- [AngularLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#angularloss) +- [CircleLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#circleloss) +- [ContrastiveLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#contrastiveloss) +- [DynamicSoftMarginLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#dynamicsoftmarginloss) +- [FastAPLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#fastaploss) +- [HistogramLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#histogramloss) +- [InstanceLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#instanceloss) +- [IntraPairVarianceLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#intrapairvarianceloss) +- [GeneralizedLiftedStructureLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#generalizedliftedstructureloss) +- [LiftedStructureLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#liftedstructureloss) +- [MarginLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#marginloss) +- [MultiSimilarityLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#multisimilarityloss) +- [NPairsLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#npairsloss) +- [NCALoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#ncaloss) +- [NTXentLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#ntxentloss) +- [PNPLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#pnploss) +- [RankedListLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#rankedlistloss) +- [SignalToNoiseRatioContrastiveLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#signaltonoisecontrastiveloss) +- [SupConLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#supconloss) +- [ThresholdConsistentMarginLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#thresholdconsistentmarginloss) +- [TripletMarginLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#tripletmarginloss) +- [TupletMarginLoss](https://kevinmusgrave.github.io/pytorch-metric-learning/losses/#tupletmarginloss) + +**Parameters:** + +For loss specific parameters, see the documentation pages linked above. In addition to the loss specific parameters, the following parameters are available: + +| Key | Type | Default value | Description | +| -------------------- | ------ | ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `miner` | `str` | `None` | Name of the miner to use with the loss. If `None`, no miner is used. All miners from [pytorch-metric-learning](https://kevinmusgrave.github.io/pytorch-metric-learning/miners/) are supported. | +| `miner_params` | `dict` | `None` | Parameters for the miner. | +| `distance` | `str` | `None` | Name of the distance metric to use with the loss. If `None`, no distance metric is used. All distance metrics from [pytorch-metric-learning](https://kevinmusgrave.github.io/pytorch-metric-learning/distances/) are supported. | +| `distance_params` | `dict` | `None` | Parameters for the distance metric. | +| `reducer` | `str` | `None` | Name of the reducer to use with the loss. If `None`, no reducer is used. All reducers from [pytorch-metric-learning](https://kevinmusgrave.github.io/pytorch-metric-learning/reducers/) are supported. | +| `reducer_params` | `dict` | `None` | Parameters for the reducer. | +| `regularizer` | `str` | `None` | Name of the regularizer to use with the loss. If `None`, no regularizer is used. All regularizers from [pytorch-metric-learning](https://kevinmusgrave.github.io/pytorch-metric-learning/regularizers/) are supported. | +| `regularizer_params` | `dict` | `None` | Parameters for the regularizer. | diff --git a/luxonis_train/attached_modules/metrics/README.md b/luxonis_train/attached_modules/metrics/README.md index 42f42fcb..59021576 100644 --- a/luxonis_train/attached_modules/metrics/README.md +++ b/luxonis_train/attached_modules/metrics/README.md @@ -8,6 +8,8 @@ List of all the available metrics. - [ObjectKeypointSimilarity](#objectkeypointsimilarity) - [MeanAveragePrecision](#meanaverageprecision) - [MeanAveragePrecisionKeypoints](#meanaverageprecisionkeypoints) +- [ClosestIsPositiveAccuracy](#closestispositiveaccuracy) +- [MedianDistances](#mediandistances) ## Torchmetrics @@ -63,3 +65,13 @@ Evaluation leverages COCO evaluation framework (COCOeval) to assess mAP performa | `area_factor` | `float` | `0.53` | Factor by which to multiply the bounding box area | | `max_dets` | `int` | `20` | Maximum number of detections per image | | `box_fotmat` | `Literal["xyxy", "xywh", "cxcywh"]` | `"xyxy"` | Format of the bounding boxes | + +## ClosestIsPositiveAccuracy + +Compute the accuracy of the closest positive sample to the query sample. +Needs to be connected to the `GhostFaceNetHead` node. + +## MedianDistances + +Compute the median distance between the query and the positive samples. +Needs to be connected to the `GhostFaceNetHead` node. diff --git a/luxonis_train/attached_modules/visualizers/README.md b/luxonis_train/attached_modules/visualizers/README.md index 03daa87f..afef5066 100644 --- a/luxonis_train/attached_modules/visualizers/README.md +++ b/luxonis_train/attached_modules/visualizers/README.md @@ -7,6 +7,8 @@ Visualizers are used to render the output of a node. They are used in the `visua - [`BBoxVisualizer`](#bboxvisualizer) - [`ClassificationVisualizer`](#classificationvisualizer) - [`KeypointVisualizer`](#keypointvisualizer) +- [`SegmentationVisualizer`](#segmentationvisualizer) +- [`EmbeddingsVisualizer`](#embeddingsvisualizer) - [`MultiVisualizer`](#multivisualizer) ## `BBoxVisualizer` @@ -72,6 +74,14 @@ Visualizer for bounding boxes. ![class_viz_example](https://github.com/luxonis/luxonis-train/blob/main/media/example_viz/class.png) +## `EmbeddingsVisualizer` + +**Parameters:** + +| Key | Type | Default value | Description | +| ------------------- | ------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------- | +| `z_score_threshold` | `float` | `3.0` | Threshold for z-score filtering. Embeddings with z-score higher than this value are considered as outliers and are not drawn. | + ## `MultiVisualizer` Special type of meta-visualizer that combines several visualizers into one. The combined visualizers share canvas. diff --git a/luxonis_train/nodes/README.md b/luxonis_train/nodes/README.md index 7e2540ee..87cc459d 100644 --- a/luxonis_train/nodes/README.md +++ b/luxonis_train/nodes/README.md @@ -18,6 +18,7 @@ arbitrarily as long as the two nodes are compatible with each other. We've group - [`DDRNet`](#ddrnet) - [`RecSubNet`](#recsubnet) - [`EfficientViT`](#efficientvit) + - [`GhostFaceNetV2`](#ghostfacenetv2) - [Necks](#necks) - [`RepPANNeck`](#reppanneck) - [Heads](#heads) @@ -29,6 +30,7 @@ arbitrarily as long as the two nodes are compatible with each other. We've group - [`DDRNetSegmentationHead`](#ddrnetsegmentationhead) - [`DiscSubNetHead`](#discsubnet) - [`FOMOHead`](#fomohead) + - [`GhostFaceNetHead`](#ghostfacenethead) Every node takes these parameters: | Key | Type | Default value | Description | @@ -186,6 +188,14 @@ Adapted from [here](https://arxiv.org/abs/2205.14756) | `expand_ratio` | `int` | `4` | Factor by which channels expand in the local module | | `dim` | `int` | `None` | Dimension size for each attention head | +### `GhostFaceNetV2` + +**Parameters:** + +| Key | Type | Default value | Description | +| --------- | --------------- | ------------- | --------------------------- | +| `variant` | `Literal["V2"]` | `"V2"` | The variant of the network. | + ## Neck ### `RepPANNeck` @@ -290,3 +300,11 @@ 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. | | `use_nms` | `bool` | `False` | If True, enable NMS. This can reduce FP, but it will also reduce TP for close neighbors. | + +### `GhostFaceNetHead` + +**Parameters:** + +| Key | Type | Default value | Description | +| ---------------- | ----- | ------------- | ---------------------------------------- | +| `embedding_size` | `int` | `512` | The size of the output embedding vector. | diff --git a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py index 66b71679..2ad0227c 100644 --- a/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py +++ b/luxonis_train/nodes/backbones/ghostfacenet/ghostfacenet.py @@ -15,11 +15,7 @@ class GhostFaceNetV2(BaseNode[Tensor, Tensor]): in_channels: int in_width: int - def __init__( - self, - variant: Literal["V2"] = "V2", - **kwargs, - ): + def __init__(self, variant: Literal["V2"] = "V2", **kwargs): """GhostFaceNetsV2 backbone. GhostFaceNetsV2 is a convolutional neural network architecture focused on face recognition, but it is From 8bc9dff26e4d9ddeb37534103d552ae304bf7f41 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Tue, 28 Jan 2025 19:20:42 -0500 Subject: [PATCH 57/57] moved colors to luxonis-ml --- .../visualizers/embeddings_visualizer.py | 40 +++---------------- .../attached_modules/visualizers/utils.py | 5 --- luxonis_train/models/luxonis_lightning.py | 2 - luxonis_train/nodes/blocks/blocks.py | 2 + 4 files changed, 7 insertions(+), 42 deletions(-) diff --git a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py index ba4d7a98..da483705 100644 --- a/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py +++ b/luxonis_train/attached_modules/visualizers/embeddings_visualizer.py @@ -1,9 +1,9 @@ -import colorsys import logging from collections.abc import Callable import numpy as np import seaborn as sns +from luxonis_ml.data.utils import ColorMap from matplotlib import pyplot as plt from sklearn.decomposition import PCA from torch import Tensor @@ -14,18 +14,12 @@ from .utils import figure_to_torch logger = logging.getLogger(__name__) -log_disable = False class EmbeddingsVisualizer(BaseVisualizer[Tensor, Tensor]): supported_tasks = [Metadata("id")] - def __init__( - self, - accumulate_n_batches: int = 2, - z_score_threshold: float = 3, - **kwargs, - ): + def __init__(self, z_score_threshold: float = 3, **kwargs): """Visualizer for embedding tasks like reID. @type accumulate_n_batches: int @@ -33,10 +27,7 @@ def __init__( before visualizing. """ super().__init__(**kwargs) - # self.memory = [] - # self.memory_size = accumulate_n_batches - self.color_dict = {} - self.gen = self._distinct_color_generator() + self.colors = ColorMap() self.z_score_threshold = z_score_threshold def forward( @@ -62,16 +53,6 @@ def forward( embeddings_np = embeddings.detach().cpu().numpy() ids_np = ids.detach().cpu().numpy().astype(int) - # if len(self.memory) < self.memory_size: - # self.memory.append((embeddings_np, ids_np)) - # return None - # - # else: - # embeddings_np = np.concatenate( - # [mem[0] for mem in self.memory], axis=0 - # ) - # ids_np = np.concatenate([mem[1] for mem in self.memory], axis=0) - # self.memory = [] pca = PCA(n_components=2) embeddings_2d = pca.fit_transform(embeddings_np) @@ -85,19 +66,8 @@ def forward( return kdeplot, scatterplot def _get_color(self, label: int) -> tuple[float, float, float]: - if label not in self.color_dict: - self.color_dict[label] = next(self.gen) - return self.color_dict[label] - - @staticmethod - def _distinct_color_generator(): - golden_ratio = 0.618033988749895 - hue = 0.0 - while True: - hue = (hue + golden_ratio) % 1 - saturation = 0.8 - value = 0.95 - yield colorsys.hsv_to_rgb(hue, saturation, value) + r, g, b = self.colors[label] + return r / 255, g / 255, b / 255 def _filter_outliers( self, points: np.ndarray, ids: np.ndarray diff --git a/luxonis_train/attached_modules/visualizers/utils.py b/luxonis_train/attached_modules/visualizers/utils.py index 28e2ebf0..ac95046b 100644 --- a/luxonis_train/attached_modules/visualizers/utils.py +++ b/luxonis_train/attached_modules/visualizers/utils.py @@ -306,15 +306,10 @@ def combine_visualizations( visualization: Tensor | tuple[Tensor, Tensor] | tuple[Tensor, list[Tensor]], - # | None, ) -> Tensor: - # ) -> Tensor | None: """Default way of combining multiple visualizations into one final image.""" - # if visualization is None: - # return None - def resize_to_match( fst: Tensor, snd: Tensor, diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index 0649baea..2b7252c8 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -790,8 +790,6 @@ def _evaluation_step( logged_images = self._logged_images for node_name, visualizations in outputs.visualizations.items(): for viz_name, viz_batch in visualizations.items(): - # if viz_batch is None: - # continue for viz in viz_batch: name = f"{mode}/visualizations/{node_name}/{viz_name}" if logged_images[name] >= self.cfg.trainer.n_log_images: diff --git a/luxonis_train/nodes/blocks/blocks.py b/luxonis_train/nodes/blocks/blocks.py index 651ae004..dfacfeab 100644 --- a/luxonis_train/nodes/blocks/blocks.py +++ b/luxonis_train/nodes/blocks/blocks.py @@ -217,6 +217,8 @@ def __init__( if activation is not False: blocks.append(activation or nn.ReLU()) + super().__init__(*blocks) + class DWConvModule(ConvModule): def __init__(