s6.app.cog.augmented_datasetΒΆ

Augmented keypoint dataset and image transforms.

Wraps s6.utils.datapipe.StructuredDatasetTorch to produce square patches around configured image/keypoint pairs with optional segmentation masks. One raw multi-camera record can expand into multiple trainable samples through data_mappings.

class s6.app.cog.augmented_dataset.AugmentationConfigBase(*, name: str, enabled: bool = True)

Bases: BaseModel

name: str
enabled: bool
class s6.app.cog.augmented_dataset.MirrorAugmentationConfig(*, name: Literal['mirror'] = 'mirror', enabled: bool = False)

Bases: AugmentationConfigBase

name: Literal['mirror']
enabled: bool
class s6.app.cog.augmented_dataset.CenterSquareCropAugmentationConfig(*, name: Literal['center_square_crop'] = 'center_square_crop', enabled: bool = True)

Bases: AugmentationConfigBase

name: Literal['center_square_crop']
enabled: bool
class s6.app.cog.augmented_dataset.CropAugmentationConfig(*, name: Literal['crop'] = 'crop', enabled: bool = True, crop_factor: float | Tuple[float, float] = 0.9, output_size: int = 256, max_rotation: float = 0.0)

Bases: AugmentationConfigBase

name: Literal['crop']
enabled: bool
crop_factor: float | Tuple[float, float]
output_size: int
max_rotation: float
class s6.app.cog.augmented_dataset.RotationAugmentationConfig(*, name: Literal['rotation'] = 'rotation', enabled: bool = True, max_rotation: float = 360.0)

Bases: AugmentationConfigBase

name: Literal['rotation']
enabled: bool
max_rotation: float
class s6.app.cog.augmented_dataset.BlurAugmentationConfig(*, name: Literal['blur'] = 'blur', enabled: bool = False, sigma_range: Tuple[float, float] = (0.0, 2.0))

Bases: AugmentationConfigBase

name: Literal['blur']
enabled: bool
sigma_range: Tuple[float, float]
class s6.app.cog.augmented_dataset.OcclusionAugmentationConfig(*, name: Literal['occlusion'] = 'occlusion', enabled: bool = False, size_range: Tuple[float, float] = (0.1, 0.3), intensity_range: Tuple[int, int] = (0, 255))

Bases: AugmentationConfigBase

name: Literal['occlusion']
enabled: bool
size_range: Tuple[float, float]
intensity_range: Tuple[int, int]
class s6.app.cog.augmented_dataset.ColorJitterAugmentationConfig(*, name: Literal['color_jitter'] = 'color_jitter', enabled: bool = True, brightness: float = 0.2, contrast: float = 0.2, gamma: float = 0.2)

Bases: AugmentationConfigBase

name: Literal['color_jitter']
enabled: bool
brightness: float
contrast: float
gamma: float
class s6.app.cog.augmented_dataset.SamplingConfig(*, enabled: bool = False, bins: int = 10, seed: int = 0)

Bases: BaseModel

Configure distance-based bucket sampling of keypoints.

enabled: bool
bins: int
seed: int
class s6.app.cog.augmented_dataset.DataMappings(*, x: List[str] = None, y: List[str] = None, mask: List[str] | None = None)

Bases: BaseModel

Logical training-target mappings to raw StructuredDataset datakeys.

x: List[str]
y: List[str]
mask: List[str] | None
property has_masks: bool
raw_datakeys() List[str]
class s6.app.cog.augmented_dataset.TrainingMappingRef(image_key: str, keypoint_key: str, mask_key: str | None = None)

Bases: object

image_key: str
keypoint_key: str
mask_key: str | None = None
class s6.app.cog.augmented_dataset.SampleRef(dataset_idx: int, source_idx: int, mapping_idx: int)

Bases: object

dataset_idx: int
source_idx: int
mapping_idx: int
class s6.app.cog.augmented_dataset.AugmentedKeypointDataset(config: Config)

Bases: Dataset

Dataset that applies keypoint-centric augmentations.

class Config(*, prefix: ConstrainedStrValue = 'model', base_dir: str | List[str] = [], data_mappings: DataMappings = None, augmentations: List[dict] = None, sd_kwargs: dict = {}, sampling: SamplingConfig = SamplingConfig(enabled=False, bins=10, seed=0), num_segmentation_classes: ConstrainedIntValue | None = None, segmentation_loss_weight: ConstrainedFloatValue = 1.0)

Bases: BaseModel

prefix: str
base_dir: str | List[str]
data_mappings: DataMappings
augmentations: List[dict]
sd_kwargs: dict
sampling: SamplingConfig
num_segmentation_classes: int | None
segmentation_loss_weight: float
classmethod default()
classmethod from_json(cfg_path: str) AugmentedKeypointDataset
property output_size: int
property has_segmentation: bool
s6.app.cog.augmented_dataset.center_square_crop(image: ndarray, keypoint: Tuple[int, int], mask: ndarray | None = None) Tuple[ndarray, ndarray, ndarray | None]
s6.app.cog.augmented_dataset.random_crop_and_resize(image: ndarray, keypoint: Tuple[int, int], crop_factor: float | Tuple[float, float] = 0.8, output_size: int = 256, max_rotation: float = 0.0, mask: ndarray | None = None) Tuple[ndarray, ndarray, ndarray | None]
s6.app.cog.augmented_dataset.random_rotate(image: ndarray, keypoint: ndarray, max_angle: float = 30.0, mask: ndarray | None = None) Tuple[ndarray, ndarray, ndarray | None]
s6.app.cog.augmented_dataset.random_color_jitter(image: ndarray, brightness: float = 0.2, contrast: float = 0.2, gamma: float = 0.2) ndarray
s6.app.cog.augmented_dataset.horizontal_mirror(image: ndarray, keypoint: Tuple[int, int], mask: ndarray | None = None) Tuple[ndarray, ndarray, ndarray | None]
s6.app.cog.augmented_dataset.random_blur(image: ndarray, sigma_range: Tuple[float, float]) ndarray
s6.app.cog.augmented_dataset.random_occlusion(image: ndarray, keypoint: ndarray, size_range: Tuple[float, float], intensity_range: Tuple[int, int], mask: ndarray | None = None) Tuple[ndarray, ndarray | None]