ContrastiveViewsGenerator#

class ContrastiveViewsGenerator(view_keys=('coord', 'color', 'normal', 'origin_coord'), view_trans_cfg=None)[source]#

Bases: object

Generate two independently-augmented views for contrastive SSL.

Copies the keys in view_keys into two sub-dicts, runs the same transform pipeline (built from view_trans_cfg) independently on each, then writes the results back under view1_<key> and view2_<key> prefixes (for every key produced by the pipeline). Registered as ContrastiveViewsGenerator — use this string as the type in a transform=[...] 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