content: from __future__ import annotations
  
  import random
  from pathlib import Path
  from typing import Optional, Tuple, Union
  
  import torch
  from PIL import Image
  from supervision.dataset.utils import train_test_split
  from torch.utils.data import Dataset
  from torchvision.transforms import Compose, ToTensor
  
  from trackers.core.reid.dataset.utils import validate_tracker_id_to_images
  
  
  class TripletsDataset(Dataset):
      """A dataset that provides triplets of images for training ReID models.
  
      This dataset is designed for training models with triplet loss, where each sample
      consists of an anchor image, a positive image (same identity as anchor),
      and a negative image (different identity from anchor).
  
      Args:
          tracker_id_to_images (dict[str, list[str]]): Dictionary mapping tracker IDs
              to lists of image paths
          transforms (Optional[Compose]): Optional image transformations to apply
  
      Attributes:
          tracker_id_to_images (dict[str, list[str]]): Dictionary mapping tracker IDs
              to lists of image paths
          transforms (Optional[Compose]): Optional image transformations to apply
          tracker_ids (list[str]): List of all unique tracker IDs in the dataset
      """
  
      def __init__(
          self,
          tracker_id_to_images: dict[str, list[str]],
          transforms: Optional[Compose] = None,
      ):
          self.tracker_id_to_images = validate_tracker_id_to_images(tracker_id_to_images)
          self.transforms = transforms or ToTensor()
          self.tracker_ids = list(self.tracker_id_to_images.keys())
  
      @classmethod
      def from_image_directories(
          cls,
          root_directory: str,
          transforms: Optional[Compose] = None,
          image_extensions: Tuple[str, ...] = (".jpg", ".jpeg", ".png"),
      ) -> TripletsDataset:
          """
          Create TripletsDataset from a directory structured by tracker IDs.
  
          Args:
              root_directory (str): Root directory with tracker folders.
              transforms (Optional[Compose]): Optional image transformations.
              image_extensions (Tuple[str, ...]): Valid image extensions to load.
  
          Returns:
              TripletsDataset: An initialized dataset.
          """
          root_path = Path(root_directory)
          tracker_id_to_images = {}
  
          for tracker_path in sorted(root_path.iterdir()):
              if not tracker_path.is_dir():
                  continue
  
              image_paths = sorted(
                  [
                      str(image_path)
                      for image_path in tracker_path.glob("*")
                      if image_path.suffix.lower() in image_extensions
                      and image_path.is_file()
                  ]
              )
  
              if image_paths:
                  tracker_id_to_images[tracker_path.name] = image_paths
  
          return cls(
              tracker_id_to_images=tracker_id_to_images,
              transforms=transforms,
          )
  
      def __len__(self) -> int:
          """
          Return the number of unique tracker IDs (identities) in the dataset.
  
          Returns:
              int: The total number of unique identities (tracker IDs) available for
                  sampling triplets.
          """
          return len(self.tracker_ids)
  
      def _load_and_transform_image(self, image_path: str) -> torch.Tensor:
          image = Image.open(image_path).convert("RGB")
          if self.transforms:
              image = self.transforms(image)
          return image
  
      def _get_triplet_image_paths(self, tracker_id: str) -> Tuple[str, str, str]:
          tracker_id_image_paths = self.tracker_id_to_images[tracker_id]
  
          anchor_image_path, positive_image_path = random.sample(  # nosec B311
              tracker_id_image_paths, 2
          )
  
          negative_candidates = [tid for tid in self.tracker_ids if tid != tracker_id]
          negative_tracker_id = random.choice(negative_candidates)  # nosec B311
  
          negative_image_path = random.choice(  # nosec B311
              self.tracker_id_to_images[negative_tracker_id]
          )
  
          return anchor_image_path, positive_image_path, negative_image_path
  
      def __getitem__(
          self, index: int
      ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
          """
          Retrieve a random triplet (anchor, positive, negative) of images for a given
          identity.
  
          For the tracker ID at the given index, samples two different images as the
          anchor and positive (same identity), and one image from a different tracker ID
          as the negative (different identity).
  
          Args:
              index (int): Index of the tracker ID (identity) to sample the triplet from.
  
          Returns:
              Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
                  A tuple containing the anchor, positive, and negative image tensors.
          """
          tracker_id = self.tracker_ids[index]
  
          anchor_image_path, positive_image_path, negative_image_path = (
              self._get_triplet_image_paths(tracker_id)
          )
  
          anchor_image = self._load_and_transform_image(anchor_image_path)
          positive_image = self._load_and_transform_image(positive_image_path)
          negative_image = self._load_and_transform_image(negative_image_path)
  
          return anchor_image, positive_image, negative_image
  
      def split(
          self,
          split_ratio: float = 0.8,
          random_state: Optional[Union[int, float, str, bytes, bytearray]] = None,
          shuffle: bool = True,
      ) -> Tuple[TripletsDataset, TripletsDataset]:
          train_tracker_id_to_images, validation_tracker_id_to_images = train_test_split(
              list(self.tracker_id_to_images.keys()),
              train_ratio=split_ratio,
              random_state=random_state,
              shuffle=shuffle,
          )
          train_tracker_id_to_images = {
              tracker_id: self.tracker_id_to_images[tracker_id]
              for tracker_id in train_tracker_id_to_images
          }
          validation_tracker_id_to_images = {
              tracker_id: self.tracker_id_to_images[tracker_id]
              for tracker_id in validation_tracker_id_to_images
          }
          return (
              TripletsDataset(train_tracker_id_to_images, self.transforms),
              TripletsDataset(validation_tracker_id_to_images, self.transforms),
          )
  