Re-Identification Model

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",
 ]

luxonis_train/attached_modules/losses/pml_loss.py

@@ -0,0 +1,122 @@
+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]
+
+        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 forward(self, inputs: Tensor, target: Tensor):
+        # miner_output = self.miner_func(inputs, target)
+
+        loss = self.loss_func(inputs, target)
+
+        return loss

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",
 ]