ContrastiveViewsGenerator#
- class ContrastiveViewsGenerator(view_keys=('coord', 'color', 'normal', 'origin_coord'), view_trans_cfg=None)[source]#
Bases:
objectGenerate two independently-augmented views for contrastive SSL.
Copies the keys in
view_keysinto two sub-dicts, runs the same transform pipeline (built fromview_trans_cfg) independently on each, then writes the results back underview1_<key>andview2_<key>prefixes (for every key produced by the pipeline). Registered asContrastiveViewsGenerator— use this string as thetypein atransform=[...]config list.- Parameters:
view_keys (tuple) – Keys copied from the source sample into each view. Defaults to
("coord", "color", "normal", "origin_coord").view_trans_cfg (list, optional) – Transform config dicts composed into the per-view pipeline applied to both views. Defaults to
None.
Example
>>> import numpy as np >>> np.random.seed(0) >>> data = {"coord": np.random.rand(10, 3).astype("f4"), ... "color": np.random.rand(10, 3).astype("f4")} >>> out = ContrastiveViewsGenerator( ... view_keys=("coord", "color"), ... view_trans_cfg=[dict(type="RandomFlip", p=1.0)])(data) >>> sorted(k for k in out if k.startswith("view")) ['view1_color', 'view1_coord', 'view2_color', 'view2_coord'] # two independently-augmented copies of each view_key