UPLOADS

densepose/__init__.py

@@ -1,20 +1,20 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-from .data.datasets import builtin  # just to register data
-from .converters import builtin as builtin_converters  # register converters
-from .config import (
-    add_densepose_config,
-    add_densepose_head_config,
-    add_hrnet_config,
-    add_dataset_category_config,
-    add_bootstrap_config,
-    load_bootstrap_config,
-)
-from .structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
-from .evaluation import DensePoseCOCOEvaluator
-from .modeling.roi_heads import DensePoseROIHeads
-from .modeling.test_time_augmentation import (
-    DensePoseGeneralizedRCNNWithTTA,
-    DensePoseDatasetMapperTTA,
-)
-from .utils.transform import load_from_cfg
-from .modeling.hrfpn import build_hrfpn_backbone
+# Copyright (c) Facebook, Inc. and its affiliates.
+from .data.datasets import builtin  # just to register data
+from .converters import builtin as builtin_converters  # register converters
+from .config import (
+    add_densepose_config,
+    add_densepose_head_config,
+    add_hrnet_config,
+    add_dataset_category_config,
+    add_bootstrap_config,
+    load_bootstrap_config,
+)
+from .structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
+from .evaluation import DensePoseCOCOEvaluator
+from .modeling.roi_heads import DensePoseROIHeads
+from .modeling.test_time_augmentation import (
+    DensePoseGeneralizedRCNNWithTTA,
+    DensePoseDatasetMapperTTA,
+)
+from .utils.transform import load_from_cfg
+from .modeling.hrfpn import build_hrfpn_backbone

densepose/config.py

@@ -1,277 +1,277 @@
-# -*- coding = utf-8 -*-
-# Copyright (c) Facebook, Inc. and its affiliates.
-# pyre-ignore-all-errors
-
-from detectron2.config import CfgNode as CN
-
-
-def add_dataset_category_config(cfg: CN) -> None:
-    """
-    Add config for additional category-related dataset options
-     - category whitelisting
-     - category mapping
-    """
-    _C = cfg
-    _C.DATASETS.CATEGORY_MAPS = CN(new_allowed=True)
-    _C.DATASETS.WHITELISTED_CATEGORIES = CN(new_allowed=True)
-    # class to mesh mapping
-    _C.DATASETS.CLASS_TO_MESH_NAME_MAPPING = CN(new_allowed=True)
-
-
-def add_evaluation_config(cfg: CN) -> None:
-    _C = cfg
-    _C.DENSEPOSE_EVALUATION = CN()
-    # evaluator type, possible values:
-    #  - "iou": evaluator for models that produce iou data
-    #  - "cse": evaluator for models that produce cse data
-    _C.DENSEPOSE_EVALUATION.TYPE = "iou"
-    # storage for DensePose results, possible values:
-    #  - "none": no explicit storage, all the results are stored in the
-    #            dictionary with predictions, memory intensive;
-    #            historically the default storage type
-    #  - "ram": RAM storage, uses per-process RAM storage, which is
-    #           reduced to a single process storage on later stages,
-    #           less memory intensive
-    #  - "file": file storage, uses per-process file-based storage,
-    #            the least memory intensive, but may create bottlenecks
-    #            on file system accesses
-    _C.DENSEPOSE_EVALUATION.STORAGE = "none"
-    # minimum threshold for IOU values: the lower its values is,
-    # the more matches are produced (and the higher the AP score)
-    _C.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD = 0.5
-    # Non-distributed inference is slower (at inference time) but can avoid RAM OOM
-    _C.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE = True
-    # evaluate mesh alignment based on vertex embeddings, only makes sense in CSE context
-    _C.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT = False
-    # meshes to compute mesh alignment for
-    _C.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES = []
-
-
-def add_bootstrap_config(cfg: CN) -> None:
-    """ """
-    _C = cfg
-    _C.BOOTSTRAP_DATASETS = []
-    _C.BOOTSTRAP_MODEL = CN()
-    _C.BOOTSTRAP_MODEL.WEIGHTS = ""
-    _C.BOOTSTRAP_MODEL.DEVICE = "cuda"
-
-
-def get_bootstrap_dataset_config() -> CN:
-    _C = CN()
-    _C.DATASET = ""
-    # ratio used to mix data loaders
-    _C.RATIO = 0.1
-    # image loader
-    _C.IMAGE_LOADER = CN(new_allowed=True)
-    _C.IMAGE_LOADER.TYPE = ""
-    _C.IMAGE_LOADER.BATCH_SIZE = 4
-    _C.IMAGE_LOADER.NUM_WORKERS = 4
-    _C.IMAGE_LOADER.CATEGORIES = []
-    _C.IMAGE_LOADER.MAX_COUNT_PER_CATEGORY = 1_000_000
-    _C.IMAGE_LOADER.CATEGORY_TO_CLASS_MAPPING = CN(new_allowed=True)
-    # inference
-    _C.INFERENCE = CN()
-    # batch size for model inputs
-    _C.INFERENCE.INPUT_BATCH_SIZE = 4
-    # batch size to group model outputs
-    _C.INFERENCE.OUTPUT_BATCH_SIZE = 2
-    # sampled data
-    _C.DATA_SAMPLER = CN(new_allowed=True)
-    _C.DATA_SAMPLER.TYPE = ""
-    _C.DATA_SAMPLER.USE_GROUND_TRUTH_CATEGORIES = False
-    # filter
-    _C.FILTER = CN(new_allowed=True)
-    _C.FILTER.TYPE = ""
-    return _C
-
-
-def load_bootstrap_config(cfg: CN) -> None:
-    """
-    Bootstrap datasets are given as a list of `dict` that are not automatically
-    converted into CfgNode. This method processes all bootstrap dataset entries
-    and ensures that they are in CfgNode format and comply with the specification
-    """
-    if not cfg.BOOTSTRAP_DATASETS:
-        return
-
-    bootstrap_datasets_cfgnodes = []
-    for dataset_cfg in cfg.BOOTSTRAP_DATASETS:
-        _C = get_bootstrap_dataset_config().clone()
-        _C.merge_from_other_cfg(CN(dataset_cfg))
-        bootstrap_datasets_cfgnodes.append(_C)
-    cfg.BOOTSTRAP_DATASETS = bootstrap_datasets_cfgnodes
-
-
-def add_densepose_head_cse_config(cfg: CN) -> None:
-    """
-    Add configuration options for Continuous Surface Embeddings (CSE)
-    """
-    _C = cfg
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE = CN()
-    # Dimensionality D of the embedding space
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE = 16
-    # Embedder specifications for various mesh IDs
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS = CN(new_allowed=True)
-    # normalization coefficient for embedding distances
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_DIST_GAUSS_SIGMA = 0.01
-    # normalization coefficient for geodesic distances
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.GEODESIC_DIST_GAUSS_SIGMA = 0.01
-    # embedding loss weight
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_WEIGHT = 0.6
-    # embedding loss name, currently the following options are supported:
-    # - EmbeddingLoss: cross-entropy on vertex labels
-    # - SoftEmbeddingLoss: cross-entropy on vertex label combined with
-    #    Gaussian penalty on distance between vertices
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_NAME = "EmbeddingLoss"
-    # optimizer hyperparameters
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR = 1.0
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR = 1.0
-    # Shape to shape cycle consistency loss parameters:
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
-    # shape to shape cycle consistency loss weight
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.025
-    # norm type used for loss computation
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
-    # normalization term for embedding similarity matrices
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.TEMPERATURE = 0.05
-    # maximum number of vertices to include into shape to shape cycle loss
-    # if negative or zero, all vertices are considered
-    # if positive, random subset of vertices of given size is considered
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.MAX_NUM_VERTICES = 4936
-    # Pixel to shape cycle consistency loss parameters:
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
-    # pixel to shape cycle consistency loss weight
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.0001
-    # norm type used for loss computation
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
-    # map images to all meshes and back (if false, use only gt meshes from the batch)
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.USE_ALL_MESHES_NOT_GT_ONLY = False
-    # Randomly select at most this number of pixels from every instance
-    # if negative or zero, all vertices are considered
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NUM_PIXELS_TO_SAMPLE = 100
-    # normalization factor for pixel to pixel distances (higher value = smoother distribution)
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.PIXEL_SIGMA = 5.0
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_PIXEL_TO_VERTEX = 0.05
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_VERTEX_TO_PIXEL = 0.05
-
-
-def add_densepose_head_config(cfg: CN) -> None:
-    """
-    Add config for densepose head.
-    """
-    _C = cfg
-
-    _C.MODEL.DENSEPOSE_ON = True
-
-    _C.MODEL.ROI_DENSEPOSE_HEAD = CN()
-    _C.MODEL.ROI_DENSEPOSE_HEAD.NAME = ""
-    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS = 8
-    # Number of parts used for point labels
-    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES = 24
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL = 4
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM = 512
-    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL = 3
-    _C.MODEL.ROI_DENSEPOSE_HEAD.UP_SCALE = 2
-    _C.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE = 112
-    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE = "ROIAlignV2"
-    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION = 28
-    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO = 2
-    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS = 2  # 15 or 2
-    # Overlap threshold for an RoI to be considered foreground (if >= FG_IOU_THRESHOLD)
-    _C.MODEL.ROI_DENSEPOSE_HEAD.FG_IOU_THRESHOLD = 0.7
-    # Loss weights for annotation masks.(14 Parts)
-    _C.MODEL.ROI_DENSEPOSE_HEAD.INDEX_WEIGHTS = 5.0
-    # Loss weights for surface parts. (24 Parts)
-    _C.MODEL.ROI_DENSEPOSE_HEAD.PART_WEIGHTS = 1.0
-    # Loss weights for UV regression.
-    _C.MODEL.ROI_DENSEPOSE_HEAD.POINT_REGRESSION_WEIGHTS = 0.01
-    # Coarse segmentation is trained using instance segmentation task data
-    _C.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS = False
-    # For Decoder
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_ON = True
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NUM_CLASSES = 256
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_CONV_DIMS = 256
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NORM = ""
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_COMMON_STRIDE = 4
-    # For DeepLab head
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB = CN()
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NORM = "GN"
-    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NONLOCAL_ON = 0
-    # Predictor class name, must be registered in DENSEPOSE_PREDICTOR_REGISTRY
-    # Some registered predictors:
-    #   "DensePoseChartPredictor": predicts segmentation and UV coordinates for predefined charts
-    #   "DensePoseChartWithConfidencePredictor": predicts segmentation, UV coordinates
-    #       and associated confidences for predefined charts (default)
-    #   "DensePoseEmbeddingWithConfidencePredictor": predicts segmentation, embeddings
-    #       and associated confidences for CSE
-    _C.MODEL.ROI_DENSEPOSE_HEAD.PREDICTOR_NAME = "DensePoseChartWithConfidencePredictor"
-    # Loss class name, must be registered in DENSEPOSE_LOSS_REGISTRY
-    # Some registered losses:
-    #   "DensePoseChartLoss": loss for chart-based models that estimate
-    #      segmentation and UV coordinates
-    #   "DensePoseChartWithConfidenceLoss": loss for chart-based models that estimate
-    #      segmentation, UV coordinates and the corresponding confidences (default)
-    _C.MODEL.ROI_DENSEPOSE_HEAD.LOSS_NAME = "DensePoseChartWithConfidenceLoss"
-    # Confidences
-    # Enable learning UV confidences (variances) along with the actual values
-    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE = CN({"ENABLED": False})
-    # UV confidence lower bound
-    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.EPSILON = 0.01
-    # Enable learning segmentation confidences (variances) along with the actual values
-    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE = CN({"ENABLED": False})
-    # Segmentation confidence lower bound
-    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE.EPSILON = 0.01
-    # Statistical model type for confidence learning, possible values:
-    # - "iid_iso": statistically independent identically distributed residuals
-    #    with isotropic covariance
-    # - "indep_aniso": statistically independent residuals with anisotropic
-    #    covariances
-    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.TYPE = "iid_iso"
-    # List of angles for rotation in data augmentation during training
-    _C.INPUT.ROTATION_ANGLES = [0]
-    _C.TEST.AUG.ROTATION_ANGLES = ()  # Rotation TTA
-
-    add_densepose_head_cse_config(cfg)
-
-
-def add_hrnet_config(cfg: CN) -> None:
-    """
-    Add config for HRNet backbone.
-    """
-    _C = cfg
-
-    # For HigherHRNet w32
-    _C.MODEL.HRNET = CN()
-    _C.MODEL.HRNET.STEM_INPLANES = 64
-    _C.MODEL.HRNET.STAGE2 = CN()
-    _C.MODEL.HRNET.STAGE2.NUM_MODULES = 1
-    _C.MODEL.HRNET.STAGE2.NUM_BRANCHES = 2
-    _C.MODEL.HRNET.STAGE2.BLOCK = "BASIC"
-    _C.MODEL.HRNET.STAGE2.NUM_BLOCKS = [4, 4]
-    _C.MODEL.HRNET.STAGE2.NUM_CHANNELS = [32, 64]
-    _C.MODEL.HRNET.STAGE2.FUSE_METHOD = "SUM"
-    _C.MODEL.HRNET.STAGE3 = CN()
-    _C.MODEL.HRNET.STAGE3.NUM_MODULES = 4
-    _C.MODEL.HRNET.STAGE3.NUM_BRANCHES = 3
-    _C.MODEL.HRNET.STAGE3.BLOCK = "BASIC"
-    _C.MODEL.HRNET.STAGE3.NUM_BLOCKS = [4, 4, 4]
-    _C.MODEL.HRNET.STAGE3.NUM_CHANNELS = [32, 64, 128]
-    _C.MODEL.HRNET.STAGE3.FUSE_METHOD = "SUM"
-    _C.MODEL.HRNET.STAGE4 = CN()
-    _C.MODEL.HRNET.STAGE4.NUM_MODULES = 3
-    _C.MODEL.HRNET.STAGE4.NUM_BRANCHES = 4
-    _C.MODEL.HRNET.STAGE4.BLOCK = "BASIC"
-    _C.MODEL.HRNET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
-    _C.MODEL.HRNET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
-    _C.MODEL.HRNET.STAGE4.FUSE_METHOD = "SUM"
-
-    _C.MODEL.HRNET.HRFPN = CN()
-    _C.MODEL.HRNET.HRFPN.OUT_CHANNELS = 256
-
-
-def add_densepose_config(cfg: CN) -> None:
-    add_densepose_head_config(cfg)
-    add_hrnet_config(cfg)
-    add_bootstrap_config(cfg)
-    add_dataset_category_config(cfg)
-    add_evaluation_config(cfg)
+# -*- coding = utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+# pyre-ignore-all-errors
+
+from detectron2.config import CfgNode as CN
+
+
+def add_dataset_category_config(cfg: CN) -> None:
+    """
+    Add config for additional category-related dataset options
+     - category whitelisting
+     - category mapping
+    """
+    _C = cfg
+    _C.DATASETS.CATEGORY_MAPS = CN(new_allowed=True)
+    _C.DATASETS.WHITELISTED_CATEGORIES = CN(new_allowed=True)
+    # class to mesh mapping
+    _C.DATASETS.CLASS_TO_MESH_NAME_MAPPING = CN(new_allowed=True)
+
+
+def add_evaluation_config(cfg: CN) -> None:
+    _C = cfg
+    _C.DENSEPOSE_EVALUATION = CN()
+    # evaluator type, possible values:
+    #  - "iou": evaluator for models that produce iou data
+    #  - "cse": evaluator for models that produce cse data
+    _C.DENSEPOSE_EVALUATION.TYPE = "iou"
+    # storage for DensePose results, possible values:
+    #  - "none": no explicit storage, all the results are stored in the
+    #            dictionary with predictions, memory intensive;
+    #            historically the default storage type
+    #  - "ram": RAM storage, uses per-process RAM storage, which is
+    #           reduced to a single process storage on later stages,
+    #           less memory intensive
+    #  - "file": file storage, uses per-process file-based storage,
+    #            the least memory intensive, but may create bottlenecks
+    #            on file system accesses
+    _C.DENSEPOSE_EVALUATION.STORAGE = "none"
+    # minimum threshold for IOU values: the lower its values is,
+    # the more matches are produced (and the higher the AP score)
+    _C.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD = 0.5
+    # Non-distributed inference is slower (at inference time) but can avoid RAM OOM
+    _C.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE = True
+    # evaluate mesh alignment based on vertex embeddings, only makes sense in CSE context
+    _C.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT = False
+    # meshes to compute mesh alignment for
+    _C.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES = []
+
+
+def add_bootstrap_config(cfg: CN) -> None:
+    """ """
+    _C = cfg
+    _C.BOOTSTRAP_DATASETS = []
+    _C.BOOTSTRAP_MODEL = CN()
+    _C.BOOTSTRAP_MODEL.WEIGHTS = ""
+    _C.BOOTSTRAP_MODEL.DEVICE = "cuda"
+
+
+def get_bootstrap_dataset_config() -> CN:
+    _C = CN()
+    _C.DATASET = ""
+    # ratio used to mix data loaders
+    _C.RATIO = 0.1
+    # image loader
+    _C.IMAGE_LOADER = CN(new_allowed=True)
+    _C.IMAGE_LOADER.TYPE = ""
+    _C.IMAGE_LOADER.BATCH_SIZE = 4
+    _C.IMAGE_LOADER.NUM_WORKERS = 4
+    _C.IMAGE_LOADER.CATEGORIES = []
+    _C.IMAGE_LOADER.MAX_COUNT_PER_CATEGORY = 1_000_000
+    _C.IMAGE_LOADER.CATEGORY_TO_CLASS_MAPPING = CN(new_allowed=True)
+    # inference
+    _C.INFERENCE = CN()
+    # batch size for model inputs
+    _C.INFERENCE.INPUT_BATCH_SIZE = 4
+    # batch size to group model outputs
+    _C.INFERENCE.OUTPUT_BATCH_SIZE = 2
+    # sampled data
+    _C.DATA_SAMPLER = CN(new_allowed=True)
+    _C.DATA_SAMPLER.TYPE = ""
+    _C.DATA_SAMPLER.USE_GROUND_TRUTH_CATEGORIES = False
+    # filter
+    _C.FILTER = CN(new_allowed=True)
+    _C.FILTER.TYPE = ""
+    return _C
+
+
+def load_bootstrap_config(cfg: CN) -> None:
+    """
+    Bootstrap datasets are given as a list of `dict` that are not automatically
+    converted into CfgNode. This method processes all bootstrap dataset entries
+    and ensures that they are in CfgNode format and comply with the specification
+    """
+    if not cfg.BOOTSTRAP_DATASETS:
+        return
+
+    bootstrap_datasets_cfgnodes = []
+    for dataset_cfg in cfg.BOOTSTRAP_DATASETS:
+        _C = get_bootstrap_dataset_config().clone()
+        _C.merge_from_other_cfg(CN(dataset_cfg))
+        bootstrap_datasets_cfgnodes.append(_C)
+    cfg.BOOTSTRAP_DATASETS = bootstrap_datasets_cfgnodes
+
+
+def add_densepose_head_cse_config(cfg: CN) -> None:
+    """
+    Add configuration options for Continuous Surface Embeddings (CSE)
+    """
+    _C = cfg
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE = CN()
+    # Dimensionality D of the embedding space
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE = 16
+    # Embedder specifications for various mesh IDs
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS = CN(new_allowed=True)
+    # normalization coefficient for embedding distances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_DIST_GAUSS_SIGMA = 0.01
+    # normalization coefficient for geodesic distances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.GEODESIC_DIST_GAUSS_SIGMA = 0.01
+    # embedding loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_WEIGHT = 0.6
+    # embedding loss name, currently the following options are supported:
+    # - EmbeddingLoss: cross-entropy on vertex labels
+    # - SoftEmbeddingLoss: cross-entropy on vertex label combined with
+    #    Gaussian penalty on distance between vertices
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_NAME = "EmbeddingLoss"
+    # optimizer hyperparameters
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR = 1.0
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR = 1.0
+    # Shape to shape cycle consistency loss parameters:
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
+    # shape to shape cycle consistency loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.025
+    # norm type used for loss computation
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
+    # normalization term for embedding similarity matrices
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.TEMPERATURE = 0.05
+    # maximum number of vertices to include into shape to shape cycle loss
+    # if negative or zero, all vertices are considered
+    # if positive, random subset of vertices of given size is considered
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.MAX_NUM_VERTICES = 4936
+    # Pixel to shape cycle consistency loss parameters:
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
+    # pixel to shape cycle consistency loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.0001
+    # norm type used for loss computation
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
+    # map images to all meshes and back (if false, use only gt meshes from the batch)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.USE_ALL_MESHES_NOT_GT_ONLY = False
+    # Randomly select at most this number of pixels from every instance
+    # if negative or zero, all vertices are considered
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NUM_PIXELS_TO_SAMPLE = 100
+    # normalization factor for pixel to pixel distances (higher value = smoother distribution)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.PIXEL_SIGMA = 5.0
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_PIXEL_TO_VERTEX = 0.05
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_VERTEX_TO_PIXEL = 0.05
+
+
+def add_densepose_head_config(cfg: CN) -> None:
+    """
+    Add config for densepose head.
+    """
+    _C = cfg
+
+    _C.MODEL.DENSEPOSE_ON = True
+
+    _C.MODEL.ROI_DENSEPOSE_HEAD = CN()
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NAME = ""
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS = 8
+    # Number of parts used for point labels
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES = 24
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL = 4
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM = 512
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL = 3
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UP_SCALE = 2
+    _C.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE = 112
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE = "ROIAlignV2"
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION = 28
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO = 2
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS = 2  # 15 or 2
+    # Overlap threshold for an RoI to be considered foreground (if >= FG_IOU_THRESHOLD)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.FG_IOU_THRESHOLD = 0.7
+    # Loss weights for annotation masks.(14 Parts)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.INDEX_WEIGHTS = 5.0
+    # Loss weights for surface parts. (24 Parts)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.PART_WEIGHTS = 1.0
+    # Loss weights for UV regression.
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POINT_REGRESSION_WEIGHTS = 0.01
+    # Coarse segmentation is trained using instance segmentation task data
+    _C.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS = False
+    # For Decoder
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_ON = True
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NUM_CLASSES = 256
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_CONV_DIMS = 256
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NORM = ""
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_COMMON_STRIDE = 4
+    # For DeepLab head
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB = CN()
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NORM = "GN"
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NONLOCAL_ON = 0
+    # Predictor class name, must be registered in DENSEPOSE_PREDICTOR_REGISTRY
+    # Some registered predictors:
+    #   "DensePoseChartPredictor": predicts segmentation and UV coordinates for predefined charts
+    #   "DensePoseChartWithConfidencePredictor": predicts segmentation, UV coordinates
+    #       and associated confidences for predefined charts (default)
+    #   "DensePoseEmbeddingWithConfidencePredictor": predicts segmentation, embeddings
+    #       and associated confidences for CSE
+    _C.MODEL.ROI_DENSEPOSE_HEAD.PREDICTOR_NAME = "DensePoseChartWithConfidencePredictor"
+    # Loss class name, must be registered in DENSEPOSE_LOSS_REGISTRY
+    # Some registered losses:
+    #   "DensePoseChartLoss": loss for chart-based models that estimate
+    #      segmentation and UV coordinates
+    #   "DensePoseChartWithConfidenceLoss": loss for chart-based models that estimate
+    #      segmentation, UV coordinates and the corresponding confidences (default)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.LOSS_NAME = "DensePoseChartWithConfidenceLoss"
+    # Confidences
+    # Enable learning UV confidences (variances) along with the actual values
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE = CN({"ENABLED": False})
+    # UV confidence lower bound
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.EPSILON = 0.01
+    # Enable learning segmentation confidences (variances) along with the actual values
+    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE = CN({"ENABLED": False})
+    # Segmentation confidence lower bound
+    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE.EPSILON = 0.01
+    # Statistical model type for confidence learning, possible values:
+    # - "iid_iso": statistically independent identically distributed residuals
+    #    with isotropic covariance
+    # - "indep_aniso": statistically independent residuals with anisotropic
+    #    covariances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.TYPE = "iid_iso"
+    # List of angles for rotation in data augmentation during training
+    _C.INPUT.ROTATION_ANGLES = [0]
+    _C.TEST.AUG.ROTATION_ANGLES = ()  # Rotation TTA
+
+    add_densepose_head_cse_config(cfg)
+
+
+def add_hrnet_config(cfg: CN) -> None:
+    """
+    Add config for HRNet backbone.
+    """
+    _C = cfg
+
+    # For HigherHRNet w32
+    _C.MODEL.HRNET = CN()
+    _C.MODEL.HRNET.STEM_INPLANES = 64
+    _C.MODEL.HRNET.STAGE2 = CN()
+    _C.MODEL.HRNET.STAGE2.NUM_MODULES = 1
+    _C.MODEL.HRNET.STAGE2.NUM_BRANCHES = 2
+    _C.MODEL.HRNET.STAGE2.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE2.NUM_BLOCKS = [4, 4]
+    _C.MODEL.HRNET.STAGE2.NUM_CHANNELS = [32, 64]
+    _C.MODEL.HRNET.STAGE2.FUSE_METHOD = "SUM"
+    _C.MODEL.HRNET.STAGE3 = CN()
+    _C.MODEL.HRNET.STAGE3.NUM_MODULES = 4
+    _C.MODEL.HRNET.STAGE3.NUM_BRANCHES = 3
+    _C.MODEL.HRNET.STAGE3.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE3.NUM_BLOCKS = [4, 4, 4]
+    _C.MODEL.HRNET.STAGE3.NUM_CHANNELS = [32, 64, 128]
+    _C.MODEL.HRNET.STAGE3.FUSE_METHOD = "SUM"
+    _C.MODEL.HRNET.STAGE4 = CN()
+    _C.MODEL.HRNET.STAGE4.NUM_MODULES = 3
+    _C.MODEL.HRNET.STAGE4.NUM_BRANCHES = 4
+    _C.MODEL.HRNET.STAGE4.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
+    _C.MODEL.HRNET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
+    _C.MODEL.HRNET.STAGE4.FUSE_METHOD = "SUM"
+
+    _C.MODEL.HRNET.HRFPN = CN()
+    _C.MODEL.HRNET.HRFPN.OUT_CHANNELS = 256
+
+
+def add_densepose_config(cfg: CN) -> None:
+    add_densepose_head_config(cfg)
+    add_hrnet_config(cfg)
+    add_bootstrap_config(cfg)
+    add_dataset_category_config(cfg)
+    add_evaluation_config(cfg)

densepose/converters/__init__.py

@@ -1,15 +1,15 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .hflip import HFlipConverter
-from .to_mask import ToMaskConverter
-from .to_chart_result import ToChartResultConverter, ToChartResultConverterWithConfidences
-from .segm_to_mask import (
-    predictor_output_with_fine_and_coarse_segm_to_mask,
-    predictor_output_with_coarse_segm_to_mask,
-    resample_fine_and_coarse_segm_to_bbox,
-)
-from .chart_output_to_chart_result import (
-    densepose_chart_predictor_output_to_result,
-    densepose_chart_predictor_output_to_result_with_confidences,
-)
-from .chart_output_hflip import densepose_chart_predictor_output_hflip
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .hflip import HFlipConverter
+from .to_mask import ToMaskConverter
+from .to_chart_result import ToChartResultConverter, ToChartResultConverterWithConfidences
+from .segm_to_mask import (
+    predictor_output_with_fine_and_coarse_segm_to_mask,
+    predictor_output_with_coarse_segm_to_mask,
+    resample_fine_and_coarse_segm_to_bbox,
+)
+from .chart_output_to_chart_result import (
+    densepose_chart_predictor_output_to_result,
+    densepose_chart_predictor_output_to_result_with_confidences,
+)
+from .chart_output_hflip import densepose_chart_predictor_output_hflip

densepose/converters/base.py

@@ -1,93 +1,93 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Any, Tuple, Type
-import torch
-
-
-class BaseConverter:
-    """
-    Converter base class to be reused by various converters.
-    Converter allows one to convert data from various source types to a particular
-    destination type. Each source type needs to register its converter. The
-    registration for each source type is valid for all descendants of that type.
-    """
-
-    @classmethod
-    def register(cls, from_type: Type, converter: Any = None):
-        """
-        Registers a converter for the specified type.
-        Can be used as a decorator (if converter is None), or called as a method.
-
-        Args:
-            from_type (type): type to register the converter for;
-                all instances of this type will use the same converter
-            converter (callable): converter to be registered for the given
-                type; if None, this method is assumed to be a decorator for the converter
-        """
-
-        if converter is not None:
-            cls._do_register(from_type, converter)
-
-        def wrapper(converter: Any) -> Any:
-            cls._do_register(from_type, converter)
-            return converter
-
-        return wrapper
-
-    @classmethod
-    def _do_register(cls, from_type: Type, converter: Any):
-        cls.registry[from_type] = converter  # pyre-ignore[16]
-
-    @classmethod
-    def _lookup_converter(cls, from_type: Type) -> Any:
-        """
-        Perform recursive lookup for the given type
-        to find registered converter. If a converter was found for some base
-        class, it gets registered for this class to save on further lookups.
-
-        Args:
-            from_type: type for which to find a converter
-        Return:
-            callable or None - registered converter or None
-                if no suitable entry was found in the registry
-        """
-        if from_type in cls.registry:  # pyre-ignore[16]
-            return cls.registry[from_type]
-        for base in from_type.__bases__:
-            converter = cls._lookup_converter(base)
-            if converter is not None:
-                cls._do_register(from_type, converter)
-                return converter
-        return None
-
-    @classmethod
-    def convert(cls, instance: Any, *args, **kwargs):
-        """
-        Convert an instance to the destination type using some registered
-        converter. Does recursive lookup for base classes, so there's no need
-        for explicit registration for derived classes.
-
-        Args:
-            instance: source instance to convert to the destination type
-        Return:
-            An instance of the destination type obtained from the source instance
-            Raises KeyError, if no suitable converter found
-        """
-        instance_type = type(instance)
-        converter = cls._lookup_converter(instance_type)
-        if converter is None:
-            if cls.dst_type is None:  # pyre-ignore[16]
-                output_type_str = "itself"
-            else:
-                output_type_str = cls.dst_type
-            raise KeyError(f"Could not find converter from {instance_type} to {output_type_str}")
-        return converter(instance, *args, **kwargs)
-
-
-IntTupleBox = Tuple[int, int, int, int]
-
-
-def make_int_box(box: torch.Tensor) -> IntTupleBox:
-    int_box = [0, 0, 0, 0]
-    int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
-    return int_box[0], int_box[1], int_box[2], int_box[3]
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Any, Tuple, Type
+import torch
+
+
+class BaseConverter:
+    """
+    Converter base class to be reused by various converters.
+    Converter allows one to convert data from various source types to a particular
+    destination type. Each source type needs to register its converter. The
+    registration for each source type is valid for all descendants of that type.
+    """
+
+    @classmethod
+    def register(cls, from_type: Type, converter: Any = None):
+        """
+        Registers a converter for the specified type.
+        Can be used as a decorator (if converter is None), or called as a method.
+
+        Args:
+            from_type (type): type to register the converter for;
+                all instances of this type will use the same converter
+            converter (callable): converter to be registered for the given
+                type; if None, this method is assumed to be a decorator for the converter
+        """
+
+        if converter is not None:
+            cls._do_register(from_type, converter)
+
+        def wrapper(converter: Any) -> Any:
+            cls._do_register(from_type, converter)
+            return converter
+
+        return wrapper
+
+    @classmethod
+    def _do_register(cls, from_type: Type, converter: Any):
+        cls.registry[from_type] = converter  # pyre-ignore[16]
+
+    @classmethod
+    def _lookup_converter(cls, from_type: Type) -> Any:
+        """
+        Perform recursive lookup for the given type
+        to find registered converter. If a converter was found for some base
+        class, it gets registered for this class to save on further lookups.
+
+        Args:
+            from_type: type for which to find a converter
+        Return:
+            callable or None - registered converter or None
+                if no suitable entry was found in the registry
+        """
+        if from_type in cls.registry:  # pyre-ignore[16]
+            return cls.registry[from_type]
+        for base in from_type.__bases__:
+            converter = cls._lookup_converter(base)
+            if converter is not None:
+                cls._do_register(from_type, converter)
+                return converter
+        return None
+
+    @classmethod
+    def convert(cls, instance: Any, *args, **kwargs):
+        """
+        Convert an instance to the destination type using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            instance: source instance to convert to the destination type
+        Return:
+            An instance of the destination type obtained from the source instance
+            Raises KeyError, if no suitable converter found
+        """
+        instance_type = type(instance)
+        converter = cls._lookup_converter(instance_type)
+        if converter is None:
+            if cls.dst_type is None:  # pyre-ignore[16]
+                output_type_str = "itself"
+            else:
+                output_type_str = cls.dst_type
+            raise KeyError(f"Could not find converter from {instance_type} to {output_type_str}")
+        return converter(instance, *args, **kwargs)
+
+
+IntTupleBox = Tuple[int, int, int, int]
+
+
+def make_int_box(box: torch.Tensor) -> IntTupleBox:
+    int_box = [0, 0, 0, 0]
+    int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
+    return int_box[0], int_box[1], int_box[2], int_box[3]

densepose/converters/builtin.py

@@ -1,31 +1,31 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from ..structures import DensePoseChartPredictorOutput, DensePoseEmbeddingPredictorOutput
-from . import (
-    HFlipConverter,
-    ToChartResultConverter,
-    ToChartResultConverterWithConfidences,
-    ToMaskConverter,
-    densepose_chart_predictor_output_hflip,
-    densepose_chart_predictor_output_to_result,
-    densepose_chart_predictor_output_to_result_with_confidences,
-    predictor_output_with_coarse_segm_to_mask,
-    predictor_output_with_fine_and_coarse_segm_to_mask,
-)
-
-ToMaskConverter.register(
-    DensePoseChartPredictorOutput, predictor_output_with_fine_and_coarse_segm_to_mask
-)
-ToMaskConverter.register(
-    DensePoseEmbeddingPredictorOutput, predictor_output_with_coarse_segm_to_mask
-)
-
-ToChartResultConverter.register(
-    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result
-)
-
-ToChartResultConverterWithConfidences.register(
-    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result_with_confidences
-)
-
-HFlipConverter.register(DensePoseChartPredictorOutput, densepose_chart_predictor_output_hflip)
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from ..structures import DensePoseChartPredictorOutput, DensePoseEmbeddingPredictorOutput
+from . import (
+    HFlipConverter,
+    ToChartResultConverter,
+    ToChartResultConverterWithConfidences,
+    ToMaskConverter,
+    densepose_chart_predictor_output_hflip,
+    densepose_chart_predictor_output_to_result,
+    densepose_chart_predictor_output_to_result_with_confidences,
+    predictor_output_with_coarse_segm_to_mask,
+    predictor_output_with_fine_and_coarse_segm_to_mask,
+)
+
+ToMaskConverter.register(
+    DensePoseChartPredictorOutput, predictor_output_with_fine_and_coarse_segm_to_mask
+)
+ToMaskConverter.register(
+    DensePoseEmbeddingPredictorOutput, predictor_output_with_coarse_segm_to_mask
+)
+
+ToChartResultConverter.register(
+    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result
+)
+
+ToChartResultConverterWithConfidences.register(
+    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result_with_confidences
+)
+
+HFlipConverter.register(DensePoseChartPredictorOutput, densepose_chart_predictor_output_hflip)

densepose/converters/chart_output_hflip.py

@@ -1,71 +1,71 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-from dataclasses import fields
-import torch
-
-from densepose.structures import DensePoseChartPredictorOutput, DensePoseTransformData
-
-
-def densepose_chart_predictor_output_hflip(
-    densepose_predictor_output: DensePoseChartPredictorOutput,
-    transform_data: DensePoseTransformData,
-) -> DensePoseChartPredictorOutput:
-    """
-    Change  to take into account a Horizontal flip.
-    """
-    if len(densepose_predictor_output) > 0:
-
-        PredictorOutput = type(densepose_predictor_output)
-        output_dict = {}
-
-        for field in fields(densepose_predictor_output):
-            field_value = getattr(densepose_predictor_output, field.name)
-            # flip tensors
-            if isinstance(field_value, torch.Tensor):
-                setattr(densepose_predictor_output, field.name, torch.flip(field_value, [3]))
-
-        densepose_predictor_output = _flip_iuv_semantics_tensor(
-            densepose_predictor_output, transform_data
-        )
-        densepose_predictor_output = _flip_segm_semantics_tensor(
-            densepose_predictor_output, transform_data
-        )
-
-        for field in fields(densepose_predictor_output):
-            output_dict[field.name] = getattr(densepose_predictor_output, field.name)
-
-        return PredictorOutput(**output_dict)
-    else:
-        return densepose_predictor_output
-
-
-def _flip_iuv_semantics_tensor(
-    densepose_predictor_output: DensePoseChartPredictorOutput,
-    dp_transform_data: DensePoseTransformData,
-) -> DensePoseChartPredictorOutput:
-    point_label_symmetries = dp_transform_data.point_label_symmetries
-    uv_symmetries = dp_transform_data.uv_symmetries
-
-    N, C, H, W = densepose_predictor_output.u.shape
-    u_loc = (densepose_predictor_output.u[:, 1:, :, :].clamp(0, 1) * 255).long()
-    v_loc = (densepose_predictor_output.v[:, 1:, :, :].clamp(0, 1) * 255).long()
-    Iindex = torch.arange(C - 1, device=densepose_predictor_output.u.device)[
-        None, :, None, None
-    ].expand(N, C - 1, H, W)
-    densepose_predictor_output.u[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc]
-    densepose_predictor_output.v[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc]
-
-    for el in ["fine_segm", "u", "v"]:
-        densepose_predictor_output.__dict__[el] = densepose_predictor_output.__dict__[el][
-            :, point_label_symmetries, :, :
-        ]
-    return densepose_predictor_output
-
-
-def _flip_segm_semantics_tensor(
-    densepose_predictor_output: DensePoseChartPredictorOutput, dp_transform_data
-):
-    if densepose_predictor_output.coarse_segm.shape[1] > 2:
-        densepose_predictor_output.coarse_segm = densepose_predictor_output.coarse_segm[
-            :, dp_transform_data.mask_label_symmetries, :, :
-        ]
-    return densepose_predictor_output
+# Copyright (c) Facebook, Inc. and its affiliates.
+from dataclasses import fields
+import torch
+
+from densepose.structures import DensePoseChartPredictorOutput, DensePoseTransformData
+
+
+def densepose_chart_predictor_output_hflip(
+    densepose_predictor_output: DensePoseChartPredictorOutput,
+    transform_data: DensePoseTransformData,
+) -> DensePoseChartPredictorOutput:
+    """
+    Change  to take into account a Horizontal flip.
+    """
+    if len(densepose_predictor_output) > 0:
+
+        PredictorOutput = type(densepose_predictor_output)
+        output_dict = {}
+
+        for field in fields(densepose_predictor_output):
+            field_value = getattr(densepose_predictor_output, field.name)
+            # flip tensors
+            if isinstance(field_value, torch.Tensor):
+                setattr(densepose_predictor_output, field.name, torch.flip(field_value, [3]))
+
+        densepose_predictor_output = _flip_iuv_semantics_tensor(
+            densepose_predictor_output, transform_data
+        )
+        densepose_predictor_output = _flip_segm_semantics_tensor(
+            densepose_predictor_output, transform_data
+        )
+
+        for field in fields(densepose_predictor_output):
+            output_dict[field.name] = getattr(densepose_predictor_output, field.name)
+
+        return PredictorOutput(**output_dict)
+    else:
+        return densepose_predictor_output
+
+
+def _flip_iuv_semantics_tensor(
+    densepose_predictor_output: DensePoseChartPredictorOutput,
+    dp_transform_data: DensePoseTransformData,
+) -> DensePoseChartPredictorOutput:
+    point_label_symmetries = dp_transform_data.point_label_symmetries
+    uv_symmetries = dp_transform_data.uv_symmetries
+
+    N, C, H, W = densepose_predictor_output.u.shape
+    u_loc = (densepose_predictor_output.u[:, 1:, :, :].clamp(0, 1) * 255).long()
+    v_loc = (densepose_predictor_output.v[:, 1:, :, :].clamp(0, 1) * 255).long()
+    Iindex = torch.arange(C - 1, device=densepose_predictor_output.u.device)[
+        None, :, None, None
+    ].expand(N, C - 1, H, W)
+    densepose_predictor_output.u[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc]
+    densepose_predictor_output.v[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc]
+
+    for el in ["fine_segm", "u", "v"]:
+        densepose_predictor_output.__dict__[el] = densepose_predictor_output.__dict__[el][
+            :, point_label_symmetries, :, :
+        ]
+    return densepose_predictor_output
+
+
+def _flip_segm_semantics_tensor(
+    densepose_predictor_output: DensePoseChartPredictorOutput, dp_transform_data
+):
+    if densepose_predictor_output.coarse_segm.shape[1] > 2:
+        densepose_predictor_output.coarse_segm = densepose_predictor_output.coarse_segm[
+            :, dp_transform_data.mask_label_symmetries, :, :
+        ]
+    return densepose_predictor_output

densepose/converters/chart_output_to_chart_result.py

@@ -1,188 +1,188 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Dict
-import torch
-from torch.nn import functional as F
-
-from detectron2.structures.boxes import Boxes, BoxMode
-
-from ..structures import (
-    DensePoseChartPredictorOutput,
-    DensePoseChartResult,
-    DensePoseChartResultWithConfidences,
-)
-from . import resample_fine_and_coarse_segm_to_bbox
-from .base import IntTupleBox, make_int_box
-
-
-def resample_uv_tensors_to_bbox(
-    u: torch.Tensor,
-    v: torch.Tensor,
-    labels: torch.Tensor,
-    box_xywh_abs: IntTupleBox,
-) -> torch.Tensor:
-    """
-    Resamples U and V coordinate estimates for the given bounding box
-
-    Args:
-        u (tensor [1, C, H, W] of float): U coordinates
-        v (tensor [1, C, H, W] of float): V coordinates
-        labels (tensor [H, W] of long): labels obtained by resampling segmentation
-            outputs for the given bounding box
-        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
-    Return:
-       Resampled U and V coordinates - a tensor [2, H, W] of float
-    """
-    x, y, w, h = box_xywh_abs
-    w = max(int(w), 1)
-    h = max(int(h), 1)
-    u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
-    v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
-    uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
-    for part_id in range(1, u_bbox.size(1)):
-        uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
-        uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
-    return uv
-
-
-def resample_uv_to_bbox(
-    predictor_output: DensePoseChartPredictorOutput,
-    labels: torch.Tensor,
-    box_xywh_abs: IntTupleBox,
-) -> torch.Tensor:
-    """
-    Resamples U and V coordinate estimates for the given bounding box
-
-    Args:
-        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
-            output to be resampled
-        labels (tensor [H, W] of long): labels obtained by resampling segmentation
-            outputs for the given bounding box
-        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
-    Return:
-       Resampled U and V coordinates - a tensor [2, H, W] of float
-    """
-    return resample_uv_tensors_to_bbox(
-        predictor_output.u,
-        predictor_output.v,
-        labels,
-        box_xywh_abs,
-    )
-
-
-def densepose_chart_predictor_output_to_result(
-    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
-) -> DensePoseChartResult:
-    """
-    Convert densepose chart predictor outputs to results
-
-    Args:
-        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
-            output to be converted to results, must contain only 1 output
-        boxes (Boxes): bounding box that corresponds to the predictor output,
-            must contain only 1 bounding box
-    Return:
-       DensePose chart-based result (DensePoseChartResult)
-    """
-    assert len(predictor_output) == 1 and len(boxes) == 1, (
-        f"Predictor output to result conversion can operate only single outputs"
-        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
-    )
-
-    boxes_xyxy_abs = boxes.tensor.clone()
-    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
-    box_xywh = make_int_box(boxes_xywh_abs[0])
-
-    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
-    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
-    return DensePoseChartResult(labels=labels, uv=uv)
-
-
-def resample_confidences_to_bbox(
-    predictor_output: DensePoseChartPredictorOutput,
-    labels: torch.Tensor,
-    box_xywh_abs: IntTupleBox,
-) -> Dict[str, torch.Tensor]:
-    """
-    Resamples confidences for the given bounding box
-
-    Args:
-        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
-            output to be resampled
-        labels (tensor [H, W] of long): labels obtained by resampling segmentation
-            outputs for the given bounding box
-        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
-    Return:
-       Resampled confidences - a dict of [H, W] tensors of float
-    """
-
-    x, y, w, h = box_xywh_abs
-    w = max(int(w), 1)
-    h = max(int(h), 1)
-
-    confidence_names = [
-        "sigma_1",
-        "sigma_2",
-        "kappa_u",
-        "kappa_v",
-        "fine_segm_confidence",
-        "coarse_segm_confidence",
-    ]
-    confidence_results = {key: None for key in confidence_names}
-    confidence_names = [
-        key for key in confidence_names if getattr(predictor_output, key) is not None
-    ]
-    confidence_base = torch.zeros([h, w], dtype=torch.float32, device=predictor_output.u.device)
-
-    # assign data from channels that correspond to the labels
-    for key in confidence_names:
-        resampled_confidence = F.interpolate(
-            getattr(predictor_output, key),
-            (h, w),
-            mode="bilinear",
-            align_corners=False,
-        )
-        result = confidence_base.clone()
-        for part_id in range(1, predictor_output.u.size(1)):
-            if resampled_confidence.size(1) != predictor_output.u.size(1):
-                # confidence is not part-based, don't try to fill it part by part
-                continue
-            result[labels == part_id] = resampled_confidence[0, part_id][labels == part_id]
-
-        if resampled_confidence.size(1) != predictor_output.u.size(1):
-            # confidence is not part-based, fill the data with the first channel
-            # (targeted for segmentation confidences that have only 1 channel)
-            result = resampled_confidence[0, 0]
-
-        confidence_results[key] = result
-
-    return confidence_results  # pyre-ignore[7]
-
-
-def densepose_chart_predictor_output_to_result_with_confidences(
-    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
-) -> DensePoseChartResultWithConfidences:
-    """
-    Convert densepose chart predictor outputs to results
-
-    Args:
-        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
-            output with confidences to be converted to results, must contain only 1 output
-        boxes (Boxes): bounding box that corresponds to the predictor output,
-            must contain only 1 bounding box
-    Return:
-       DensePose chart-based result with confidences (DensePoseChartResultWithConfidences)
-    """
-    assert len(predictor_output) == 1 and len(boxes) == 1, (
-        f"Predictor output to result conversion can operate only single outputs"
-        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
-    )
-
-    boxes_xyxy_abs = boxes.tensor.clone()
-    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
-    box_xywh = make_int_box(boxes_xywh_abs[0])
-
-    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
-    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
-    confidences = resample_confidences_to_bbox(predictor_output, labels, box_xywh)
-    return DensePoseChartResultWithConfidences(labels=labels, uv=uv, **confidences)
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Dict
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures.boxes import Boxes, BoxMode
+
+from ..structures import (
+    DensePoseChartPredictorOutput,
+    DensePoseChartResult,
+    DensePoseChartResultWithConfidences,
+)
+from . import resample_fine_and_coarse_segm_to_bbox
+from .base import IntTupleBox, make_int_box
+
+
+def resample_uv_tensors_to_bbox(
+    u: torch.Tensor,
+    v: torch.Tensor,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> torch.Tensor:
+    """
+    Resamples U and V coordinate estimates for the given bounding box
+
+    Args:
+        u (tensor [1, C, H, W] of float): U coordinates
+        v (tensor [1, C, H, W] of float): V coordinates
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled U and V coordinates - a tensor [2, H, W] of float
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
+    v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
+    uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
+    for part_id in range(1, u_bbox.size(1)):
+        uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
+        uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
+    return uv
+
+
+def resample_uv_to_bbox(
+    predictor_output: DensePoseChartPredictorOutput,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> torch.Tensor:
+    """
+    Resamples U and V coordinate estimates for the given bounding box
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be resampled
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled U and V coordinates - a tensor [2, H, W] of float
+    """
+    return resample_uv_tensors_to_bbox(
+        predictor_output.u,
+        predictor_output.v,
+        labels,
+        box_xywh_abs,
+    )
+
+
+def densepose_chart_predictor_output_to_result(
+    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
+) -> DensePoseChartResult:
+    """
+    Convert densepose chart predictor outputs to results
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be converted to results, must contain only 1 output
+        boxes (Boxes): bounding box that corresponds to the predictor output,
+            must contain only 1 bounding box
+    Return:
+       DensePose chart-based result (DensePoseChartResult)
+    """
+    assert len(predictor_output) == 1 and len(boxes) == 1, (
+        f"Predictor output to result conversion can operate only single outputs"
+        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
+    )
+
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    box_xywh = make_int_box(boxes_xywh_abs[0])
+
+    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
+    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
+    return DensePoseChartResult(labels=labels, uv=uv)
+
+
+def resample_confidences_to_bbox(
+    predictor_output: DensePoseChartPredictorOutput,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> Dict[str, torch.Tensor]:
+    """
+    Resamples confidences for the given bounding box
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be resampled
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled confidences - a dict of [H, W] tensors of float
+    """
+
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+
+    confidence_names = [
+        "sigma_1",
+        "sigma_2",
+        "kappa_u",
+        "kappa_v",
+        "fine_segm_confidence",
+        "coarse_segm_confidence",
+    ]
+    confidence_results = {key: None for key in confidence_names}
+    confidence_names = [
+        key for key in confidence_names if getattr(predictor_output, key) is not None
+    ]
+    confidence_base = torch.zeros([h, w], dtype=torch.float32, device=predictor_output.u.device)
+
+    # assign data from channels that correspond to the labels
+    for key in confidence_names:
+        resampled_confidence = F.interpolate(
+            getattr(predictor_output, key),
+            (h, w),
+            mode="bilinear",
+            align_corners=False,
+        )
+        result = confidence_base.clone()
+        for part_id in range(1, predictor_output.u.size(1)):
+            if resampled_confidence.size(1) != predictor_output.u.size(1):
+                # confidence is not part-based, don't try to fill it part by part
+                continue
+            result[labels == part_id] = resampled_confidence[0, part_id][labels == part_id]
+
+        if resampled_confidence.size(1) != predictor_output.u.size(1):
+            # confidence is not part-based, fill the data with the first channel
+            # (targeted for segmentation confidences that have only 1 channel)
+            result = resampled_confidence[0, 0]
+
+        confidence_results[key] = result
+
+    return confidence_results  # pyre-ignore[7]
+
+
+def densepose_chart_predictor_output_to_result_with_confidences(
+    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
+) -> DensePoseChartResultWithConfidences:
+    """
+    Convert densepose chart predictor outputs to results
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output with confidences to be converted to results, must contain only 1 output
+        boxes (Boxes): bounding box that corresponds to the predictor output,
+            must contain only 1 bounding box
+    Return:
+       DensePose chart-based result with confidences (DensePoseChartResultWithConfidences)
+    """
+    assert len(predictor_output) == 1 and len(boxes) == 1, (
+        f"Predictor output to result conversion can operate only single outputs"
+        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
+    )
+
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    box_xywh = make_int_box(boxes_xywh_abs[0])
+
+    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
+    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
+    confidences = resample_confidences_to_bbox(predictor_output, labels, box_xywh)
+    return DensePoseChartResultWithConfidences(labels=labels, uv=uv, **confidences)

densepose/converters/hflip.py

@@ -1,34 +1,34 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Any
-
-from .base import BaseConverter
-
-
-class HFlipConverter(BaseConverter):
-    """
-    Converts various DensePose predictor outputs to DensePose results.
-    Each DensePose predictor output type has to register its convertion strategy.
-    """
-
-    registry = {}
-    dst_type = None
-
-    @classmethod
-    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
-    #  inconsistently.
-    def convert(cls, predictor_outputs: Any, transform_data: Any, *args, **kwargs):
-        """
-        Performs an horizontal flip on DensePose predictor outputs.
-        Does recursive lookup for base classes, so there's no need
-        for explicit registration for derived classes.
-
-        Args:
-            predictor_outputs: DensePose predictor output to be converted to BitMasks
-            transform_data: Anything useful for the flip
-        Return:
-            An instance of the same type as predictor_outputs
-        """
-        return super(HFlipConverter, cls).convert(
-            predictor_outputs, transform_data, *args, **kwargs
-        )
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Any
+
+from .base import BaseConverter
+
+
+class HFlipConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = None
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(cls, predictor_outputs: Any, transform_data: Any, *args, **kwargs):
+        """
+        Performs an horizontal flip on DensePose predictor outputs.
+        Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            predictor_outputs: DensePose predictor output to be converted to BitMasks
+            transform_data: Anything useful for the flip
+        Return:
+            An instance of the same type as predictor_outputs
+        """
+        return super(HFlipConverter, cls).convert(
+            predictor_outputs, transform_data, *args, **kwargs
+        )

densepose/converters/segm_to_mask.py

@@ -1,150 +1,150 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Any
-import torch
-from torch.nn import functional as F
-
-from detectron2.structures import BitMasks, Boxes, BoxMode
-
-from .base import IntTupleBox, make_int_box
-from .to_mask import ImageSizeType
-
-
-def resample_coarse_segm_tensor_to_bbox(coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox):
-    """
-    Resample coarse segmentation tensor to the given
-    bounding box and derive labels for each pixel of the bounding box
-
-    Args:
-        coarse_segm: float tensor of shape [1, K, Hout, Wout]
-        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
-            corner coordinates, width (W) and height (H)
-    Return:
-        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
-    """
-    x, y, w, h = box_xywh_abs
-    w = max(int(w), 1)
-    h = max(int(h), 1)
-    labels = F.interpolate(coarse_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
-    return labels
-
-
-def resample_fine_and_coarse_segm_tensors_to_bbox(
-    fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
-):
-    """
-    Resample fine and coarse segmentation tensors to the given
-    bounding box and derive labels for each pixel of the bounding box
-
-    Args:
-        fine_segm: float tensor of shape [1, C, Hout, Wout]
-        coarse_segm: float tensor of shape [1, K, Hout, Wout]
-        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
-            corner coordinates, width (W) and height (H)
-    Return:
-        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
-    """
-    x, y, w, h = box_xywh_abs
-    w = max(int(w), 1)
-    h = max(int(h), 1)
-    # coarse segmentation
-    coarse_segm_bbox = F.interpolate(
-        coarse_segm,
-        (h, w),
-        mode="bilinear",
-        align_corners=False,
-    ).argmax(dim=1)
-    # combined coarse and fine segmentation
-    labels = (
-        F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
-        * (coarse_segm_bbox > 0).long()
-    )
-    return labels
-
-
-def resample_fine_and_coarse_segm_to_bbox(predictor_output: Any, box_xywh_abs: IntTupleBox):
-    """
-    Resample fine and coarse segmentation outputs from a predictor to the given
-    bounding box and derive labels for each pixel of the bounding box
-
-    Args:
-        predictor_output: DensePose predictor output that contains segmentation
-            results to be resampled
-        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
-            corner coordinates, width (W) and height (H)
-    Return:
-        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
-    """
-    return resample_fine_and_coarse_segm_tensors_to_bbox(
-        predictor_output.fine_segm,
-        predictor_output.coarse_segm,
-        box_xywh_abs,
-    )
-
-
-def predictor_output_with_coarse_segm_to_mask(
-    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
-) -> BitMasks:
-    """
-    Convert predictor output with coarse and fine segmentation to a mask.
-    Assumes that predictor output has the following attributes:
-     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
-         unnormalized scores for N instances; D is the number of coarse
-         segmentation labels, H and W is the resolution of the estimate
-
-    Args:
-        predictor_output: DensePose predictor output to be converted to mask
-        boxes (Boxes): bounding boxes that correspond to the DensePose
-            predictor outputs
-        image_size_hw (tuple [int, int]): image height Himg and width Wimg
-    Return:
-        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
-        a mask of the size of the image for each instance
-    """
-    H, W = image_size_hw
-    boxes_xyxy_abs = boxes.tensor.clone()
-    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
-    N = len(boxes_xywh_abs)
-    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
-    for i in range(len(boxes_xywh_abs)):
-        box_xywh = make_int_box(boxes_xywh_abs[i])
-        box_mask = resample_coarse_segm_tensor_to_bbox(predictor_output[i].coarse_segm, box_xywh)
-        x, y, w, h = box_xywh
-        masks[i, y : y + h, x : x + w] = box_mask
-
-    return BitMasks(masks)
-
-
-def predictor_output_with_fine_and_coarse_segm_to_mask(
-    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
-) -> BitMasks:
-    """
-    Convert predictor output with coarse and fine segmentation to a mask.
-    Assumes that predictor output has the following attributes:
-     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
-         unnormalized scores for N instances; D is the number of coarse
-         segmentation labels, H and W is the resolution of the estimate
-     - fine_segm (tensor of size [N, C, H, W]): fine segmentation
-         unnormalized scores for N instances; C is the number of fine
-         segmentation labels, H and W is the resolution of the estimate
-
-    Args:
-        predictor_output: DensePose predictor output to be converted to mask
-        boxes (Boxes): bounding boxes that correspond to the DensePose
-            predictor outputs
-        image_size_hw (tuple [int, int]): image height Himg and width Wimg
-    Return:
-        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
-        a mask of the size of the image for each instance
-    """
-    H, W = image_size_hw
-    boxes_xyxy_abs = boxes.tensor.clone()
-    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
-    N = len(boxes_xywh_abs)
-    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
-    for i in range(len(boxes_xywh_abs)):
-        box_xywh = make_int_box(boxes_xywh_abs[i])
-        labels_i = resample_fine_and_coarse_segm_to_bbox(predictor_output[i], box_xywh)
-        x, y, w, h = box_xywh
-        masks[i, y : y + h, x : x + w] = labels_i > 0
-    return BitMasks(masks)
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Any
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures import BitMasks, Boxes, BoxMode
+
+from .base import IntTupleBox, make_int_box
+from .to_mask import ImageSizeType
+
+
+def resample_coarse_segm_tensor_to_bbox(coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox):
+    """
+    Resample coarse segmentation tensor to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        coarse_segm: float tensor of shape [1, K, Hout, Wout]
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    labels = F.interpolate(coarse_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
+    return labels
+
+
+def resample_fine_and_coarse_segm_tensors_to_bbox(
+    fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
+):
+    """
+    Resample fine and coarse segmentation tensors to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        fine_segm: float tensor of shape [1, C, Hout, Wout]
+        coarse_segm: float tensor of shape [1, K, Hout, Wout]
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    # coarse segmentation
+    coarse_segm_bbox = F.interpolate(
+        coarse_segm,
+        (h, w),
+        mode="bilinear",
+        align_corners=False,
+    ).argmax(dim=1)
+    # combined coarse and fine segmentation
+    labels = (
+        F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
+        * (coarse_segm_bbox > 0).long()
+    )
+    return labels
+
+
+def resample_fine_and_coarse_segm_to_bbox(predictor_output: Any, box_xywh_abs: IntTupleBox):
+    """
+    Resample fine and coarse segmentation outputs from a predictor to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        predictor_output: DensePose predictor output that contains segmentation
+            results to be resampled
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    return resample_fine_and_coarse_segm_tensors_to_bbox(
+        predictor_output.fine_segm,
+        predictor_output.coarse_segm,
+        box_xywh_abs,
+    )
+
+
+def predictor_output_with_coarse_segm_to_mask(
+    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
+) -> BitMasks:
+    """
+    Convert predictor output with coarse and fine segmentation to a mask.
+    Assumes that predictor output has the following attributes:
+     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
+         unnormalized scores for N instances; D is the number of coarse
+         segmentation labels, H and W is the resolution of the estimate
+
+    Args:
+        predictor_output: DensePose predictor output to be converted to mask
+        boxes (Boxes): bounding boxes that correspond to the DensePose
+            predictor outputs
+        image_size_hw (tuple [int, int]): image height Himg and width Wimg
+    Return:
+        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
+        a mask of the size of the image for each instance
+    """
+    H, W = image_size_hw
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    N = len(boxes_xywh_abs)
+    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
+    for i in range(len(boxes_xywh_abs)):
+        box_xywh = make_int_box(boxes_xywh_abs[i])
+        box_mask = resample_coarse_segm_tensor_to_bbox(predictor_output[i].coarse_segm, box_xywh)
+        x, y, w, h = box_xywh
+        masks[i, y : y + h, x : x + w] = box_mask
+
+    return BitMasks(masks)
+
+
+def predictor_output_with_fine_and_coarse_segm_to_mask(
+    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
+) -> BitMasks:
+    """
+    Convert predictor output with coarse and fine segmentation to a mask.
+    Assumes that predictor output has the following attributes:
+     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
+         unnormalized scores for N instances; D is the number of coarse
+         segmentation labels, H and W is the resolution of the estimate
+     - fine_segm (tensor of size [N, C, H, W]): fine segmentation
+         unnormalized scores for N instances; C is the number of fine
+         segmentation labels, H and W is the resolution of the estimate
+
+    Args:
+        predictor_output: DensePose predictor output to be converted to mask
+        boxes (Boxes): bounding boxes that correspond to the DensePose
+            predictor outputs
+        image_size_hw (tuple [int, int]): image height Himg and width Wimg
+    Return:
+        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
+        a mask of the size of the image for each instance
+    """
+    H, W = image_size_hw
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    N = len(boxes_xywh_abs)
+    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
+    for i in range(len(boxes_xywh_abs)):
+        box_xywh = make_int_box(boxes_xywh_abs[i])
+        labels_i = resample_fine_and_coarse_segm_to_bbox(predictor_output[i], box_xywh)
+        x, y, w, h = box_xywh
+        masks[i, y : y + h, x : x + w] = labels_i > 0
+    return BitMasks(masks)

densepose/converters/to_chart_result.py

@@ -1,70 +1,70 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Any
-
-from detectron2.structures import Boxes
-
-from ..structures import DensePoseChartResult, DensePoseChartResultWithConfidences
-from .base import BaseConverter
-
-
-class ToChartResultConverter(BaseConverter):
-    """
-    Converts various DensePose predictor outputs to DensePose results.
-    Each DensePose predictor output type has to register its convertion strategy.
-    """
-
-    registry = {}
-    dst_type = DensePoseChartResult
-
-    @classmethod
-    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
-    #  inconsistently.
-    def convert(cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs) -> DensePoseChartResult:
-        """
-        Convert DensePose predictor outputs to DensePoseResult using some registered
-        converter. Does recursive lookup for base classes, so there's no need
-        for explicit registration for derived classes.
-
-        Args:
-            densepose_predictor_outputs: DensePose predictor output to be
-                converted to BitMasks
-            boxes (Boxes): bounding boxes that correspond to the DensePose
-                predictor outputs
-        Return:
-            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
-        """
-        return super(ToChartResultConverter, cls).convert(predictor_outputs, boxes, *args, **kwargs)
-
-
-class ToChartResultConverterWithConfidences(BaseConverter):
-    """
-    Converts various DensePose predictor outputs to DensePose results.
-    Each DensePose predictor output type has to register its convertion strategy.
-    """
-
-    registry = {}
-    dst_type = DensePoseChartResultWithConfidences
-
-    @classmethod
-    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
-    #  inconsistently.
-    def convert(
-        cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs
-    ) -> DensePoseChartResultWithConfidences:
-        """
-        Convert DensePose predictor outputs to DensePoseResult with confidences
-        using some registered converter. Does recursive lookup for base classes,
-        so there's no need for explicit registration for derived classes.
-
-        Args:
-            densepose_predictor_outputs: DensePose predictor output with confidences
-                to be converted to BitMasks
-            boxes (Boxes): bounding boxes that correspond to the DensePose
-                predictor outputs
-        Return:
-            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
-        """
-        return super(ToChartResultConverterWithConfidences, cls).convert(
-            predictor_outputs, boxes, *args, **kwargs
-        )
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Any
+
+from detectron2.structures import Boxes
+
+from ..structures import DensePoseChartResult, DensePoseChartResultWithConfidences
+from .base import BaseConverter
+
+
+class ToChartResultConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = DensePoseChartResult
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs) -> DensePoseChartResult:
+        """
+        Convert DensePose predictor outputs to DensePoseResult using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output to be
+                converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+        Return:
+            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
+        """
+        return super(ToChartResultConverter, cls).convert(predictor_outputs, boxes, *args, **kwargs)
+
+
+class ToChartResultConverterWithConfidences(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = DensePoseChartResultWithConfidences
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(
+        cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs
+    ) -> DensePoseChartResultWithConfidences:
+        """
+        Convert DensePose predictor outputs to DensePoseResult with confidences
+        using some registered converter. Does recursive lookup for base classes,
+        so there's no need for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output with confidences
+                to be converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+        Return:
+            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
+        """
+        return super(ToChartResultConverterWithConfidences, cls).convert(
+            predictor_outputs, boxes, *args, **kwargs
+        )

densepose/converters/to_mask.py

@@ -1,49 +1,49 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Any, Tuple
-
-from detectron2.structures import BitMasks, Boxes
-
-from .base import BaseConverter
-
-ImageSizeType = Tuple[int, int]
-
-
-class ToMaskConverter(BaseConverter):
-    """
-    Converts various DensePose predictor outputs to masks
-    in bit mask format (see `BitMasks`). Each DensePose predictor output type
-    has to register its convertion strategy.
-    """
-
-    registry = {}
-    dst_type = BitMasks
-
-    @classmethod
-    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
-    #  inconsistently.
-    def convert(
-        cls,
-        densepose_predictor_outputs: Any,
-        boxes: Boxes,
-        image_size_hw: ImageSizeType,
-        *args,
-        **kwargs
-    ) -> BitMasks:
-        """
-        Convert DensePose predictor outputs to BitMasks using some registered
-        converter. Does recursive lookup for base classes, so there's no need
-        for explicit registration for derived classes.
-
-        Args:
-            densepose_predictor_outputs: DensePose predictor output to be
-                converted to BitMasks
-            boxes (Boxes): bounding boxes that correspond to the DensePose
-                predictor outputs
-            image_size_hw (tuple [int, int]): image height and width
-        Return:
-            An instance of `BitMasks`. If no suitable converter was found, raises KeyError
-        """
-        return super(ToMaskConverter, cls).convert(
-            densepose_predictor_outputs, boxes, image_size_hw, *args, **kwargs
-        )
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Any, Tuple
+
+from detectron2.structures import BitMasks, Boxes
+
+from .base import BaseConverter
+
+ImageSizeType = Tuple[int, int]
+
+
+class ToMaskConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to masks
+    in bit mask format (see `BitMasks`). Each DensePose predictor output type
+    has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = BitMasks
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(
+        cls,
+        densepose_predictor_outputs: Any,
+        boxes: Boxes,
+        image_size_hw: ImageSizeType,
+        *args,
+        **kwargs
+    ) -> BitMasks:
+        """
+        Convert DensePose predictor outputs to BitMasks using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output to be
+                converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+            image_size_hw (tuple [int, int]): image height and width
+        Return:
+            An instance of `BitMasks`. If no suitable converter was found, raises KeyError
+        """
+        return super(ToMaskConverter, cls).convert(
+            densepose_predictor_outputs, boxes, image_size_hw, *args, **kwargs
+        )

densepose/data/__init__.py

@@ -1,25 +1,25 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .meshes import builtin
-from .build import (
-    build_detection_test_loader,
-    build_detection_train_loader,
-    build_combined_loader,
-    build_frame_selector,
-    build_inference_based_loaders,
-    has_inference_based_loaders,
-    BootstrapDatasetFactoryCatalog,
-)
-from .combined_loader import CombinedDataLoader
-from .dataset_mapper import DatasetMapper
-from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
-from .image_list_dataset import ImageListDataset
-from .utils import is_relative_local_path, maybe_prepend_base_path
-
-# ensure the builtin datasets are registered
-from . import datasets
-
-# ensure the bootstrap datasets builders are registered
-from . import build
-
-__all__ = [k for k in globals().keys() if not k.startswith("_")]
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .meshes import builtin
+from .build import (
+    build_detection_test_loader,
+    build_detection_train_loader,
+    build_combined_loader,
+    build_frame_selector,
+    build_inference_based_loaders,
+    has_inference_based_loaders,
+    BootstrapDatasetFactoryCatalog,
+)
+from .combined_loader import CombinedDataLoader
+from .dataset_mapper import DatasetMapper
+from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
+from .image_list_dataset import ImageListDataset
+from .utils import is_relative_local_path, maybe_prepend_base_path
+
+# ensure the builtin datasets are registered
+from . import datasets
+
+# ensure the bootstrap datasets builders are registered
+from . import build
+
+__all__ = [k for k in globals().keys() if not k.startswith("_")]

densepose/data/build.py

@@ -1,736 +1,736 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import itertools
-import logging
-import numpy as np
-from collections import UserDict, defaultdict
-from dataclasses import dataclass
-from typing import Any, Callable, Collection, Dict, Iterable, List, Optional, Sequence, Tuple
-import torch
-from torch.utils.data.dataset import Dataset
-
-from detectron2.config import CfgNode
-from detectron2.data.build import build_detection_test_loader as d2_build_detection_test_loader
-from detectron2.data.build import build_detection_train_loader as d2_build_detection_train_loader
-from detectron2.data.build import (
-    load_proposals_into_dataset,
-    print_instances_class_histogram,
-    trivial_batch_collator,
-    worker_init_reset_seed,
-)
-from detectron2.data.catalog import DatasetCatalog, Metadata, MetadataCatalog
-from detectron2.data.samplers import TrainingSampler
-from detectron2.utils.comm import get_world_size
-
-from densepose.config import get_bootstrap_dataset_config
-from densepose.modeling import build_densepose_embedder
-
-from .combined_loader import CombinedDataLoader, Loader
-from .dataset_mapper import DatasetMapper
-from .datasets.coco import DENSEPOSE_CSE_KEYS_WITHOUT_MASK, DENSEPOSE_IUV_KEYS_WITHOUT_MASK
-from .datasets.dataset_type import DatasetType
-from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
-from .samplers import (
-    DensePoseConfidenceBasedSampler,
-    DensePoseCSEConfidenceBasedSampler,
-    DensePoseCSEUniformSampler,
-    DensePoseUniformSampler,
-    MaskFromDensePoseSampler,
-    PredictionToGroundTruthSampler,
-)
-from .transform import ImageResizeTransform
-from .utils import get_category_to_class_mapping, get_class_to_mesh_name_mapping
-from .video import (
-    FirstKFramesSelector,
-    FrameSelectionStrategy,
-    LastKFramesSelector,
-    RandomKFramesSelector,
-    VideoKeyframeDataset,
-    video_list_from_file,
-)
-
-__all__ = ["build_detection_train_loader", "build_detection_test_loader"]
-
-
-Instance = Dict[str, Any]
-InstancePredicate = Callable[[Instance], bool]
-
-
-def _compute_num_images_per_worker(cfg: CfgNode) -> int:
-    num_workers = get_world_size()
-    images_per_batch = cfg.SOLVER.IMS_PER_BATCH
-    assert (
-        images_per_batch % num_workers == 0
-    ), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of workers ({}).".format(
-        images_per_batch, num_workers
-    )
-    assert (
-        images_per_batch >= num_workers
-    ), "SOLVER.IMS_PER_BATCH ({}) must be larger than the number of workers ({}).".format(
-        images_per_batch, num_workers
-    )
-    images_per_worker = images_per_batch // num_workers
-    return images_per_worker
-
-
-def _map_category_id_to_contiguous_id(dataset_name: str, dataset_dicts: Iterable[Instance]) -> None:
-    meta = MetadataCatalog.get(dataset_name)
-    for dataset_dict in dataset_dicts:
-        for ann in dataset_dict["annotations"]:
-            ann["category_id"] = meta.thing_dataset_id_to_contiguous_id[ann["category_id"]]
-
-
-@dataclass
-class _DatasetCategory:
-    """
-    Class representing category data in a dataset:
-     - id: category ID, as specified in the dataset annotations file
-     - name: category name, as specified in the dataset annotations file
-     - mapped_id: category ID after applying category maps (DATASETS.CATEGORY_MAPS config option)
-     - mapped_name: category name after applying category maps
-     - dataset_name: dataset in which the category is defined
-
-    For example, when training models in a class-agnostic manner, one could take LVIS 1.0
-    dataset and map the animal categories to the same category as human data from COCO:
-     id = 225
-     name = "cat"
-     mapped_id = 1
-     mapped_name = "person"
-     dataset_name = "lvis_v1_animals_dp_train"
-    """
-
-    id: int
-    name: str
-    mapped_id: int
-    mapped_name: str
-    dataset_name: str
-
-
-_MergedCategoriesT = Dict[int, List[_DatasetCategory]]
-
-
-def _add_category_id_to_contiguous_id_maps_to_metadata(
-    merged_categories: _MergedCategoriesT,
-) -> None:
-    merged_categories_per_dataset = {}
-    for contiguous_cat_id, cat_id in enumerate(sorted(merged_categories.keys())):
-        for cat in merged_categories[cat_id]:
-            if cat.dataset_name not in merged_categories_per_dataset:
-                merged_categories_per_dataset[cat.dataset_name] = defaultdict(list)
-            merged_categories_per_dataset[cat.dataset_name][cat_id].append(
-                (
-                    contiguous_cat_id,
-                    cat,
-                )
-            )
-
-    logger = logging.getLogger(__name__)
-    for dataset_name, merged_categories in merged_categories_per_dataset.items():
-        meta = MetadataCatalog.get(dataset_name)
-        if not hasattr(meta, "thing_classes"):
-            meta.thing_classes = []
-            meta.thing_dataset_id_to_contiguous_id = {}
-            meta.thing_dataset_id_to_merged_id = {}
-        else:
-            meta.thing_classes.clear()
-            meta.thing_dataset_id_to_contiguous_id.clear()
-            meta.thing_dataset_id_to_merged_id.clear()
-        logger.info(f"Dataset {dataset_name}: category ID to contiguous ID mapping:")
-        for _cat_id, categories in sorted(merged_categories.items()):
-            added_to_thing_classes = False
-            for contiguous_cat_id, cat in categories:
-                if not added_to_thing_classes:
-                    meta.thing_classes.append(cat.mapped_name)
-                    added_to_thing_classes = True
-                meta.thing_dataset_id_to_contiguous_id[cat.id] = contiguous_cat_id
-                meta.thing_dataset_id_to_merged_id[cat.id] = cat.mapped_id
-                logger.info(f"{cat.id} ({cat.name}) -> {contiguous_cat_id}")
-
-
-def _maybe_create_general_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
-    def has_annotations(instance: Instance) -> bool:
-        return "annotations" in instance
-
-    def has_only_crowd_anotations(instance: Instance) -> bool:
-        for ann in instance["annotations"]:
-            if ann.get("is_crowd", 0) == 0:
-                return False
-        return True
-
-    def general_keep_instance_predicate(instance: Instance) -> bool:
-        return has_annotations(instance) and not has_only_crowd_anotations(instance)
-
-    if not cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS:
-        return None
-    return general_keep_instance_predicate
-
-
-def _maybe_create_keypoints_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
-
-    min_num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
-
-    def has_sufficient_num_keypoints(instance: Instance) -> bool:
-        num_kpts = sum(
-            (np.array(ann["keypoints"][2::3]) > 0).sum()
-            for ann in instance["annotations"]
-            if "keypoints" in ann
-        )
-        return num_kpts >= min_num_keypoints
-
-    if cfg.MODEL.KEYPOINT_ON and (min_num_keypoints > 0):
-        return has_sufficient_num_keypoints
-    return None
-
-
-def _maybe_create_mask_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
-    if not cfg.MODEL.MASK_ON:
-        return None
-
-    def has_mask_annotations(instance: Instance) -> bool:
-        return any("segmentation" in ann for ann in instance["annotations"])
-
-    return has_mask_annotations
-
-
-def _maybe_create_densepose_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
-    if not cfg.MODEL.DENSEPOSE_ON:
-        return None
-
-    use_masks = cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS
-
-    def has_densepose_annotations(instance: Instance) -> bool:
-        for ann in instance["annotations"]:
-            if all(key in ann for key in DENSEPOSE_IUV_KEYS_WITHOUT_MASK) or all(
-                key in ann for key in DENSEPOSE_CSE_KEYS_WITHOUT_MASK
-            ):
-                return True
-            if use_masks and "segmentation" in ann:
-                return True
-        return False
-
-    return has_densepose_annotations
-
-
-def _maybe_create_specific_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
-    specific_predicate_creators = [
-        _maybe_create_keypoints_keep_instance_predicate,
-        _maybe_create_mask_keep_instance_predicate,
-        _maybe_create_densepose_keep_instance_predicate,
-    ]
-    predicates = [creator(cfg) for creator in specific_predicate_creators]
-    predicates = [p for p in predicates if p is not None]
-    if not predicates:
-        return None
-
-    def combined_predicate(instance: Instance) -> bool:
-        return any(p(instance) for p in predicates)
-
-    return combined_predicate
-
-
-def _get_train_keep_instance_predicate(cfg: CfgNode):
-    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
-    combined_specific_keep_predicate = _maybe_create_specific_keep_instance_predicate(cfg)
-
-    def combined_general_specific_keep_predicate(instance: Instance) -> bool:
-        return general_keep_predicate(instance) and combined_specific_keep_predicate(instance)
-
-    if (general_keep_predicate is None) and (combined_specific_keep_predicate is None):
-        return None
-    if general_keep_predicate is None:
-        return combined_specific_keep_predicate
-    if combined_specific_keep_predicate is None:
-        return general_keep_predicate
-    return combined_general_specific_keep_predicate
-
-
-def _get_test_keep_instance_predicate(cfg: CfgNode):
-    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
-    return general_keep_predicate
-
-
-def _maybe_filter_and_map_categories(
-    dataset_name: str, dataset_dicts: List[Instance]
-) -> List[Instance]:
-    meta = MetadataCatalog.get(dataset_name)
-    category_id_map = meta.thing_dataset_id_to_contiguous_id
-    filtered_dataset_dicts = []
-    for dataset_dict in dataset_dicts:
-        anns = []
-        for ann in dataset_dict["annotations"]:
-            cat_id = ann["category_id"]
-            if cat_id not in category_id_map:
-                continue
-            ann["category_id"] = category_id_map[cat_id]
-            anns.append(ann)
-        dataset_dict["annotations"] = anns
-        filtered_dataset_dicts.append(dataset_dict)
-    return filtered_dataset_dicts
-
-
-def _add_category_whitelists_to_metadata(cfg: CfgNode) -> None:
-    for dataset_name, whitelisted_cat_ids in cfg.DATASETS.WHITELISTED_CATEGORIES.items():
-        meta = MetadataCatalog.get(dataset_name)
-        meta.whitelisted_categories = whitelisted_cat_ids
-        logger = logging.getLogger(__name__)
-        logger.info(
-            "Whitelisted categories for dataset {}: {}".format(
-                dataset_name, meta.whitelisted_categories
-            )
-        )
-
-
-def _add_category_maps_to_metadata(cfg: CfgNode) -> None:
-    for dataset_name, category_map in cfg.DATASETS.CATEGORY_MAPS.items():
-        category_map = {
-            int(cat_id_src): int(cat_id_dst) for cat_id_src, cat_id_dst in category_map.items()
-        }
-        meta = MetadataCatalog.get(dataset_name)
-        meta.category_map = category_map
-        logger = logging.getLogger(__name__)
-        logger.info("Category maps for dataset {}: {}".format(dataset_name, meta.category_map))
-
-
-def _add_category_info_to_bootstrapping_metadata(dataset_name: str, dataset_cfg: CfgNode) -> None:
-    meta = MetadataCatalog.get(dataset_name)
-    meta.category_to_class_mapping = get_category_to_class_mapping(dataset_cfg)
-    meta.categories = dataset_cfg.CATEGORIES
-    meta.max_count_per_category = dataset_cfg.MAX_COUNT_PER_CATEGORY
-    logger = logging.getLogger(__name__)
-    logger.info(
-        "Category to class mapping for dataset {}: {}".format(
-            dataset_name, meta.category_to_class_mapping
-        )
-    )
-
-
-def _maybe_add_class_to_mesh_name_map_to_metadata(dataset_names: List[str], cfg: CfgNode) -> None:
-    for dataset_name in dataset_names:
-        meta = MetadataCatalog.get(dataset_name)
-        if not hasattr(meta, "class_to_mesh_name"):
-            meta.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
-
-
-def _merge_categories(dataset_names: Collection[str]) -> _MergedCategoriesT:
-    merged_categories = defaultdict(list)
-    category_names = {}
-    for dataset_name in dataset_names:
-        meta = MetadataCatalog.get(dataset_name)
-        whitelisted_categories = meta.get("whitelisted_categories")
-        category_map = meta.get("category_map", {})
-        cat_ids = (
-            whitelisted_categories if whitelisted_categories is not None else meta.categories.keys()
-        )
-        for cat_id in cat_ids:
-            cat_name = meta.categories[cat_id]
-            cat_id_mapped = category_map.get(cat_id, cat_id)
-            if cat_id_mapped == cat_id or cat_id_mapped in cat_ids:
-                category_names[cat_id] = cat_name
-            else:
-                category_names[cat_id] = str(cat_id_mapped)
-            # assign temporary mapped category name, this name can be changed
-            # during the second pass, since mapped ID can correspond to a category
-            # from a different dataset
-            cat_name_mapped = meta.categories[cat_id_mapped]
-            merged_categories[cat_id_mapped].append(
-                _DatasetCategory(
-                    id=cat_id,
-                    name=cat_name,
-                    mapped_id=cat_id_mapped,
-                    mapped_name=cat_name_mapped,
-                    dataset_name=dataset_name,
-                )
-            )
-    # second pass to assign proper mapped category names
-    for cat_id, categories in merged_categories.items():
-        for cat in categories:
-            if cat_id in category_names and cat.mapped_name != category_names[cat_id]:
-                cat.mapped_name = category_names[cat_id]
-
-    return merged_categories
-
-
-def _warn_if_merged_different_categories(merged_categories: _MergedCategoriesT) -> None:
-    logger = logging.getLogger(__name__)
-    for cat_id in merged_categories:
-        merged_categories_i = merged_categories[cat_id]
-        first_cat_name = merged_categories_i[0].name
-        if len(merged_categories_i) > 1 and not all(
-            cat.name == first_cat_name for cat in merged_categories_i[1:]
-        ):
-            cat_summary_str = ", ".join(
-                [f"{cat.id} ({cat.name}) from {cat.dataset_name}" for cat in merged_categories_i]
-            )
-            logger.warning(
-                f"Merged category {cat_id} corresponds to the following categories: "
-                f"{cat_summary_str}"
-            )
-
-
-def combine_detection_dataset_dicts(
-    dataset_names: Collection[str],
-    keep_instance_predicate: Optional[InstancePredicate] = None,
-    proposal_files: Optional[Collection[str]] = None,
-) -> List[Instance]:
-    """
-    Load and prepare dataset dicts for training / testing
-
-    Args:
-        dataset_names (Collection[str]): a list of dataset names
-        keep_instance_predicate (Callable: Dict[str, Any] -> bool): predicate
-            applied to instance dicts which defines whether to keep the instance
-        proposal_files (Collection[str]): if given, a list of object proposal files
-            that match each dataset in `dataset_names`.
-    """
-    assert len(dataset_names)
-    if proposal_files is None:
-        proposal_files = [None] * len(dataset_names)
-    assert len(dataset_names) == len(proposal_files)
-    # load datasets and metadata
-    dataset_name_to_dicts = {}
-    for dataset_name in dataset_names:
-        dataset_name_to_dicts[dataset_name] = DatasetCatalog.get(dataset_name)
-        assert len(dataset_name_to_dicts), f"Dataset '{dataset_name}' is empty!"
-    # merge categories, requires category metadata to be loaded
-    # cat_id -> [(orig_cat_id, cat_name, dataset_name)]
-    merged_categories = _merge_categories(dataset_names)
-    _warn_if_merged_different_categories(merged_categories)
-    merged_category_names = [
-        merged_categories[cat_id][0].mapped_name for cat_id in sorted(merged_categories)
-    ]
-    # map to contiguous category IDs
-    _add_category_id_to_contiguous_id_maps_to_metadata(merged_categories)
-    # load annotations and dataset metadata
-    for dataset_name, proposal_file in zip(dataset_names, proposal_files):
-        dataset_dicts = dataset_name_to_dicts[dataset_name]
-        assert len(dataset_dicts), f"Dataset '{dataset_name}' is empty!"
-        if proposal_file is not None:
-            dataset_dicts = load_proposals_into_dataset(dataset_dicts, proposal_file)
-        dataset_dicts = _maybe_filter_and_map_categories(dataset_name, dataset_dicts)
-        print_instances_class_histogram(dataset_dicts, merged_category_names)
-        dataset_name_to_dicts[dataset_name] = dataset_dicts
-
-    if keep_instance_predicate is not None:
-        all_datasets_dicts_plain = [
-            d
-            for d in itertools.chain.from_iterable(dataset_name_to_dicts.values())
-            if keep_instance_predicate(d)
-        ]
-    else:
-        all_datasets_dicts_plain = list(
-            itertools.chain.from_iterable(dataset_name_to_dicts.values())
-        )
-    return all_datasets_dicts_plain
-
-
-def build_detection_train_loader(cfg: CfgNode, mapper=None):
-    """
-    A data loader is created in a way similar to that of Detectron2.
-    The main differences are:
-     - it allows to combine datasets with different but compatible object category sets
-
-    The data loader is created by the following steps:
-    1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
-    2. Start workers to work on the dicts. Each worker will:
-        * Map each metadata dict into another format to be consumed by the model.
-        * Batch them by simply putting dicts into a list.
-    The batched ``list[mapped_dict]`` is what this dataloader will return.
-
-    Args:
-        cfg (CfgNode): the config
-        mapper (callable): a callable which takes a sample (dict) from dataset and
-            returns the format to be consumed by the model.
-            By default it will be `DatasetMapper(cfg, True)`.
-
-    Returns:
-        an infinite iterator of training data
-    """
-
-    _add_category_whitelists_to_metadata(cfg)
-    _add_category_maps_to_metadata(cfg)
-    _maybe_add_class_to_mesh_name_map_to_metadata(cfg.DATASETS.TRAIN, cfg)
-    dataset_dicts = combine_detection_dataset_dicts(
-        cfg.DATASETS.TRAIN,
-        keep_instance_predicate=_get_train_keep_instance_predicate(cfg),
-        proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
-    )
-    if mapper is None:
-        mapper = DatasetMapper(cfg, True)
-    return d2_build_detection_train_loader(cfg, dataset=dataset_dicts, mapper=mapper)
-
-
-def build_detection_test_loader(cfg, dataset_name, mapper=None):
-    """
-    Similar to `build_detection_train_loader`.
-    But this function uses the given `dataset_name` argument (instead of the names in cfg),
-    and uses batch size 1.
-
-    Args:
-        cfg: a detectron2 CfgNode
-        dataset_name (str): a name of the dataset that's available in the DatasetCatalog
-        mapper (callable): a callable which takes a sample (dict) from dataset
-            and returns the format to be consumed by the model.
-            By default it will be `DatasetMapper(cfg, False)`.
-
-    Returns:
-        DataLoader: a torch DataLoader, that loads the given detection
-            dataset, with test-time transformation and batching.
-    """
-    _add_category_whitelists_to_metadata(cfg)
-    _add_category_maps_to_metadata(cfg)
-    _maybe_add_class_to_mesh_name_map_to_metadata([dataset_name], cfg)
-    dataset_dicts = combine_detection_dataset_dicts(
-        [dataset_name],
-        keep_instance_predicate=_get_test_keep_instance_predicate(cfg),
-        proposal_files=[
-            cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(dataset_name)]
-        ]
-        if cfg.MODEL.LOAD_PROPOSALS
-        else None,
-    )
-    sampler = None
-    if not cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE:
-        sampler = torch.utils.data.SequentialSampler(dataset_dicts)
-    if mapper is None:
-        mapper = DatasetMapper(cfg, False)
-    return d2_build_detection_test_loader(
-        dataset_dicts, mapper=mapper, num_workers=cfg.DATALOADER.NUM_WORKERS, sampler=sampler
-    )
-
-
-def build_frame_selector(cfg: CfgNode):
-    strategy = FrameSelectionStrategy(cfg.STRATEGY)
-    if strategy == FrameSelectionStrategy.RANDOM_K:
-        frame_selector = RandomKFramesSelector(cfg.NUM_IMAGES)
-    elif strategy == FrameSelectionStrategy.FIRST_K:
-        frame_selector = FirstKFramesSelector(cfg.NUM_IMAGES)
-    elif strategy == FrameSelectionStrategy.LAST_K:
-        frame_selector = LastKFramesSelector(cfg.NUM_IMAGES)
-    elif strategy == FrameSelectionStrategy.ALL:
-        frame_selector = None
-    # pyre-fixme[61]: `frame_selector` may not be initialized here.
-    return frame_selector
-
-
-def build_transform(cfg: CfgNode, data_type: str):
-    if cfg.TYPE == "resize":
-        if data_type == "image":
-            return ImageResizeTransform(cfg.MIN_SIZE, cfg.MAX_SIZE)
-    raise ValueError(f"Unknown transform {cfg.TYPE} for data type {data_type}")
-
-
-def build_combined_loader(cfg: CfgNode, loaders: Collection[Loader], ratios: Sequence[float]):
-    images_per_worker = _compute_num_images_per_worker(cfg)
-    return CombinedDataLoader(loaders, images_per_worker, ratios)
-
-
-def build_bootstrap_dataset(dataset_name: str, cfg: CfgNode) -> Sequence[torch.Tensor]:
-    """
-    Build dataset that provides data to bootstrap on
-
-    Args:
-        dataset_name (str): Name of the dataset, needs to have associated metadata
-            to load the data
-        cfg (CfgNode): bootstrapping config
-    Returns:
-        Sequence[Tensor] - dataset that provides image batches, Tensors of size
-            [N, C, H, W] of type float32
-    """
-    logger = logging.getLogger(__name__)
-    _add_category_info_to_bootstrapping_metadata(dataset_name, cfg)
-    meta = MetadataCatalog.get(dataset_name)
-    factory = BootstrapDatasetFactoryCatalog.get(meta.dataset_type)
-    dataset = None
-    if factory is not None:
-        dataset = factory(meta, cfg)
-    if dataset is None:
-        logger.warning(f"Failed to create dataset {dataset_name} of type {meta.dataset_type}")
-    return dataset
-
-
-def build_data_sampler(cfg: CfgNode, sampler_cfg: CfgNode, embedder: Optional[torch.nn.Module]):
-    if sampler_cfg.TYPE == "densepose_uniform":
-        data_sampler = PredictionToGroundTruthSampler()
-        # transform densepose pred -> gt
-        data_sampler.register_sampler(
-            "pred_densepose",
-            "gt_densepose",
-            DensePoseUniformSampler(count_per_class=sampler_cfg.COUNT_PER_CLASS),
-        )
-        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
-        return data_sampler
-    elif sampler_cfg.TYPE == "densepose_UV_confidence":
-        data_sampler = PredictionToGroundTruthSampler()
-        # transform densepose pred -> gt
-        data_sampler.register_sampler(
-            "pred_densepose",
-            "gt_densepose",
-            DensePoseConfidenceBasedSampler(
-                confidence_channel="sigma_2",
-                count_per_class=sampler_cfg.COUNT_PER_CLASS,
-                search_proportion=0.5,
-            ),
-        )
-        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
-        return data_sampler
-    elif sampler_cfg.TYPE == "densepose_fine_segm_confidence":
-        data_sampler = PredictionToGroundTruthSampler()
-        # transform densepose pred -> gt
-        data_sampler.register_sampler(
-            "pred_densepose",
-            "gt_densepose",
-            DensePoseConfidenceBasedSampler(
-                confidence_channel="fine_segm_confidence",
-                count_per_class=sampler_cfg.COUNT_PER_CLASS,
-                search_proportion=0.5,
-            ),
-        )
-        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
-        return data_sampler
-    elif sampler_cfg.TYPE == "densepose_coarse_segm_confidence":
-        data_sampler = PredictionToGroundTruthSampler()
-        # transform densepose pred -> gt
-        data_sampler.register_sampler(
-            "pred_densepose",
-            "gt_densepose",
-            DensePoseConfidenceBasedSampler(
-                confidence_channel="coarse_segm_confidence",
-                count_per_class=sampler_cfg.COUNT_PER_CLASS,
-                search_proportion=0.5,
-            ),
-        )
-        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
-        return data_sampler
-    elif sampler_cfg.TYPE == "densepose_cse_uniform":
-        assert embedder is not None
-        data_sampler = PredictionToGroundTruthSampler()
-        # transform densepose pred -> gt
-        data_sampler.register_sampler(
-            "pred_densepose",
-            "gt_densepose",
-            DensePoseCSEUniformSampler(
-                cfg=cfg,
-                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
-                embedder=embedder,
-                count_per_class=sampler_cfg.COUNT_PER_CLASS,
-            ),
-        )
-        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
-        return data_sampler
-    elif sampler_cfg.TYPE == "densepose_cse_coarse_segm_confidence":
-        assert embedder is not None
-        data_sampler = PredictionToGroundTruthSampler()
-        # transform densepose pred -> gt
-        data_sampler.register_sampler(
-            "pred_densepose",
-            "gt_densepose",
-            DensePoseCSEConfidenceBasedSampler(
-                cfg=cfg,
-                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
-                embedder=embedder,
-                confidence_channel="coarse_segm_confidence",
-                count_per_class=sampler_cfg.COUNT_PER_CLASS,
-                search_proportion=0.5,
-            ),
-        )
-        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
-        return data_sampler
-
-    raise ValueError(f"Unknown data sampler type {sampler_cfg.TYPE}")
-
-
-def build_data_filter(cfg: CfgNode):
-    if cfg.TYPE == "detection_score":
-        min_score = cfg.MIN_VALUE
-        return ScoreBasedFilter(min_score=min_score)
-    raise ValueError(f"Unknown data filter type {cfg.TYPE}")
-
-
-def build_inference_based_loader(
-    cfg: CfgNode,
-    dataset_cfg: CfgNode,
-    model: torch.nn.Module,
-    embedder: Optional[torch.nn.Module] = None,
-) -> InferenceBasedLoader:
-    """
-    Constructs data loader based on inference results of a model.
-    """
-    dataset = build_bootstrap_dataset(dataset_cfg.DATASET, dataset_cfg.IMAGE_LOADER)
-    meta = MetadataCatalog.get(dataset_cfg.DATASET)
-    training_sampler = TrainingSampler(len(dataset))
-    data_loader = torch.utils.data.DataLoader(
-        dataset,  # pyre-ignore[6]
-        batch_size=dataset_cfg.IMAGE_LOADER.BATCH_SIZE,
-        sampler=training_sampler,
-        num_workers=dataset_cfg.IMAGE_LOADER.NUM_WORKERS,
-        collate_fn=trivial_batch_collator,
-        worker_init_fn=worker_init_reset_seed,
-    )
-    return InferenceBasedLoader(
-        model,
-        data_loader=data_loader,
-        data_sampler=build_data_sampler(cfg, dataset_cfg.DATA_SAMPLER, embedder),
-        data_filter=build_data_filter(dataset_cfg.FILTER),
-        shuffle=True,
-        batch_size=dataset_cfg.INFERENCE.OUTPUT_BATCH_SIZE,
-        inference_batch_size=dataset_cfg.INFERENCE.INPUT_BATCH_SIZE,
-        category_to_class_mapping=meta.category_to_class_mapping,
-    )
-
-
-def has_inference_based_loaders(cfg: CfgNode) -> bool:
-    """
-    Returns True, if at least one inferense-based loader must
-    be instantiated for training
-    """
-    return len(cfg.BOOTSTRAP_DATASETS) > 0
-
-
-def build_inference_based_loaders(
-    cfg: CfgNode, model: torch.nn.Module
-) -> Tuple[List[InferenceBasedLoader], List[float]]:
-    loaders = []
-    ratios = []
-    embedder = build_densepose_embedder(cfg).to(device=model.device)  # pyre-ignore[16]
-    for dataset_spec in cfg.BOOTSTRAP_DATASETS:
-        dataset_cfg = get_bootstrap_dataset_config().clone()
-        dataset_cfg.merge_from_other_cfg(CfgNode(dataset_spec))
-        loader = build_inference_based_loader(cfg, dataset_cfg, model, embedder)
-        loaders.append(loader)
-        ratios.append(dataset_cfg.RATIO)
-    return loaders, ratios
-
-
-def build_video_list_dataset(meta: Metadata, cfg: CfgNode):
-    video_list_fpath = meta.video_list_fpath
-    video_base_path = meta.video_base_path
-    category = meta.category
-    if cfg.TYPE == "video_keyframe":
-        frame_selector = build_frame_selector(cfg.SELECT)
-        transform = build_transform(cfg.TRANSFORM, data_type="image")
-        video_list = video_list_from_file(video_list_fpath, video_base_path)
-        keyframe_helper_fpath = getattr(cfg, "KEYFRAME_HELPER", None)
-        return VideoKeyframeDataset(
-            video_list, category, frame_selector, transform, keyframe_helper_fpath
-        )
-
-
-class _BootstrapDatasetFactoryCatalog(UserDict):
-    """
-    A global dictionary that stores information about bootstrapped datasets creation functions
-    from metadata and config, for diverse DatasetType
-    """
-
-    def register(self, dataset_type: DatasetType, factory: Callable[[Metadata, CfgNode], Dataset]):
-        """
-        Args:
-            dataset_type (DatasetType): a DatasetType e.g. DatasetType.VIDEO_LIST
-            factory (Callable[Metadata, CfgNode]): a callable which takes Metadata and cfg
-            arguments and returns a dataset object.
-        """
-        assert dataset_type not in self, "Dataset '{}' is already registered!".format(dataset_type)
-        self[dataset_type] = factory
-
-
-BootstrapDatasetFactoryCatalog = _BootstrapDatasetFactoryCatalog()
-BootstrapDatasetFactoryCatalog.register(DatasetType.VIDEO_LIST, build_video_list_dataset)
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import itertools
+import logging
+import numpy as np
+from collections import UserDict, defaultdict
+from dataclasses import dataclass
+from typing import Any, Callable, Collection, Dict, Iterable, List, Optional, Sequence, Tuple
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.config import CfgNode
+from detectron2.data.build import build_detection_test_loader as d2_build_detection_test_loader
+from detectron2.data.build import build_detection_train_loader as d2_build_detection_train_loader
+from detectron2.data.build import (
+    load_proposals_into_dataset,
+    print_instances_class_histogram,
+    trivial_batch_collator,
+    worker_init_reset_seed,
+)
+from detectron2.data.catalog import DatasetCatalog, Metadata, MetadataCatalog
+from detectron2.data.samplers import TrainingSampler
+from detectron2.utils.comm import get_world_size
+
+from densepose.config import get_bootstrap_dataset_config
+from densepose.modeling import build_densepose_embedder
+
+from .combined_loader import CombinedDataLoader, Loader
+from .dataset_mapper import DatasetMapper
+from .datasets.coco import DENSEPOSE_CSE_KEYS_WITHOUT_MASK, DENSEPOSE_IUV_KEYS_WITHOUT_MASK
+from .datasets.dataset_type import DatasetType
+from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
+from .samplers import (
+    DensePoseConfidenceBasedSampler,
+    DensePoseCSEConfidenceBasedSampler,
+    DensePoseCSEUniformSampler,
+    DensePoseUniformSampler,
+    MaskFromDensePoseSampler,
+    PredictionToGroundTruthSampler,
+)
+from .transform import ImageResizeTransform
+from .utils import get_category_to_class_mapping, get_class_to_mesh_name_mapping
+from .video import (
+    FirstKFramesSelector,
+    FrameSelectionStrategy,
+    LastKFramesSelector,
+    RandomKFramesSelector,
+    VideoKeyframeDataset,
+    video_list_from_file,
+)
+
+__all__ = ["build_detection_train_loader", "build_detection_test_loader"]
+
+
+Instance = Dict[str, Any]
+InstancePredicate = Callable[[Instance], bool]
+
+
+def _compute_num_images_per_worker(cfg: CfgNode) -> int:
+    num_workers = get_world_size()
+    images_per_batch = cfg.SOLVER.IMS_PER_BATCH
+    assert (
+        images_per_batch % num_workers == 0
+    ), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of workers ({}).".format(
+        images_per_batch, num_workers
+    )
+    assert (
+        images_per_batch >= num_workers
+    ), "SOLVER.IMS_PER_BATCH ({}) must be larger than the number of workers ({}).".format(
+        images_per_batch, num_workers
+    )
+    images_per_worker = images_per_batch // num_workers
+    return images_per_worker
+
+
+def _map_category_id_to_contiguous_id(dataset_name: str, dataset_dicts: Iterable[Instance]) -> None:
+    meta = MetadataCatalog.get(dataset_name)
+    for dataset_dict in dataset_dicts:
+        for ann in dataset_dict["annotations"]:
+            ann["category_id"] = meta.thing_dataset_id_to_contiguous_id[ann["category_id"]]
+
+
+@dataclass
+class _DatasetCategory:
+    """
+    Class representing category data in a dataset:
+     - id: category ID, as specified in the dataset annotations file
+     - name: category name, as specified in the dataset annotations file
+     - mapped_id: category ID after applying category maps (DATASETS.CATEGORY_MAPS config option)
+     - mapped_name: category name after applying category maps
+     - dataset_name: dataset in which the category is defined
+
+    For example, when training models in a class-agnostic manner, one could take LVIS 1.0
+    dataset and map the animal categories to the same category as human data from COCO:
+     id = 225
+     name = "cat"
+     mapped_id = 1
+     mapped_name = "person"
+     dataset_name = "lvis_v1_animals_dp_train"
+    """
+
+    id: int
+    name: str
+    mapped_id: int
+    mapped_name: str
+    dataset_name: str
+
+
+_MergedCategoriesT = Dict[int, List[_DatasetCategory]]
+
+
+def _add_category_id_to_contiguous_id_maps_to_metadata(
+    merged_categories: _MergedCategoriesT,
+) -> None:
+    merged_categories_per_dataset = {}
+    for contiguous_cat_id, cat_id in enumerate(sorted(merged_categories.keys())):
+        for cat in merged_categories[cat_id]:
+            if cat.dataset_name not in merged_categories_per_dataset:
+                merged_categories_per_dataset[cat.dataset_name] = defaultdict(list)
+            merged_categories_per_dataset[cat.dataset_name][cat_id].append(
+                (
+                    contiguous_cat_id,
+                    cat,
+                )
+            )
+
+    logger = logging.getLogger(__name__)
+    for dataset_name, merged_categories in merged_categories_per_dataset.items():
+        meta = MetadataCatalog.get(dataset_name)
+        if not hasattr(meta, "thing_classes"):
+            meta.thing_classes = []
+            meta.thing_dataset_id_to_contiguous_id = {}
+            meta.thing_dataset_id_to_merged_id = {}
+        else:
+            meta.thing_classes.clear()
+            meta.thing_dataset_id_to_contiguous_id.clear()
+            meta.thing_dataset_id_to_merged_id.clear()
+        logger.info(f"Dataset {dataset_name}: category ID to contiguous ID mapping:")
+        for _cat_id, categories in sorted(merged_categories.items()):
+            added_to_thing_classes = False
+            for contiguous_cat_id, cat in categories:
+                if not added_to_thing_classes:
+                    meta.thing_classes.append(cat.mapped_name)
+                    added_to_thing_classes = True
+                meta.thing_dataset_id_to_contiguous_id[cat.id] = contiguous_cat_id
+                meta.thing_dataset_id_to_merged_id[cat.id] = cat.mapped_id
+                logger.info(f"{cat.id} ({cat.name}) -> {contiguous_cat_id}")
+
+
+def _maybe_create_general_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    def has_annotations(instance: Instance) -> bool:
+        return "annotations" in instance
+
+    def has_only_crowd_anotations(instance: Instance) -> bool:
+        for ann in instance["annotations"]:
+            if ann.get("is_crowd", 0) == 0:
+                return False
+        return True
+
+    def general_keep_instance_predicate(instance: Instance) -> bool:
+        return has_annotations(instance) and not has_only_crowd_anotations(instance)
+
+    if not cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS:
+        return None
+    return general_keep_instance_predicate
+
+
+def _maybe_create_keypoints_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+
+    min_num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
+
+    def has_sufficient_num_keypoints(instance: Instance) -> bool:
+        num_kpts = sum(
+            (np.array(ann["keypoints"][2::3]) > 0).sum()
+            for ann in instance["annotations"]
+            if "keypoints" in ann
+        )
+        return num_kpts >= min_num_keypoints
+
+    if cfg.MODEL.KEYPOINT_ON and (min_num_keypoints > 0):
+        return has_sufficient_num_keypoints
+    return None
+
+
+def _maybe_create_mask_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    if not cfg.MODEL.MASK_ON:
+        return None
+
+    def has_mask_annotations(instance: Instance) -> bool:
+        return any("segmentation" in ann for ann in instance["annotations"])
+
+    return has_mask_annotations
+
+
+def _maybe_create_densepose_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    if not cfg.MODEL.DENSEPOSE_ON:
+        return None
+
+    use_masks = cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS
+
+    def has_densepose_annotations(instance: Instance) -> bool:
+        for ann in instance["annotations"]:
+            if all(key in ann for key in DENSEPOSE_IUV_KEYS_WITHOUT_MASK) or all(
+                key in ann for key in DENSEPOSE_CSE_KEYS_WITHOUT_MASK
+            ):
+                return True
+            if use_masks and "segmentation" in ann:
+                return True
+        return False
+
+    return has_densepose_annotations
+
+
+def _maybe_create_specific_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    specific_predicate_creators = [
+        _maybe_create_keypoints_keep_instance_predicate,
+        _maybe_create_mask_keep_instance_predicate,
+        _maybe_create_densepose_keep_instance_predicate,
+    ]
+    predicates = [creator(cfg) for creator in specific_predicate_creators]
+    predicates = [p for p in predicates if p is not None]
+    if not predicates:
+        return None
+
+    def combined_predicate(instance: Instance) -> bool:
+        return any(p(instance) for p in predicates)
+
+    return combined_predicate
+
+
+def _get_train_keep_instance_predicate(cfg: CfgNode):
+    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
+    combined_specific_keep_predicate = _maybe_create_specific_keep_instance_predicate(cfg)
+
+    def combined_general_specific_keep_predicate(instance: Instance) -> bool:
+        return general_keep_predicate(instance) and combined_specific_keep_predicate(instance)
+
+    if (general_keep_predicate is None) and (combined_specific_keep_predicate is None):
+        return None
+    if general_keep_predicate is None:
+        return combined_specific_keep_predicate
+    if combined_specific_keep_predicate is None:
+        return general_keep_predicate
+    return combined_general_specific_keep_predicate
+
+
+def _get_test_keep_instance_predicate(cfg: CfgNode):
+    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
+    return general_keep_predicate
+
+
+def _maybe_filter_and_map_categories(
+    dataset_name: str, dataset_dicts: List[Instance]
+) -> List[Instance]:
+    meta = MetadataCatalog.get(dataset_name)
+    category_id_map = meta.thing_dataset_id_to_contiguous_id
+    filtered_dataset_dicts = []
+    for dataset_dict in dataset_dicts:
+        anns = []
+        for ann in dataset_dict["annotations"]:
+            cat_id = ann["category_id"]
+            if cat_id not in category_id_map:
+                continue
+            ann["category_id"] = category_id_map[cat_id]
+            anns.append(ann)
+        dataset_dict["annotations"] = anns
+        filtered_dataset_dicts.append(dataset_dict)
+    return filtered_dataset_dicts
+
+
+def _add_category_whitelists_to_metadata(cfg: CfgNode) -> None:
+    for dataset_name, whitelisted_cat_ids in cfg.DATASETS.WHITELISTED_CATEGORIES.items():
+        meta = MetadataCatalog.get(dataset_name)
+        meta.whitelisted_categories = whitelisted_cat_ids
+        logger = logging.getLogger(__name__)
+        logger.info(
+            "Whitelisted categories for dataset {}: {}".format(
+                dataset_name, meta.whitelisted_categories
+            )
+        )
+
+
+def _add_category_maps_to_metadata(cfg: CfgNode) -> None:
+    for dataset_name, category_map in cfg.DATASETS.CATEGORY_MAPS.items():
+        category_map = {
+            int(cat_id_src): int(cat_id_dst) for cat_id_src, cat_id_dst in category_map.items()
+        }
+        meta = MetadataCatalog.get(dataset_name)
+        meta.category_map = category_map
+        logger = logging.getLogger(__name__)
+        logger.info("Category maps for dataset {}: {}".format(dataset_name, meta.category_map))
+
+
+def _add_category_info_to_bootstrapping_metadata(dataset_name: str, dataset_cfg: CfgNode) -> None:
+    meta = MetadataCatalog.get(dataset_name)
+    meta.category_to_class_mapping = get_category_to_class_mapping(dataset_cfg)
+    meta.categories = dataset_cfg.CATEGORIES
+    meta.max_count_per_category = dataset_cfg.MAX_COUNT_PER_CATEGORY
+    logger = logging.getLogger(__name__)
+    logger.info(
+        "Category to class mapping for dataset {}: {}".format(
+            dataset_name, meta.category_to_class_mapping
+        )
+    )
+
+
+def _maybe_add_class_to_mesh_name_map_to_metadata(dataset_names: List[str], cfg: CfgNode) -> None:
+    for dataset_name in dataset_names:
+        meta = MetadataCatalog.get(dataset_name)
+        if not hasattr(meta, "class_to_mesh_name"):
+            meta.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
+
+
+def _merge_categories(dataset_names: Collection[str]) -> _MergedCategoriesT:
+    merged_categories = defaultdict(list)
+    category_names = {}
+    for dataset_name in dataset_names:
+        meta = MetadataCatalog.get(dataset_name)
+        whitelisted_categories = meta.get("whitelisted_categories")
+        category_map = meta.get("category_map", {})
+        cat_ids = (
+            whitelisted_categories if whitelisted_categories is not None else meta.categories.keys()
+        )
+        for cat_id in cat_ids:
+            cat_name = meta.categories[cat_id]
+            cat_id_mapped = category_map.get(cat_id, cat_id)
+            if cat_id_mapped == cat_id or cat_id_mapped in cat_ids:
+                category_names[cat_id] = cat_name
+            else:
+                category_names[cat_id] = str(cat_id_mapped)
+            # assign temporary mapped category name, this name can be changed
+            # during the second pass, since mapped ID can correspond to a category
+            # from a different dataset
+            cat_name_mapped = meta.categories[cat_id_mapped]
+            merged_categories[cat_id_mapped].append(
+                _DatasetCategory(
+                    id=cat_id,
+                    name=cat_name,
+                    mapped_id=cat_id_mapped,
+                    mapped_name=cat_name_mapped,
+                    dataset_name=dataset_name,
+                )
+            )
+    # second pass to assign proper mapped category names
+    for cat_id, categories in merged_categories.items():
+        for cat in categories:
+            if cat_id in category_names and cat.mapped_name != category_names[cat_id]:
+                cat.mapped_name = category_names[cat_id]
+
+    return merged_categories
+
+
+def _warn_if_merged_different_categories(merged_categories: _MergedCategoriesT) -> None:
+    logger = logging.getLogger(__name__)
+    for cat_id in merged_categories:
+        merged_categories_i = merged_categories[cat_id]
+        first_cat_name = merged_categories_i[0].name
+        if len(merged_categories_i) > 1 and not all(
+            cat.name == first_cat_name for cat in merged_categories_i[1:]
+        ):
+            cat_summary_str = ", ".join(
+                [f"{cat.id} ({cat.name}) from {cat.dataset_name}" for cat in merged_categories_i]
+            )
+            logger.warning(
+                f"Merged category {cat_id} corresponds to the following categories: "
+                f"{cat_summary_str}"
+            )
+
+
+def combine_detection_dataset_dicts(
+    dataset_names: Collection[str],
+    keep_instance_predicate: Optional[InstancePredicate] = None,
+    proposal_files: Optional[Collection[str]] = None,
+) -> List[Instance]:
+    """
+    Load and prepare dataset dicts for training / testing
+
+    Args:
+        dataset_names (Collection[str]): a list of dataset names
+        keep_instance_predicate (Callable: Dict[str, Any] -> bool): predicate
+            applied to instance dicts which defines whether to keep the instance
+        proposal_files (Collection[str]): if given, a list of object proposal files
+            that match each dataset in `dataset_names`.
+    """
+    assert len(dataset_names)
+    if proposal_files is None:
+        proposal_files = [None] * len(dataset_names)
+    assert len(dataset_names) == len(proposal_files)
+    # load datasets and metadata
+    dataset_name_to_dicts = {}
+    for dataset_name in dataset_names:
+        dataset_name_to_dicts[dataset_name] = DatasetCatalog.get(dataset_name)
+        assert len(dataset_name_to_dicts), f"Dataset '{dataset_name}' is empty!"
+    # merge categories, requires category metadata to be loaded
+    # cat_id -> [(orig_cat_id, cat_name, dataset_name)]
+    merged_categories = _merge_categories(dataset_names)
+    _warn_if_merged_different_categories(merged_categories)
+    merged_category_names = [
+        merged_categories[cat_id][0].mapped_name for cat_id in sorted(merged_categories)
+    ]
+    # map to contiguous category IDs
+    _add_category_id_to_contiguous_id_maps_to_metadata(merged_categories)
+    # load annotations and dataset metadata
+    for dataset_name, proposal_file in zip(dataset_names, proposal_files):
+        dataset_dicts = dataset_name_to_dicts[dataset_name]
+        assert len(dataset_dicts), f"Dataset '{dataset_name}' is empty!"
+        if proposal_file is not None:
+            dataset_dicts = load_proposals_into_dataset(dataset_dicts, proposal_file)
+        dataset_dicts = _maybe_filter_and_map_categories(dataset_name, dataset_dicts)
+        print_instances_class_histogram(dataset_dicts, merged_category_names)
+        dataset_name_to_dicts[dataset_name] = dataset_dicts
+
+    if keep_instance_predicate is not None:
+        all_datasets_dicts_plain = [
+            d
+            for d in itertools.chain.from_iterable(dataset_name_to_dicts.values())
+            if keep_instance_predicate(d)
+        ]
+    else:
+        all_datasets_dicts_plain = list(
+            itertools.chain.from_iterable(dataset_name_to_dicts.values())
+        )
+    return all_datasets_dicts_plain
+
+
+def build_detection_train_loader(cfg: CfgNode, mapper=None):
+    """
+    A data loader is created in a way similar to that of Detectron2.
+    The main differences are:
+     - it allows to combine datasets with different but compatible object category sets
+
+    The data loader is created by the following steps:
+    1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
+    2. Start workers to work on the dicts. Each worker will:
+        * Map each metadata dict into another format to be consumed by the model.
+        * Batch them by simply putting dicts into a list.
+    The batched ``list[mapped_dict]`` is what this dataloader will return.
+
+    Args:
+        cfg (CfgNode): the config
+        mapper (callable): a callable which takes a sample (dict) from dataset and
+            returns the format to be consumed by the model.
+            By default it will be `DatasetMapper(cfg, True)`.
+
+    Returns:
+        an infinite iterator of training data
+    """
+
+    _add_category_whitelists_to_metadata(cfg)
+    _add_category_maps_to_metadata(cfg)
+    _maybe_add_class_to_mesh_name_map_to_metadata(cfg.DATASETS.TRAIN, cfg)
+    dataset_dicts = combine_detection_dataset_dicts(
+        cfg.DATASETS.TRAIN,
+        keep_instance_predicate=_get_train_keep_instance_predicate(cfg),
+        proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
+    )
+    if mapper is None:
+        mapper = DatasetMapper(cfg, True)
+    return d2_build_detection_train_loader(cfg, dataset=dataset_dicts, mapper=mapper)
+
+
+def build_detection_test_loader(cfg, dataset_name, mapper=None):
+    """
+    Similar to `build_detection_train_loader`.
+    But this function uses the given `dataset_name` argument (instead of the names in cfg),
+    and uses batch size 1.
+
+    Args:
+        cfg: a detectron2 CfgNode
+        dataset_name (str): a name of the dataset that's available in the DatasetCatalog
+        mapper (callable): a callable which takes a sample (dict) from dataset
+            and returns the format to be consumed by the model.
+            By default it will be `DatasetMapper(cfg, False)`.
+
+    Returns:
+        DataLoader: a torch DataLoader, that loads the given detection
+            dataset, with test-time transformation and batching.
+    """
+    _add_category_whitelists_to_metadata(cfg)
+    _add_category_maps_to_metadata(cfg)
+    _maybe_add_class_to_mesh_name_map_to_metadata([dataset_name], cfg)
+    dataset_dicts = combine_detection_dataset_dicts(
+        [dataset_name],
+        keep_instance_predicate=_get_test_keep_instance_predicate(cfg),
+        proposal_files=[
+            cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(dataset_name)]
+        ]
+        if cfg.MODEL.LOAD_PROPOSALS
+        else None,
+    )
+    sampler = None
+    if not cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE:
+        sampler = torch.utils.data.SequentialSampler(dataset_dicts)
+    if mapper is None:
+        mapper = DatasetMapper(cfg, False)
+    return d2_build_detection_test_loader(
+        dataset_dicts, mapper=mapper, num_workers=cfg.DATALOADER.NUM_WORKERS, sampler=sampler
+    )
+
+
+def build_frame_selector(cfg: CfgNode):
+    strategy = FrameSelectionStrategy(cfg.STRATEGY)
+    if strategy == FrameSelectionStrategy.RANDOM_K:
+        frame_selector = RandomKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.FIRST_K:
+        frame_selector = FirstKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.LAST_K:
+        frame_selector = LastKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.ALL:
+        frame_selector = None
+    # pyre-fixme[61]: `frame_selector` may not be initialized here.
+    return frame_selector
+
+
+def build_transform(cfg: CfgNode, data_type: str):
+    if cfg.TYPE == "resize":
+        if data_type == "image":
+            return ImageResizeTransform(cfg.MIN_SIZE, cfg.MAX_SIZE)
+    raise ValueError(f"Unknown transform {cfg.TYPE} for data type {data_type}")
+
+
+def build_combined_loader(cfg: CfgNode, loaders: Collection[Loader], ratios: Sequence[float]):
+    images_per_worker = _compute_num_images_per_worker(cfg)
+    return CombinedDataLoader(loaders, images_per_worker, ratios)
+
+
+def build_bootstrap_dataset(dataset_name: str, cfg: CfgNode) -> Sequence[torch.Tensor]:
+    """
+    Build dataset that provides data to bootstrap on
+
+    Args:
+        dataset_name (str): Name of the dataset, needs to have associated metadata
+            to load the data
+        cfg (CfgNode): bootstrapping config
+    Returns:
+        Sequence[Tensor] - dataset that provides image batches, Tensors of size
+            [N, C, H, W] of type float32
+    """
+    logger = logging.getLogger(__name__)
+    _add_category_info_to_bootstrapping_metadata(dataset_name, cfg)
+    meta = MetadataCatalog.get(dataset_name)
+    factory = BootstrapDatasetFactoryCatalog.get(meta.dataset_type)
+    dataset = None
+    if factory is not None:
+        dataset = factory(meta, cfg)
+    if dataset is None:
+        logger.warning(f"Failed to create dataset {dataset_name} of type {meta.dataset_type}")
+    return dataset
+
+
+def build_data_sampler(cfg: CfgNode, sampler_cfg: CfgNode, embedder: Optional[torch.nn.Module]):
+    if sampler_cfg.TYPE == "densepose_uniform":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseUniformSampler(count_per_class=sampler_cfg.COUNT_PER_CLASS),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_UV_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="sigma_2",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_fine_segm_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="fine_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_coarse_segm_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="coarse_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_cse_uniform":
+        assert embedder is not None
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseCSEUniformSampler(
+                cfg=cfg,
+                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
+                embedder=embedder,
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_cse_coarse_segm_confidence":
+        assert embedder is not None
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseCSEConfidenceBasedSampler(
+                cfg=cfg,
+                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
+                embedder=embedder,
+                confidence_channel="coarse_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+
+    raise ValueError(f"Unknown data sampler type {sampler_cfg.TYPE}")
+
+
+def build_data_filter(cfg: CfgNode):
+    if cfg.TYPE == "detection_score":
+        min_score = cfg.MIN_VALUE
+        return ScoreBasedFilter(min_score=min_score)
+    raise ValueError(f"Unknown data filter type {cfg.TYPE}")
+
+
+def build_inference_based_loader(
+    cfg: CfgNode,
+    dataset_cfg: CfgNode,
+    model: torch.nn.Module,
+    embedder: Optional[torch.nn.Module] = None,
+) -> InferenceBasedLoader:
+    """
+    Constructs data loader based on inference results of a model.
+    """
+    dataset = build_bootstrap_dataset(dataset_cfg.DATASET, dataset_cfg.IMAGE_LOADER)
+    meta = MetadataCatalog.get(dataset_cfg.DATASET)
+    training_sampler = TrainingSampler(len(dataset))
+    data_loader = torch.utils.data.DataLoader(
+        dataset,  # pyre-ignore[6]
+        batch_size=dataset_cfg.IMAGE_LOADER.BATCH_SIZE,
+        sampler=training_sampler,
+        num_workers=dataset_cfg.IMAGE_LOADER.NUM_WORKERS,
+        collate_fn=trivial_batch_collator,
+        worker_init_fn=worker_init_reset_seed,
+    )
+    return InferenceBasedLoader(
+        model,
+        data_loader=data_loader,
+        data_sampler=build_data_sampler(cfg, dataset_cfg.DATA_SAMPLER, embedder),
+        data_filter=build_data_filter(dataset_cfg.FILTER),
+        shuffle=True,
+        batch_size=dataset_cfg.INFERENCE.OUTPUT_BATCH_SIZE,
+        inference_batch_size=dataset_cfg.INFERENCE.INPUT_BATCH_SIZE,
+        category_to_class_mapping=meta.category_to_class_mapping,
+    )
+
+
+def has_inference_based_loaders(cfg: CfgNode) -> bool:
+    """
+    Returns True, if at least one inferense-based loader must
+    be instantiated for training
+    """
+    return len(cfg.BOOTSTRAP_DATASETS) > 0
+
+
+def build_inference_based_loaders(
+    cfg: CfgNode, model: torch.nn.Module
+) -> Tuple[List[InferenceBasedLoader], List[float]]:
+    loaders = []
+    ratios = []
+    embedder = build_densepose_embedder(cfg).to(device=model.device)  # pyre-ignore[16]
+    for dataset_spec in cfg.BOOTSTRAP_DATASETS:
+        dataset_cfg = get_bootstrap_dataset_config().clone()
+        dataset_cfg.merge_from_other_cfg(CfgNode(dataset_spec))
+        loader = build_inference_based_loader(cfg, dataset_cfg, model, embedder)
+        loaders.append(loader)
+        ratios.append(dataset_cfg.RATIO)
+    return loaders, ratios
+
+
+def build_video_list_dataset(meta: Metadata, cfg: CfgNode):
+    video_list_fpath = meta.video_list_fpath
+    video_base_path = meta.video_base_path
+    category = meta.category
+    if cfg.TYPE == "video_keyframe":
+        frame_selector = build_frame_selector(cfg.SELECT)
+        transform = build_transform(cfg.TRANSFORM, data_type="image")
+        video_list = video_list_from_file(video_list_fpath, video_base_path)
+        keyframe_helper_fpath = getattr(cfg, "KEYFRAME_HELPER", None)
+        return VideoKeyframeDataset(
+            video_list, category, frame_selector, transform, keyframe_helper_fpath
+        )
+
+
+class _BootstrapDatasetFactoryCatalog(UserDict):
+    """
+    A global dictionary that stores information about bootstrapped datasets creation functions
+    from metadata and config, for diverse DatasetType
+    """
+
+    def register(self, dataset_type: DatasetType, factory: Callable[[Metadata, CfgNode], Dataset]):
+        """
+        Args:
+            dataset_type (DatasetType): a DatasetType e.g. DatasetType.VIDEO_LIST
+            factory (Callable[Metadata, CfgNode]): a callable which takes Metadata and cfg
+            arguments and returns a dataset object.
+        """
+        assert dataset_type not in self, "Dataset '{}' is already registered!".format(dataset_type)
+        self[dataset_type] = factory
+
+
+BootstrapDatasetFactoryCatalog = _BootstrapDatasetFactoryCatalog()
+BootstrapDatasetFactoryCatalog.register(DatasetType.VIDEO_LIST, build_video_list_dataset)

densepose/data/combined_loader.py

@@ -1,44 +1,44 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import random
-from collections import deque
-from typing import Any, Collection, Deque, Iterable, Iterator, List, Sequence
-
-Loader = Iterable[Any]
-
-
-def _pooled_next(iterator: Iterator[Any], pool: Deque[Any]):
-    if not pool:
-        pool.extend(next(iterator))
-    return pool.popleft()
-
-
-class CombinedDataLoader:
-    """
-    Combines data loaders using the provided sampling ratios
-    """
-
-    BATCH_COUNT = 100
-
-    def __init__(self, loaders: Collection[Loader], batch_size: int, ratios: Sequence[float]):
-        self.loaders = loaders
-        self.batch_size = batch_size
-        self.ratios = ratios
-
-    def __iter__(self) -> Iterator[List[Any]]:
-        iters = [iter(loader) for loader in self.loaders]
-        indices = []
-        pool = [deque()] * len(iters)
-        # infinite iterator, as in D2
-        while True:
-            if not indices:
-                # just a buffer of indices, its size doesn't matter
-                # as long as it's a multiple of batch_size
-                k = self.batch_size * self.BATCH_COUNT
-                indices = random.choices(range(len(self.loaders)), self.ratios, k=k)
-            try:
-                batch = [_pooled_next(iters[i], pool[i]) for i in indices[: self.batch_size]]
-            except StopIteration:
-                break
-            indices = indices[self.batch_size :]
-            yield batch
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import random
+from collections import deque
+from typing import Any, Collection, Deque, Iterable, Iterator, List, Sequence
+
+Loader = Iterable[Any]
+
+
+def _pooled_next(iterator: Iterator[Any], pool: Deque[Any]):
+    if not pool:
+        pool.extend(next(iterator))
+    return pool.popleft()
+
+
+class CombinedDataLoader:
+    """
+    Combines data loaders using the provided sampling ratios
+    """
+
+    BATCH_COUNT = 100
+
+    def __init__(self, loaders: Collection[Loader], batch_size: int, ratios: Sequence[float]):
+        self.loaders = loaders
+        self.batch_size = batch_size
+        self.ratios = ratios
+
+    def __iter__(self) -> Iterator[List[Any]]:
+        iters = [iter(loader) for loader in self.loaders]
+        indices = []
+        pool = [deque()] * len(iters)
+        # infinite iterator, as in D2
+        while True:
+            if not indices:
+                # just a buffer of indices, its size doesn't matter
+                # as long as it's a multiple of batch_size
+                k = self.batch_size * self.BATCH_COUNT
+                indices = random.choices(range(len(self.loaders)), self.ratios, k=k)
+            try:
+                batch = [_pooled_next(iters[i], pool[i]) for i in indices[: self.batch_size]]
+            except StopIteration:
+                break
+            indices = indices[self.batch_size :]
+            yield batch

densepose/data/dataset_mapper.py

@@ -1,168 +1,168 @@
-# -*- coding: utf-8 -*-
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import copy
-import logging
-from typing import Any, Dict, List, Tuple
-import torch
-
-from detectron2.data import MetadataCatalog
-from detectron2.data import detection_utils as utils
-from detectron2.data import transforms as T
-from detectron2.layers import ROIAlign
-from detectron2.structures import BoxMode
-from detectron2.utils.file_io import PathManager
-
-from densepose.structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
-
-
-def build_augmentation(cfg, is_train):
-    logger = logging.getLogger(__name__)
-    result = utils.build_augmentation(cfg, is_train)
-    if is_train:
-        random_rotation = T.RandomRotation(
-            cfg.INPUT.ROTATION_ANGLES, expand=False, sample_style="choice"
-        )
-        result.append(random_rotation)
-        logger.info("DensePose-specific augmentation used in training: " + str(random_rotation))
-    return result
-
-
-class DatasetMapper:
-    """
-    A customized version of `detectron2.data.DatasetMapper`
-    """
-
-    def __init__(self, cfg, is_train=True):
-        self.augmentation = build_augmentation(cfg, is_train)
-
-        # fmt: off
-        self.img_format     = cfg.INPUT.FORMAT
-        self.mask_on        = (
-            cfg.MODEL.MASK_ON or (
-                cfg.MODEL.DENSEPOSE_ON
-                and cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS)
-        )
-        self.keypoint_on    = cfg.MODEL.KEYPOINT_ON
-        self.densepose_on   = cfg.MODEL.DENSEPOSE_ON
-        assert not cfg.MODEL.LOAD_PROPOSALS, "not supported yet"
-        # fmt: on
-        if self.keypoint_on and is_train:
-            # Flip only makes sense in training
-            self.keypoint_hflip_indices = utils.create_keypoint_hflip_indices(cfg.DATASETS.TRAIN)
-        else:
-            self.keypoint_hflip_indices = None
-
-        if self.densepose_on:
-            densepose_transform_srcs = [
-                MetadataCatalog.get(ds).densepose_transform_src
-                for ds in cfg.DATASETS.TRAIN + cfg.DATASETS.TEST
-            ]
-            assert len(densepose_transform_srcs) > 0
-            # TODO: check that DensePose transformation data is the same for
-            # all the datasets. Otherwise one would have to pass DB ID with
-            # each entry to select proper transformation data. For now, since
-            # all DensePose annotated data uses the same data semantics, we
-            # omit this check.
-            densepose_transform_data_fpath = PathManager.get_local_path(densepose_transform_srcs[0])
-            self.densepose_transform_data = DensePoseTransformData.load(
-                densepose_transform_data_fpath
-            )
-
-        self.is_train = is_train
-
-    def __call__(self, dataset_dict):
-        """
-        Args:
-            dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
-
-        Returns:
-            dict: a format that builtin models in detectron2 accept
-        """
-        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
-        image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
-        utils.check_image_size(dataset_dict, image)
-
-        image, transforms = T.apply_transform_gens(self.augmentation, image)
-        image_shape = image.shape[:2]  # h, w
-        dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))
-
-        if not self.is_train:
-            dataset_dict.pop("annotations", None)
-            return dataset_dict
-
-        for anno in dataset_dict["annotations"]:
-            if not self.mask_on:
-                anno.pop("segmentation", None)
-            if not self.keypoint_on:
-                anno.pop("keypoints", None)
-
-        # USER: Implement additional transformations if you have other types of data
-        # USER: Don't call transpose_densepose if you don't need
-        annos = [
-            self._transform_densepose(
-                utils.transform_instance_annotations(
-                    obj, transforms, image_shape, keypoint_hflip_indices=self.keypoint_hflip_indices
-                ),
-                transforms,
-            )
-            for obj in dataset_dict.pop("annotations")
-            if obj.get("iscrowd", 0) == 0
-        ]
-
-        if self.mask_on:
-            self._add_densepose_masks_as_segmentation(annos, image_shape)
-
-        instances = utils.annotations_to_instances(annos, image_shape, mask_format="bitmask")
-        densepose_annotations = [obj.get("densepose") for obj in annos]
-        if densepose_annotations and not all(v is None for v in densepose_annotations):
-            instances.gt_densepose = DensePoseList(
-                densepose_annotations, instances.gt_boxes, image_shape
-            )
-
-        dataset_dict["instances"] = instances[instances.gt_boxes.nonempty()]
-        return dataset_dict
-
-    def _transform_densepose(self, annotation, transforms):
-        if not self.densepose_on:
-            return annotation
-
-        # Handle densepose annotations
-        is_valid, reason_not_valid = DensePoseDataRelative.validate_annotation(annotation)
-        if is_valid:
-            densepose_data = DensePoseDataRelative(annotation, cleanup=True)
-            densepose_data.apply_transform(transforms, self.densepose_transform_data)
-            annotation["densepose"] = densepose_data
-        else:
-            # logger = logging.getLogger(__name__)
-            # logger.debug("Could not load DensePose annotation: {}".format(reason_not_valid))
-            DensePoseDataRelative.cleanup_annotation(annotation)
-            # NOTE: annotations for certain instances may be unavailable.
-            # 'None' is accepted by the DensePostList data structure.
-            annotation["densepose"] = None
-        return annotation
-
-    def _add_densepose_masks_as_segmentation(
-        self, annotations: List[Dict[str, Any]], image_shape_hw: Tuple[int, int]
-    ):
-        for obj in annotations:
-            if ("densepose" not in obj) or ("segmentation" in obj):
-                continue
-            # DP segmentation: torch.Tensor [S, S] of float32, S=256
-            segm_dp = torch.zeros_like(obj["densepose"].segm)
-            segm_dp[obj["densepose"].segm > 0] = 1
-            segm_h, segm_w = segm_dp.shape
-            bbox_segm_dp = torch.tensor((0, 0, segm_h - 1, segm_w - 1), dtype=torch.float32)
-            # image bbox
-            x0, y0, x1, y1 = (
-                v.item() for v in BoxMode.convert(obj["bbox"], obj["bbox_mode"], BoxMode.XYXY_ABS)
-            )
-            segm_aligned = (
-                ROIAlign((y1 - y0, x1 - x0), 1.0, 0, aligned=True)
-                .forward(segm_dp.view(1, 1, *segm_dp.shape), bbox_segm_dp)
-                .squeeze()
-            )
-            image_mask = torch.zeros(*image_shape_hw, dtype=torch.float32)
-            image_mask[y0:y1, x0:x1] = segm_aligned
-            # segmentation for BitMask: np.array [H, W] of bool
-            obj["segmentation"] = image_mask >= 0.5
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import copy
+import logging
+from typing import Any, Dict, List, Tuple
+import torch
+
+from detectron2.data import MetadataCatalog
+from detectron2.data import detection_utils as utils
+from detectron2.data import transforms as T
+from detectron2.layers import ROIAlign
+from detectron2.structures import BoxMode
+from detectron2.utils.file_io import PathManager
+
+from densepose.structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
+
+
+def build_augmentation(cfg, is_train):
+    logger = logging.getLogger(__name__)
+    result = utils.build_augmentation(cfg, is_train)
+    if is_train:
+        random_rotation = T.RandomRotation(
+            cfg.INPUT.ROTATION_ANGLES, expand=False, sample_style="choice"
+        )
+        result.append(random_rotation)
+        logger.info("DensePose-specific augmentation used in training: " + str(random_rotation))
+    return result
+
+
+class DatasetMapper:
+    """
+    A customized version of `detectron2.data.DatasetMapper`
+    """
+
+    def __init__(self, cfg, is_train=True):
+        self.augmentation = build_augmentation(cfg, is_train)
+
+        # fmt: off
+        self.img_format     = cfg.INPUT.FORMAT
+        self.mask_on        = (
+            cfg.MODEL.MASK_ON or (
+                cfg.MODEL.DENSEPOSE_ON
+                and cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS)
+        )
+        self.keypoint_on    = cfg.MODEL.KEYPOINT_ON
+        self.densepose_on   = cfg.MODEL.DENSEPOSE_ON
+        assert not cfg.MODEL.LOAD_PROPOSALS, "not supported yet"
+        # fmt: on
+        if self.keypoint_on and is_train:
+            # Flip only makes sense in training
+            self.keypoint_hflip_indices = utils.create_keypoint_hflip_indices(cfg.DATASETS.TRAIN)
+        else:
+            self.keypoint_hflip_indices = None
+
+        if self.densepose_on:
+            densepose_transform_srcs = [
+                MetadataCatalog.get(ds).densepose_transform_src
+                for ds in cfg.DATASETS.TRAIN + cfg.DATASETS.TEST
+            ]
+            assert len(densepose_transform_srcs) > 0
+            # TODO: check that DensePose transformation data is the same for
+            # all the datasets. Otherwise one would have to pass DB ID with
+            # each entry to select proper transformation data. For now, since
+            # all DensePose annotated data uses the same data semantics, we
+            # omit this check.
+            densepose_transform_data_fpath = PathManager.get_local_path(densepose_transform_srcs[0])
+            self.densepose_transform_data = DensePoseTransformData.load(
+                densepose_transform_data_fpath
+            )
+
+        self.is_train = is_train
+
+    def __call__(self, dataset_dict):
+        """
+        Args:
+            dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
+
+        Returns:
+            dict: a format that builtin models in detectron2 accept
+        """
+        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
+        image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
+        utils.check_image_size(dataset_dict, image)
+
+        image, transforms = T.apply_transform_gens(self.augmentation, image)
+        image_shape = image.shape[:2]  # h, w
+        dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))
+
+        if not self.is_train:
+            dataset_dict.pop("annotations", None)
+            return dataset_dict
+
+        for anno in dataset_dict["annotations"]:
+            if not self.mask_on:
+                anno.pop("segmentation", None)
+            if not self.keypoint_on:
+                anno.pop("keypoints", None)
+
+        # USER: Implement additional transformations if you have other types of data
+        # USER: Don't call transpose_densepose if you don't need
+        annos = [
+            self._transform_densepose(
+                utils.transform_instance_annotations(
+                    obj, transforms, image_shape, keypoint_hflip_indices=self.keypoint_hflip_indices
+                ),
+                transforms,
+            )
+            for obj in dataset_dict.pop("annotations")
+            if obj.get("iscrowd", 0) == 0
+        ]
+
+        if self.mask_on:
+            self._add_densepose_masks_as_segmentation(annos, image_shape)
+
+        instances = utils.annotations_to_instances(annos, image_shape, mask_format="bitmask")
+        densepose_annotations = [obj.get("densepose") for obj in annos]
+        if densepose_annotations and not all(v is None for v in densepose_annotations):
+            instances.gt_densepose = DensePoseList(
+                densepose_annotations, instances.gt_boxes, image_shape
+            )
+
+        dataset_dict["instances"] = instances[instances.gt_boxes.nonempty()]
+        return dataset_dict
+
+    def _transform_densepose(self, annotation, transforms):
+        if not self.densepose_on:
+            return annotation
+
+        # Handle densepose annotations
+        is_valid, reason_not_valid = DensePoseDataRelative.validate_annotation(annotation)
+        if is_valid:
+            densepose_data = DensePoseDataRelative(annotation, cleanup=True)
+            densepose_data.apply_transform(transforms, self.densepose_transform_data)
+            annotation["densepose"] = densepose_data
+        else:
+            # logger = logging.getLogger(__name__)
+            # logger.debug("Could not load DensePose annotation: {}".format(reason_not_valid))
+            DensePoseDataRelative.cleanup_annotation(annotation)
+            # NOTE: annotations for certain instances may be unavailable.
+            # 'None' is accepted by the DensePostList data structure.
+            annotation["densepose"] = None
+        return annotation
+
+    def _add_densepose_masks_as_segmentation(
+        self, annotations: List[Dict[str, Any]], image_shape_hw: Tuple[int, int]
+    ):
+        for obj in annotations:
+            if ("densepose" not in obj) or ("segmentation" in obj):
+                continue
+            # DP segmentation: torch.Tensor [S, S] of float32, S=256
+            segm_dp = torch.zeros_like(obj["densepose"].segm)
+            segm_dp[obj["densepose"].segm > 0] = 1
+            segm_h, segm_w = segm_dp.shape
+            bbox_segm_dp = torch.tensor((0, 0, segm_h - 1, segm_w - 1), dtype=torch.float32)
+            # image bbox
+            x0, y0, x1, y1 = (
+                v.item() for v in BoxMode.convert(obj["bbox"], obj["bbox_mode"], BoxMode.XYXY_ABS)
+            )
+            segm_aligned = (
+                ROIAlign((y1 - y0, x1 - x0), 1.0, 0, aligned=True)
+                .forward(segm_dp.view(1, 1, *segm_dp.shape), bbox_segm_dp)
+                .squeeze()
+            )
+            image_mask = torch.zeros(*image_shape_hw, dtype=torch.float32)
+            image_mask[y0:y1, x0:x1] = segm_aligned
+            # segmentation for BitMask: np.array [H, W] of bool
+            obj["segmentation"] = image_mask >= 0.5

densepose/data/datasets/__init__.py

@@ -1,5 +1,5 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from . import builtin  # ensure the builtin datasets are registered
-
-__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from . import builtin  # ensure the builtin datasets are registered
+
+__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]

densepose/data/datasets/builtin.py

@@ -1,16 +1,16 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-from .chimpnsee import register_dataset as register_chimpnsee_dataset
-from .coco import BASE_DATASETS as BASE_COCO_DATASETS
-from .coco import DATASETS as COCO_DATASETS
-from .coco import register_datasets as register_coco_datasets
-from .lvis import DATASETS as LVIS_DATASETS
-from .lvis import register_datasets as register_lvis_datasets
-
-DEFAULT_DATASETS_ROOT = "datasets"
-
-
-register_coco_datasets(COCO_DATASETS, DEFAULT_DATASETS_ROOT)
-register_coco_datasets(BASE_COCO_DATASETS, DEFAULT_DATASETS_ROOT)
-register_lvis_datasets(LVIS_DATASETS, DEFAULT_DATASETS_ROOT)
-
-register_chimpnsee_dataset(DEFAULT_DATASETS_ROOT)  # pyre-ignore[19]
+# Copyright (c) Facebook, Inc. and its affiliates.
+from .chimpnsee import register_dataset as register_chimpnsee_dataset
+from .coco import BASE_DATASETS as BASE_COCO_DATASETS
+from .coco import DATASETS as COCO_DATASETS
+from .coco import register_datasets as register_coco_datasets
+from .lvis import DATASETS as LVIS_DATASETS
+from .lvis import register_datasets as register_lvis_datasets
+
+DEFAULT_DATASETS_ROOT = "datasets"
+
+
+register_coco_datasets(COCO_DATASETS, DEFAULT_DATASETS_ROOT)
+register_coco_datasets(BASE_COCO_DATASETS, DEFAULT_DATASETS_ROOT)
+register_lvis_datasets(LVIS_DATASETS, DEFAULT_DATASETS_ROOT)
+
+register_chimpnsee_dataset(DEFAULT_DATASETS_ROOT)  # pyre-ignore[19]

densepose/data/datasets/chimpnsee.py

@@ -1,29 +1,29 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Optional
-
-from detectron2.data import DatasetCatalog, MetadataCatalog
-
-from ..utils import maybe_prepend_base_path
-from .dataset_type import DatasetType
-
-CHIMPNSEE_DATASET_NAME = "chimpnsee"
-
-
-def register_dataset(datasets_root: Optional[str] = None) -> None:
-    def empty_load_callback():
-        pass
-
-    video_list_fpath = maybe_prepend_base_path(
-        datasets_root,
-        "chimpnsee/cdna.eva.mpg.de/video_list.txt",
-    )
-    video_base_path = maybe_prepend_base_path(datasets_root, "chimpnsee/cdna.eva.mpg.de")
-
-    DatasetCatalog.register(CHIMPNSEE_DATASET_NAME, empty_load_callback)
-    MetadataCatalog.get(CHIMPNSEE_DATASET_NAME).set(
-        dataset_type=DatasetType.VIDEO_LIST,
-        video_list_fpath=video_list_fpath,
-        video_base_path=video_base_path,
-        category="chimpanzee",
-    )
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Optional
+
+from detectron2.data import DatasetCatalog, MetadataCatalog
+
+from ..utils import maybe_prepend_base_path
+from .dataset_type import DatasetType
+
+CHIMPNSEE_DATASET_NAME = "chimpnsee"
+
+
+def register_dataset(datasets_root: Optional[str] = None) -> None:
+    def empty_load_callback():
+        pass
+
+    video_list_fpath = maybe_prepend_base_path(
+        datasets_root,
+        "chimpnsee/cdna.eva.mpg.de/video_list.txt",
+    )
+    video_base_path = maybe_prepend_base_path(datasets_root, "chimpnsee/cdna.eva.mpg.de")
+
+    DatasetCatalog.register(CHIMPNSEE_DATASET_NAME, empty_load_callback)
+    MetadataCatalog.get(CHIMPNSEE_DATASET_NAME).set(
+        dataset_type=DatasetType.VIDEO_LIST,
+        video_list_fpath=video_list_fpath,
+        video_base_path=video_base_path,
+        category="chimpanzee",
+    )

densepose/data/datasets/coco.py

@@ -1,432 +1,432 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-import contextlib
-import io
-import logging
-import os
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Optional
-from fvcore.common.timer import Timer
-
-from detectron2.data import DatasetCatalog, MetadataCatalog
-from detectron2.structures import BoxMode
-from detectron2.utils.file_io import PathManager
-
-from ..utils import maybe_prepend_base_path
-
-DENSEPOSE_MASK_KEY = "dp_masks"
-DENSEPOSE_IUV_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_I", "dp_U", "dp_V"]
-DENSEPOSE_CSE_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_vertex", "ref_model"]
-DENSEPOSE_ALL_POSSIBLE_KEYS = set(
-    DENSEPOSE_IUV_KEYS_WITHOUT_MASK + DENSEPOSE_CSE_KEYS_WITHOUT_MASK + [DENSEPOSE_MASK_KEY]
-)
-DENSEPOSE_METADATA_URL_PREFIX = "https://dl.fbaipublicfiles.com/densepose/data/"
-
-
-@dataclass
-class CocoDatasetInfo:
-    name: str
-    images_root: str
-    annotations_fpath: str
-
-
-DATASETS = [
-    CocoDatasetInfo(
-        name="densepose_coco_2014_train",
-        images_root="coco/train2014",
-        annotations_fpath="coco/annotations/densepose_train2014.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_coco_2014_minival",
-        images_root="coco/val2014",
-        annotations_fpath="coco/annotations/densepose_minival2014.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_coco_2014_minival_100",
-        images_root="coco/val2014",
-        annotations_fpath="coco/annotations/densepose_minival2014_100.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_coco_2014_valminusminival",
-        images_root="coco/val2014",
-        annotations_fpath="coco/annotations/densepose_valminusminival2014.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_coco_2014_train_cse",
-        images_root="coco/train2014",
-        annotations_fpath="coco_cse/densepose_train2014_cse.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_coco_2014_minival_cse",
-        images_root="coco/val2014",
-        annotations_fpath="coco_cse/densepose_minival2014_cse.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_coco_2014_minival_100_cse",
-        images_root="coco/val2014",
-        annotations_fpath="coco_cse/densepose_minival2014_100_cse.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_coco_2014_valminusminival_cse",
-        images_root="coco/val2014",
-        annotations_fpath="coco_cse/densepose_valminusminival2014_cse.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_chimps",
-        images_root="densepose_chimps/images",
-        annotations_fpath="densepose_chimps/densepose_chimps_densepose.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_chimps_cse_train",
-        images_root="densepose_chimps/images",
-        annotations_fpath="densepose_chimps/densepose_chimps_cse_train.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_chimps_cse_val",
-        images_root="densepose_chimps/images",
-        annotations_fpath="densepose_chimps/densepose_chimps_cse_val.json",
-    ),
-    CocoDatasetInfo(
-        name="posetrack2017_train",
-        images_root="posetrack2017/posetrack_data_2017",
-        annotations_fpath="posetrack2017/densepose_posetrack_train2017.json",
-    ),
-    CocoDatasetInfo(
-        name="posetrack2017_val",
-        images_root="posetrack2017/posetrack_data_2017",
-        annotations_fpath="posetrack2017/densepose_posetrack_val2017.json",
-    ),
-    CocoDatasetInfo(
-        name="lvis_v05_train",
-        images_root="coco/train2017",
-        annotations_fpath="lvis/lvis_v0.5_plus_dp_train.json",
-    ),
-    CocoDatasetInfo(
-        name="lvis_v05_val",
-        images_root="coco/val2017",
-        annotations_fpath="lvis/lvis_v0.5_plus_dp_val.json",
-    ),
-]
-
-
-BASE_DATASETS = [
-    CocoDatasetInfo(
-        name="base_coco_2017_train",
-        images_root="coco/train2017",
-        annotations_fpath="coco/annotations/instances_train2017.json",
-    ),
-    CocoDatasetInfo(
-        name="base_coco_2017_val",
-        images_root="coco/val2017",
-        annotations_fpath="coco/annotations/instances_val2017.json",
-    ),
-    CocoDatasetInfo(
-        name="base_coco_2017_val_100",
-        images_root="coco/val2017",
-        annotations_fpath="coco/annotations/instances_val2017_100.json",
-    ),
-]
-
-
-def get_metadata(base_path: Optional[str]) -> Dict[str, Any]:
-    """
-    Returns metadata associated with COCO DensePose datasets
-
-    Args:
-    base_path: Optional[str]
-        Base path used to load metadata from
-
-    Returns:
-    Dict[str, Any]
-        Metadata in the form of a dictionary
-    """
-    meta = {
-        "densepose_transform_src": maybe_prepend_base_path(base_path, "UV_symmetry_transforms.mat"),
-        "densepose_smpl_subdiv": maybe_prepend_base_path(base_path, "SMPL_subdiv.mat"),
-        "densepose_smpl_subdiv_transform": maybe_prepend_base_path(
-            base_path,
-            "SMPL_SUBDIV_TRANSFORM.mat",
-        ),
-    }
-    return meta
-
-
-def _load_coco_annotations(json_file: str):
-    """
-    Load COCO annotations from a JSON file
-
-    Args:
-        json_file: str
-            Path to the file to load annotations from
-    Returns:
-        Instance of `pycocotools.coco.COCO` that provides access to annotations
-        data
-    """
-    from pycocotools.coco import COCO
-
-    logger = logging.getLogger(__name__)
-    timer = Timer()
-    with contextlib.redirect_stdout(io.StringIO()):
-        coco_api = COCO(json_file)
-    if timer.seconds() > 1:
-        logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
-    return coco_api
-
-
-def _add_categories_metadata(dataset_name: str, categories: List[Dict[str, Any]]):
-    meta = MetadataCatalog.get(dataset_name)
-    meta.categories = {c["id"]: c["name"] for c in categories}
-    logger = logging.getLogger(__name__)
-    logger.info("Dataset {} categories: {}".format(dataset_name, meta.categories))
-
-
-def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]):
-    if "minival" in json_file:
-        # Skip validation on COCO2014 valminusminival and minival annotations
-        # The ratio of buggy annotations there is tiny and does not affect accuracy
-        # Therefore we explicitly white-list them
-        return
-    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
-    assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
-        json_file
-    )
-
-
-def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
-    if "bbox" not in ann_dict:
-        return
-    obj["bbox"] = ann_dict["bbox"]
-    obj["bbox_mode"] = BoxMode.XYWH_ABS
-
-
-def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
-    if "segmentation" not in ann_dict:
-        return
-    segm = ann_dict["segmentation"]
-    if not isinstance(segm, dict):
-        # filter out invalid polygons (< 3 points)
-        segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
-        if len(segm) == 0:
-            return
-    obj["segmentation"] = segm
-
-
-def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
-    if "keypoints" not in ann_dict:
-        return
-    keypts = ann_dict["keypoints"]  # list[int]
-    for idx, v in enumerate(keypts):
-        if idx % 3 != 2:
-            # COCO's segmentation coordinates are floating points in [0, H or W],
-            # but keypoint coordinates are integers in [0, H-1 or W-1]
-            # Therefore we assume the coordinates are "pixel indices" and
-            # add 0.5 to convert to floating point coordinates.
-            keypts[idx] = v + 0.5
-    obj["keypoints"] = keypts
-
-
-def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
-    for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
-        if key in ann_dict:
-            obj[key] = ann_dict[key]
-
-
-def _combine_images_with_annotations(
-    dataset_name: str,
-    image_root: str,
-    img_datas: Iterable[Dict[str, Any]],
-    ann_datas: Iterable[Iterable[Dict[str, Any]]],
-):
-
-    ann_keys = ["iscrowd", "category_id"]
-    dataset_dicts = []
-    contains_video_frame_info = False
-
-    for img_dict, ann_dicts in zip(img_datas, ann_datas):
-        record = {}
-        record["file_name"] = os.path.join(image_root, img_dict["file_name"])
-        record["height"] = img_dict["height"]
-        record["width"] = img_dict["width"]
-        record["image_id"] = img_dict["id"]
-        record["dataset"] = dataset_name
-        if "frame_id" in img_dict:
-            record["frame_id"] = img_dict["frame_id"]
-            record["video_id"] = img_dict.get("vid_id", None)
-            contains_video_frame_info = True
-        objs = []
-        for ann_dict in ann_dicts:
-            assert ann_dict["image_id"] == record["image_id"]
-            assert ann_dict.get("ignore", 0) == 0
-            obj = {key: ann_dict[key] for key in ann_keys if key in ann_dict}
-            _maybe_add_bbox(obj, ann_dict)
-            _maybe_add_segm(obj, ann_dict)
-            _maybe_add_keypoints(obj, ann_dict)
-            _maybe_add_densepose(obj, ann_dict)
-            objs.append(obj)
-        record["annotations"] = objs
-        dataset_dicts.append(record)
-    if contains_video_frame_info:
-        create_video_frame_mapping(dataset_name, dataset_dicts)
-    return dataset_dicts
-
-
-def get_contiguous_id_to_category_id_map(metadata):
-    cat_id_2_cont_id = metadata.thing_dataset_id_to_contiguous_id
-    cont_id_2_cat_id = {}
-    for cat_id, cont_id in cat_id_2_cont_id.items():
-        if cont_id in cont_id_2_cat_id:
-            continue
-        cont_id_2_cat_id[cont_id] = cat_id
-    return cont_id_2_cat_id
-
-
-def maybe_filter_categories_cocoapi(dataset_name, coco_api):
-    meta = MetadataCatalog.get(dataset_name)
-    cont_id_2_cat_id = get_contiguous_id_to_category_id_map(meta)
-    cat_id_2_cont_id = meta.thing_dataset_id_to_contiguous_id
-    # filter categories
-    cats = []
-    for cat in coco_api.dataset["categories"]:
-        cat_id = cat["id"]
-        if cat_id not in cat_id_2_cont_id:
-            continue
-        cont_id = cat_id_2_cont_id[cat_id]
-        if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
-            cats.append(cat)
-    coco_api.dataset["categories"] = cats
-    # filter annotations, if multiple categories are mapped to a single
-    # contiguous ID, use only one category ID and map all annotations to that category ID
-    anns = []
-    for ann in coco_api.dataset["annotations"]:
-        cat_id = ann["category_id"]
-        if cat_id not in cat_id_2_cont_id:
-            continue
-        cont_id = cat_id_2_cont_id[cat_id]
-        ann["category_id"] = cont_id_2_cat_id[cont_id]
-        anns.append(ann)
-    coco_api.dataset["annotations"] = anns
-    # recreate index
-    coco_api.createIndex()
-
-
-def maybe_filter_and_map_categories_cocoapi(dataset_name, coco_api):
-    meta = MetadataCatalog.get(dataset_name)
-    category_id_map = meta.thing_dataset_id_to_contiguous_id
-    # map categories
-    cats = []
-    for cat in coco_api.dataset["categories"]:
-        cat_id = cat["id"]
-        if cat_id not in category_id_map:
-            continue
-        cat["id"] = category_id_map[cat_id]
-        cats.append(cat)
-    coco_api.dataset["categories"] = cats
-    # map annotation categories
-    anns = []
-    for ann in coco_api.dataset["annotations"]:
-        cat_id = ann["category_id"]
-        if cat_id not in category_id_map:
-            continue
-        ann["category_id"] = category_id_map[cat_id]
-        anns.append(ann)
-    coco_api.dataset["annotations"] = anns
-    # recreate index
-    coco_api.createIndex()
-
-
-def create_video_frame_mapping(dataset_name, dataset_dicts):
-    mapping = defaultdict(dict)
-    for d in dataset_dicts:
-        video_id = d.get("video_id")
-        if video_id is None:
-            continue
-        mapping[video_id].update({d["frame_id"]: d["file_name"]})
-    MetadataCatalog.get(dataset_name).set(video_frame_mapping=mapping)
-
-
-def load_coco_json(annotations_json_file: str, image_root: str, dataset_name: str):
-    """
-    Loads a JSON file with annotations in COCO instances format.
-    Replaces `detectron2.data.datasets.coco.load_coco_json` to handle metadata
-    in a more flexible way. Postpones category mapping to a later stage to be
-    able to combine several datasets with different (but coherent) sets of
-    categories.
-
-    Args:
-
-    annotations_json_file: str
-        Path to the JSON file with annotations in COCO instances format.
-    image_root: str
-        directory that contains all the images
-    dataset_name: str
-        the name that identifies a dataset, e.g. "densepose_coco_2014_train"
-    extra_annotation_keys: Optional[List[str]]
-        If provided, these keys are used to extract additional data from
-        the annotations.
-    """
-    coco_api = _load_coco_annotations(PathManager.get_local_path(annotations_json_file))
-    _add_categories_metadata(dataset_name, coco_api.loadCats(coco_api.getCatIds()))
-    # sort indices for reproducible results
-    img_ids = sorted(coco_api.imgs.keys())
-    # imgs is a list of dicts, each looks something like:
-    # {'license': 4,
-    #  'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
-    #  'file_name': 'COCO_val2014_000000001268.jpg',
-    #  'height': 427,
-    #  'width': 640,
-    #  'date_captured': '2013-11-17 05:57:24',
-    #  'id': 1268}
-    imgs = coco_api.loadImgs(img_ids)
-    logger = logging.getLogger(__name__)
-    logger.info("Loaded {} images in COCO format from {}".format(len(imgs), annotations_json_file))
-    # anns is a list[list[dict]], where each dict is an annotation
-    # record for an object. The inner list enumerates the objects in an image
-    # and the outer list enumerates over images.
-    anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
-    _verify_annotations_have_unique_ids(annotations_json_file, anns)
-    dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
-    return dataset_records
-
-
-def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None):
-    """
-    Registers provided COCO DensePose dataset
-
-    Args:
-    dataset_data: CocoDatasetInfo
-        Dataset data
-    datasets_root: Optional[str]
-        Datasets root folder (default: None)
-    """
-    annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
-    images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
-
-    def load_annotations():
-        return load_coco_json(
-            annotations_json_file=annotations_fpath,
-            image_root=images_root,
-            dataset_name=dataset_data.name,
-        )
-
-    DatasetCatalog.register(dataset_data.name, load_annotations)
-    MetadataCatalog.get(dataset_data.name).set(
-        json_file=annotations_fpath,
-        image_root=images_root,
-        **get_metadata(DENSEPOSE_METADATA_URL_PREFIX)
-    )
-
-
-def register_datasets(
-    datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
-):
-    """
-    Registers provided COCO DensePose datasets
-
-    Args:
-    datasets_data: Iterable[CocoDatasetInfo]
-        An iterable of dataset datas
-    datasets_root: Optional[str]
-        Datasets root folder (default: None)
-    """
-    for dataset_data in datasets_data:
-        register_dataset(dataset_data, datasets_root)
+# Copyright (c) Facebook, Inc. and its affiliates.
+import contextlib
+import io
+import logging
+import os
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Optional
+from fvcore.common.timer import Timer
+
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from detectron2.structures import BoxMode
+from detectron2.utils.file_io import PathManager
+
+from ..utils import maybe_prepend_base_path
+
+DENSEPOSE_MASK_KEY = "dp_masks"
+DENSEPOSE_IUV_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_I", "dp_U", "dp_V"]
+DENSEPOSE_CSE_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_vertex", "ref_model"]
+DENSEPOSE_ALL_POSSIBLE_KEYS = set(
+    DENSEPOSE_IUV_KEYS_WITHOUT_MASK + DENSEPOSE_CSE_KEYS_WITHOUT_MASK + [DENSEPOSE_MASK_KEY]
+)
+DENSEPOSE_METADATA_URL_PREFIX = "https://dl.fbaipublicfiles.com/densepose/data/"
+
+
+@dataclass
+class CocoDatasetInfo:
+    name: str
+    images_root: str
+    annotations_fpath: str
+
+
+DATASETS = [
+    CocoDatasetInfo(
+        name="densepose_coco_2014_train",
+        images_root="coco/train2014",
+        annotations_fpath="coco/annotations/densepose_train2014.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival",
+        images_root="coco/val2014",
+        annotations_fpath="coco/annotations/densepose_minival2014.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival_100",
+        images_root="coco/val2014",
+        annotations_fpath="coco/annotations/densepose_minival2014_100.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_valminusminival",
+        images_root="coco/val2014",
+        annotations_fpath="coco/annotations/densepose_valminusminival2014.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_train_cse",
+        images_root="coco/train2014",
+        annotations_fpath="coco_cse/densepose_train2014_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival_cse",
+        images_root="coco/val2014",
+        annotations_fpath="coco_cse/densepose_minival2014_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival_100_cse",
+        images_root="coco/val2014",
+        annotations_fpath="coco_cse/densepose_minival2014_100_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_valminusminival_cse",
+        images_root="coco/val2014",
+        annotations_fpath="coco_cse/densepose_valminusminival2014_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_chimps",
+        images_root="densepose_chimps/images",
+        annotations_fpath="densepose_chimps/densepose_chimps_densepose.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_chimps_cse_train",
+        images_root="densepose_chimps/images",
+        annotations_fpath="densepose_chimps/densepose_chimps_cse_train.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_chimps_cse_val",
+        images_root="densepose_chimps/images",
+        annotations_fpath="densepose_chimps/densepose_chimps_cse_val.json",
+    ),
+    CocoDatasetInfo(
+        name="posetrack2017_train",
+        images_root="posetrack2017/posetrack_data_2017",
+        annotations_fpath="posetrack2017/densepose_posetrack_train2017.json",
+    ),
+    CocoDatasetInfo(
+        name="posetrack2017_val",
+        images_root="posetrack2017/posetrack_data_2017",
+        annotations_fpath="posetrack2017/densepose_posetrack_val2017.json",
+    ),
+    CocoDatasetInfo(
+        name="lvis_v05_train",
+        images_root="coco/train2017",
+        annotations_fpath="lvis/lvis_v0.5_plus_dp_train.json",
+    ),
+    CocoDatasetInfo(
+        name="lvis_v05_val",
+        images_root="coco/val2017",
+        annotations_fpath="lvis/lvis_v0.5_plus_dp_val.json",
+    ),
+]
+
+
+BASE_DATASETS = [
+    CocoDatasetInfo(
+        name="base_coco_2017_train",
+        images_root="coco/train2017",
+        annotations_fpath="coco/annotations/instances_train2017.json",
+    ),
+    CocoDatasetInfo(
+        name="base_coco_2017_val",
+        images_root="coco/val2017",
+        annotations_fpath="coco/annotations/instances_val2017.json",
+    ),
+    CocoDatasetInfo(
+        name="base_coco_2017_val_100",
+        images_root="coco/val2017",
+        annotations_fpath="coco/annotations/instances_val2017_100.json",
+    ),
+]
+
+
+def get_metadata(base_path: Optional[str]) -> Dict[str, Any]:
+    """
+    Returns metadata associated with COCO DensePose datasets
+
+    Args:
+    base_path: Optional[str]
+        Base path used to load metadata from
+
+    Returns:
+    Dict[str, Any]
+        Metadata in the form of a dictionary
+    """
+    meta = {
+        "densepose_transform_src": maybe_prepend_base_path(base_path, "UV_symmetry_transforms.mat"),
+        "densepose_smpl_subdiv": maybe_prepend_base_path(base_path, "SMPL_subdiv.mat"),
+        "densepose_smpl_subdiv_transform": maybe_prepend_base_path(
+            base_path,
+            "SMPL_SUBDIV_TRANSFORM.mat",
+        ),
+    }
+    return meta
+
+
+def _load_coco_annotations(json_file: str):
+    """
+    Load COCO annotations from a JSON file
+
+    Args:
+        json_file: str
+            Path to the file to load annotations from
+    Returns:
+        Instance of `pycocotools.coco.COCO` that provides access to annotations
+        data
+    """
+    from pycocotools.coco import COCO
+
+    logger = logging.getLogger(__name__)
+    timer = Timer()
+    with contextlib.redirect_stdout(io.StringIO()):
+        coco_api = COCO(json_file)
+    if timer.seconds() > 1:
+        logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
+    return coco_api
+
+
+def _add_categories_metadata(dataset_name: str, categories: List[Dict[str, Any]]):
+    meta = MetadataCatalog.get(dataset_name)
+    meta.categories = {c["id"]: c["name"] for c in categories}
+    logger = logging.getLogger(__name__)
+    logger.info("Dataset {} categories: {}".format(dataset_name, meta.categories))
+
+
+def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]):
+    if "minival" in json_file:
+        # Skip validation on COCO2014 valminusminival and minival annotations
+        # The ratio of buggy annotations there is tiny and does not affect accuracy
+        # Therefore we explicitly white-list them
+        return
+    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+    assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
+        json_file
+    )
+
+
+def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    if "bbox" not in ann_dict:
+        return
+    obj["bbox"] = ann_dict["bbox"]
+    obj["bbox_mode"] = BoxMode.XYWH_ABS
+
+
+def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    if "segmentation" not in ann_dict:
+        return
+    segm = ann_dict["segmentation"]
+    if not isinstance(segm, dict):
+        # filter out invalid polygons (< 3 points)
+        segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
+        if len(segm) == 0:
+            return
+    obj["segmentation"] = segm
+
+
+def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    if "keypoints" not in ann_dict:
+        return
+    keypts = ann_dict["keypoints"]  # list[int]
+    for idx, v in enumerate(keypts):
+        if idx % 3 != 2:
+            # COCO's segmentation coordinates are floating points in [0, H or W],
+            # but keypoint coordinates are integers in [0, H-1 or W-1]
+            # Therefore we assume the coordinates are "pixel indices" and
+            # add 0.5 to convert to floating point coordinates.
+            keypts[idx] = v + 0.5
+    obj["keypoints"] = keypts
+
+
+def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
+        if key in ann_dict:
+            obj[key] = ann_dict[key]
+
+
+def _combine_images_with_annotations(
+    dataset_name: str,
+    image_root: str,
+    img_datas: Iterable[Dict[str, Any]],
+    ann_datas: Iterable[Iterable[Dict[str, Any]]],
+):
+
+    ann_keys = ["iscrowd", "category_id"]
+    dataset_dicts = []
+    contains_video_frame_info = False
+
+    for img_dict, ann_dicts in zip(img_datas, ann_datas):
+        record = {}
+        record["file_name"] = os.path.join(image_root, img_dict["file_name"])
+        record["height"] = img_dict["height"]
+        record["width"] = img_dict["width"]
+        record["image_id"] = img_dict["id"]
+        record["dataset"] = dataset_name
+        if "frame_id" in img_dict:
+            record["frame_id"] = img_dict["frame_id"]
+            record["video_id"] = img_dict.get("vid_id", None)
+            contains_video_frame_info = True
+        objs = []
+        for ann_dict in ann_dicts:
+            assert ann_dict["image_id"] == record["image_id"]
+            assert ann_dict.get("ignore", 0) == 0
+            obj = {key: ann_dict[key] for key in ann_keys if key in ann_dict}
+            _maybe_add_bbox(obj, ann_dict)
+            _maybe_add_segm(obj, ann_dict)
+            _maybe_add_keypoints(obj, ann_dict)
+            _maybe_add_densepose(obj, ann_dict)
+            objs.append(obj)
+        record["annotations"] = objs
+        dataset_dicts.append(record)
+    if contains_video_frame_info:
+        create_video_frame_mapping(dataset_name, dataset_dicts)
+    return dataset_dicts
+
+
+def get_contiguous_id_to_category_id_map(metadata):
+    cat_id_2_cont_id = metadata.thing_dataset_id_to_contiguous_id
+    cont_id_2_cat_id = {}
+    for cat_id, cont_id in cat_id_2_cont_id.items():
+        if cont_id in cont_id_2_cat_id:
+            continue
+        cont_id_2_cat_id[cont_id] = cat_id
+    return cont_id_2_cat_id
+
+
+def maybe_filter_categories_cocoapi(dataset_name, coco_api):
+    meta = MetadataCatalog.get(dataset_name)
+    cont_id_2_cat_id = get_contiguous_id_to_category_id_map(meta)
+    cat_id_2_cont_id = meta.thing_dataset_id_to_contiguous_id
+    # filter categories
+    cats = []
+    for cat in coco_api.dataset["categories"]:
+        cat_id = cat["id"]
+        if cat_id not in cat_id_2_cont_id:
+            continue
+        cont_id = cat_id_2_cont_id[cat_id]
+        if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
+            cats.append(cat)
+    coco_api.dataset["categories"] = cats
+    # filter annotations, if multiple categories are mapped to a single
+    # contiguous ID, use only one category ID and map all annotations to that category ID
+    anns = []
+    for ann in coco_api.dataset["annotations"]:
+        cat_id = ann["category_id"]
+        if cat_id not in cat_id_2_cont_id:
+            continue
+        cont_id = cat_id_2_cont_id[cat_id]
+        ann["category_id"] = cont_id_2_cat_id[cont_id]
+        anns.append(ann)
+    coco_api.dataset["annotations"] = anns
+    # recreate index
+    coco_api.createIndex()
+
+
+def maybe_filter_and_map_categories_cocoapi(dataset_name, coco_api):
+    meta = MetadataCatalog.get(dataset_name)
+    category_id_map = meta.thing_dataset_id_to_contiguous_id
+    # map categories
+    cats = []
+    for cat in coco_api.dataset["categories"]:
+        cat_id = cat["id"]
+        if cat_id not in category_id_map:
+            continue
+        cat["id"] = category_id_map[cat_id]
+        cats.append(cat)
+    coco_api.dataset["categories"] = cats
+    # map annotation categories
+    anns = []
+    for ann in coco_api.dataset["annotations"]:
+        cat_id = ann["category_id"]
+        if cat_id not in category_id_map:
+            continue
+        ann["category_id"] = category_id_map[cat_id]
+        anns.append(ann)
+    coco_api.dataset["annotations"] = anns
+    # recreate index
+    coco_api.createIndex()
+
+
+def create_video_frame_mapping(dataset_name, dataset_dicts):
+    mapping = defaultdict(dict)
+    for d in dataset_dicts:
+        video_id = d.get("video_id")
+        if video_id is None:
+            continue
+        mapping[video_id].update({d["frame_id"]: d["file_name"]})
+    MetadataCatalog.get(dataset_name).set(video_frame_mapping=mapping)
+
+
+def load_coco_json(annotations_json_file: str, image_root: str, dataset_name: str):
+    """
+    Loads a JSON file with annotations in COCO instances format.
+    Replaces `detectron2.data.datasets.coco.load_coco_json` to handle metadata
+    in a more flexible way. Postpones category mapping to a later stage to be
+    able to combine several datasets with different (but coherent) sets of
+    categories.
+
+    Args:
+
+    annotations_json_file: str
+        Path to the JSON file with annotations in COCO instances format.
+    image_root: str
+        directory that contains all the images
+    dataset_name: str
+        the name that identifies a dataset, e.g. "densepose_coco_2014_train"
+    extra_annotation_keys: Optional[List[str]]
+        If provided, these keys are used to extract additional data from
+        the annotations.
+    """
+    coco_api = _load_coco_annotations(PathManager.get_local_path(annotations_json_file))
+    _add_categories_metadata(dataset_name, coco_api.loadCats(coco_api.getCatIds()))
+    # sort indices for reproducible results
+    img_ids = sorted(coco_api.imgs.keys())
+    # imgs is a list of dicts, each looks something like:
+    # {'license': 4,
+    #  'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
+    #  'file_name': 'COCO_val2014_000000001268.jpg',
+    #  'height': 427,
+    #  'width': 640,
+    #  'date_captured': '2013-11-17 05:57:24',
+    #  'id': 1268}
+    imgs = coco_api.loadImgs(img_ids)
+    logger = logging.getLogger(__name__)
+    logger.info("Loaded {} images in COCO format from {}".format(len(imgs), annotations_json_file))
+    # anns is a list[list[dict]], where each dict is an annotation
+    # record for an object. The inner list enumerates the objects in an image
+    # and the outer list enumerates over images.
+    anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
+    _verify_annotations_have_unique_ids(annotations_json_file, anns)
+    dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
+    return dataset_records
+
+
+def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None):
+    """
+    Registers provided COCO DensePose dataset
+
+    Args:
+    dataset_data: CocoDatasetInfo
+        Dataset data
+    datasets_root: Optional[str]
+        Datasets root folder (default: None)
+    """
+    annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
+    images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
+
+    def load_annotations():
+        return load_coco_json(
+            annotations_json_file=annotations_fpath,
+            image_root=images_root,
+            dataset_name=dataset_data.name,
+        )
+
+    DatasetCatalog.register(dataset_data.name, load_annotations)
+    MetadataCatalog.get(dataset_data.name).set(
+        json_file=annotations_fpath,
+        image_root=images_root,
+        **get_metadata(DENSEPOSE_METADATA_URL_PREFIX)
+    )
+
+
+def register_datasets(
+    datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
+):
+    """
+    Registers provided COCO DensePose datasets
+
+    Args:
+    datasets_data: Iterable[CocoDatasetInfo]
+        An iterable of dataset datas
+    datasets_root: Optional[str]
+        Datasets root folder (default: None)
+    """
+    for dataset_data in datasets_data:
+        register_dataset(dataset_data, datasets_root)

densepose/data/datasets/dataset_type.py

@@ -1,11 +1,11 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from enum import Enum
-
-
-class DatasetType(Enum):
-    """
-    Dataset type, mostly used for datasets that contain data to bootstrap models on
-    """
-
-    VIDEO_LIST = "video_list"
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from enum import Enum
+
+
+class DatasetType(Enum):
+    """
+    Dataset type, mostly used for datasets that contain data to bootstrap models on
+    """
+
+    VIDEO_LIST = "video_list"

densepose/data/datasets/lvis.py

@@ -1,257 +1,257 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-import logging
-import os
-from typing import Any, Dict, Iterable, List, Optional
-from fvcore.common.timer import Timer
-
-from detectron2.data import DatasetCatalog, MetadataCatalog
-from detectron2.data.datasets.lvis import get_lvis_instances_meta
-from detectron2.structures import BoxMode
-from detectron2.utils.file_io import PathManager
-
-from ..utils import maybe_prepend_base_path
-from .coco import (
-    DENSEPOSE_ALL_POSSIBLE_KEYS,
-    DENSEPOSE_METADATA_URL_PREFIX,
-    CocoDatasetInfo,
-    get_metadata,
-)
-
-DATASETS = [
-    CocoDatasetInfo(
-        name="densepose_lvis_v1_ds1_train_v1",
-        images_root="coco_",
-        annotations_fpath="lvis/densepose_lvis_v1_ds1_train_v1.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_lvis_v1_ds1_val_v1",
-        images_root="coco_",
-        annotations_fpath="lvis/densepose_lvis_v1_ds1_val_v1.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_lvis_v1_ds2_train_v1",
-        images_root="coco_",
-        annotations_fpath="lvis/densepose_lvis_v1_ds2_train_v1.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_lvis_v1_ds2_val_v1",
-        images_root="coco_",
-        annotations_fpath="lvis/densepose_lvis_v1_ds2_val_v1.json",
-    ),
-    CocoDatasetInfo(
-        name="densepose_lvis_v1_ds1_val_animals_100",
-        images_root="coco_",
-        annotations_fpath="lvis/densepose_lvis_v1_val_animals_100_v2.json",
-    ),
-]
-
-
-def _load_lvis_annotations(json_file: str):
-    """
-    Load COCO annotations from a JSON file
-
-    Args:
-        json_file: str
-            Path to the file to load annotations from
-    Returns:
-        Instance of `pycocotools.coco.COCO` that provides access to annotations
-        data
-    """
-    from lvis import LVIS
-
-    json_file = PathManager.get_local_path(json_file)
-    logger = logging.getLogger(__name__)
-    timer = Timer()
-    lvis_api = LVIS(json_file)
-    if timer.seconds() > 1:
-        logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
-    return lvis_api
-
-
-def _add_categories_metadata(dataset_name: str) -> None:
-    metadict = get_lvis_instances_meta(dataset_name)
-    categories = metadict["thing_classes"]
-    metadata = MetadataCatalog.get(dataset_name)
-    metadata.categories = {i + 1: categories[i] for i in range(len(categories))}
-    logger = logging.getLogger(__name__)
-    logger.info(f"Dataset {dataset_name} has {len(categories)} categories")
-
-
-def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]) -> None:
-    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
-    assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
-        json_file
-    )
-
-
-def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
-    if "bbox" not in ann_dict:
-        return
-    obj["bbox"] = ann_dict["bbox"]
-    obj["bbox_mode"] = BoxMode.XYWH_ABS
-
-
-def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
-    if "segmentation" not in ann_dict:
-        return
-    segm = ann_dict["segmentation"]
-    if not isinstance(segm, dict):
-        # filter out invalid polygons (< 3 points)
-        segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
-        if len(segm) == 0:
-            return
-    obj["segmentation"] = segm
-
-
-def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
-    if "keypoints" not in ann_dict:
-        return
-    keypts = ann_dict["keypoints"]  # list[int]
-    for idx, v in enumerate(keypts):
-        if idx % 3 != 2:
-            # COCO's segmentation coordinates are floating points in [0, H or W],
-            # but keypoint coordinates are integers in [0, H-1 or W-1]
-            # Therefore we assume the coordinates are "pixel indices" and
-            # add 0.5 to convert to floating point coordinates.
-            keypts[idx] = v + 0.5
-    obj["keypoints"] = keypts
-
-
-def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
-    for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
-        if key in ann_dict:
-            obj[key] = ann_dict[key]
-
-
-def _combine_images_with_annotations(
-    dataset_name: str,
-    image_root: str,
-    img_datas: Iterable[Dict[str, Any]],
-    ann_datas: Iterable[Iterable[Dict[str, Any]]],
-):
-
-    dataset_dicts = []
-
-    def get_file_name(img_root, img_dict):
-        # Determine the path including the split folder ("train2017", "val2017", "test2017") from
-        # the coco_url field. Example:
-        #   'coco_url': 'http://images.cocodataset.org/train2017/000000155379.jpg'
-        split_folder, file_name = img_dict["coco_url"].split("/")[-2:]
-        return os.path.join(img_root + split_folder, file_name)
-
-    for img_dict, ann_dicts in zip(img_datas, ann_datas):
-        record = {}
-        record["file_name"] = get_file_name(image_root, img_dict)
-        record["height"] = img_dict["height"]
-        record["width"] = img_dict["width"]
-        record["not_exhaustive_category_ids"] = img_dict.get("not_exhaustive_category_ids", [])
-        record["neg_category_ids"] = img_dict.get("neg_category_ids", [])
-        record["image_id"] = img_dict["id"]
-        record["dataset"] = dataset_name
-
-        objs = []
-        for ann_dict in ann_dicts:
-            assert ann_dict["image_id"] == record["image_id"]
-            obj = {}
-            _maybe_add_bbox(obj, ann_dict)
-            obj["iscrowd"] = ann_dict.get("iscrowd", 0)
-            obj["category_id"] = ann_dict["category_id"]
-            _maybe_add_segm(obj, ann_dict)
-            _maybe_add_keypoints(obj, ann_dict)
-            _maybe_add_densepose(obj, ann_dict)
-            objs.append(obj)
-        record["annotations"] = objs
-        dataset_dicts.append(record)
-    return dataset_dicts
-
-
-def load_lvis_json(annotations_json_file: str, image_root: str, dataset_name: str):
-    """
-    Loads a JSON file with annotations in LVIS instances format.
-    Replaces `detectron2.data.datasets.coco.load_lvis_json` to handle metadata
-    in a more flexible way. Postpones category mapping to a later stage to be
-    able to combine several datasets with different (but coherent) sets of
-    categories.
-
-    Args:
-
-    annotations_json_file: str
-        Path to the JSON file with annotations in COCO instances format.
-    image_root: str
-        directory that contains all the images
-    dataset_name: str
-        the name that identifies a dataset, e.g. "densepose_coco_2014_train"
-    extra_annotation_keys: Optional[List[str]]
-        If provided, these keys are used to extract additional data from
-        the annotations.
-    """
-    lvis_api = _load_lvis_annotations(PathManager.get_local_path(annotations_json_file))
-
-    _add_categories_metadata(dataset_name)
-
-    # sort indices for reproducible results
-    img_ids = sorted(lvis_api.imgs.keys())
-    # imgs is a list of dicts, each looks something like:
-    # {'license': 4,
-    #  'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
-    #  'file_name': 'COCO_val2014_000000001268.jpg',
-    #  'height': 427,
-    #  'width': 640,
-    #  'date_captured': '2013-11-17 05:57:24',
-    #  'id': 1268}
-    imgs = lvis_api.load_imgs(img_ids)
-    logger = logging.getLogger(__name__)
-    logger.info("Loaded {} images in LVIS format from {}".format(len(imgs), annotations_json_file))
-    # anns is a list[list[dict]], where each dict is an annotation
-    # record for an object. The inner list enumerates the objects in an image
-    # and the outer list enumerates over images.
-    anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
-
-    _verify_annotations_have_unique_ids(annotations_json_file, anns)
-    dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
-    return dataset_records
-
-
-def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None) -> None:
-    """
-    Registers provided LVIS DensePose dataset
-
-    Args:
-        dataset_data: CocoDatasetInfo
-            Dataset data
-        datasets_root: Optional[str]
-            Datasets root folder (default: None)
-    """
-    annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
-    images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
-
-    def load_annotations():
-        return load_lvis_json(
-            annotations_json_file=annotations_fpath,
-            image_root=images_root,
-            dataset_name=dataset_data.name,
-        )
-
-    DatasetCatalog.register(dataset_data.name, load_annotations)
-    MetadataCatalog.get(dataset_data.name).set(
-        json_file=annotations_fpath,
-        image_root=images_root,
-        evaluator_type="lvis",
-        **get_metadata(DENSEPOSE_METADATA_URL_PREFIX),
-    )
-
-
-def register_datasets(
-    datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
-) -> None:
-    """
-    Registers provided LVIS DensePose datasets
-
-    Args:
-        datasets_data: Iterable[CocoDatasetInfo]
-            An iterable of dataset datas
-        datasets_root: Optional[str]
-            Datasets root folder (default: None)
-    """
-    for dataset_data in datasets_data:
-        register_dataset(dataset_data, datasets_root)
+# Copyright (c) Facebook, Inc. and its affiliates.
+import logging
+import os
+from typing import Any, Dict, Iterable, List, Optional
+from fvcore.common.timer import Timer
+
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from detectron2.data.datasets.lvis import get_lvis_instances_meta
+from detectron2.structures import BoxMode
+from detectron2.utils.file_io import PathManager
+
+from ..utils import maybe_prepend_base_path
+from .coco import (
+    DENSEPOSE_ALL_POSSIBLE_KEYS,
+    DENSEPOSE_METADATA_URL_PREFIX,
+    CocoDatasetInfo,
+    get_metadata,
+)
+
+DATASETS = [
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds1_train_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds1_train_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds1_val_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds1_val_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds2_train_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds2_train_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds2_val_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds2_val_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds1_val_animals_100",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_val_animals_100_v2.json",
+    ),
+]
+
+
+def _load_lvis_annotations(json_file: str):
+    """
+    Load COCO annotations from a JSON file
+
+    Args:
+        json_file: str
+            Path to the file to load annotations from
+    Returns:
+        Instance of `pycocotools.coco.COCO` that provides access to annotations
+        data
+    """
+    from lvis import LVIS
+
+    json_file = PathManager.get_local_path(json_file)
+    logger = logging.getLogger(__name__)
+    timer = Timer()
+    lvis_api = LVIS(json_file)
+    if timer.seconds() > 1:
+        logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
+    return lvis_api
+
+
+def _add_categories_metadata(dataset_name: str) -> None:
+    metadict = get_lvis_instances_meta(dataset_name)
+    categories = metadict["thing_classes"]
+    metadata = MetadataCatalog.get(dataset_name)
+    metadata.categories = {i + 1: categories[i] for i in range(len(categories))}
+    logger = logging.getLogger(__name__)
+    logger.info(f"Dataset {dataset_name} has {len(categories)} categories")
+
+
+def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]) -> None:
+    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+    assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
+        json_file
+    )
+
+
+def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    if "bbox" not in ann_dict:
+        return
+    obj["bbox"] = ann_dict["bbox"]
+    obj["bbox_mode"] = BoxMode.XYWH_ABS
+
+
+def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    if "segmentation" not in ann_dict:
+        return
+    segm = ann_dict["segmentation"]
+    if not isinstance(segm, dict):
+        # filter out invalid polygons (< 3 points)
+        segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
+        if len(segm) == 0:
+            return
+    obj["segmentation"] = segm
+
+
+def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    if "keypoints" not in ann_dict:
+        return
+    keypts = ann_dict["keypoints"]  # list[int]
+    for idx, v in enumerate(keypts):
+        if idx % 3 != 2:
+            # COCO's segmentation coordinates are floating points in [0, H or W],
+            # but keypoint coordinates are integers in [0, H-1 or W-1]
+            # Therefore we assume the coordinates are "pixel indices" and
+            # add 0.5 to convert to floating point coordinates.
+            keypts[idx] = v + 0.5
+    obj["keypoints"] = keypts
+
+
+def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
+        if key in ann_dict:
+            obj[key] = ann_dict[key]
+
+
+def _combine_images_with_annotations(
+    dataset_name: str,
+    image_root: str,
+    img_datas: Iterable[Dict[str, Any]],
+    ann_datas: Iterable[Iterable[Dict[str, Any]]],
+):
+
+    dataset_dicts = []
+
+    def get_file_name(img_root, img_dict):
+        # Determine the path including the split folder ("train2017", "val2017", "test2017") from
+        # the coco_url field. Example:
+        #   'coco_url': 'http://images.cocodataset.org/train2017/000000155379.jpg'
+        split_folder, file_name = img_dict["coco_url"].split("/")[-2:]
+        return os.path.join(img_root + split_folder, file_name)
+
+    for img_dict, ann_dicts in zip(img_datas, ann_datas):
+        record = {}
+        record["file_name"] = get_file_name(image_root, img_dict)
+        record["height"] = img_dict["height"]
+        record["width"] = img_dict["width"]
+        record["not_exhaustive_category_ids"] = img_dict.get("not_exhaustive_category_ids", [])
+        record["neg_category_ids"] = img_dict.get("neg_category_ids", [])
+        record["image_id"] = img_dict["id"]
+        record["dataset"] = dataset_name
+
+        objs = []
+        for ann_dict in ann_dicts:
+            assert ann_dict["image_id"] == record["image_id"]
+            obj = {}
+            _maybe_add_bbox(obj, ann_dict)
+            obj["iscrowd"] = ann_dict.get("iscrowd", 0)
+            obj["category_id"] = ann_dict["category_id"]
+            _maybe_add_segm(obj, ann_dict)
+            _maybe_add_keypoints(obj, ann_dict)
+            _maybe_add_densepose(obj, ann_dict)
+            objs.append(obj)
+        record["annotations"] = objs
+        dataset_dicts.append(record)
+    return dataset_dicts
+
+
+def load_lvis_json(annotations_json_file: str, image_root: str, dataset_name: str):
+    """
+    Loads a JSON file with annotations in LVIS instances format.
+    Replaces `detectron2.data.datasets.coco.load_lvis_json` to handle metadata
+    in a more flexible way. Postpones category mapping to a later stage to be
+    able to combine several datasets with different (but coherent) sets of
+    categories.
+
+    Args:
+
+    annotations_json_file: str
+        Path to the JSON file with annotations in COCO instances format.
+    image_root: str
+        directory that contains all the images
+    dataset_name: str
+        the name that identifies a dataset, e.g. "densepose_coco_2014_train"
+    extra_annotation_keys: Optional[List[str]]
+        If provided, these keys are used to extract additional data from
+        the annotations.
+    """
+    lvis_api = _load_lvis_annotations(PathManager.get_local_path(annotations_json_file))
+
+    _add_categories_metadata(dataset_name)
+
+    # sort indices for reproducible results
+    img_ids = sorted(lvis_api.imgs.keys())
+    # imgs is a list of dicts, each looks something like:
+    # {'license': 4,
+    #  'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
+    #  'file_name': 'COCO_val2014_000000001268.jpg',
+    #  'height': 427,
+    #  'width': 640,
+    #  'date_captured': '2013-11-17 05:57:24',
+    #  'id': 1268}
+    imgs = lvis_api.load_imgs(img_ids)
+    logger = logging.getLogger(__name__)
+    logger.info("Loaded {} images in LVIS format from {}".format(len(imgs), annotations_json_file))
+    # anns is a list[list[dict]], where each dict is an annotation
+    # record for an object. The inner list enumerates the objects in an image
+    # and the outer list enumerates over images.
+    anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+    _verify_annotations_have_unique_ids(annotations_json_file, anns)
+    dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
+    return dataset_records
+
+
+def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None) -> None:
+    """
+    Registers provided LVIS DensePose dataset
+
+    Args:
+        dataset_data: CocoDatasetInfo
+            Dataset data
+        datasets_root: Optional[str]
+            Datasets root folder (default: None)
+    """
+    annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
+    images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
+
+    def load_annotations():
+        return load_lvis_json(
+            annotations_json_file=annotations_fpath,
+            image_root=images_root,
+            dataset_name=dataset_data.name,
+        )
+
+    DatasetCatalog.register(dataset_data.name, load_annotations)
+    MetadataCatalog.get(dataset_data.name).set(
+        json_file=annotations_fpath,
+        image_root=images_root,
+        evaluator_type="lvis",
+        **get_metadata(DENSEPOSE_METADATA_URL_PREFIX),
+    )
+
+
+def register_datasets(
+    datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
+) -> None:
+    """
+    Registers provided LVIS DensePose datasets
+
+    Args:
+        datasets_data: Iterable[CocoDatasetInfo]
+            An iterable of dataset datas
+        datasets_root: Optional[str]
+            Datasets root folder (default: None)
+    """
+    for dataset_data in datasets_data:
+        register_dataset(dataset_data, datasets_root)

densepose/data/image_list_dataset.py

@@ -1,72 +1,72 @@
-# -*- coding: utf-8 -*-
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import logging
-import numpy as np
-from typing import Any, Callable, Dict, List, Optional, Union
-import torch
-from torch.utils.data.dataset import Dataset
-
-from detectron2.data.detection_utils import read_image
-
-ImageTransform = Callable[[torch.Tensor], torch.Tensor]
-
-
-class ImageListDataset(Dataset):
-    """
-    Dataset that provides images from a list.
-    """
-
-    _EMPTY_IMAGE = torch.empty((0, 3, 1, 1))
-
-    def __init__(
-        self,
-        image_list: List[str],
-        category_list: Union[str, List[str], None] = None,
-        transform: Optional[ImageTransform] = None,
-    ):
-        """
-        Args:
-            image_list (List[str]): list of paths to image files
-            category_list (Union[str, List[str], None]): list of animal categories for
-                each image. If it is a string, or None, this applies to all images
-        """
-        if type(category_list) == list:
-            self.category_list = category_list
-        else:
-            self.category_list = [category_list] * len(image_list)
-        assert len(image_list) == len(
-            self.category_list
-        ), "length of image and category lists must be equal"
-        self.image_list = image_list
-        self.transform = transform
-
-    def __getitem__(self, idx: int) -> Dict[str, Any]:
-        """
-        Gets selected images from the list
-
-        Args:
-            idx (int): video index in the video list file
-        Returns:
-            A dictionary containing two keys:
-                images (torch.Tensor): tensor of size [N, 3, H, W] (N = 1, or 0 for _EMPTY_IMAGE)
-                categories (List[str]): categories of the frames
-        """
-        categories = [self.category_list[idx]]
-        fpath = self.image_list[idx]
-        transform = self.transform
-
-        try:
-            image = torch.from_numpy(np.ascontiguousarray(read_image(fpath, format="BGR")))
-            image = image.permute(2, 0, 1).unsqueeze(0).float()  # HWC -> NCHW
-            if transform is not None:
-                image = transform(image)
-            return {"images": image, "categories": categories}
-        except (OSError, RuntimeError) as e:
-            logger = logging.getLogger(__name__)
-            logger.warning(f"Error opening image file container {fpath}: {e}")
-
-        return {"images": self._EMPTY_IMAGE, "categories": []}
-
-    def __len__(self):
-        return len(self.image_list)
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import logging
+import numpy as np
+from typing import Any, Callable, Dict, List, Optional, Union
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.data.detection_utils import read_image
+
+ImageTransform = Callable[[torch.Tensor], torch.Tensor]
+
+
+class ImageListDataset(Dataset):
+    """
+    Dataset that provides images from a list.
+    """
+
+    _EMPTY_IMAGE = torch.empty((0, 3, 1, 1))
+
+    def __init__(
+        self,
+        image_list: List[str],
+        category_list: Union[str, List[str], None] = None,
+        transform: Optional[ImageTransform] = None,
+    ):
+        """
+        Args:
+            image_list (List[str]): list of paths to image files
+            category_list (Union[str, List[str], None]): list of animal categories for
+                each image. If it is a string, or None, this applies to all images
+        """
+        if type(category_list) == list:
+            self.category_list = category_list
+        else:
+            self.category_list = [category_list] * len(image_list)
+        assert len(image_list) == len(
+            self.category_list
+        ), "length of image and category lists must be equal"
+        self.image_list = image_list
+        self.transform = transform
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """
+        Gets selected images from the list
+
+        Args:
+            idx (int): video index in the video list file
+        Returns:
+            A dictionary containing two keys:
+                images (torch.Tensor): tensor of size [N, 3, H, W] (N = 1, or 0 for _EMPTY_IMAGE)
+                categories (List[str]): categories of the frames
+        """
+        categories = [self.category_list[idx]]
+        fpath = self.image_list[idx]
+        transform = self.transform
+
+        try:
+            image = torch.from_numpy(np.ascontiguousarray(read_image(fpath, format="BGR")))
+            image = image.permute(2, 0, 1).unsqueeze(0).float()  # HWC -> NCHW
+            if transform is not None:
+                image = transform(image)
+            return {"images": image, "categories": categories}
+        except (OSError, RuntimeError) as e:
+            logger = logging.getLogger(__name__)
+            logger.warning(f"Error opening image file container {fpath}: {e}")
+
+        return {"images": self._EMPTY_IMAGE, "categories": []}
+
+    def __len__(self):
+        return len(self.image_list)

densepose/data/inference_based_loader.py

@@ -1,172 +1,172 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import random
-from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple
-import torch
-from torch import nn
-
-SampledData = Any
-ModelOutput = Any
-
-
-def _grouper(iterable: Iterable[Any], n: int, fillvalue=None) -> Iterator[Tuple[Any]]:
-    """
-    Group elements of an iterable by chunks of size `n`, e.g.
-    grouper(range(9), 4) ->
-        (0, 1, 2, 3), (4, 5, 6, 7), (8, None, None, None)
-    """
-    it = iter(iterable)
-    while True:
-        values = []
-        for _ in range(n):
-            try:
-                value = next(it)
-            except StopIteration:
-                if values:
-                    values.extend([fillvalue] * (n - len(values)))
-                    yield tuple(values)
-                return
-            values.append(value)
-        yield tuple(values)
-
-
-class ScoreBasedFilter:
-    """
-    Filters entries in model output based on their scores
-    Discards all entries with score less than the specified minimum
-    """
-
-    def __init__(self, min_score: float = 0.8):
-        self.min_score = min_score
-
-    def __call__(self, model_output: ModelOutput) -> ModelOutput:
-        for model_output_i in model_output:
-            instances = model_output_i["instances"]
-            if not instances.has("scores"):
-                continue
-            instances_filtered = instances[instances.scores >= self.min_score]
-            model_output_i["instances"] = instances_filtered
-        return model_output
-
-
-class InferenceBasedLoader:
-    """
-    Data loader based on results inferred by a model. Consists of:
-     - a data loader that provides batches of images
-     - a model that is used to infer the results
-     - a data sampler that converts inferred results to annotations
-    """
-
-    def __init__(
-        self,
-        model: nn.Module,
-        data_loader: Iterable[List[Dict[str, Any]]],
-        data_sampler: Optional[Callable[[ModelOutput], List[SampledData]]] = None,
-        data_filter: Optional[Callable[[ModelOutput], ModelOutput]] = None,
-        shuffle: bool = True,
-        batch_size: int = 4,
-        inference_batch_size: int = 4,
-        drop_last: bool = False,
-        category_to_class_mapping: Optional[dict] = None,
-    ):
-        """
-        Constructor
-
-        Args:
-          model (torch.nn.Module): model used to produce data
-          data_loader (Iterable[List[Dict[str, Any]]]): iterable that provides
-            dictionaries with "images" and "categories" fields to perform inference on
-          data_sampler (Callable: ModelOutput -> SampledData): functor
-              that produces annotation data from inference results;
-              (optional, default: None)
-          data_filter (Callable: ModelOutput -> ModelOutput): filter
-              that selects model outputs for further processing
-              (optional, default: None)
-          shuffle (bool): if True, the input images get shuffled
-          batch_size (int): batch size for the produced annotation data
-          inference_batch_size (int): batch size for input images
-          drop_last (bool): if True, drop the last batch if it is undersized
-          category_to_class_mapping (dict): category to class mapping
-        """
-        self.model = model
-        self.model.eval()
-        self.data_loader = data_loader
-        self.data_sampler = data_sampler
-        self.data_filter = data_filter
-        self.shuffle = shuffle
-        self.batch_size = batch_size
-        self.inference_batch_size = inference_batch_size
-        self.drop_last = drop_last
-        if category_to_class_mapping is not None:
-            self.category_to_class_mapping = category_to_class_mapping
-        else:
-            self.category_to_class_mapping = {}
-
-    def __iter__(self) -> Iterator[List[SampledData]]:
-        for batch in self.data_loader:
-            # batch : List[Dict[str: Tensor[N, C, H, W], str: Optional[str]]]
-            # images_batch : Tensor[N, C, H, W]
-            # image : Tensor[C, H, W]
-            images_and_categories = [
-                {"image": image, "category": category}
-                for element in batch
-                for image, category in zip(element["images"], element["categories"])
-            ]
-            if not images_and_categories:
-                continue
-            if self.shuffle:
-                random.shuffle(images_and_categories)
-            yield from self._produce_data(images_and_categories)  # pyre-ignore[6]
-
-    def _produce_data(
-        self, images_and_categories: List[Tuple[torch.Tensor, Optional[str]]]
-    ) -> Iterator[List[SampledData]]:
-        """
-        Produce batches of data from images
-
-        Args:
-          images_and_categories (List[Tuple[torch.Tensor, Optional[str]]]):
-            list of images and corresponding categories to process
-
-        Returns:
-          Iterator over batches of data sampled from model outputs
-        """
-        data_batches: List[SampledData] = []
-        category_to_class_mapping = self.category_to_class_mapping
-        batched_images_and_categories = _grouper(images_and_categories, self.inference_batch_size)
-        for batch in batched_images_and_categories:
-            batch = [
-                {
-                    "image": image_and_category["image"].to(self.model.device),
-                    "category": image_and_category["category"],
-                }
-                for image_and_category in batch
-                if image_and_category is not None
-            ]
-            if not batch:
-                continue
-            with torch.no_grad():
-                model_output = self.model(batch)
-            for model_output_i, batch_i in zip(model_output, batch):
-                assert len(batch_i["image"].shape) == 3
-                model_output_i["image"] = batch_i["image"]
-                instance_class = category_to_class_mapping.get(batch_i["category"], 0)
-                model_output_i["instances"].dataset_classes = torch.tensor(
-                    [instance_class] * len(model_output_i["instances"])
-                )
-            model_output_filtered = (
-                model_output if self.data_filter is None else self.data_filter(model_output)
-            )
-            data = (
-                model_output_filtered
-                if self.data_sampler is None
-                else self.data_sampler(model_output_filtered)
-            )
-            for data_i in data:
-                if len(data_i["instances"]):
-                    data_batches.append(data_i)
-            if len(data_batches) >= self.batch_size:
-                yield data_batches[: self.batch_size]
-                data_batches = data_batches[self.batch_size :]
-        if not self.drop_last and data_batches:
-            yield data_batches
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import random
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple
+import torch
+from torch import nn
+
+SampledData = Any
+ModelOutput = Any
+
+
+def _grouper(iterable: Iterable[Any], n: int, fillvalue=None) -> Iterator[Tuple[Any]]:
+    """
+    Group elements of an iterable by chunks of size `n`, e.g.
+    grouper(range(9), 4) ->
+        (0, 1, 2, 3), (4, 5, 6, 7), (8, None, None, None)
+    """
+    it = iter(iterable)
+    while True:
+        values = []
+        for _ in range(n):
+            try:
+                value = next(it)
+            except StopIteration:
+                if values:
+                    values.extend([fillvalue] * (n - len(values)))
+                    yield tuple(values)
+                return
+            values.append(value)
+        yield tuple(values)
+
+
+class ScoreBasedFilter:
+    """
+    Filters entries in model output based on their scores
+    Discards all entries with score less than the specified minimum
+    """
+
+    def __init__(self, min_score: float = 0.8):
+        self.min_score = min_score
+
+    def __call__(self, model_output: ModelOutput) -> ModelOutput:
+        for model_output_i in model_output:
+            instances = model_output_i["instances"]
+            if not instances.has("scores"):
+                continue
+            instances_filtered = instances[instances.scores >= self.min_score]
+            model_output_i["instances"] = instances_filtered
+        return model_output
+
+
+class InferenceBasedLoader:
+    """
+    Data loader based on results inferred by a model. Consists of:
+     - a data loader that provides batches of images
+     - a model that is used to infer the results
+     - a data sampler that converts inferred results to annotations
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        data_loader: Iterable[List[Dict[str, Any]]],
+        data_sampler: Optional[Callable[[ModelOutput], List[SampledData]]] = None,
+        data_filter: Optional[Callable[[ModelOutput], ModelOutput]] = None,
+        shuffle: bool = True,
+        batch_size: int = 4,
+        inference_batch_size: int = 4,
+        drop_last: bool = False,
+        category_to_class_mapping: Optional[dict] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          model (torch.nn.Module): model used to produce data
+          data_loader (Iterable[List[Dict[str, Any]]]): iterable that provides
+            dictionaries with "images" and "categories" fields to perform inference on
+          data_sampler (Callable: ModelOutput -> SampledData): functor
+              that produces annotation data from inference results;
+              (optional, default: None)
+          data_filter (Callable: ModelOutput -> ModelOutput): filter
+              that selects model outputs for further processing
+              (optional, default: None)
+          shuffle (bool): if True, the input images get shuffled
+          batch_size (int): batch size for the produced annotation data
+          inference_batch_size (int): batch size for input images
+          drop_last (bool): if True, drop the last batch if it is undersized
+          category_to_class_mapping (dict): category to class mapping
+        """
+        self.model = model
+        self.model.eval()
+        self.data_loader = data_loader
+        self.data_sampler = data_sampler
+        self.data_filter = data_filter
+        self.shuffle = shuffle
+        self.batch_size = batch_size
+        self.inference_batch_size = inference_batch_size
+        self.drop_last = drop_last
+        if category_to_class_mapping is not None:
+            self.category_to_class_mapping = category_to_class_mapping
+        else:
+            self.category_to_class_mapping = {}
+
+    def __iter__(self) -> Iterator[List[SampledData]]:
+        for batch in self.data_loader:
+            # batch : List[Dict[str: Tensor[N, C, H, W], str: Optional[str]]]
+            # images_batch : Tensor[N, C, H, W]
+            # image : Tensor[C, H, W]
+            images_and_categories = [
+                {"image": image, "category": category}
+                for element in batch
+                for image, category in zip(element["images"], element["categories"])
+            ]
+            if not images_and_categories:
+                continue
+            if self.shuffle:
+                random.shuffle(images_and_categories)
+            yield from self._produce_data(images_and_categories)  # pyre-ignore[6]
+
+    def _produce_data(
+        self, images_and_categories: List[Tuple[torch.Tensor, Optional[str]]]
+    ) -> Iterator[List[SampledData]]:
+        """
+        Produce batches of data from images
+
+        Args:
+          images_and_categories (List[Tuple[torch.Tensor, Optional[str]]]):
+            list of images and corresponding categories to process
+
+        Returns:
+          Iterator over batches of data sampled from model outputs
+        """
+        data_batches: List[SampledData] = []
+        category_to_class_mapping = self.category_to_class_mapping
+        batched_images_and_categories = _grouper(images_and_categories, self.inference_batch_size)
+        for batch in batched_images_and_categories:
+            batch = [
+                {
+                    "image": image_and_category["image"].to(self.model.device),
+                    "category": image_and_category["category"],
+                }
+                for image_and_category in batch
+                if image_and_category is not None
+            ]
+            if not batch:
+                continue
+            with torch.no_grad():
+                model_output = self.model(batch)
+            for model_output_i, batch_i in zip(model_output, batch):
+                assert len(batch_i["image"].shape) == 3
+                model_output_i["image"] = batch_i["image"]
+                instance_class = category_to_class_mapping.get(batch_i["category"], 0)
+                model_output_i["instances"].dataset_classes = torch.tensor(
+                    [instance_class] * len(model_output_i["instances"])
+                )
+            model_output_filtered = (
+                model_output if self.data_filter is None else self.data_filter(model_output)
+            )
+            data = (
+                model_output_filtered
+                if self.data_sampler is None
+                else self.data_sampler(model_output_filtered)
+            )
+            for data_i in data:
+                if len(data_i["instances"]):
+                    data_batches.append(data_i)
+            if len(data_batches) >= self.batch_size:
+                yield data_batches[: self.batch_size]
+                data_batches = data_batches[self.batch_size :]
+        if not self.drop_last and data_batches:
+            yield data_batches

densepose/data/meshes/__init__.py

@@ -1,5 +1,5 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-from . import builtin
-
-__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from . import builtin
+
+__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]

densepose/data/meshes/builtin.py

@@ -1,101 +1,101 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-from .catalog import MeshInfo, register_meshes
-
-DENSEPOSE_MESHES_DIR = "https://dl.fbaipublicfiles.com/densepose/meshes/"
-
-MESHES = [
-    MeshInfo(
-        name="smpl_27554",
-        data="smpl_27554.pkl",
-        geodists="geodists/geodists_smpl_27554.pkl",
-        symmetry="symmetry/symmetry_smpl_27554.pkl",
-        texcoords="texcoords/texcoords_smpl_27554.pkl",
-    ),
-    MeshInfo(
-        name="chimp_5029",
-        data="chimp_5029.pkl",
-        geodists="geodists/geodists_chimp_5029.pkl",
-        symmetry="symmetry/symmetry_chimp_5029.pkl",
-        texcoords="texcoords/texcoords_chimp_5029.pkl",
-    ),
-    MeshInfo(
-        name="cat_5001",
-        data="cat_5001.pkl",
-        geodists="geodists/geodists_cat_5001.pkl",
-        symmetry="symmetry/symmetry_cat_5001.pkl",
-        texcoords="texcoords/texcoords_cat_5001.pkl",
-    ),
-    MeshInfo(
-        name="cat_7466",
-        data="cat_7466.pkl",
-        geodists="geodists/geodists_cat_7466.pkl",
-        symmetry="symmetry/symmetry_cat_7466.pkl",
-        texcoords="texcoords/texcoords_cat_7466.pkl",
-    ),
-    MeshInfo(
-        name="sheep_5004",
-        data="sheep_5004.pkl",
-        geodists="geodists/geodists_sheep_5004.pkl",
-        symmetry="symmetry/symmetry_sheep_5004.pkl",
-        texcoords="texcoords/texcoords_sheep_5004.pkl",
-    ),
-    MeshInfo(
-        name="zebra_5002",
-        data="zebra_5002.pkl",
-        geodists="geodists/geodists_zebra_5002.pkl",
-        symmetry="symmetry/symmetry_zebra_5002.pkl",
-        texcoords="texcoords/texcoords_zebra_5002.pkl",
-    ),
-    MeshInfo(
-        name="horse_5004",
-        data="horse_5004.pkl",
-        geodists="geodists/geodists_horse_5004.pkl",
-        symmetry="symmetry/symmetry_horse_5004.pkl",
-        texcoords="texcoords/texcoords_zebra_5002.pkl",
-    ),
-    MeshInfo(
-        name="giraffe_5002",
-        data="giraffe_5002.pkl",
-        geodists="geodists/geodists_giraffe_5002.pkl",
-        symmetry="symmetry/symmetry_giraffe_5002.pkl",
-        texcoords="texcoords/texcoords_giraffe_5002.pkl",
-    ),
-    MeshInfo(
-        name="elephant_5002",
-        data="elephant_5002.pkl",
-        geodists="geodists/geodists_elephant_5002.pkl",
-        symmetry="symmetry/symmetry_elephant_5002.pkl",
-        texcoords="texcoords/texcoords_elephant_5002.pkl",
-    ),
-    MeshInfo(
-        name="dog_5002",
-        data="dog_5002.pkl",
-        geodists="geodists/geodists_dog_5002.pkl",
-        symmetry="symmetry/symmetry_dog_5002.pkl",
-        texcoords="texcoords/texcoords_dog_5002.pkl",
-    ),
-    MeshInfo(
-        name="dog_7466",
-        data="dog_7466.pkl",
-        geodists="geodists/geodists_dog_7466.pkl",
-        symmetry="symmetry/symmetry_dog_7466.pkl",
-        texcoords="texcoords/texcoords_dog_7466.pkl",
-    ),
-    MeshInfo(
-        name="cow_5002",
-        data="cow_5002.pkl",
-        geodists="geodists/geodists_cow_5002.pkl",
-        symmetry="symmetry/symmetry_cow_5002.pkl",
-        texcoords="texcoords/texcoords_cow_5002.pkl",
-    ),
-    MeshInfo(
-        name="bear_4936",
-        data="bear_4936.pkl",
-        geodists="geodists/geodists_bear_4936.pkl",
-        symmetry="symmetry/symmetry_bear_4936.pkl",
-        texcoords="texcoords/texcoords_bear_4936.pkl",
-    ),
-]
-
-register_meshes(MESHES, DENSEPOSE_MESHES_DIR)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from .catalog import MeshInfo, register_meshes
+
+DENSEPOSE_MESHES_DIR = "https://dl.fbaipublicfiles.com/densepose/meshes/"
+
+MESHES = [
+    MeshInfo(
+        name="smpl_27554",
+        data="smpl_27554.pkl",
+        geodists="geodists/geodists_smpl_27554.pkl",
+        symmetry="symmetry/symmetry_smpl_27554.pkl",
+        texcoords="texcoords/texcoords_smpl_27554.pkl",
+    ),
+    MeshInfo(
+        name="chimp_5029",
+        data="chimp_5029.pkl",
+        geodists="geodists/geodists_chimp_5029.pkl",
+        symmetry="symmetry/symmetry_chimp_5029.pkl",
+        texcoords="texcoords/texcoords_chimp_5029.pkl",
+    ),
+    MeshInfo(
+        name="cat_5001",
+        data="cat_5001.pkl",
+        geodists="geodists/geodists_cat_5001.pkl",
+        symmetry="symmetry/symmetry_cat_5001.pkl",
+        texcoords="texcoords/texcoords_cat_5001.pkl",
+    ),
+    MeshInfo(
+        name="cat_7466",
+        data="cat_7466.pkl",
+        geodists="geodists/geodists_cat_7466.pkl",
+        symmetry="symmetry/symmetry_cat_7466.pkl",
+        texcoords="texcoords/texcoords_cat_7466.pkl",
+    ),
+    MeshInfo(
+        name="sheep_5004",
+        data="sheep_5004.pkl",
+        geodists="geodists/geodists_sheep_5004.pkl",
+        symmetry="symmetry/symmetry_sheep_5004.pkl",
+        texcoords="texcoords/texcoords_sheep_5004.pkl",
+    ),
+    MeshInfo(
+        name="zebra_5002",
+        data="zebra_5002.pkl",
+        geodists="geodists/geodists_zebra_5002.pkl",
+        symmetry="symmetry/symmetry_zebra_5002.pkl",
+        texcoords="texcoords/texcoords_zebra_5002.pkl",
+    ),
+    MeshInfo(
+        name="horse_5004",
+        data="horse_5004.pkl",
+        geodists="geodists/geodists_horse_5004.pkl",
+        symmetry="symmetry/symmetry_horse_5004.pkl",
+        texcoords="texcoords/texcoords_zebra_5002.pkl",
+    ),
+    MeshInfo(
+        name="giraffe_5002",
+        data="giraffe_5002.pkl",
+        geodists="geodists/geodists_giraffe_5002.pkl",
+        symmetry="symmetry/symmetry_giraffe_5002.pkl",
+        texcoords="texcoords/texcoords_giraffe_5002.pkl",
+    ),
+    MeshInfo(
+        name="elephant_5002",
+        data="elephant_5002.pkl",
+        geodists="geodists/geodists_elephant_5002.pkl",
+        symmetry="symmetry/symmetry_elephant_5002.pkl",
+        texcoords="texcoords/texcoords_elephant_5002.pkl",
+    ),
+    MeshInfo(
+        name="dog_5002",
+        data="dog_5002.pkl",
+        geodists="geodists/geodists_dog_5002.pkl",
+        symmetry="symmetry/symmetry_dog_5002.pkl",
+        texcoords="texcoords/texcoords_dog_5002.pkl",
+    ),
+    MeshInfo(
+        name="dog_7466",
+        data="dog_7466.pkl",
+        geodists="geodists/geodists_dog_7466.pkl",
+        symmetry="symmetry/symmetry_dog_7466.pkl",
+        texcoords="texcoords/texcoords_dog_7466.pkl",
+    ),
+    MeshInfo(
+        name="cow_5002",
+        data="cow_5002.pkl",
+        geodists="geodists/geodists_cow_5002.pkl",
+        symmetry="symmetry/symmetry_cow_5002.pkl",
+        texcoords="texcoords/texcoords_cow_5002.pkl",
+    ),
+    MeshInfo(
+        name="bear_4936",
+        data="bear_4936.pkl",
+        geodists="geodists/geodists_bear_4936.pkl",
+        symmetry="symmetry/symmetry_bear_4936.pkl",
+        texcoords="texcoords/texcoords_bear_4936.pkl",
+    ),
+]
+
+register_meshes(MESHES, DENSEPOSE_MESHES_DIR)

densepose/data/meshes/catalog.py

@@ -1,71 +1,71 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-import logging
-from collections import UserDict
-from dataclasses import dataclass
-from typing import Iterable, Optional
-
-from ..utils import maybe_prepend_base_path
-
-
-@dataclass
-class MeshInfo:
-    name: str
-    data: str
-    geodists: Optional[str] = None
-    symmetry: Optional[str] = None
-    texcoords: Optional[str] = None
-
-
-class _MeshCatalog(UserDict):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.mesh_ids = {}
-        self.mesh_names = {}
-        self.max_mesh_id = -1
-
-    def __setitem__(self, key, value):
-        if key in self:
-            logger = logging.getLogger(__name__)
-            logger.warning(
-                f"Overwriting mesh catalog entry '{key}': old value {self[key]}"
-                f", new value {value}"
-            )
-            mesh_id = self.mesh_ids[key]
-        else:
-            self.max_mesh_id += 1
-            mesh_id = self.max_mesh_id
-        super().__setitem__(key, value)
-        self.mesh_ids[key] = mesh_id
-        self.mesh_names[mesh_id] = key
-
-    def get_mesh_id(self, shape_name: str) -> int:
-        return self.mesh_ids[shape_name]
-
-    def get_mesh_name(self, mesh_id: int) -> str:
-        return self.mesh_names[mesh_id]
-
-
-MeshCatalog = _MeshCatalog()
-
-
-def register_mesh(mesh_info: MeshInfo, base_path: Optional[str]) -> None:
-    geodists, symmetry, texcoords = mesh_info.geodists, mesh_info.symmetry, mesh_info.texcoords
-    if geodists:
-        geodists = maybe_prepend_base_path(base_path, geodists)
-    if symmetry:
-        symmetry = maybe_prepend_base_path(base_path, symmetry)
-    if texcoords:
-        texcoords = maybe_prepend_base_path(base_path, texcoords)
-    MeshCatalog[mesh_info.name] = MeshInfo(
-        name=mesh_info.name,
-        data=maybe_prepend_base_path(base_path, mesh_info.data),
-        geodists=geodists,
-        symmetry=symmetry,
-        texcoords=texcoords,
-    )
-
-
-def register_meshes(mesh_infos: Iterable[MeshInfo], base_path: Optional[str]) -> None:
-    for mesh_info in mesh_infos:
-        register_mesh(mesh_info, base_path)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import logging
+from collections import UserDict
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+from ..utils import maybe_prepend_base_path
+
+
+@dataclass
+class MeshInfo:
+    name: str
+    data: str
+    geodists: Optional[str] = None
+    symmetry: Optional[str] = None
+    texcoords: Optional[str] = None
+
+
+class _MeshCatalog(UserDict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mesh_ids = {}
+        self.mesh_names = {}
+        self.max_mesh_id = -1
+
+    def __setitem__(self, key, value):
+        if key in self:
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                f"Overwriting mesh catalog entry '{key}': old value {self[key]}"
+                f", new value {value}"
+            )
+            mesh_id = self.mesh_ids[key]
+        else:
+            self.max_mesh_id += 1
+            mesh_id = self.max_mesh_id
+        super().__setitem__(key, value)
+        self.mesh_ids[key] = mesh_id
+        self.mesh_names[mesh_id] = key
+
+    def get_mesh_id(self, shape_name: str) -> int:
+        return self.mesh_ids[shape_name]
+
+    def get_mesh_name(self, mesh_id: int) -> str:
+        return self.mesh_names[mesh_id]
+
+
+MeshCatalog = _MeshCatalog()
+
+
+def register_mesh(mesh_info: MeshInfo, base_path: Optional[str]) -> None:
+    geodists, symmetry, texcoords = mesh_info.geodists, mesh_info.symmetry, mesh_info.texcoords
+    if geodists:
+        geodists = maybe_prepend_base_path(base_path, geodists)
+    if symmetry:
+        symmetry = maybe_prepend_base_path(base_path, symmetry)
+    if texcoords:
+        texcoords = maybe_prepend_base_path(base_path, texcoords)
+    MeshCatalog[mesh_info.name] = MeshInfo(
+        name=mesh_info.name,
+        data=maybe_prepend_base_path(base_path, mesh_info.data),
+        geodists=geodists,
+        symmetry=symmetry,
+        texcoords=texcoords,
+    )
+
+
+def register_meshes(mesh_infos: Iterable[MeshInfo], base_path: Optional[str]) -> None:
+    for mesh_info in mesh_infos:
+        register_mesh(mesh_info, base_path)

densepose/data/samplers/__init__.py

@@ -1,8 +1,8 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .densepose_uniform import DensePoseUniformSampler
-from .densepose_confidence_based import DensePoseConfidenceBasedSampler
-from .densepose_cse_uniform import DensePoseCSEUniformSampler
-from .densepose_cse_confidence_based import DensePoseCSEConfidenceBasedSampler
-from .mask_from_densepose import MaskFromDensePoseSampler
-from .prediction_to_gt import PredictionToGroundTruthSampler
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .densepose_uniform import DensePoseUniformSampler
+from .densepose_confidence_based import DensePoseConfidenceBasedSampler
+from .densepose_cse_uniform import DensePoseCSEUniformSampler
+from .densepose_cse_confidence_based import DensePoseCSEConfidenceBasedSampler
+from .mask_from_densepose import MaskFromDensePoseSampler
+from .prediction_to_gt import PredictionToGroundTruthSampler

densepose/data/samplers/densepose_base.py

@@ -1,203 +1,203 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Any, Dict, List, Tuple
-import torch
-from torch.nn import functional as F
-
-from detectron2.structures import BoxMode, Instances
-
-from densepose.converters import ToChartResultConverter
-from densepose.converters.base import IntTupleBox, make_int_box
-from densepose.structures import DensePoseDataRelative, DensePoseList
-
-
-class DensePoseBaseSampler:
-    """
-    Base DensePose sampler to produce DensePose data from DensePose predictions.
-    Samples for each class are drawn according to some distribution over all pixels estimated
-    to belong to that class.
-    """
-
-    def __init__(self, count_per_class: int = 8):
-        """
-        Constructor
-
-        Args:
-          count_per_class (int): the sampler produces at most `count_per_class`
-              samples for each category
-        """
-        self.count_per_class = count_per_class
-
-    def __call__(self, instances: Instances) -> DensePoseList:
-        """
-        Convert DensePose predictions (an instance of `DensePoseChartPredictorOutput`)
-        into DensePose annotations data (an instance of `DensePoseList`)
-        """
-        boxes_xyxy_abs = instances.pred_boxes.tensor.clone().cpu()
-        boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
-        dp_datas = []
-        for i in range(len(boxes_xywh_abs)):
-            annotation_i = self._sample(instances[i], make_int_box(boxes_xywh_abs[i]))
-            annotation_i[DensePoseDataRelative.S_KEY] = self._resample_mask(  # pyre-ignore[6]
-                instances[i].pred_densepose
-            )
-            dp_datas.append(DensePoseDataRelative(annotation_i))
-        # create densepose annotations on CPU
-        dp_list = DensePoseList(dp_datas, boxes_xyxy_abs, instances.image_size)
-        return dp_list
-
-    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
-        """
-        Sample DensPoseDataRelative from estimation results
-        """
-        labels, dp_result = self._produce_labels_and_results(instance)
-        annotation = {
-            DensePoseDataRelative.X_KEY: [],
-            DensePoseDataRelative.Y_KEY: [],
-            DensePoseDataRelative.U_KEY: [],
-            DensePoseDataRelative.V_KEY: [],
-            DensePoseDataRelative.I_KEY: [],
-        }
-        n, h, w = dp_result.shape
-        for part_id in range(1, DensePoseDataRelative.N_PART_LABELS + 1):
-            # indices - tuple of 3 1D tensors of size k
-            # 0: index along the first dimension N
-            # 1: index along H dimension
-            # 2: index along W dimension
-            indices = torch.nonzero(labels.expand(n, h, w) == part_id, as_tuple=True)
-            # values - an array of size [n, k]
-            # n: number of channels (U, V, confidences)
-            # k: number of points labeled with part_id
-            values = dp_result[indices].view(n, -1)
-            k = values.shape[1]
-            count = min(self.count_per_class, k)
-            if count <= 0:
-                continue
-            index_sample = self._produce_index_sample(values, count)
-            sampled_values = values[:, index_sample]
-            sampled_y = indices[1][index_sample] + 0.5
-            sampled_x = indices[2][index_sample] + 0.5
-            # prepare / normalize data
-            x = (sampled_x / w * 256.0).cpu().tolist()
-            y = (sampled_y / h * 256.0).cpu().tolist()
-            u = sampled_values[0].clamp(0, 1).cpu().tolist()
-            v = sampled_values[1].clamp(0, 1).cpu().tolist()
-            fine_segm_labels = [part_id] * count
-            # extend annotations
-            annotation[DensePoseDataRelative.X_KEY].extend(x)
-            annotation[DensePoseDataRelative.Y_KEY].extend(y)
-            annotation[DensePoseDataRelative.U_KEY].extend(u)
-            annotation[DensePoseDataRelative.V_KEY].extend(v)
-            annotation[DensePoseDataRelative.I_KEY].extend(fine_segm_labels)
-        return annotation
-
-    def _produce_index_sample(self, values: torch.Tensor, count: int):
-        """
-        Abstract method to produce a sample of indices to select data
-        To be implemented in descendants
-
-        Args:
-            values (torch.Tensor): an array of size [n, k] that contains
-                estimated values (U, V, confidences);
-                n: number of channels (U, V, confidences)
-                k: number of points labeled with part_id
-            count (int): number of samples to produce, should be positive and <= k
-
-        Return:
-            list(int): indices of values (along axis 1) selected as a sample
-        """
-        raise NotImplementedError
-
-    def _produce_labels_and_results(self, instance: Instances) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Method to get labels and DensePose results from an instance
-
-        Args:
-            instance (Instances): an instance of `DensePoseChartPredictorOutput`
-
-        Return:
-            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
-            dp_result (torch.Tensor): shape [2, H, W], stacked DensePose results u and v
-        """
-        converter = ToChartResultConverter
-        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
-        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
-        return labels, dp_result
-
-    def _resample_mask(self, output: Any) -> torch.Tensor:
-        """
-        Convert DensePose predictor output to segmentation annotation - tensors of size
-        (256, 256) and type `int64`.
-
-        Args:
-            output: DensePose predictor output with the following attributes:
-             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
-               segmentation scores
-             - fine_segm: tensor of size [N, C, H, W] with unnormalized fine
-               segmentation scores
-        Return:
-            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
-            where S = DensePoseDataRelative.MASK_SIZE
-        """
-        sz = DensePoseDataRelative.MASK_SIZE
-        S = (
-            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
-            .argmax(dim=1)
-            .long()
-        )
-        I = (
-            (
-                F.interpolate(
-                    output.fine_segm,
-                    (sz, sz),
-                    mode="bilinear",
-                    align_corners=False,
-                ).argmax(dim=1)
-                * (S > 0).long()
-            )
-            .squeeze()
-            .cpu()
-        )
-        # Map fine segmentation results to coarse segmentation ground truth
-        # TODO: extract this into separate classes
-        # coarse segmentation: 1 = Torso, 2 = Right Hand, 3 = Left Hand,
-        # 4 = Left Foot, 5 = Right Foot, 6 = Upper Leg Right, 7 = Upper Leg Left,
-        # 8 = Lower Leg Right, 9 = Lower Leg Left, 10 = Upper Arm Left,
-        # 11 = Upper Arm Right, 12 = Lower Arm Left, 13 = Lower Arm Right,
-        # 14 = Head
-        # fine segmentation: 1, 2 = Torso, 3 = Right Hand, 4 = Left Hand,
-        # 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right,
-        # 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right,
-        # 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left,
-        # 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left,
-        # 20, 22 = Lower Arm Right, 23, 24 = Head
-        FINE_TO_COARSE_SEGMENTATION = {
-            1: 1,
-            2: 1,
-            3: 2,
-            4: 3,
-            5: 4,
-            6: 5,
-            7: 6,
-            8: 7,
-            9: 6,
-            10: 7,
-            11: 8,
-            12: 9,
-            13: 8,
-            14: 9,
-            15: 10,
-            16: 11,
-            17: 10,
-            18: 11,
-            19: 12,
-            20: 13,
-            21: 12,
-            22: 13,
-            23: 14,
-            24: 14,
-        }
-        mask = torch.zeros((sz, sz), dtype=torch.int64, device=torch.device("cpu"))
-        for i in range(DensePoseDataRelative.N_PART_LABELS):
-            mask[I == i + 1] = FINE_TO_COARSE_SEGMENTATION[i + 1]
-        return mask
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Any, Dict, List, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures import BoxMode, Instances
+
+from densepose.converters import ToChartResultConverter
+from densepose.converters.base import IntTupleBox, make_int_box
+from densepose.structures import DensePoseDataRelative, DensePoseList
+
+
+class DensePoseBaseSampler:
+    """
+    Base DensePose sampler to produce DensePose data from DensePose predictions.
+    Samples for each class are drawn according to some distribution over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(self, count_per_class: int = 8):
+        """
+        Constructor
+
+        Args:
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        self.count_per_class = count_per_class
+
+    def __call__(self, instances: Instances) -> DensePoseList:
+        """
+        Convert DensePose predictions (an instance of `DensePoseChartPredictorOutput`)
+        into DensePose annotations data (an instance of `DensePoseList`)
+        """
+        boxes_xyxy_abs = instances.pred_boxes.tensor.clone().cpu()
+        boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+        dp_datas = []
+        for i in range(len(boxes_xywh_abs)):
+            annotation_i = self._sample(instances[i], make_int_box(boxes_xywh_abs[i]))
+            annotation_i[DensePoseDataRelative.S_KEY] = self._resample_mask(  # pyre-ignore[6]
+                instances[i].pred_densepose
+            )
+            dp_datas.append(DensePoseDataRelative(annotation_i))
+        # create densepose annotations on CPU
+        dp_list = DensePoseList(dp_datas, boxes_xyxy_abs, instances.image_size)
+        return dp_list
+
+    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
+        """
+        Sample DensPoseDataRelative from estimation results
+        """
+        labels, dp_result = self._produce_labels_and_results(instance)
+        annotation = {
+            DensePoseDataRelative.X_KEY: [],
+            DensePoseDataRelative.Y_KEY: [],
+            DensePoseDataRelative.U_KEY: [],
+            DensePoseDataRelative.V_KEY: [],
+            DensePoseDataRelative.I_KEY: [],
+        }
+        n, h, w = dp_result.shape
+        for part_id in range(1, DensePoseDataRelative.N_PART_LABELS + 1):
+            # indices - tuple of 3 1D tensors of size k
+            # 0: index along the first dimension N
+            # 1: index along H dimension
+            # 2: index along W dimension
+            indices = torch.nonzero(labels.expand(n, h, w) == part_id, as_tuple=True)
+            # values - an array of size [n, k]
+            # n: number of channels (U, V, confidences)
+            # k: number of points labeled with part_id
+            values = dp_result[indices].view(n, -1)
+            k = values.shape[1]
+            count = min(self.count_per_class, k)
+            if count <= 0:
+                continue
+            index_sample = self._produce_index_sample(values, count)
+            sampled_values = values[:, index_sample]
+            sampled_y = indices[1][index_sample] + 0.5
+            sampled_x = indices[2][index_sample] + 0.5
+            # prepare / normalize data
+            x = (sampled_x / w * 256.0).cpu().tolist()
+            y = (sampled_y / h * 256.0).cpu().tolist()
+            u = sampled_values[0].clamp(0, 1).cpu().tolist()
+            v = sampled_values[1].clamp(0, 1).cpu().tolist()
+            fine_segm_labels = [part_id] * count
+            # extend annotations
+            annotation[DensePoseDataRelative.X_KEY].extend(x)
+            annotation[DensePoseDataRelative.Y_KEY].extend(y)
+            annotation[DensePoseDataRelative.U_KEY].extend(u)
+            annotation[DensePoseDataRelative.V_KEY].extend(v)
+            annotation[DensePoseDataRelative.I_KEY].extend(fine_segm_labels)
+        return annotation
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Abstract method to produce a sample of indices to select data
+        To be implemented in descendants
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        raise NotImplementedError
+
+    def _produce_labels_and_results(self, instance: Instances) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of `DensePoseChartPredictorOutput`
+
+        Return:
+            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
+            dp_result (torch.Tensor): shape [2, H, W], stacked DensePose results u and v
+        """
+        converter = ToChartResultConverter
+        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
+        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
+        return labels, dp_result
+
+    def _resample_mask(self, output: Any) -> torch.Tensor:
+        """
+        Convert DensePose predictor output to segmentation annotation - tensors of size
+        (256, 256) and type `int64`.
+
+        Args:
+            output: DensePose predictor output with the following attributes:
+             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
+               segmentation scores
+             - fine_segm: tensor of size [N, C, H, W] with unnormalized fine
+               segmentation scores
+        Return:
+            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
+            where S = DensePoseDataRelative.MASK_SIZE
+        """
+        sz = DensePoseDataRelative.MASK_SIZE
+        S = (
+            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
+            .argmax(dim=1)
+            .long()
+        )
+        I = (
+            (
+                F.interpolate(
+                    output.fine_segm,
+                    (sz, sz),
+                    mode="bilinear",
+                    align_corners=False,
+                ).argmax(dim=1)
+                * (S > 0).long()
+            )
+            .squeeze()
+            .cpu()
+        )
+        # Map fine segmentation results to coarse segmentation ground truth
+        # TODO: extract this into separate classes
+        # coarse segmentation: 1 = Torso, 2 = Right Hand, 3 = Left Hand,
+        # 4 = Left Foot, 5 = Right Foot, 6 = Upper Leg Right, 7 = Upper Leg Left,
+        # 8 = Lower Leg Right, 9 = Lower Leg Left, 10 = Upper Arm Left,
+        # 11 = Upper Arm Right, 12 = Lower Arm Left, 13 = Lower Arm Right,
+        # 14 = Head
+        # fine segmentation: 1, 2 = Torso, 3 = Right Hand, 4 = Left Hand,
+        # 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right,
+        # 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right,
+        # 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left,
+        # 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left,
+        # 20, 22 = Lower Arm Right, 23, 24 = Head
+        FINE_TO_COARSE_SEGMENTATION = {
+            1: 1,
+            2: 1,
+            3: 2,
+            4: 3,
+            5: 4,
+            6: 5,
+            7: 6,
+            8: 7,
+            9: 6,
+            10: 7,
+            11: 8,
+            12: 9,
+            13: 8,
+            14: 9,
+            15: 10,
+            16: 11,
+            17: 10,
+            18: 11,
+            19: 12,
+            20: 13,
+            21: 12,
+            22: 13,
+            23: 14,
+            24: 14,
+        }
+        mask = torch.zeros((sz, sz), dtype=torch.int64, device=torch.device("cpu"))
+        for i in range(DensePoseDataRelative.N_PART_LABELS):
+            mask[I == i + 1] = FINE_TO_COARSE_SEGMENTATION[i + 1]
+        return mask

densepose/data/samplers/densepose_confidence_based.py

@@ -1,108 +1,108 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import random
-from typing import Optional, Tuple
-import torch
-
-from densepose.converters import ToChartResultConverterWithConfidences
-
-from .densepose_base import DensePoseBaseSampler
-
-
-class DensePoseConfidenceBasedSampler(DensePoseBaseSampler):
-    """
-    Samples DensePose data from DensePose predictions.
-    Samples for each class are drawn using confidence value estimates.
-    """
-
-    def __init__(
-        self,
-        confidence_channel: str,
-        count_per_class: int = 8,
-        search_count_multiplier: Optional[float] = None,
-        search_proportion: Optional[float] = None,
-    ):
-        """
-        Constructor
-
-        Args:
-          confidence_channel (str): confidence channel to use for sampling;
-            possible values:
-              "sigma_2": confidences for UV values
-              "fine_segm_confidence": confidences for fine segmentation
-              "coarse_segm_confidence": confidences for coarse segmentation
-            (default: "sigma_2")
-          count_per_class (int): the sampler produces at most `count_per_class`
-              samples for each category (default: 8)
-          search_count_multiplier (float or None): if not None, the total number
-              of the most confident estimates of a given class to consider is
-              defined as `min(search_count_multiplier * count_per_class, N)`,
-              where `N` is the total number of estimates of the class; cannot be
-              specified together with `search_proportion` (default: None)
-          search_proportion (float or None): if not None, the total number of the
-              of the most confident estimates of a given class to consider is
-              defined as `min(max(search_proportion * N, count_per_class), N)`,
-              where `N` is the total number of estimates of the class; cannot be
-              specified together with `search_count_multiplier` (default: None)
-        """
-        super().__init__(count_per_class)
-        self.confidence_channel = confidence_channel
-        self.search_count_multiplier = search_count_multiplier
-        self.search_proportion = search_proportion
-        assert (search_count_multiplier is None) or (search_proportion is None), (
-            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
-            f"and search_proportion (={search_proportion})"
-        )
-
-    def _produce_index_sample(self, values: torch.Tensor, count: int):
-        """
-        Produce a sample of indices to select data based on confidences
-
-        Args:
-            values (torch.Tensor): an array of size [n, k] that contains
-                estimated values (U, V, confidences);
-                n: number of channels (U, V, confidences)
-                k: number of points labeled with part_id
-            count (int): number of samples to produce, should be positive and <= k
-
-        Return:
-            list(int): indices of values (along axis 1) selected as a sample
-        """
-        k = values.shape[1]
-        if k == count:
-            index_sample = list(range(k))
-        else:
-            # take the best count * search_count_multiplier pixels,
-            # sample from them uniformly
-            # (here best = smallest variance)
-            _, sorted_confidence_indices = torch.sort(values[2])
-            if self.search_count_multiplier is not None:
-                search_count = min(int(count * self.search_count_multiplier), k)
-            elif self.search_proportion is not None:
-                search_count = min(max(int(k * self.search_proportion), count), k)
-            else:
-                search_count = min(count, k)
-            sample_from_top = random.sample(range(search_count), count)
-            index_sample = sorted_confidence_indices[:search_count][sample_from_top]
-        return index_sample
-
-    def _produce_labels_and_results(self, instance) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Method to get labels and DensePose results from an instance, with confidences
-
-        Args:
-            instance (Instances): an instance of `DensePoseChartPredictorOutputWithConfidences`
-
-        Return:
-            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
-            dp_result (torch.Tensor): shape [3, H, W], DensePose results u and v
-                stacked with the confidence channel
-        """
-        converter = ToChartResultConverterWithConfidences
-        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
-        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
-        dp_result = torch.cat(
-            (dp_result, getattr(chart_result, self.confidence_channel)[None].cpu())
-        )
-
-        return labels, dp_result
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import random
+from typing import Optional, Tuple
+import torch
+
+from densepose.converters import ToChartResultConverterWithConfidences
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseConfidenceBasedSampler(DensePoseBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn using confidence value estimates.
+    """
+
+    def __init__(
+        self,
+        confidence_channel: str,
+        count_per_class: int = 8,
+        search_count_multiplier: Optional[float] = None,
+        search_proportion: Optional[float] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          confidence_channel (str): confidence channel to use for sampling;
+            possible values:
+              "sigma_2": confidences for UV values
+              "fine_segm_confidence": confidences for fine segmentation
+              "coarse_segm_confidence": confidences for coarse segmentation
+            (default: "sigma_2")
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category (default: 8)
+          search_count_multiplier (float or None): if not None, the total number
+              of the most confident estimates of a given class to consider is
+              defined as `min(search_count_multiplier * count_per_class, N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_proportion` (default: None)
+          search_proportion (float or None): if not None, the total number of the
+              of the most confident estimates of a given class to consider is
+              defined as `min(max(search_proportion * N, count_per_class), N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_count_multiplier` (default: None)
+        """
+        super().__init__(count_per_class)
+        self.confidence_channel = confidence_channel
+        self.search_count_multiplier = search_count_multiplier
+        self.search_proportion = search_proportion
+        assert (search_count_multiplier is None) or (search_proportion is None), (
+            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
+            f"and search_proportion (={search_proportion})"
+        )
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a sample of indices to select data based on confidences
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        if k == count:
+            index_sample = list(range(k))
+        else:
+            # take the best count * search_count_multiplier pixels,
+            # sample from them uniformly
+            # (here best = smallest variance)
+            _, sorted_confidence_indices = torch.sort(values[2])
+            if self.search_count_multiplier is not None:
+                search_count = min(int(count * self.search_count_multiplier), k)
+            elif self.search_proportion is not None:
+                search_count = min(max(int(k * self.search_proportion), count), k)
+            else:
+                search_count = min(count, k)
+            sample_from_top = random.sample(range(search_count), count)
+            index_sample = sorted_confidence_indices[:search_count][sample_from_top]
+        return index_sample
+
+    def _produce_labels_and_results(self, instance) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance, with confidences
+
+        Args:
+            instance (Instances): an instance of `DensePoseChartPredictorOutputWithConfidences`
+
+        Return:
+            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
+            dp_result (torch.Tensor): shape [3, H, W], DensePose results u and v
+                stacked with the confidence channel
+        """
+        converter = ToChartResultConverterWithConfidences
+        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
+        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
+        dp_result = torch.cat(
+            (dp_result, getattr(chart_result, self.confidence_channel)[None].cpu())
+        )
+
+        return labels, dp_result

densepose/data/samplers/densepose_cse_base.py

@@ -1,139 +1,139 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from typing import Any, Dict, List, Tuple
-import torch
-from torch.nn import functional as F
-
-from detectron2.config import CfgNode
-from detectron2.structures import Instances
-
-from densepose.converters.base import IntTupleBox
-from densepose.data.utils import get_class_to_mesh_name_mapping
-from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
-from densepose.structures import DensePoseDataRelative
-
-from .densepose_base import DensePoseBaseSampler
-
-
-class DensePoseCSEBaseSampler(DensePoseBaseSampler):
-    """
-    Base DensePose sampler to produce DensePose data from DensePose predictions.
-    Samples for each class are drawn according to some distribution over all pixels estimated
-    to belong to that class.
-    """
-
-    def __init__(
-        self,
-        cfg: CfgNode,
-        use_gt_categories: bool,
-        embedder: torch.nn.Module,
-        count_per_class: int = 8,
-    ):
-        """
-        Constructor
-
-        Args:
-          cfg (CfgNode): the config of the model
-          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
-          count_per_class (int): the sampler produces at most `count_per_class`
-              samples for each category
-        """
-        super().__init__(count_per_class)
-        self.embedder = embedder
-        self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
-        self.use_gt_categories = use_gt_categories
-
-    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
-        """
-        Sample DensPoseDataRelative from estimation results
-        """
-        if self.use_gt_categories:
-            instance_class = instance.dataset_classes.tolist()[0]
-        else:
-            instance_class = instance.pred_classes.tolist()[0]
-        mesh_name = self.class_to_mesh_name[instance_class]
-
-        annotation = {
-            DensePoseDataRelative.X_KEY: [],
-            DensePoseDataRelative.Y_KEY: [],
-            DensePoseDataRelative.VERTEX_IDS_KEY: [],
-            DensePoseDataRelative.MESH_NAME_KEY: mesh_name,
-        }
-
-        mask, embeddings, other_values = self._produce_mask_and_results(instance, bbox_xywh)
-        indices = torch.nonzero(mask, as_tuple=True)
-        selected_embeddings = embeddings.permute(1, 2, 0)[indices].cpu()
-        values = other_values[:, indices[0], indices[1]]
-        k = values.shape[1]
-
-        count = min(self.count_per_class, k)
-        if count <= 0:
-            return annotation
-
-        index_sample = self._produce_index_sample(values, count)
-        closest_vertices = squared_euclidean_distance_matrix(
-            selected_embeddings[index_sample], self.embedder(mesh_name)
-        )
-        closest_vertices = torch.argmin(closest_vertices, dim=1)
-
-        sampled_y = indices[0][index_sample] + 0.5
-        sampled_x = indices[1][index_sample] + 0.5
-        # prepare / normalize data
-        _, _, w, h = bbox_xywh
-        x = (sampled_x / w * 256.0).cpu().tolist()
-        y = (sampled_y / h * 256.0).cpu().tolist()
-        # extend annotations
-        annotation[DensePoseDataRelative.X_KEY].extend(x)
-        annotation[DensePoseDataRelative.Y_KEY].extend(y)
-        annotation[DensePoseDataRelative.VERTEX_IDS_KEY].extend(closest_vertices.cpu().tolist())
-        return annotation
-
-    def _produce_mask_and_results(
-        self, instance: Instances, bbox_xywh: IntTupleBox
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Method to get labels and DensePose results from an instance
-
-        Args:
-            instance (Instances): an instance of `DensePoseEmbeddingPredictorOutput`
-            bbox_xywh (IntTupleBox): the corresponding bounding box
-
-        Return:
-            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
-            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W],
-                DensePose CSE Embeddings
-            other_values (Tuple[torch.Tensor]): a tensor of shape [0, H, W],
-                for potential other values
-        """
-        densepose_output = instance.pred_densepose
-        S = densepose_output.coarse_segm
-        E = densepose_output.embedding
-        _, _, w, h = bbox_xywh
-        embeddings = F.interpolate(E, size=(h, w), mode="bilinear")[0]
-        coarse_segm_resized = F.interpolate(S, size=(h, w), mode="bilinear")[0]
-        mask = coarse_segm_resized.argmax(0) > 0
-        other_values = torch.empty((0, h, w), device=E.device)
-        return mask, embeddings, other_values
-
-    def _resample_mask(self, output: Any) -> torch.Tensor:
-        """
-        Convert DensePose predictor output to segmentation annotation - tensors of size
-        (256, 256) and type `int64`.
-
-        Args:
-            output: DensePose predictor output with the following attributes:
-             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
-               segmentation scores
-        Return:
-            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
-            where S = DensePoseDataRelative.MASK_SIZE
-        """
-        sz = DensePoseDataRelative.MASK_SIZE
-        mask = (
-            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
-            .argmax(dim=1)
-            .long()
-            .squeeze()
-            .cpu()
-        )
-        return mask
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from typing import Any, Dict, List, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.config import CfgNode
+from detectron2.structures import Instances
+
+from densepose.converters.base import IntTupleBox
+from densepose.data.utils import get_class_to_mesh_name_mapping
+from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
+from densepose.structures import DensePoseDataRelative
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseCSEBaseSampler(DensePoseBaseSampler):
+    """
+    Base DensePose sampler to produce DensePose data from DensePose predictions.
+    Samples for each class are drawn according to some distribution over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(
+        self,
+        cfg: CfgNode,
+        use_gt_categories: bool,
+        embedder: torch.nn.Module,
+        count_per_class: int = 8,
+    ):
+        """
+        Constructor
+
+        Args:
+          cfg (CfgNode): the config of the model
+          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        super().__init__(count_per_class)
+        self.embedder = embedder
+        self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
+        self.use_gt_categories = use_gt_categories
+
+    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
+        """
+        Sample DensPoseDataRelative from estimation results
+        """
+        if self.use_gt_categories:
+            instance_class = instance.dataset_classes.tolist()[0]
+        else:
+            instance_class = instance.pred_classes.tolist()[0]
+        mesh_name = self.class_to_mesh_name[instance_class]
+
+        annotation = {
+            DensePoseDataRelative.X_KEY: [],
+            DensePoseDataRelative.Y_KEY: [],
+            DensePoseDataRelative.VERTEX_IDS_KEY: [],
+            DensePoseDataRelative.MESH_NAME_KEY: mesh_name,
+        }
+
+        mask, embeddings, other_values = self._produce_mask_and_results(instance, bbox_xywh)
+        indices = torch.nonzero(mask, as_tuple=True)
+        selected_embeddings = embeddings.permute(1, 2, 0)[indices].cpu()
+        values = other_values[:, indices[0], indices[1]]
+        k = values.shape[1]
+
+        count = min(self.count_per_class, k)
+        if count <= 0:
+            return annotation
+
+        index_sample = self._produce_index_sample(values, count)
+        closest_vertices = squared_euclidean_distance_matrix(
+            selected_embeddings[index_sample], self.embedder(mesh_name)
+        )
+        closest_vertices = torch.argmin(closest_vertices, dim=1)
+
+        sampled_y = indices[0][index_sample] + 0.5
+        sampled_x = indices[1][index_sample] + 0.5
+        # prepare / normalize data
+        _, _, w, h = bbox_xywh
+        x = (sampled_x / w * 256.0).cpu().tolist()
+        y = (sampled_y / h * 256.0).cpu().tolist()
+        # extend annotations
+        annotation[DensePoseDataRelative.X_KEY].extend(x)
+        annotation[DensePoseDataRelative.Y_KEY].extend(y)
+        annotation[DensePoseDataRelative.VERTEX_IDS_KEY].extend(closest_vertices.cpu().tolist())
+        return annotation
+
+    def _produce_mask_and_results(
+        self, instance: Instances, bbox_xywh: IntTupleBox
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of `DensePoseEmbeddingPredictorOutput`
+            bbox_xywh (IntTupleBox): the corresponding bounding box
+
+        Return:
+            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
+            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W],
+                DensePose CSE Embeddings
+            other_values (Tuple[torch.Tensor]): a tensor of shape [0, H, W],
+                for potential other values
+        """
+        densepose_output = instance.pred_densepose
+        S = densepose_output.coarse_segm
+        E = densepose_output.embedding
+        _, _, w, h = bbox_xywh
+        embeddings = F.interpolate(E, size=(h, w), mode="bilinear")[0]
+        coarse_segm_resized = F.interpolate(S, size=(h, w), mode="bilinear")[0]
+        mask = coarse_segm_resized.argmax(0) > 0
+        other_values = torch.empty((0, h, w), device=E.device)
+        return mask, embeddings, other_values
+
+    def _resample_mask(self, output: Any) -> torch.Tensor:
+        """
+        Convert DensePose predictor output to segmentation annotation - tensors of size
+        (256, 256) and type `int64`.
+
+        Args:
+            output: DensePose predictor output with the following attributes:
+             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
+               segmentation scores
+        Return:
+            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
+            where S = DensePoseDataRelative.MASK_SIZE
+        """
+        sz = DensePoseDataRelative.MASK_SIZE
+        mask = (
+            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
+            .argmax(dim=1)
+            .long()
+            .squeeze()
+            .cpu()
+        )
+        return mask

densepose/data/samplers/densepose_cse_confidence_based.py

@@ -1,119 +1,119 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import random
-from typing import Optional, Tuple
-import torch
-from torch.nn import functional as F
-
-from detectron2.config import CfgNode
-from detectron2.structures import Instances
-
-from densepose.converters.base import IntTupleBox
-
-from .densepose_cse_base import DensePoseCSEBaseSampler
-
-
-class DensePoseCSEConfidenceBasedSampler(DensePoseCSEBaseSampler):
-    """
-    Samples DensePose data from DensePose predictions.
-    Samples for each class are drawn using confidence value estimates.
-    """
-
-    def __init__(
-        self,
-        cfg: CfgNode,
-        use_gt_categories: bool,
-        embedder: torch.nn.Module,
-        confidence_channel: str,
-        count_per_class: int = 8,
-        search_count_multiplier: Optional[float] = None,
-        search_proportion: Optional[float] = None,
-    ):
-        """
-        Constructor
-
-        Args:
-          cfg (CfgNode): the config of the model
-          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
-          confidence_channel (str): confidence channel to use for sampling;
-            possible values:
-              "coarse_segm_confidence": confidences for coarse segmentation
-            (default: "coarse_segm_confidence")
-          count_per_class (int): the sampler produces at most `count_per_class`
-              samples for each category (default: 8)
-          search_count_multiplier (float or None): if not None, the total number
-              of the most confident estimates of a given class to consider is
-              defined as `min(search_count_multiplier * count_per_class, N)`,
-              where `N` is the total number of estimates of the class; cannot be
-              specified together with `search_proportion` (default: None)
-          search_proportion (float or None): if not None, the total number of the
-              of the most confident estimates of a given class to consider is
-              defined as `min(max(search_proportion * N, count_per_class), N)`,
-              where `N` is the total number of estimates of the class; cannot be
-              specified together with `search_count_multiplier` (default: None)
-        """
-        super().__init__(cfg, use_gt_categories, embedder, count_per_class)
-        self.confidence_channel = confidence_channel
-        self.search_count_multiplier = search_count_multiplier
-        self.search_proportion = search_proportion
-        assert (search_count_multiplier is None) or (search_proportion is None), (
-            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
-            f"and search_proportion (={search_proportion})"
-        )
-
-    def _produce_index_sample(self, values: torch.Tensor, count: int):
-        """
-        Produce a sample of indices to select data based on confidences
-
-        Args:
-            values (torch.Tensor): a tensor of length k that contains confidences
-                k: number of points labeled with part_id
-            count (int): number of samples to produce, should be positive and <= k
-
-        Return:
-            list(int): indices of values (along axis 1) selected as a sample
-        """
-        k = values.shape[1]
-        if k == count:
-            index_sample = list(range(k))
-        else:
-            # take the best count * search_count_multiplier pixels,
-            # sample from them uniformly
-            # (here best = smallest variance)
-            _, sorted_confidence_indices = torch.sort(values[0])
-            if self.search_count_multiplier is not None:
-                search_count = min(int(count * self.search_count_multiplier), k)
-            elif self.search_proportion is not None:
-                search_count = min(max(int(k * self.search_proportion), count), k)
-            else:
-                search_count = min(count, k)
-            sample_from_top = random.sample(range(search_count), count)
-            index_sample = sorted_confidence_indices[-search_count:][sample_from_top]
-        return index_sample
-
-    def _produce_mask_and_results(
-        self, instance: Instances, bbox_xywh: IntTupleBox
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Method to get labels and DensePose results from an instance
-
-        Args:
-            instance (Instances): an instance of
-                `DensePoseEmbeddingPredictorOutputWithConfidences`
-            bbox_xywh (IntTupleBox): the corresponding bounding box
-
-        Return:
-            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
-            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W]
-                DensePose CSE Embeddings
-            other_values: a tensor of shape [1, H, W], DensePose CSE confidence
-        """
-        _, _, w, h = bbox_xywh
-        densepose_output = instance.pred_densepose
-        mask, embeddings, _ = super()._produce_mask_and_results(instance, bbox_xywh)
-        other_values = F.interpolate(
-            getattr(densepose_output, self.confidence_channel),
-            size=(h, w),
-            mode="bilinear",
-        )[0].cpu()
-        return mask, embeddings, other_values
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import random
+from typing import Optional, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.config import CfgNode
+from detectron2.structures import Instances
+
+from densepose.converters.base import IntTupleBox
+
+from .densepose_cse_base import DensePoseCSEBaseSampler
+
+
+class DensePoseCSEConfidenceBasedSampler(DensePoseCSEBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn using confidence value estimates.
+    """
+
+    def __init__(
+        self,
+        cfg: CfgNode,
+        use_gt_categories: bool,
+        embedder: torch.nn.Module,
+        confidence_channel: str,
+        count_per_class: int = 8,
+        search_count_multiplier: Optional[float] = None,
+        search_proportion: Optional[float] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          cfg (CfgNode): the config of the model
+          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
+          confidence_channel (str): confidence channel to use for sampling;
+            possible values:
+              "coarse_segm_confidence": confidences for coarse segmentation
+            (default: "coarse_segm_confidence")
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category (default: 8)
+          search_count_multiplier (float or None): if not None, the total number
+              of the most confident estimates of a given class to consider is
+              defined as `min(search_count_multiplier * count_per_class, N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_proportion` (default: None)
+          search_proportion (float or None): if not None, the total number of the
+              of the most confident estimates of a given class to consider is
+              defined as `min(max(search_proportion * N, count_per_class), N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_count_multiplier` (default: None)
+        """
+        super().__init__(cfg, use_gt_categories, embedder, count_per_class)
+        self.confidence_channel = confidence_channel
+        self.search_count_multiplier = search_count_multiplier
+        self.search_proportion = search_proportion
+        assert (search_count_multiplier is None) or (search_proportion is None), (
+            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
+            f"and search_proportion (={search_proportion})"
+        )
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a sample of indices to select data based on confidences
+
+        Args:
+            values (torch.Tensor): a tensor of length k that contains confidences
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        if k == count:
+            index_sample = list(range(k))
+        else:
+            # take the best count * search_count_multiplier pixels,
+            # sample from them uniformly
+            # (here best = smallest variance)
+            _, sorted_confidence_indices = torch.sort(values[0])
+            if self.search_count_multiplier is not None:
+                search_count = min(int(count * self.search_count_multiplier), k)
+            elif self.search_proportion is not None:
+                search_count = min(max(int(k * self.search_proportion), count), k)
+            else:
+                search_count = min(count, k)
+            sample_from_top = random.sample(range(search_count), count)
+            index_sample = sorted_confidence_indices[-search_count:][sample_from_top]
+        return index_sample
+
+    def _produce_mask_and_results(
+        self, instance: Instances, bbox_xywh: IntTupleBox
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of
+                `DensePoseEmbeddingPredictorOutputWithConfidences`
+            bbox_xywh (IntTupleBox): the corresponding bounding box
+
+        Return:
+            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
+            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W]
+                DensePose CSE Embeddings
+            other_values: a tensor of shape [1, H, W], DensePose CSE confidence
+        """
+        _, _, w, h = bbox_xywh
+        densepose_output = instance.pred_densepose
+        mask, embeddings, _ = super()._produce_mask_and_results(instance, bbox_xywh)
+        other_values = F.interpolate(
+            getattr(densepose_output, self.confidence_channel),
+            size=(h, w),
+            mode="bilinear",
+        )[0].cpu()
+        return mask, embeddings, other_values

densepose/data/samplers/densepose_cse_uniform.py

@@ -1,12 +1,12 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .densepose_cse_base import DensePoseCSEBaseSampler
-from .densepose_uniform import DensePoseUniformSampler
-
-
-class DensePoseCSEUniformSampler(DensePoseCSEBaseSampler, DensePoseUniformSampler):
-    """
-    Uniform Sampler for CSE
-    """
-
-    pass
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .densepose_cse_base import DensePoseCSEBaseSampler
+from .densepose_uniform import DensePoseUniformSampler
+
+
+class DensePoseCSEUniformSampler(DensePoseCSEBaseSampler, DensePoseUniformSampler):
+    """
+    Uniform Sampler for CSE
+    """
+
+    pass

densepose/data/samplers/densepose_uniform.py

@@ -1,41 +1,41 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import random
-import torch
-
-from .densepose_base import DensePoseBaseSampler
-
-
-class DensePoseUniformSampler(DensePoseBaseSampler):
-    """
-    Samples DensePose data from DensePose predictions.
-    Samples for each class are drawn uniformly over all pixels estimated
-    to belong to that class.
-    """
-
-    def __init__(self, count_per_class: int = 8):
-        """
-        Constructor
-
-        Args:
-          count_per_class (int): the sampler produces at most `count_per_class`
-              samples for each category
-        """
-        super().__init__(count_per_class)
-
-    def _produce_index_sample(self, values: torch.Tensor, count: int):
-        """
-        Produce a uniform sample of indices to select data
-
-        Args:
-            values (torch.Tensor): an array of size [n, k] that contains
-                estimated values (U, V, confidences);
-                n: number of channels (U, V, confidences)
-                k: number of points labeled with part_id
-            count (int): number of samples to produce, should be positive and <= k
-
-        Return:
-            list(int): indices of values (along axis 1) selected as a sample
-        """
-        k = values.shape[1]
-        return random.sample(range(k), count)
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import random
+import torch
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseUniformSampler(DensePoseBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn uniformly over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(self, count_per_class: int = 8):
+        """
+        Constructor
+
+        Args:
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        super().__init__(count_per_class)
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a uniform sample of indices to select data
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        return random.sample(range(k), count)

densepose/data/samplers/mask_from_densepose.py

@@ -1,28 +1,28 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from detectron2.structures import BitMasks, Instances
-
-from densepose.converters import ToMaskConverter
-
-
-class MaskFromDensePoseSampler:
-    """
-    Produce mask GT from DensePose predictions
-    This sampler simply converts DensePose predictions to BitMasks
-    that a contain a bool tensor of the size of the input image
-    """
-
-    def __call__(self, instances: Instances) -> BitMasks:
-        """
-        Converts predicted data from `instances` into the GT mask data
-
-        Args:
-            instances (Instances): predicted results, expected to have `pred_densepose` field
-
-        Returns:
-            Boolean Tensor of the size of the input image that has non-zero
-            values at pixels that are estimated to belong to the detected object
-        """
-        return ToMaskConverter.convert(
-            instances.pred_densepose, instances.pred_boxes, instances.image_size
-        )
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from detectron2.structures import BitMasks, Instances
+
+from densepose.converters import ToMaskConverter
+
+
+class MaskFromDensePoseSampler:
+    """
+    Produce mask GT from DensePose predictions
+    This sampler simply converts DensePose predictions to BitMasks
+    that a contain a bool tensor of the size of the input image
+    """
+
+    def __call__(self, instances: Instances) -> BitMasks:
+        """
+        Converts predicted data from `instances` into the GT mask data
+
+        Args:
+            instances (Instances): predicted results, expected to have `pred_densepose` field
+
+        Returns:
+            Boolean Tensor of the size of the input image that has non-zero
+            values at pixels that are estimated to belong to the detected object
+        """
+        return ToMaskConverter.convert(
+            instances.pred_densepose, instances.pred_boxes, instances.image_size
+        )

densepose/data/samplers/prediction_to_gt.py

@@ -1,98 +1,98 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional
-
-from detectron2.structures import Instances
-
-ModelOutput = Dict[str, Any]
-SampledData = Dict[str, Any]
-
-
-@dataclass
-class _Sampler:
-    """
-    Sampler registry entry that contains:
-     - src (str): source field to sample from (deleted after sampling)
-     - dst (Optional[str]): destination field to sample to, if not None
-     - func (Optional[Callable: Any -> Any]): function that performs sampling,
-         if None, reference copy is performed
-    """
-
-    src: str
-    dst: Optional[str]
-    func: Optional[Callable[[Any], Any]]
-
-
-class PredictionToGroundTruthSampler:
-    """
-    Sampler implementation that converts predictions to GT using registered
-    samplers for different fields of `Instances`.
-    """
-
-    def __init__(self, dataset_name: str = ""):
-        self.dataset_name = dataset_name
-        self._samplers = {}
-        self.register_sampler("pred_boxes", "gt_boxes", None)
-        self.register_sampler("pred_classes", "gt_classes", None)
-        # delete scores
-        self.register_sampler("scores")
-
-    def __call__(self, model_output: List[ModelOutput]) -> List[SampledData]:
-        """
-        Transform model output into ground truth data through sampling
-
-        Args:
-          model_output (Dict[str, Any]): model output
-        Returns:
-          Dict[str, Any]: sampled data
-        """
-        for model_output_i in model_output:
-            instances: Instances = model_output_i["instances"]
-            # transform data in each field
-            for _, sampler in self._samplers.items():
-                if not instances.has(sampler.src) or sampler.dst is None:
-                    continue
-                if sampler.func is None:
-                    instances.set(sampler.dst, instances.get(sampler.src))
-                else:
-                    instances.set(sampler.dst, sampler.func(instances))
-            # delete model output data that was transformed
-            for _, sampler in self._samplers.items():
-                if sampler.src != sampler.dst and instances.has(sampler.src):
-                    instances.remove(sampler.src)
-            model_output_i["dataset"] = self.dataset_name
-        return model_output
-
-    def register_sampler(
-        self,
-        prediction_attr: str,
-        gt_attr: Optional[str] = None,
-        func: Optional[Callable[[Any], Any]] = None,
-    ):
-        """
-        Register sampler for a field
-
-        Args:
-          prediction_attr (str): field to replace with a sampled value
-          gt_attr (Optional[str]): field to store the sampled value to, if not None
-          func (Optional[Callable: Any -> Any]): sampler function
-        """
-        self._samplers[(prediction_attr, gt_attr)] = _Sampler(
-            src=prediction_attr, dst=gt_attr, func=func
-        )
-
-    def remove_sampler(
-        self,
-        prediction_attr: str,
-        gt_attr: Optional[str] = None,
-    ):
-        """
-        Remove sampler for a field
-
-        Args:
-          prediction_attr (str): field to replace with a sampled value
-          gt_attr (Optional[str]): field to store the sampled value to, if not None
-        """
-        assert (prediction_attr, gt_attr) in self._samplers
-        del self._samplers[(prediction_attr, gt_attr)]
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+from detectron2.structures import Instances
+
+ModelOutput = Dict[str, Any]
+SampledData = Dict[str, Any]
+
+
+@dataclass
+class _Sampler:
+    """
+    Sampler registry entry that contains:
+     - src (str): source field to sample from (deleted after sampling)
+     - dst (Optional[str]): destination field to sample to, if not None
+     - func (Optional[Callable: Any -> Any]): function that performs sampling,
+         if None, reference copy is performed
+    """
+
+    src: str
+    dst: Optional[str]
+    func: Optional[Callable[[Any], Any]]
+
+
+class PredictionToGroundTruthSampler:
+    """
+    Sampler implementation that converts predictions to GT using registered
+    samplers for different fields of `Instances`.
+    """
+
+    def __init__(self, dataset_name: str = ""):
+        self.dataset_name = dataset_name
+        self._samplers = {}
+        self.register_sampler("pred_boxes", "gt_boxes", None)
+        self.register_sampler("pred_classes", "gt_classes", None)
+        # delete scores
+        self.register_sampler("scores")
+
+    def __call__(self, model_output: List[ModelOutput]) -> List[SampledData]:
+        """
+        Transform model output into ground truth data through sampling
+
+        Args:
+          model_output (Dict[str, Any]): model output
+        Returns:
+          Dict[str, Any]: sampled data
+        """
+        for model_output_i in model_output:
+            instances: Instances = model_output_i["instances"]
+            # transform data in each field
+            for _, sampler in self._samplers.items():
+                if not instances.has(sampler.src) or sampler.dst is None:
+                    continue
+                if sampler.func is None:
+                    instances.set(sampler.dst, instances.get(sampler.src))
+                else:
+                    instances.set(sampler.dst, sampler.func(instances))
+            # delete model output data that was transformed
+            for _, sampler in self._samplers.items():
+                if sampler.src != sampler.dst and instances.has(sampler.src):
+                    instances.remove(sampler.src)
+            model_output_i["dataset"] = self.dataset_name
+        return model_output
+
+    def register_sampler(
+        self,
+        prediction_attr: str,
+        gt_attr: Optional[str] = None,
+        func: Optional[Callable[[Any], Any]] = None,
+    ):
+        """
+        Register sampler for a field
+
+        Args:
+          prediction_attr (str): field to replace with a sampled value
+          gt_attr (Optional[str]): field to store the sampled value to, if not None
+          func (Optional[Callable: Any -> Any]): sampler function
+        """
+        self._samplers[(prediction_attr, gt_attr)] = _Sampler(
+            src=prediction_attr, dst=gt_attr, func=func
+        )
+
+    def remove_sampler(
+        self,
+        prediction_attr: str,
+        gt_attr: Optional[str] = None,
+    ):
+        """
+        Remove sampler for a field
+
+        Args:
+          prediction_attr (str): field to replace with a sampled value
+          gt_attr (Optional[str]): field to store the sampled value to, if not None
+        """
+        assert (prediction_attr, gt_attr) in self._samplers
+        del self._samplers[(prediction_attr, gt_attr)]

densepose/data/transform/__init__.py

@@ -1,3 +1,3 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .image import ImageResizeTransform
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .image import ImageResizeTransform

densepose/data/transform/image.py

@@ -1,39 +1,39 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import torch
-
-
-class ImageResizeTransform:
-    """
-    Transform that resizes images loaded from a dataset
-    (BGR data in NCHW channel order, typically uint8) to a format ready to be
-    consumed by DensePose training (BGR float32 data in NCHW channel order)
-    """
-
-    def __init__(self, min_size: int = 800, max_size: int = 1333):
-        self.min_size = min_size
-        self.max_size = max_size
-
-    def __call__(self, images: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            images (torch.Tensor): tensor of size [N, 3, H, W] that contains
-                BGR data (typically in uint8)
-        Returns:
-            images (torch.Tensor): tensor of size [N, 3, H1, W1] where
-                H1 and W1 are chosen to respect the specified min and max sizes
-                and preserve the original aspect ratio, the data channels
-                follow BGR order and the data type is `torch.float32`
-        """
-        # resize with min size
-        images = images.float()
-        min_size = min(images.shape[-2:])
-        max_size = max(images.shape[-2:])
-        scale = min(self.min_size / min_size, self.max_size / max_size)
-        images = torch.nn.functional.interpolate(
-            images,
-            scale_factor=scale,
-            mode="bilinear",
-            align_corners=False,
-        )
-        return images
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import torch
+
+
+class ImageResizeTransform:
+    """
+    Transform that resizes images loaded from a dataset
+    (BGR data in NCHW channel order, typically uint8) to a format ready to be
+    consumed by DensePose training (BGR float32 data in NCHW channel order)
+    """
+
+    def __init__(self, min_size: int = 800, max_size: int = 1333):
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            images (torch.Tensor): tensor of size [N, 3, H, W] that contains
+                BGR data (typically in uint8)
+        Returns:
+            images (torch.Tensor): tensor of size [N, 3, H1, W1] where
+                H1 and W1 are chosen to respect the specified min and max sizes
+                and preserve the original aspect ratio, the data channels
+                follow BGR order and the data type is `torch.float32`
+        """
+        # resize with min size
+        images = images.float()
+        min_size = min(images.shape[-2:])
+        max_size = max(images.shape[-2:])
+        scale = min(self.min_size / min_size, self.max_size / max_size)
+        images = torch.nn.functional.interpolate(
+            images,
+            scale_factor=scale,
+            mode="bilinear",
+            align_corners=False,
+        )
+        return images

densepose/data/utils.py

@@ -1,38 +1,38 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import os
-from typing import Dict, Optional
-
-from detectron2.config import CfgNode
-
-
-def is_relative_local_path(path: str) -> bool:
-    path_str = os.fsdecode(path)
-    return ("://" not in path_str) and not os.path.isabs(path)
-
-
-def maybe_prepend_base_path(base_path: Optional[str], path: str):
-    """
-    Prepends the provided path with a base path prefix if:
-    1) base path is not None;
-    2) path is a local path
-    """
-    if base_path is None:
-        return path
-    if is_relative_local_path(path):
-        return os.path.join(base_path, path)
-    return path
-
-
-def get_class_to_mesh_name_mapping(cfg: CfgNode) -> Dict[int, str]:
-    return {
-        int(class_id): mesh_name
-        for class_id, mesh_name in cfg.DATASETS.CLASS_TO_MESH_NAME_MAPPING.items()
-    }
-
-
-def get_category_to_class_mapping(dataset_cfg: CfgNode) -> Dict[str, int]:
-    return {
-        category: int(class_id)
-        for category, class_id in dataset_cfg.CATEGORY_TO_CLASS_MAPPING.items()
-    }
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import os
+from typing import Dict, Optional
+
+from detectron2.config import CfgNode
+
+
+def is_relative_local_path(path: str) -> bool:
+    path_str = os.fsdecode(path)
+    return ("://" not in path_str) and not os.path.isabs(path)
+
+
+def maybe_prepend_base_path(base_path: Optional[str], path: str):
+    """
+    Prepends the provided path with a base path prefix if:
+    1) base path is not None;
+    2) path is a local path
+    """
+    if base_path is None:
+        return path
+    if is_relative_local_path(path):
+        return os.path.join(base_path, path)
+    return path
+
+
+def get_class_to_mesh_name_mapping(cfg: CfgNode) -> Dict[int, str]:
+    return {
+        int(class_id): mesh_name
+        for class_id, mesh_name in cfg.DATASETS.CLASS_TO_MESH_NAME_MAPPING.items()
+    }
+
+
+def get_category_to_class_mapping(dataset_cfg: CfgNode) -> Dict[str, int]:
+    return {
+        category: int(class_id)
+        for category, class_id in dataset_cfg.CATEGORY_TO_CLASS_MAPPING.items()
+    }

densepose/data/video/__init__.py

@@ -1,17 +1,17 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .frame_selector import (
-    FrameSelectionStrategy,
-    RandomKFramesSelector,
-    FirstKFramesSelector,
-    LastKFramesSelector,
-    FrameTsList,
-    FrameSelector,
-)
-
-from .video_keyframe_dataset import (
-    VideoKeyframeDataset,
-    video_list_from_file,
-    list_keyframes,
-    read_keyframes,
-)
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .frame_selector import (
+    FrameSelectionStrategy,
+    RandomKFramesSelector,
+    FirstKFramesSelector,
+    LastKFramesSelector,
+    FrameTsList,
+    FrameSelector,
+)
+
+from .video_keyframe_dataset import (
+    VideoKeyframeDataset,
+    video_list_from_file,
+    list_keyframes,
+    read_keyframes,
+)

densepose/data/video/frame_selector.py

@@ -1,87 +1,87 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import random
-from collections.abc import Callable
-from enum import Enum
-from typing import Callable as TCallable
-from typing import List
-
-FrameTsList = List[int]
-FrameSelector = TCallable[[FrameTsList], FrameTsList]
-
-
-class FrameSelectionStrategy(Enum):
-    """
-    Frame selection strategy used with videos:
-     - "random_k": select k random frames
-     - "first_k": select k first frames
-     - "last_k": select k last frames
-     - "all": select all frames
-    """
-
-    # fmt: off
-    RANDOM_K = "random_k"
-    FIRST_K  = "first_k"
-    LAST_K   = "last_k"
-    ALL      = "all"
-    # fmt: on
-
-
-class RandomKFramesSelector(Callable):  # pyre-ignore[39]
-    """
-    Selector that retains at most `k` random frames
-    """
-
-    def __init__(self, k: int):
-        self.k = k
-
-    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
-        """
-        Select `k` random frames
-
-        Args:
-          frames_tss (List[int]): timestamps of input frames
-        Returns:
-          List[int]: timestamps of selected frames
-        """
-        return random.sample(frame_tss, min(self.k, len(frame_tss)))
-
-
-class FirstKFramesSelector(Callable):  # pyre-ignore[39]
-    """
-    Selector that retains at most `k` first frames
-    """
-
-    def __init__(self, k: int):
-        self.k = k
-
-    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
-        """
-        Select `k` first frames
-
-        Args:
-          frames_tss (List[int]): timestamps of input frames
-        Returns:
-          List[int]: timestamps of selected frames
-        """
-        return frame_tss[: self.k]
-
-
-class LastKFramesSelector(Callable):  # pyre-ignore[39]
-    """
-    Selector that retains at most `k` last frames from video data
-    """
-
-    def __init__(self, k: int):
-        self.k = k
-
-    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
-        """
-        Select `k` last frames
-
-        Args:
-          frames_tss (List[int]): timestamps of input frames
-        Returns:
-          List[int]: timestamps of selected frames
-        """
-        return frame_tss[-self.k :]
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import random
+from collections.abc import Callable
+from enum import Enum
+from typing import Callable as TCallable
+from typing import List
+
+FrameTsList = List[int]
+FrameSelector = TCallable[[FrameTsList], FrameTsList]
+
+
+class FrameSelectionStrategy(Enum):
+    """
+    Frame selection strategy used with videos:
+     - "random_k": select k random frames
+     - "first_k": select k first frames
+     - "last_k": select k last frames
+     - "all": select all frames
+    """
+
+    # fmt: off
+    RANDOM_K = "random_k"
+    FIRST_K  = "first_k"
+    LAST_K   = "last_k"
+    ALL      = "all"
+    # fmt: on
+
+
+class RandomKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` random frames
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` random frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return random.sample(frame_tss, min(self.k, len(frame_tss)))
+
+
+class FirstKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` first frames
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` first frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return frame_tss[: self.k]
+
+
+class LastKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` last frames from video data
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` last frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return frame_tss[-self.k :]

densepose/data/video/video_keyframe_dataset.py

@@ -1,300 +1,300 @@
-# -*- coding: utf-8 -*-
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import csv
-import logging
-import numpy as np
-from typing import Any, Callable, Dict, List, Optional, Union
-import av
-import torch
-from torch.utils.data.dataset import Dataset
-
-from detectron2.utils.file_io import PathManager
-
-from ..utils import maybe_prepend_base_path
-from .frame_selector import FrameSelector, FrameTsList
-
-FrameList = List[av.frame.Frame]  # pyre-ignore[16]
-FrameTransform = Callable[[torch.Tensor], torch.Tensor]
-
-
-def list_keyframes(video_fpath: str, video_stream_idx: int = 0) -> FrameTsList:
-    """
-    Traverses all keyframes of a video file. Returns a list of keyframe
-    timestamps. Timestamps are counts in timebase units.
-
-    Args:
-       video_fpath (str): Video file path
-       video_stream_idx (int): Video stream index (default: 0)
-    Returns:
-       List[int]: list of keyframe timestaps (timestamp is a count in timebase
-           units)
-    """
-    try:
-        with PathManager.open(video_fpath, "rb") as io:
-            container = av.open(io, mode="r")
-            stream = container.streams.video[video_stream_idx]
-            keyframes = []
-            pts = -1
-            # Note: even though we request forward seeks for keyframes, sometimes
-            # a keyframe in backwards direction is returned. We introduce tolerance
-            # as a max count of ignored backward seeks
-            tolerance_backward_seeks = 2
-            while True:
-                try:
-                    container.seek(pts + 1, backward=False, any_frame=False, stream=stream)
-                except av.AVError as e:
-                    # the exception occurs when the video length is exceeded,
-                    # we then return whatever data we've already collected
-                    logger = logging.getLogger(__name__)
-                    logger.debug(
-                        f"List keyframes: Error seeking video file {video_fpath}, "
-                        f"video stream {video_stream_idx}, pts {pts + 1}, AV error: {e}"
-                    )
-                    return keyframes
-                except OSError as e:
-                    logger = logging.getLogger(__name__)
-                    logger.warning(
-                        f"List keyframes: Error seeking video file {video_fpath}, "
-                        f"video stream {video_stream_idx}, pts {pts + 1}, OS error: {e}"
-                    )
-                    return []
-                packet = next(container.demux(video=video_stream_idx))
-                if packet.pts is not None and packet.pts <= pts:
-                    logger = logging.getLogger(__name__)
-                    logger.warning(
-                        f"Video file {video_fpath}, stream {video_stream_idx}: "
-                        f"bad seek for packet {pts + 1} (got packet {packet.pts}), "
-                        f"tolerance {tolerance_backward_seeks}."
-                    )
-                    tolerance_backward_seeks -= 1
-                    if tolerance_backward_seeks == 0:
-                        return []
-                    pts += 1
-                    continue
-                tolerance_backward_seeks = 2
-                pts = packet.pts
-                if pts is None:
-                    return keyframes
-                if packet.is_keyframe:
-                    keyframes.append(pts)
-            return keyframes
-    except OSError as e:
-        logger = logging.getLogger(__name__)
-        logger.warning(
-            f"List keyframes: Error opening video file container {video_fpath}, " f"OS error: {e}"
-        )
-    except RuntimeError as e:
-        logger = logging.getLogger(__name__)
-        logger.warning(
-            f"List keyframes: Error opening video file container {video_fpath}, "
-            f"Runtime error: {e}"
-        )
-    return []
-
-
-def read_keyframes(
-    video_fpath: str, keyframes: FrameTsList, video_stream_idx: int = 0
-) -> FrameList:  # pyre-ignore[11]
-    """
-    Reads keyframe data from a video file.
-
-    Args:
-        video_fpath (str): Video file path
-        keyframes (List[int]): List of keyframe timestamps (as counts in
-            timebase units to be used in container seek operations)
-        video_stream_idx (int): Video stream index (default: 0)
-    Returns:
-        List[Frame]: list of frames that correspond to the specified timestamps
-    """
-    try:
-        with PathManager.open(video_fpath, "rb") as io:
-            container = av.open(io)
-            stream = container.streams.video[video_stream_idx]
-            frames = []
-            for pts in keyframes:
-                try:
-                    container.seek(pts, any_frame=False, stream=stream)
-                    frame = next(container.decode(video=0))
-                    frames.append(frame)
-                except av.AVError as e:
-                    logger = logging.getLogger(__name__)
-                    logger.warning(
-                        f"Read keyframes: Error seeking video file {video_fpath}, "
-                        f"video stream {video_stream_idx}, pts {pts}, AV error: {e}"
-                    )
-                    container.close()
-                    return frames
-                except OSError as e:
-                    logger = logging.getLogger(__name__)
-                    logger.warning(
-                        f"Read keyframes: Error seeking video file {video_fpath}, "
-                        f"video stream {video_stream_idx}, pts {pts}, OS error: {e}"
-                    )
-                    container.close()
-                    return frames
-                except StopIteration:
-                    logger = logging.getLogger(__name__)
-                    logger.warning(
-                        f"Read keyframes: Error decoding frame from {video_fpath}, "
-                        f"video stream {video_stream_idx}, pts {pts}"
-                    )
-                    container.close()
-                    return frames
-
-            container.close()
-            return frames
-    except OSError as e:
-        logger = logging.getLogger(__name__)
-        logger.warning(
-            f"Read keyframes: Error opening video file container {video_fpath}, OS error: {e}"
-        )
-    except RuntimeError as e:
-        logger = logging.getLogger(__name__)
-        logger.warning(
-            f"Read keyframes: Error opening video file container {video_fpath}, Runtime error: {e}"
-        )
-    return []
-
-
-def video_list_from_file(video_list_fpath: str, base_path: Optional[str] = None):
-    """
-    Create a list of paths to video files from a text file.
-
-    Args:
-        video_list_fpath (str): path to a plain text file with the list of videos
-        base_path (str): base path for entries from the video list (default: None)
-    """
-    video_list = []
-    with PathManager.open(video_list_fpath, "r") as io:
-        for line in io:
-            video_list.append(maybe_prepend_base_path(base_path, str(line.strip())))
-    return video_list
-
-
-def read_keyframe_helper_data(fpath: str):
-    """
-    Read keyframe data from a file in CSV format: the header should contain
-    "video_id" and "keyframes" fields. Value specifications are:
-      video_id: int
-      keyframes: list(int)
-    Example of contents:
-      video_id,keyframes
-      2,"[1,11,21,31,41,51,61,71,81]"
-
-    Args:
-        fpath (str): File containing keyframe data
-
-    Return:
-        video_id_to_keyframes (dict: int -> list(int)): for a given video ID it
-          contains a list of keyframes for that video
-    """
-    video_id_to_keyframes = {}
-    try:
-        with PathManager.open(fpath, "r") as io:
-            csv_reader = csv.reader(io)
-            header = next(csv_reader)
-            video_id_idx = header.index("video_id")
-            keyframes_idx = header.index("keyframes")
-            for row in csv_reader:
-                video_id = int(row[video_id_idx])
-                assert (
-                    video_id not in video_id_to_keyframes
-                ), f"Duplicate keyframes entry for video {fpath}"
-                video_id_to_keyframes[video_id] = (
-                    [int(v) for v in row[keyframes_idx][1:-1].split(",")]
-                    if len(row[keyframes_idx]) > 2
-                    else []
-                )
-    except Exception as e:
-        logger = logging.getLogger(__name__)
-        logger.warning(f"Error reading keyframe helper data from {fpath}: {e}")
-    return video_id_to_keyframes
-
-
-class VideoKeyframeDataset(Dataset):
-    """
-    Dataset that provides keyframes for a set of videos.
-    """
-
-    _EMPTY_FRAMES = torch.empty((0, 3, 1, 1))
-
-    def __init__(
-        self,
-        video_list: List[str],
-        category_list: Union[str, List[str], None] = None,
-        frame_selector: Optional[FrameSelector] = None,
-        transform: Optional[FrameTransform] = None,
-        keyframe_helper_fpath: Optional[str] = None,
-    ):
-        """
-        Dataset constructor
-
-        Args:
-            video_list (List[str]): list of paths to video files
-            category_list (Union[str, List[str], None]): list of animal categories for each
-                video file. If it is a string, or None, this applies to all videos
-            frame_selector (Callable: KeyFrameList -> KeyFrameList):
-                selects keyframes to process, keyframes are given by
-                packet timestamps in timebase counts. If None, all keyframes
-                are selected (default: None)
-            transform (Callable: torch.Tensor -> torch.Tensor):
-                transforms a batch of RGB images (tensors of size [B, 3, H, W]),
-                returns a tensor of the same size. If None, no transform is
-                applied (default: None)
-
-        """
-        if type(category_list) == list:
-            self.category_list = category_list
-        else:
-            self.category_list = [category_list] * len(video_list)
-        assert len(video_list) == len(
-            self.category_list
-        ), "length of video and category lists must be equal"
-        self.video_list = video_list
-        self.frame_selector = frame_selector
-        self.transform = transform
-        self.keyframe_helper_data = (
-            read_keyframe_helper_data(keyframe_helper_fpath)
-            if keyframe_helper_fpath is not None
-            else None
-        )
-
-    def __getitem__(self, idx: int) -> Dict[str, Any]:
-        """
-        Gets selected keyframes from a given video
-
-        Args:
-            idx (int): video index in the video list file
-        Returns:
-            A dictionary containing two keys:
-                images (torch.Tensor): tensor of size [N, H, W, 3] or of size
-                    defined by the transform that contains keyframes data
-                categories (List[str]): categories of the frames
-        """
-        categories = [self.category_list[idx]]
-        fpath = self.video_list[idx]
-        keyframes = (
-            list_keyframes(fpath)
-            if self.keyframe_helper_data is None or idx not in self.keyframe_helper_data
-            else self.keyframe_helper_data[idx]
-        )
-        transform = self.transform
-        frame_selector = self.frame_selector
-        if not keyframes:
-            return {"images": self._EMPTY_FRAMES, "categories": []}
-        if frame_selector is not None:
-            keyframes = frame_selector(keyframes)
-        frames = read_keyframes(fpath, keyframes)
-        if not frames:
-            return {"images": self._EMPTY_FRAMES, "categories": []}
-        frames = np.stack([frame.to_rgb().to_ndarray() for frame in frames])
-        frames = torch.as_tensor(frames, device=torch.device("cpu"))
-        frames = frames[..., [2, 1, 0]]  # RGB -> BGR
-        frames = frames.permute(0, 3, 1, 2).float()  # NHWC -> NCHW
-        if transform is not None:
-            frames = transform(frames)
-        return {"images": frames, "categories": categories}
-
-    def __len__(self):
-        return len(self.video_list)
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import csv
+import logging
+import numpy as np
+from typing import Any, Callable, Dict, List, Optional, Union
+import av
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.utils.file_io import PathManager
+
+from ..utils import maybe_prepend_base_path
+from .frame_selector import FrameSelector, FrameTsList
+
+FrameList = List[av.frame.Frame]  # pyre-ignore[16]
+FrameTransform = Callable[[torch.Tensor], torch.Tensor]
+
+
+def list_keyframes(video_fpath: str, video_stream_idx: int = 0) -> FrameTsList:
+    """
+    Traverses all keyframes of a video file. Returns a list of keyframe
+    timestamps. Timestamps are counts in timebase units.
+
+    Args:
+       video_fpath (str): Video file path
+       video_stream_idx (int): Video stream index (default: 0)
+    Returns:
+       List[int]: list of keyframe timestaps (timestamp is a count in timebase
+           units)
+    """
+    try:
+        with PathManager.open(video_fpath, "rb") as io:
+            container = av.open(io, mode="r")
+            stream = container.streams.video[video_stream_idx]
+            keyframes = []
+            pts = -1
+            # Note: even though we request forward seeks for keyframes, sometimes
+            # a keyframe in backwards direction is returned. We introduce tolerance
+            # as a max count of ignored backward seeks
+            tolerance_backward_seeks = 2
+            while True:
+                try:
+                    container.seek(pts + 1, backward=False, any_frame=False, stream=stream)
+                except av.AVError as e:
+                    # the exception occurs when the video length is exceeded,
+                    # we then return whatever data we've already collected
+                    logger = logging.getLogger(__name__)
+                    logger.debug(
+                        f"List keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts + 1}, AV error: {e}"
+                    )
+                    return keyframes
+                except OSError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"List keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts + 1}, OS error: {e}"
+                    )
+                    return []
+                packet = next(container.demux(video=video_stream_idx))
+                if packet.pts is not None and packet.pts <= pts:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Video file {video_fpath}, stream {video_stream_idx}: "
+                        f"bad seek for packet {pts + 1} (got packet {packet.pts}), "
+                        f"tolerance {tolerance_backward_seeks}."
+                    )
+                    tolerance_backward_seeks -= 1
+                    if tolerance_backward_seeks == 0:
+                        return []
+                    pts += 1
+                    continue
+                tolerance_backward_seeks = 2
+                pts = packet.pts
+                if pts is None:
+                    return keyframes
+                if packet.is_keyframe:
+                    keyframes.append(pts)
+            return keyframes
+    except OSError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"List keyframes: Error opening video file container {video_fpath}, " f"OS error: {e}"
+        )
+    except RuntimeError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"List keyframes: Error opening video file container {video_fpath}, "
+            f"Runtime error: {e}"
+        )
+    return []
+
+
+def read_keyframes(
+    video_fpath: str, keyframes: FrameTsList, video_stream_idx: int = 0
+) -> FrameList:  # pyre-ignore[11]
+    """
+    Reads keyframe data from a video file.
+
+    Args:
+        video_fpath (str): Video file path
+        keyframes (List[int]): List of keyframe timestamps (as counts in
+            timebase units to be used in container seek operations)
+        video_stream_idx (int): Video stream index (default: 0)
+    Returns:
+        List[Frame]: list of frames that correspond to the specified timestamps
+    """
+    try:
+        with PathManager.open(video_fpath, "rb") as io:
+            container = av.open(io)
+            stream = container.streams.video[video_stream_idx]
+            frames = []
+            for pts in keyframes:
+                try:
+                    container.seek(pts, any_frame=False, stream=stream)
+                    frame = next(container.decode(video=0))
+                    frames.append(frame)
+                except av.AVError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}, AV error: {e}"
+                    )
+                    container.close()
+                    return frames
+                except OSError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}, OS error: {e}"
+                    )
+                    container.close()
+                    return frames
+                except StopIteration:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error decoding frame from {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}"
+                    )
+                    container.close()
+                    return frames
+
+            container.close()
+            return frames
+    except OSError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"Read keyframes: Error opening video file container {video_fpath}, OS error: {e}"
+        )
+    except RuntimeError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"Read keyframes: Error opening video file container {video_fpath}, Runtime error: {e}"
+        )
+    return []
+
+
+def video_list_from_file(video_list_fpath: str, base_path: Optional[str] = None):
+    """
+    Create a list of paths to video files from a text file.
+
+    Args:
+        video_list_fpath (str): path to a plain text file with the list of videos
+        base_path (str): base path for entries from the video list (default: None)
+    """
+    video_list = []
+    with PathManager.open(video_list_fpath, "r") as io:
+        for line in io:
+            video_list.append(maybe_prepend_base_path(base_path, str(line.strip())))
+    return video_list
+
+
+def read_keyframe_helper_data(fpath: str):
+    """
+    Read keyframe data from a file in CSV format: the header should contain
+    "video_id" and "keyframes" fields. Value specifications are:
+      video_id: int
+      keyframes: list(int)
+    Example of contents:
+      video_id,keyframes
+      2,"[1,11,21,31,41,51,61,71,81]"
+
+    Args:
+        fpath (str): File containing keyframe data
+
+    Return:
+        video_id_to_keyframes (dict: int -> list(int)): for a given video ID it
+          contains a list of keyframes for that video
+    """
+    video_id_to_keyframes = {}
+    try:
+        with PathManager.open(fpath, "r") as io:
+            csv_reader = csv.reader(io)
+            header = next(csv_reader)
+            video_id_idx = header.index("video_id")
+            keyframes_idx = header.index("keyframes")
+            for row in csv_reader:
+                video_id = int(row[video_id_idx])
+                assert (
+                    video_id not in video_id_to_keyframes
+                ), f"Duplicate keyframes entry for video {fpath}"
+                video_id_to_keyframes[video_id] = (
+                    [int(v) for v in row[keyframes_idx][1:-1].split(",")]
+                    if len(row[keyframes_idx]) > 2
+                    else []
+                )
+    except Exception as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(f"Error reading keyframe helper data from {fpath}: {e}")
+    return video_id_to_keyframes
+
+
+class VideoKeyframeDataset(Dataset):
+    """
+    Dataset that provides keyframes for a set of videos.
+    """
+
+    _EMPTY_FRAMES = torch.empty((0, 3, 1, 1))
+
+    def __init__(
+        self,
+        video_list: List[str],
+        category_list: Union[str, List[str], None] = None,
+        frame_selector: Optional[FrameSelector] = None,
+        transform: Optional[FrameTransform] = None,
+        keyframe_helper_fpath: Optional[str] = None,
+    ):
+        """
+        Dataset constructor
+
+        Args:
+            video_list (List[str]): list of paths to video files
+            category_list (Union[str, List[str], None]): list of animal categories for each
+                video file. If it is a string, or None, this applies to all videos
+            frame_selector (Callable: KeyFrameList -> KeyFrameList):
+                selects keyframes to process, keyframes are given by
+                packet timestamps in timebase counts. If None, all keyframes
+                are selected (default: None)
+            transform (Callable: torch.Tensor -> torch.Tensor):
+                transforms a batch of RGB images (tensors of size [B, 3, H, W]),
+                returns a tensor of the same size. If None, no transform is
+                applied (default: None)
+
+        """
+        if type(category_list) == list:
+            self.category_list = category_list
+        else:
+            self.category_list = [category_list] * len(video_list)
+        assert len(video_list) == len(
+            self.category_list
+        ), "length of video and category lists must be equal"
+        self.video_list = video_list
+        self.frame_selector = frame_selector
+        self.transform = transform
+        self.keyframe_helper_data = (
+            read_keyframe_helper_data(keyframe_helper_fpath)
+            if keyframe_helper_fpath is not None
+            else None
+        )
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """
+        Gets selected keyframes from a given video
+
+        Args:
+            idx (int): video index in the video list file
+        Returns:
+            A dictionary containing two keys:
+                images (torch.Tensor): tensor of size [N, H, W, 3] or of size
+                    defined by the transform that contains keyframes data
+                categories (List[str]): categories of the frames
+        """
+        categories = [self.category_list[idx]]
+        fpath = self.video_list[idx]
+        keyframes = (
+            list_keyframes(fpath)
+            if self.keyframe_helper_data is None or idx not in self.keyframe_helper_data
+            else self.keyframe_helper_data[idx]
+        )
+        transform = self.transform
+        frame_selector = self.frame_selector
+        if not keyframes:
+            return {"images": self._EMPTY_FRAMES, "categories": []}
+        if frame_selector is not None:
+            keyframes = frame_selector(keyframes)
+        frames = read_keyframes(fpath, keyframes)
+        if not frames:
+            return {"images": self._EMPTY_FRAMES, "categories": []}
+        frames = np.stack([frame.to_rgb().to_ndarray() for frame in frames])
+        frames = torch.as_tensor(frames, device=torch.device("cpu"))
+        frames = frames[..., [2, 1, 0]]  # RGB -> BGR
+        frames = frames.permute(0, 3, 1, 2).float()  # NHWC -> NCHW
+        if transform is not None:
+            frames = transform(frames)
+        return {"images": frames, "categories": categories}
+
+    def __len__(self):
+        return len(self.video_list)

densepose/engine/__init__.py

@@ -1,3 +1,3 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .trainer import Trainer
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .trainer import Trainer

densepose/engine/trainer.py

@@ -1,258 +1,258 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-import logging
-import os
-from collections import OrderedDict
-from typing import List, Optional, Union
-import torch
-from torch import nn
-
-from detectron2.checkpoint import DetectionCheckpointer
-from detectron2.config import CfgNode
-from detectron2.engine import DefaultTrainer
-from detectron2.evaluation import (
-    DatasetEvaluator,
-    DatasetEvaluators,
-    inference_on_dataset,
-    print_csv_format,
-)
-from detectron2.solver.build import get_default_optimizer_params, maybe_add_gradient_clipping
-from detectron2.utils import comm
-from detectron2.utils.events import EventWriter, get_event_storage
-
-from densepose import DensePoseDatasetMapperTTA, DensePoseGeneralizedRCNNWithTTA, load_from_cfg
-from densepose.data import (
-    DatasetMapper,
-    build_combined_loader,
-    build_detection_test_loader,
-    build_detection_train_loader,
-    build_inference_based_loaders,
-    has_inference_based_loaders,
-)
-from densepose.evaluation.d2_evaluator_adapter import Detectron2COCOEvaluatorAdapter
-from densepose.evaluation.evaluator import DensePoseCOCOEvaluator, build_densepose_evaluator_storage
-from densepose.modeling.cse import Embedder
-
-
-class SampleCountingLoader:
-    def __init__(self, loader):
-        self.loader = loader
-
-    def __iter__(self):
-        it = iter(self.loader)
-        storage = get_event_storage()
-        while True:
-            try:
-                batch = next(it)
-                num_inst_per_dataset = {}
-                for data in batch:
-                    dataset_name = data["dataset"]
-                    if dataset_name not in num_inst_per_dataset:
-                        num_inst_per_dataset[dataset_name] = 0
-                    num_inst = len(data["instances"])
-                    num_inst_per_dataset[dataset_name] += num_inst
-                for dataset_name in num_inst_per_dataset:
-                    storage.put_scalar(f"batch/{dataset_name}", num_inst_per_dataset[dataset_name])
-                yield batch
-            except StopIteration:
-                break
-
-
-class SampleCountMetricPrinter(EventWriter):
-    def __init__(self):
-        self.logger = logging.getLogger(__name__)
-
-    def write(self):
-        storage = get_event_storage()
-        batch_stats_strs = []
-        for key, buf in storage.histories().items():
-            if key.startswith("batch/"):
-                batch_stats_strs.append(f"{key} {buf.avg(20)}")
-        self.logger.info(", ".join(batch_stats_strs))
-
-
-class Trainer(DefaultTrainer):
-    @classmethod
-    def extract_embedder_from_model(cls, model: nn.Module) -> Optional[Embedder]:
-        if isinstance(model, nn.parallel.DistributedDataParallel):
-            model = model.module
-        if hasattr(model, "roi_heads") and hasattr(model.roi_heads, "embedder"):
-            return model.roi_heads.embedder
-        return None
-
-    # TODO: the only reason to copy the base class code here is to pass the embedder from
-    # the model to the evaluator; that should be refactored to avoid unnecessary copy-pasting
-    @classmethod
-    def test(
-        cls,
-        cfg: CfgNode,
-        model: nn.Module,
-        evaluators: Optional[Union[DatasetEvaluator, List[DatasetEvaluator]]] = None,
-    ):
-        """
-        Args:
-            cfg (CfgNode):
-            model (nn.Module):
-            evaluators (DatasetEvaluator, list[DatasetEvaluator] or None): if None, will call
-                :meth:`build_evaluator`. Otherwise, must have the same length as
-                ``cfg.DATASETS.TEST``.
-
-        Returns:
-            dict: a dict of result metrics
-        """
-        logger = logging.getLogger(__name__)
-        if isinstance(evaluators, DatasetEvaluator):
-            evaluators = [evaluators]
-        if evaluators is not None:
-            assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
-                len(cfg.DATASETS.TEST), len(evaluators)
-            )
-
-        results = OrderedDict()
-        for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
-            data_loader = cls.build_test_loader(cfg, dataset_name)
-            # When evaluators are passed in as arguments,
-            # implicitly assume that evaluators can be created before data_loader.
-            if evaluators is not None:
-                evaluator = evaluators[idx]
-            else:
-                try:
-                    embedder = cls.extract_embedder_from_model(model)
-                    evaluator = cls.build_evaluator(cfg, dataset_name, embedder=embedder)
-                except NotImplementedError:
-                    logger.warn(
-                        "No evaluator found. Use `DefaultTrainer.test(evaluators=)`, "
-                        "or implement its `build_evaluator` method."
-                    )
-                    results[dataset_name] = {}
-                    continue
-            if cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE or comm.is_main_process():
-                results_i = inference_on_dataset(model, data_loader, evaluator)
-            else:
-                results_i = {}
-            results[dataset_name] = results_i
-            if comm.is_main_process():
-                assert isinstance(
-                    results_i, dict
-                ), "Evaluator must return a dict on the main process. Got {} instead.".format(
-                    results_i
-                )
-                logger.info("Evaluation results for {} in csv format:".format(dataset_name))
-                print_csv_format(results_i)
-
-        if len(results) == 1:
-            results = list(results.values())[0]
-        return results
-
-    @classmethod
-    def build_evaluator(
-        cls,
-        cfg: CfgNode,
-        dataset_name: str,
-        output_folder: Optional[str] = None,
-        embedder: Optional[Embedder] = None,
-    ) -> DatasetEvaluators:
-        if output_folder is None:
-            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
-        evaluators = []
-        distributed = cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE
-        # Note: we currently use COCO evaluator for both COCO and LVIS datasets
-        # to have compatible metrics. LVIS bbox evaluator could also be used
-        # with an adapter to properly handle filtered / mapped categories
-        # evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
-        # if evaluator_type == "coco":
-        #     evaluators.append(COCOEvaluator(dataset_name, output_dir=output_folder))
-        # elif evaluator_type == "lvis":
-        #     evaluators.append(LVISEvaluator(dataset_name, output_dir=output_folder))
-        evaluators.append(
-            Detectron2COCOEvaluatorAdapter(
-                dataset_name, output_dir=output_folder, distributed=distributed
-            )
-        )
-        if cfg.MODEL.DENSEPOSE_ON:
-            storage = build_densepose_evaluator_storage(cfg, output_folder)
-            evaluators.append(
-                DensePoseCOCOEvaluator(
-                    dataset_name,
-                    distributed,
-                    output_folder,
-                    evaluator_type=cfg.DENSEPOSE_EVALUATION.TYPE,
-                    min_iou_threshold=cfg.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD,
-                    storage=storage,
-                    embedder=embedder,
-                    should_evaluate_mesh_alignment=cfg.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT,
-                    mesh_alignment_mesh_names=cfg.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES,
-                )
-            )
-        return DatasetEvaluators(evaluators)
-
-    @classmethod
-    def build_optimizer(cls, cfg: CfgNode, model: nn.Module):
-        params = get_default_optimizer_params(
-            model,
-            base_lr=cfg.SOLVER.BASE_LR,
-            weight_decay_norm=cfg.SOLVER.WEIGHT_DECAY_NORM,
-            bias_lr_factor=cfg.SOLVER.BIAS_LR_FACTOR,
-            weight_decay_bias=cfg.SOLVER.WEIGHT_DECAY_BIAS,
-            overrides={
-                "features": {
-                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR,
-                },
-                "embeddings": {
-                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR,
-                },
-            },
-        )
-        optimizer = torch.optim.SGD(
-            params,
-            cfg.SOLVER.BASE_LR,
-            momentum=cfg.SOLVER.MOMENTUM,
-            nesterov=cfg.SOLVER.NESTEROV,
-            weight_decay=cfg.SOLVER.WEIGHT_DECAY,
-        )
-        # pyre-fixme[6]: For 2nd param expected `Type[Optimizer]` but got `SGD`.
-        return maybe_add_gradient_clipping(cfg, optimizer)
-
-    @classmethod
-    def build_test_loader(cls, cfg: CfgNode, dataset_name):
-        return build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, False))
-
-    @classmethod
-    def build_train_loader(cls, cfg: CfgNode):
-        data_loader = build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True))
-        if not has_inference_based_loaders(cfg):
-            return data_loader
-        model = cls.build_model(cfg)
-        model.to(cfg.BOOTSTRAP_MODEL.DEVICE)
-        DetectionCheckpointer(model).resume_or_load(cfg.BOOTSTRAP_MODEL.WEIGHTS, resume=False)
-        inference_based_loaders, ratios = build_inference_based_loaders(cfg, model)
-        loaders = [data_loader] + inference_based_loaders
-        ratios = [1.0] + ratios
-        combined_data_loader = build_combined_loader(cfg, loaders, ratios)
-        sample_counting_loader = SampleCountingLoader(combined_data_loader)
-        return sample_counting_loader
-
-    def build_writers(self):
-        writers = super().build_writers()
-        writers.append(SampleCountMetricPrinter())
-        return writers
-
-    @classmethod
-    def test_with_TTA(cls, cfg: CfgNode, model):
-        logger = logging.getLogger("detectron2.trainer")
-        # In the end of training, run an evaluation with TTA
-        # Only support some R-CNN models.
-        logger.info("Running inference with test-time augmentation ...")
-        transform_data = load_from_cfg(cfg)
-        model = DensePoseGeneralizedRCNNWithTTA(
-            cfg, model, transform_data, DensePoseDatasetMapperTTA(cfg)
-        )
-        evaluators = [
-            cls.build_evaluator(
-                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
-            )
-            for name in cfg.DATASETS.TEST
-        ]
-        res = cls.test(cfg, model, evaluators)  # pyre-ignore[6]
-        res = OrderedDict({k + "_TTA": v for k, v in res.items()})
-        return res
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import logging
+import os
+from collections import OrderedDict
+from typing import List, Optional, Union
+import torch
+from torch import nn
+
+from detectron2.checkpoint import DetectionCheckpointer
+from detectron2.config import CfgNode
+from detectron2.engine import DefaultTrainer
+from detectron2.evaluation import (
+    DatasetEvaluator,
+    DatasetEvaluators,
+    inference_on_dataset,
+    print_csv_format,
+)
+from detectron2.solver.build import get_default_optimizer_params, maybe_add_gradient_clipping
+from detectron2.utils import comm
+from detectron2.utils.events import EventWriter, get_event_storage
+
+from densepose import DensePoseDatasetMapperTTA, DensePoseGeneralizedRCNNWithTTA, load_from_cfg
+from densepose.data import (
+    DatasetMapper,
+    build_combined_loader,
+    build_detection_test_loader,
+    build_detection_train_loader,
+    build_inference_based_loaders,
+    has_inference_based_loaders,
+)
+from densepose.evaluation.d2_evaluator_adapter import Detectron2COCOEvaluatorAdapter
+from densepose.evaluation.evaluator import DensePoseCOCOEvaluator, build_densepose_evaluator_storage
+from densepose.modeling.cse import Embedder
+
+
+class SampleCountingLoader:
+    def __init__(self, loader):
+        self.loader = loader
+
+    def __iter__(self):
+        it = iter(self.loader)
+        storage = get_event_storage()
+        while True:
+            try:
+                batch = next(it)
+                num_inst_per_dataset = {}
+                for data in batch:
+                    dataset_name = data["dataset"]
+                    if dataset_name not in num_inst_per_dataset:
+                        num_inst_per_dataset[dataset_name] = 0
+                    num_inst = len(data["instances"])
+                    num_inst_per_dataset[dataset_name] += num_inst
+                for dataset_name in num_inst_per_dataset:
+                    storage.put_scalar(f"batch/{dataset_name}", num_inst_per_dataset[dataset_name])
+                yield batch
+            except StopIteration:
+                break
+
+
+class SampleCountMetricPrinter(EventWriter):
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+
+    def write(self):
+        storage = get_event_storage()
+        batch_stats_strs = []
+        for key, buf in storage.histories().items():
+            if key.startswith("batch/"):
+                batch_stats_strs.append(f"{key} {buf.avg(20)}")
+        self.logger.info(", ".join(batch_stats_strs))
+
+
+class Trainer(DefaultTrainer):
+    @classmethod
+    def extract_embedder_from_model(cls, model: nn.Module) -> Optional[Embedder]:
+        if isinstance(model, nn.parallel.DistributedDataParallel):
+            model = model.module
+        if hasattr(model, "roi_heads") and hasattr(model.roi_heads, "embedder"):
+            return model.roi_heads.embedder
+        return None
+
+    # TODO: the only reason to copy the base class code here is to pass the embedder from
+    # the model to the evaluator; that should be refactored to avoid unnecessary copy-pasting
+    @classmethod
+    def test(
+        cls,
+        cfg: CfgNode,
+        model: nn.Module,
+        evaluators: Optional[Union[DatasetEvaluator, List[DatasetEvaluator]]] = None,
+    ):
+        """
+        Args:
+            cfg (CfgNode):
+            model (nn.Module):
+            evaluators (DatasetEvaluator, list[DatasetEvaluator] or None): if None, will call
+                :meth:`build_evaluator`. Otherwise, must have the same length as
+                ``cfg.DATASETS.TEST``.
+
+        Returns:
+            dict: a dict of result metrics
+        """
+        logger = logging.getLogger(__name__)
+        if isinstance(evaluators, DatasetEvaluator):
+            evaluators = [evaluators]
+        if evaluators is not None:
+            assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
+                len(cfg.DATASETS.TEST), len(evaluators)
+            )
+
+        results = OrderedDict()
+        for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
+            data_loader = cls.build_test_loader(cfg, dataset_name)
+            # When evaluators are passed in as arguments,
+            # implicitly assume that evaluators can be created before data_loader.
+            if evaluators is not None:
+                evaluator = evaluators[idx]
+            else:
+                try:
+                    embedder = cls.extract_embedder_from_model(model)
+                    evaluator = cls.build_evaluator(cfg, dataset_name, embedder=embedder)
+                except NotImplementedError:
+                    logger.warn(
+                        "No evaluator found. Use `DefaultTrainer.test(evaluators=)`, "
+                        "or implement its `build_evaluator` method."
+                    )
+                    results[dataset_name] = {}
+                    continue
+            if cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE or comm.is_main_process():
+                results_i = inference_on_dataset(model, data_loader, evaluator)
+            else:
+                results_i = {}
+            results[dataset_name] = results_i
+            if comm.is_main_process():
+                assert isinstance(
+                    results_i, dict
+                ), "Evaluator must return a dict on the main process. Got {} instead.".format(
+                    results_i
+                )
+                logger.info("Evaluation results for {} in csv format:".format(dataset_name))
+                print_csv_format(results_i)
+
+        if len(results) == 1:
+            results = list(results.values())[0]
+        return results
+
+    @classmethod
+    def build_evaluator(
+        cls,
+        cfg: CfgNode,
+        dataset_name: str,
+        output_folder: Optional[str] = None,
+        embedder: Optional[Embedder] = None,
+    ) -> DatasetEvaluators:
+        if output_folder is None:
+            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
+        evaluators = []
+        distributed = cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE
+        # Note: we currently use COCO evaluator for both COCO and LVIS datasets
+        # to have compatible metrics. LVIS bbox evaluator could also be used
+        # with an adapter to properly handle filtered / mapped categories
+        # evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
+        # if evaluator_type == "coco":
+        #     evaluators.append(COCOEvaluator(dataset_name, output_dir=output_folder))
+        # elif evaluator_type == "lvis":
+        #     evaluators.append(LVISEvaluator(dataset_name, output_dir=output_folder))
+        evaluators.append(
+            Detectron2COCOEvaluatorAdapter(
+                dataset_name, output_dir=output_folder, distributed=distributed
+            )
+        )
+        if cfg.MODEL.DENSEPOSE_ON:
+            storage = build_densepose_evaluator_storage(cfg, output_folder)
+            evaluators.append(
+                DensePoseCOCOEvaluator(
+                    dataset_name,
+                    distributed,
+                    output_folder,
+                    evaluator_type=cfg.DENSEPOSE_EVALUATION.TYPE,
+                    min_iou_threshold=cfg.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD,
+                    storage=storage,
+                    embedder=embedder,
+                    should_evaluate_mesh_alignment=cfg.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT,
+                    mesh_alignment_mesh_names=cfg.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES,
+                )
+            )
+        return DatasetEvaluators(evaluators)
+
+    @classmethod
+    def build_optimizer(cls, cfg: CfgNode, model: nn.Module):
+        params = get_default_optimizer_params(
+            model,
+            base_lr=cfg.SOLVER.BASE_LR,
+            weight_decay_norm=cfg.SOLVER.WEIGHT_DECAY_NORM,
+            bias_lr_factor=cfg.SOLVER.BIAS_LR_FACTOR,
+            weight_decay_bias=cfg.SOLVER.WEIGHT_DECAY_BIAS,
+            overrides={
+                "features": {
+                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR,
+                },
+                "embeddings": {
+                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR,
+                },
+            },
+        )
+        optimizer = torch.optim.SGD(
+            params,
+            cfg.SOLVER.BASE_LR,
+            momentum=cfg.SOLVER.MOMENTUM,
+            nesterov=cfg.SOLVER.NESTEROV,
+            weight_decay=cfg.SOLVER.WEIGHT_DECAY,
+        )
+        # pyre-fixme[6]: For 2nd param expected `Type[Optimizer]` but got `SGD`.
+        return maybe_add_gradient_clipping(cfg, optimizer)
+
+    @classmethod
+    def build_test_loader(cls, cfg: CfgNode, dataset_name):
+        return build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, False))
+
+    @classmethod
+    def build_train_loader(cls, cfg: CfgNode):
+        data_loader = build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True))
+        if not has_inference_based_loaders(cfg):
+            return data_loader
+        model = cls.build_model(cfg)
+        model.to(cfg.BOOTSTRAP_MODEL.DEVICE)
+        DetectionCheckpointer(model).resume_or_load(cfg.BOOTSTRAP_MODEL.WEIGHTS, resume=False)
+        inference_based_loaders, ratios = build_inference_based_loaders(cfg, model)
+        loaders = [data_loader] + inference_based_loaders
+        ratios = [1.0] + ratios
+        combined_data_loader = build_combined_loader(cfg, loaders, ratios)
+        sample_counting_loader = SampleCountingLoader(combined_data_loader)
+        return sample_counting_loader
+
+    def build_writers(self):
+        writers = super().build_writers()
+        writers.append(SampleCountMetricPrinter())
+        return writers
+
+    @classmethod
+    def test_with_TTA(cls, cfg: CfgNode, model):
+        logger = logging.getLogger("detectron2.trainer")
+        # In the end of training, run an evaluation with TTA
+        # Only support some R-CNN models.
+        logger.info("Running inference with test-time augmentation ...")
+        transform_data = load_from_cfg(cfg)
+        model = DensePoseGeneralizedRCNNWithTTA(
+            cfg, model, transform_data, DensePoseDatasetMapperTTA(cfg)
+        )
+        evaluators = [
+            cls.build_evaluator(
+                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
+            )
+            for name in cfg.DATASETS.TEST
+        ]
+        res = cls.test(cfg, model, evaluators)  # pyre-ignore[6]
+        res = OrderedDict({k + "_TTA": v for k, v in res.items()})
+        return res

densepose/evaluation/__init__.py

@@ -1,3 +1,3 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .evaluator import DensePoseCOCOEvaluator
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .evaluator import DensePoseCOCOEvaluator

densepose/evaluation/d2_evaluator_adapter.py

@@ -1,50 +1,50 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from detectron2.data.catalog import Metadata
-from detectron2.evaluation import COCOEvaluator
-
-from densepose.data.datasets.coco import (
-    get_contiguous_id_to_category_id_map,
-    maybe_filter_categories_cocoapi,
-)
-
-
-def _maybe_add_iscrowd_annotations(cocoapi) -> None:
-    for ann in cocoapi.dataset["annotations"]:
-        if "iscrowd" not in ann:
-            ann["iscrowd"] = 0
-
-
-class Detectron2COCOEvaluatorAdapter(COCOEvaluator):
-    def __init__(
-        self,
-        dataset_name,
-        output_dir=None,
-        distributed=True,
-    ):
-        super().__init__(dataset_name, output_dir=output_dir, distributed=distributed)
-        maybe_filter_categories_cocoapi(dataset_name, self._coco_api)
-        _maybe_add_iscrowd_annotations(self._coco_api)
-        # substitute category metadata to account for categories
-        # that are mapped to the same contiguous id
-        if hasattr(self._metadata, "thing_dataset_id_to_contiguous_id"):
-            self._maybe_substitute_metadata()
-
-    def _maybe_substitute_metadata(self):
-        cont_id_2_cat_id = get_contiguous_id_to_category_id_map(self._metadata)
-        cat_id_2_cont_id = self._metadata.thing_dataset_id_to_contiguous_id
-        if len(cont_id_2_cat_id) == len(cat_id_2_cont_id):
-            return
-
-        cat_id_2_cont_id_injective = {}
-        for cat_id, cont_id in cat_id_2_cont_id.items():
-            if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
-                cat_id_2_cont_id_injective[cat_id] = cont_id
-
-        metadata_new = Metadata(name=self._metadata.name)
-        for key, value in self._metadata.__dict__.items():
-            if key == "thing_dataset_id_to_contiguous_id":
-                setattr(metadata_new, key, cat_id_2_cont_id_injective)
-            else:
-                setattr(metadata_new, key, value)
-        self._metadata = metadata_new
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from detectron2.data.catalog import Metadata
+from detectron2.evaluation import COCOEvaluator
+
+from densepose.data.datasets.coco import (
+    get_contiguous_id_to_category_id_map,
+    maybe_filter_categories_cocoapi,
+)
+
+
+def _maybe_add_iscrowd_annotations(cocoapi) -> None:
+    for ann in cocoapi.dataset["annotations"]:
+        if "iscrowd" not in ann:
+            ann["iscrowd"] = 0
+
+
+class Detectron2COCOEvaluatorAdapter(COCOEvaluator):
+    def __init__(
+        self,
+        dataset_name,
+        output_dir=None,
+        distributed=True,
+    ):
+        super().__init__(dataset_name, output_dir=output_dir, distributed=distributed)
+        maybe_filter_categories_cocoapi(dataset_name, self._coco_api)
+        _maybe_add_iscrowd_annotations(self._coco_api)
+        # substitute category metadata to account for categories
+        # that are mapped to the same contiguous id
+        if hasattr(self._metadata, "thing_dataset_id_to_contiguous_id"):
+            self._maybe_substitute_metadata()
+
+    def _maybe_substitute_metadata(self):
+        cont_id_2_cat_id = get_contiguous_id_to_category_id_map(self._metadata)
+        cat_id_2_cont_id = self._metadata.thing_dataset_id_to_contiguous_id
+        if len(cont_id_2_cat_id) == len(cat_id_2_cont_id):
+            return
+
+        cat_id_2_cont_id_injective = {}
+        for cat_id, cont_id in cat_id_2_cont_id.items():
+            if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
+                cat_id_2_cont_id_injective[cat_id] = cont_id
+
+        metadata_new = Metadata(name=self._metadata.name)
+        for key, value in self._metadata.__dict__.items():
+            if key == "thing_dataset_id_to_contiguous_id":
+                setattr(metadata_new, key, cat_id_2_cont_id_injective)
+            else:
+                setattr(metadata_new, key, value)
+        self._metadata = metadata_new

densepose/evaluation/evaluator.py

@@ -1,421 +1,421 @@
-# -*- coding: utf-8 -*-
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import contextlib
-import copy
-import io
-import itertools
-import logging
-import numpy as np
-import os
-from collections import OrderedDict
-from typing import Dict, Iterable, List, Optional
-import pycocotools.mask as mask_utils
-import torch
-from pycocotools.coco import COCO
-from tabulate import tabulate
-
-from detectron2.config import CfgNode
-from detectron2.data import MetadataCatalog
-from detectron2.evaluation import DatasetEvaluator
-from detectron2.structures import BoxMode
-from detectron2.utils.comm import gather, get_rank, is_main_process, synchronize
-from detectron2.utils.file_io import PathManager
-from detectron2.utils.logger import create_small_table
-
-from densepose.converters import ToChartResultConverter, ToMaskConverter
-from densepose.data.datasets.coco import maybe_filter_and_map_categories_cocoapi
-from densepose.structures import (
-    DensePoseChartPredictorOutput,
-    DensePoseEmbeddingPredictorOutput,
-    quantize_densepose_chart_result,
-)
-
-from .densepose_coco_evaluation import DensePoseCocoEval, DensePoseEvalMode
-from .mesh_alignment_evaluator import MeshAlignmentEvaluator
-from .tensor_storage import (
-    SingleProcessFileTensorStorage,
-    SingleProcessRamTensorStorage,
-    SingleProcessTensorStorage,
-    SizeData,
-    storage_gather,
-)
-
-
-class DensePoseCOCOEvaluator(DatasetEvaluator):
-    def __init__(
-        self,
-        dataset_name,
-        distributed,
-        output_dir=None,
-        evaluator_type: str = "iuv",
-        min_iou_threshold: float = 0.5,
-        storage: Optional[SingleProcessTensorStorage] = None,
-        embedder=None,
-        should_evaluate_mesh_alignment: bool = False,
-        mesh_alignment_mesh_names: Optional[List[str]] = None,
-    ):
-        self._embedder = embedder
-        self._distributed = distributed
-        self._output_dir = output_dir
-        self._evaluator_type = evaluator_type
-        self._storage = storage
-        self._should_evaluate_mesh_alignment = should_evaluate_mesh_alignment
-
-        assert not (
-            should_evaluate_mesh_alignment and embedder is None
-        ), "Mesh alignment evaluation is activated, but no vertex embedder provided!"
-        if should_evaluate_mesh_alignment:
-            self._mesh_alignment_evaluator = MeshAlignmentEvaluator(
-                embedder,
-                mesh_alignment_mesh_names,
-            )
-
-        self._cpu_device = torch.device("cpu")
-        self._logger = logging.getLogger(__name__)
-
-        self._metadata = MetadataCatalog.get(dataset_name)
-        self._min_threshold = min_iou_threshold
-        json_file = PathManager.get_local_path(self._metadata.json_file)
-        with contextlib.redirect_stdout(io.StringIO()):
-            self._coco_api = COCO(json_file)
-        maybe_filter_and_map_categories_cocoapi(dataset_name, self._coco_api)
-
-    def reset(self):
-        self._predictions = []
-
-    def process(self, inputs, outputs):
-        """
-        Args:
-            inputs: the inputs to a COCO model (e.g., GeneralizedRCNN).
-                It is a list of dict. Each dict corresponds to an image and
-                contains keys like "height", "width", "file_name", "image_id".
-            outputs: the outputs of a COCO model. It is a list of dicts with key
-                "instances" that contains :class:`Instances`.
-                The :class:`Instances` object needs to have `densepose` field.
-        """
-        for input, output in zip(inputs, outputs):
-            instances = output["instances"].to(self._cpu_device)
-            if not instances.has("pred_densepose"):
-                continue
-            prediction_list = prediction_to_dict(
-                instances,
-                input["image_id"],
-                self._embedder,
-                self._metadata.class_to_mesh_name,
-                self._storage is not None,
-            )
-            if self._storage is not None:
-                for prediction_dict in prediction_list:
-                    dict_to_store = {}
-                    for field_name in self._storage.data_schema:
-                        dict_to_store[field_name] = prediction_dict[field_name]
-                    record_id = self._storage.put(dict_to_store)
-                    prediction_dict["record_id"] = record_id
-                    prediction_dict["rank"] = get_rank()
-                    for field_name in self._storage.data_schema:
-                        del prediction_dict[field_name]
-            self._predictions.extend(prediction_list)
-
-    def evaluate(self, img_ids=None):
-        if self._distributed:
-            synchronize()
-            predictions = gather(self._predictions)
-            predictions = list(itertools.chain(*predictions))
-        else:
-            predictions = self._predictions
-
-        multi_storage = storage_gather(self._storage) if self._storage is not None else None
-
-        if not is_main_process():
-            return
-        return copy.deepcopy(self._eval_predictions(predictions, multi_storage, img_ids))
-
-    def _eval_predictions(self, predictions, multi_storage=None, img_ids=None):
-        """
-        Evaluate predictions on densepose.
-        Return results with the metrics of the tasks.
-        """
-        self._logger.info("Preparing results for COCO format ...")
-
-        if self._output_dir:
-            PathManager.mkdirs(self._output_dir)
-            file_path = os.path.join(self._output_dir, "coco_densepose_predictions.pth")
-            with PathManager.open(file_path, "wb") as f:
-                torch.save(predictions, f)
-
-        self._logger.info("Evaluating predictions ...")
-        res = OrderedDict()
-        results_gps, results_gpsm, results_segm = _evaluate_predictions_on_coco(
-            self._coco_api,
-            predictions,
-            multi_storage,
-            self._embedder,
-            class_names=self._metadata.get("thing_classes"),
-            min_threshold=self._min_threshold,
-            img_ids=img_ids,
-        )
-        res["densepose_gps"] = results_gps
-        res["densepose_gpsm"] = results_gpsm
-        res["densepose_segm"] = results_segm
-        if self._should_evaluate_mesh_alignment:
-            res["densepose_mesh_alignment"] = self._evaluate_mesh_alignment()
-        return res
-
-    def _evaluate_mesh_alignment(self):
-        self._logger.info("Mesh alignment evaluation ...")
-        mean_ge, mean_gps, per_mesh_metrics = self._mesh_alignment_evaluator.evaluate()
-        results = {
-            "GE": mean_ge * 100,
-            "GPS": mean_gps * 100,
-        }
-        mesh_names = set()
-        for metric_name in per_mesh_metrics:
-            for mesh_name, value in per_mesh_metrics[metric_name].items():
-                results[f"{metric_name}-{mesh_name}"] = value * 100
-                mesh_names.add(mesh_name)
-        self._print_mesh_alignment_results(results, mesh_names)
-        return results
-
-    def _print_mesh_alignment_results(self, results: Dict[str, float], mesh_names: Iterable[str]):
-        self._logger.info("Evaluation results for densepose, mesh alignment:")
-        self._logger.info(f'| {"Mesh":13s} | {"GErr":7s} | {"GPS":7s} |')
-        self._logger.info("| :-----------: | :-----: | :-----: |")
-        for mesh_name in mesh_names:
-            ge_key = f"GE-{mesh_name}"
-            ge_str = f"{results[ge_key]:.4f}" if ge_key in results else " "
-            gps_key = f"GPS-{mesh_name}"
-            gps_str = f"{results[gps_key]:.4f}" if gps_key in results else " "
-            self._logger.info(f"| {mesh_name:13s} | {ge_str:7s} | {gps_str:7s} |")
-        self._logger.info("| :-------------------------------: |")
-        ge_key = "GE"
-        ge_str = f"{results[ge_key]:.4f}" if ge_key in results else " "
-        gps_key = "GPS"
-        gps_str = f"{results[gps_key]:.4f}" if gps_key in results else " "
-        self._logger.info(f'| {"MEAN":13s} | {ge_str:7s} | {gps_str:7s} |')
-
-
-def prediction_to_dict(instances, img_id, embedder, class_to_mesh_name, use_storage):
-    """
-    Args:
-        instances (Instances): the output of the model
-        img_id (str): the image id in COCO
-
-    Returns:
-        list[dict]: the results in densepose evaluation format
-    """
-    scores = instances.scores.tolist()
-    classes = instances.pred_classes.tolist()
-    raw_boxes_xywh = BoxMode.convert(
-        instances.pred_boxes.tensor.clone(), BoxMode.XYXY_ABS, BoxMode.XYWH_ABS
-    )
-
-    if isinstance(instances.pred_densepose, DensePoseEmbeddingPredictorOutput):
-        results_densepose = densepose_cse_predictions_to_dict(
-            instances, embedder, class_to_mesh_name, use_storage
-        )
-    elif isinstance(instances.pred_densepose, DensePoseChartPredictorOutput):
-        if not use_storage:
-            results_densepose = densepose_chart_predictions_to_dict(instances)
-        else:
-            results_densepose = densepose_chart_predictions_to_storage_dict(instances)
-
-    results = []
-    for k in range(len(instances)):
-        result = {
-            "image_id": img_id,
-            "category_id": classes[k],
-            "bbox": raw_boxes_xywh[k].tolist(),
-            "score": scores[k],
-        }
-        results.append({**result, **results_densepose[k]})
-    return results
-
-
-def densepose_chart_predictions_to_dict(instances):
-    segmentations = ToMaskConverter.convert(
-        instances.pred_densepose, instances.pred_boxes, instances.image_size
-    )
-
-    results = []
-    for k in range(len(instances)):
-        densepose_results_quantized = quantize_densepose_chart_result(
-            ToChartResultConverter.convert(instances.pred_densepose[k], instances.pred_boxes[k])
-        )
-        densepose_results_quantized.labels_uv_uint8 = (
-            densepose_results_quantized.labels_uv_uint8.cpu()
-        )
-        segmentation = segmentations.tensor[k]
-        segmentation_encoded = mask_utils.encode(
-            np.require(segmentation.numpy(), dtype=np.uint8, requirements=["F"])
-        )
-        segmentation_encoded["counts"] = segmentation_encoded["counts"].decode("utf-8")
-        result = {
-            "densepose": densepose_results_quantized,
-            "segmentation": segmentation_encoded,
-        }
-        results.append(result)
-    return results
-
-
-def densepose_chart_predictions_to_storage_dict(instances):
-    results = []
-    for k in range(len(instances)):
-        densepose_predictor_output = instances.pred_densepose[k]
-        result = {
-            "coarse_segm": densepose_predictor_output.coarse_segm.squeeze(0).cpu(),
-            "fine_segm": densepose_predictor_output.fine_segm.squeeze(0).cpu(),
-            "u": densepose_predictor_output.u.squeeze(0).cpu(),
-            "v": densepose_predictor_output.v.squeeze(0).cpu(),
-        }
-        results.append(result)
-    return results
-
-
-def densepose_cse_predictions_to_dict(instances, embedder, class_to_mesh_name, use_storage):
-    results = []
-    for k in range(len(instances)):
-        cse = instances.pred_densepose[k]
-        results.append(
-            {
-                "coarse_segm": cse.coarse_segm[0].cpu(),
-                "embedding": cse.embedding[0].cpu(),
-            }
-        )
-    return results
-
-
-def _evaluate_predictions_on_coco(
-    coco_gt,
-    coco_results,
-    multi_storage=None,
-    embedder=None,
-    class_names=None,
-    min_threshold: float = 0.5,
-    img_ids=None,
-):
-    logger = logging.getLogger(__name__)
-
-    densepose_metrics = _get_densepose_metrics(min_threshold)
-    if len(coco_results) == 0:  # cocoapi does not handle empty results very well
-        logger.warn("No predictions from the model! Set scores to -1")
-        results_gps = {metric: -1 for metric in densepose_metrics}
-        results_gpsm = {metric: -1 for metric in densepose_metrics}
-        results_segm = {metric: -1 for metric in densepose_metrics}
-        return results_gps, results_gpsm, results_segm
-
-    coco_dt = coco_gt.loadRes(coco_results)
-
-    results = []
-    for eval_mode_name in ["GPS", "GPSM", "IOU"]:
-        eval_mode = getattr(DensePoseEvalMode, eval_mode_name)
-        coco_eval = DensePoseCocoEval(
-            coco_gt, coco_dt, "densepose", multi_storage, embedder, dpEvalMode=eval_mode
-        )
-        result = _derive_results_from_coco_eval(
-            coco_eval, eval_mode_name, densepose_metrics, class_names, min_threshold, img_ids
-        )
-        results.append(result)
-    return results
-
-
-def _get_densepose_metrics(min_threshold: float = 0.5):
-    metrics = ["AP"]
-    if min_threshold <= 0.201:
-        metrics += ["AP20"]
-    if min_threshold <= 0.301:
-        metrics += ["AP30"]
-    if min_threshold <= 0.401:
-        metrics += ["AP40"]
-    metrics.extend(["AP50", "AP75", "APm", "APl", "AR", "AR50", "AR75", "ARm", "ARl"])
-    return metrics
-
-
-def _derive_results_from_coco_eval(
-    coco_eval, eval_mode_name, metrics, class_names, min_threshold: float, img_ids
-):
-    if img_ids is not None:
-        coco_eval.params.imgIds = img_ids
-    coco_eval.params.iouThrs = np.linspace(
-        min_threshold, 0.95, int(np.round((0.95 - min_threshold) / 0.05)) + 1, endpoint=True
-    )
-    coco_eval.evaluate()
-    coco_eval.accumulate()
-    coco_eval.summarize()
-    results = {metric: float(coco_eval.stats[idx] * 100) for idx, metric in enumerate(metrics)}
-    logger = logging.getLogger(__name__)
-    logger.info(
-        f"Evaluation results for densepose, {eval_mode_name} metric: \n"
-        + create_small_table(results)
-    )
-    if class_names is None or len(class_names) <= 1:
-        return results
-
-    # Compute per-category AP, the same way as it is done in D2
-    # (see detectron2/evaluation/coco_evaluation.py):
-    precisions = coco_eval.eval["precision"]
-    # precision has dims (iou, recall, cls, area range, max dets)
-    assert len(class_names) == precisions.shape[2]
-
-    results_per_category = []
-    for idx, name in enumerate(class_names):
-        # area range index 0: all area ranges
-        # max dets index -1: typically 100 per image
-        precision = precisions[:, :, idx, 0, -1]
-        precision = precision[precision > -1]
-        ap = np.mean(precision) if precision.size else float("nan")
-        results_per_category.append((f"{name}", float(ap * 100)))
-
-    # tabulate it
-    n_cols = min(6, len(results_per_category) * 2)
-    results_flatten = list(itertools.chain(*results_per_category))
-    results_2d = itertools.zip_longest(*[results_flatten[i::n_cols] for i in range(n_cols)])
-    table = tabulate(
-        results_2d,
-        tablefmt="pipe",
-        floatfmt=".3f",
-        headers=["category", "AP"] * (n_cols // 2),
-        numalign="left",
-    )
-    logger.info(f"Per-category {eval_mode_name} AP: \n" + table)
-
-    results.update({"AP-" + name: ap for name, ap in results_per_category})
-    return results
-
-
-def build_densepose_evaluator_storage(cfg: CfgNode, output_folder: str):
-    storage_spec = cfg.DENSEPOSE_EVALUATION.STORAGE
-    if storage_spec == "none":
-        return None
-    evaluator_type = cfg.DENSEPOSE_EVALUATION.TYPE
-    # common output tensor sizes
-    hout = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE
-    wout = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE
-    n_csc = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS
-    # specific output tensors
-    if evaluator_type == "iuv":
-        n_fsc = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES + 1
-        schema = {
-            "coarse_segm": SizeData(dtype="float32", shape=(n_csc, hout, wout)),
-            "fine_segm": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
-            "u": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
-            "v": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
-        }
-    elif evaluator_type == "cse":
-        embed_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE
-        schema = {
-            "coarse_segm": SizeData(dtype="float32", shape=(n_csc, hout, wout)),
-            "embedding": SizeData(dtype="float32", shape=(embed_size, hout, wout)),
-        }
-    else:
-        raise ValueError(f"Unknown evaluator type: {evaluator_type}")
-    # storage types
-    if storage_spec == "ram":
-        storage = SingleProcessRamTensorStorage(schema, io.BytesIO())
-    elif storage_spec == "file":
-        fpath = os.path.join(output_folder, f"DensePoseEvaluatorStorage.{get_rank()}.bin")
-        PathManager.mkdirs(output_folder)
-        storage = SingleProcessFileTensorStorage(schema, fpath, "wb")
-    else:
-        raise ValueError(f"Unknown storage specification: {storage_spec}")
-    return storage
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import contextlib
+import copy
+import io
+import itertools
+import logging
+import numpy as np
+import os
+from collections import OrderedDict
+from typing import Dict, Iterable, List, Optional
+import pycocotools.mask as mask_utils
+import torch
+from pycocotools.coco import COCO
+from tabulate import tabulate
+
+from detectron2.config import CfgNode
+from detectron2.data import MetadataCatalog
+from detectron2.evaluation import DatasetEvaluator
+from detectron2.structures import BoxMode
+from detectron2.utils.comm import gather, get_rank, is_main_process, synchronize
+from detectron2.utils.file_io import PathManager
+from detectron2.utils.logger import create_small_table
+
+from densepose.converters import ToChartResultConverter, ToMaskConverter
+from densepose.data.datasets.coco import maybe_filter_and_map_categories_cocoapi
+from densepose.structures import (
+    DensePoseChartPredictorOutput,
+    DensePoseEmbeddingPredictorOutput,
+    quantize_densepose_chart_result,
+)
+
+from .densepose_coco_evaluation import DensePoseCocoEval, DensePoseEvalMode
+from .mesh_alignment_evaluator import MeshAlignmentEvaluator
+from .tensor_storage import (
+    SingleProcessFileTensorStorage,
+    SingleProcessRamTensorStorage,
+    SingleProcessTensorStorage,
+    SizeData,
+    storage_gather,
+)
+
+
+class DensePoseCOCOEvaluator(DatasetEvaluator):
+    def __init__(
+        self,
+        dataset_name,
+        distributed,
+        output_dir=None,
+        evaluator_type: str = "iuv",
+        min_iou_threshold: float = 0.5,
+        storage: Optional[SingleProcessTensorStorage] = None,
+        embedder=None,
+        should_evaluate_mesh_alignment: bool = False,
+        mesh_alignment_mesh_names: Optional[List[str]] = None,
+    ):
+        self._embedder = embedder
+        self._distributed = distributed
+        self._output_dir = output_dir
+        self._evaluator_type = evaluator_type
+        self._storage = storage
+        self._should_evaluate_mesh_alignment = should_evaluate_mesh_alignment
+
+        assert not (
+            should_evaluate_mesh_alignment and embedder is None
+        ), "Mesh alignment evaluation is activated, but no vertex embedder provided!"
+        if should_evaluate_mesh_alignment:
+            self._mesh_alignment_evaluator = MeshAlignmentEvaluator(
+                embedder,
+                mesh_alignment_mesh_names,
+            )
+
+        self._cpu_device = torch.device("cpu")
+        self._logger = logging.getLogger(__name__)
+
+        self._metadata = MetadataCatalog.get(dataset_name)
+        self._min_threshold = min_iou_threshold
+        json_file = PathManager.get_local_path(self._metadata.json_file)
+        with contextlib.redirect_stdout(io.StringIO()):
+            self._coco_api = COCO(json_file)
+        maybe_filter_and_map_categories_cocoapi(dataset_name, self._coco_api)
+
+    def reset(self):
+        self._predictions = []
+
+    def process(self, inputs, outputs):
+        """
+        Args:
+            inputs: the inputs to a COCO model (e.g., GeneralizedRCNN).
+                It is a list of dict. Each dict corresponds to an image and
+                contains keys like "height", "width", "file_name", "image_id".
+            outputs: the outputs of a COCO model. It is a list of dicts with key
+                "instances" that contains :class:`Instances`.
+                The :class:`Instances` object needs to have `densepose` field.
+        """
+        for input, output in zip(inputs, outputs):
+            instances = output["instances"].to(self._cpu_device)
+            if not instances.has("pred_densepose"):
+                continue
+            prediction_list = prediction_to_dict(
+                instances,
+                input["image_id"],
+                self._embedder,
+                self._metadata.class_to_mesh_name,
+                self._storage is not None,
+            )
+            if self._storage is not None:
+                for prediction_dict in prediction_list:
+                    dict_to_store = {}
+                    for field_name in self._storage.data_schema:
+                        dict_to_store[field_name] = prediction_dict[field_name]
+                    record_id = self._storage.put(dict_to_store)
+                    prediction_dict["record_id"] = record_id
+                    prediction_dict["rank"] = get_rank()
+                    for field_name in self._storage.data_schema:
+                        del prediction_dict[field_name]
+            self._predictions.extend(prediction_list)
+
+    def evaluate(self, img_ids=None):
+        if self._distributed:
+            synchronize()
+            predictions = gather(self._predictions)
+            predictions = list(itertools.chain(*predictions))
+        else:
+            predictions = self._predictions
+
+        multi_storage = storage_gather(self._storage) if self._storage is not None else None
+
+        if not is_main_process():
+            return
+        return copy.deepcopy(self._eval_predictions(predictions, multi_storage, img_ids))
+
+    def _eval_predictions(self, predictions, multi_storage=None, img_ids=None):
+        """
+        Evaluate predictions on densepose.
+        Return results with the metrics of the tasks.
+        """
+        self._logger.info("Preparing results for COCO format ...")
+
+        if self._output_dir:
+            PathManager.mkdirs(self._output_dir)
+            file_path = os.path.join(self._output_dir, "coco_densepose_predictions.pth")
+            with PathManager.open(file_path, "wb") as f:
+                torch.save(predictions, f)
+
+        self._logger.info("Evaluating predictions ...")
+        res = OrderedDict()
+        results_gps, results_gpsm, results_segm = _evaluate_predictions_on_coco(
+            self._coco_api,
+            predictions,
+            multi_storage,
+            self._embedder,
+            class_names=self._metadata.get("thing_classes"),
+            min_threshold=self._min_threshold,
+            img_ids=img_ids,
+        )
+        res["densepose_gps"] = results_gps
+        res["densepose_gpsm"] = results_gpsm
+        res["densepose_segm"] = results_segm
+        if self._should_evaluate_mesh_alignment:
+            res["densepose_mesh_alignment"] = self._evaluate_mesh_alignment()
+        return res
+
+    def _evaluate_mesh_alignment(self):
+        self._logger.info("Mesh alignment evaluation ...")
+        mean_ge, mean_gps, per_mesh_metrics = self._mesh_alignment_evaluator.evaluate()
+        results = {
+            "GE": mean_ge * 100,
+            "GPS": mean_gps * 100,
+        }
+        mesh_names = set()
+        for metric_name in per_mesh_metrics:
+            for mesh_name, value in per_mesh_metrics[metric_name].items():
+                results[f"{metric_name}-{mesh_name}"] = value * 100
+                mesh_names.add(mesh_name)
+        self._print_mesh_alignment_results(results, mesh_names)
+        return results
+
+    def _print_mesh_alignment_results(self, results: Dict[str, float], mesh_names: Iterable[str]):
+        self._logger.info("Evaluation results for densepose, mesh alignment:")
+        self._logger.info(f'| {"Mesh":13s} | {"GErr":7s} | {"GPS":7s} |')
+        self._logger.info("| :-----------: | :-----: | :-----: |")
+        for mesh_name in mesh_names:
+            ge_key = f"GE-{mesh_name}"
+            ge_str = f"{results[ge_key]:.4f}" if ge_key in results else " "
+            gps_key = f"GPS-{mesh_name}"
+            gps_str = f"{results[gps_key]:.4f}" if gps_key in results else " "
+            self._logger.info(f"| {mesh_name:13s} | {ge_str:7s} | {gps_str:7s} |")
+        self._logger.info("| :-------------------------------: |")
+        ge_key = "GE"
+        ge_str = f"{results[ge_key]:.4f}" if ge_key in results else " "
+        gps_key = "GPS"
+        gps_str = f"{results[gps_key]:.4f}" if gps_key in results else " "
+        self._logger.info(f'| {"MEAN":13s} | {ge_str:7s} | {gps_str:7s} |')
+
+
+def prediction_to_dict(instances, img_id, embedder, class_to_mesh_name, use_storage):
+    """
+    Args:
+        instances (Instances): the output of the model
+        img_id (str): the image id in COCO
+
+    Returns:
+        list[dict]: the results in densepose evaluation format
+    """
+    scores = instances.scores.tolist()
+    classes = instances.pred_classes.tolist()
+    raw_boxes_xywh = BoxMode.convert(
+        instances.pred_boxes.tensor.clone(), BoxMode.XYXY_ABS, BoxMode.XYWH_ABS
+    )
+
+    if isinstance(instances.pred_densepose, DensePoseEmbeddingPredictorOutput):
+        results_densepose = densepose_cse_predictions_to_dict(
+            instances, embedder, class_to_mesh_name, use_storage
+        )
+    elif isinstance(instances.pred_densepose, DensePoseChartPredictorOutput):
+        if not use_storage:
+            results_densepose = densepose_chart_predictions_to_dict(instances)
+        else:
+            results_densepose = densepose_chart_predictions_to_storage_dict(instances)
+
+    results = []
+    for k in range(len(instances)):
+        result = {
+            "image_id": img_id,
+            "category_id": classes[k],
+            "bbox": raw_boxes_xywh[k].tolist(),
+            "score": scores[k],
+        }
+        results.append({**result, **results_densepose[k]})
+    return results
+
+
+def densepose_chart_predictions_to_dict(instances):
+    segmentations = ToMaskConverter.convert(
+        instances.pred_densepose, instances.pred_boxes, instances.image_size
+    )
+
+    results = []
+    for k in range(len(instances)):
+        densepose_results_quantized = quantize_densepose_chart_result(
+            ToChartResultConverter.convert(instances.pred_densepose[k], instances.pred_boxes[k])
+        )
+        densepose_results_quantized.labels_uv_uint8 = (
+            densepose_results_quantized.labels_uv_uint8.cpu()
+        )
+        segmentation = segmentations.tensor[k]
+        segmentation_encoded = mask_utils.encode(
+            np.require(segmentation.numpy(), dtype=np.uint8, requirements=["F"])
+        )
+        segmentation_encoded["counts"] = segmentation_encoded["counts"].decode("utf-8")
+        result = {
+            "densepose": densepose_results_quantized,
+            "segmentation": segmentation_encoded,
+        }
+        results.append(result)
+    return results
+
+
+def densepose_chart_predictions_to_storage_dict(instances):
+    results = []
+    for k in range(len(instances)):
+        densepose_predictor_output = instances.pred_densepose[k]
+        result = {
+            "coarse_segm": densepose_predictor_output.coarse_segm.squeeze(0).cpu(),
+            "fine_segm": densepose_predictor_output.fine_segm.squeeze(0).cpu(),
+            "u": densepose_predictor_output.u.squeeze(0).cpu(),
+            "v": densepose_predictor_output.v.squeeze(0).cpu(),
+        }
+        results.append(result)
+    return results
+
+
+def densepose_cse_predictions_to_dict(instances, embedder, class_to_mesh_name, use_storage):
+    results = []
+    for k in range(len(instances)):
+        cse = instances.pred_densepose[k]
+        results.append(
+            {
+                "coarse_segm": cse.coarse_segm[0].cpu(),
+                "embedding": cse.embedding[0].cpu(),
+            }
+        )
+    return results
+
+
+def _evaluate_predictions_on_coco(
+    coco_gt,
+    coco_results,
+    multi_storage=None,
+    embedder=None,
+    class_names=None,
+    min_threshold: float = 0.5,
+    img_ids=None,
+):
+    logger = logging.getLogger(__name__)
+
+    densepose_metrics = _get_densepose_metrics(min_threshold)
+    if len(coco_results) == 0:  # cocoapi does not handle empty results very well
+        logger.warn("No predictions from the model! Set scores to -1")
+        results_gps = {metric: -1 for metric in densepose_metrics}
+        results_gpsm = {metric: -1 for metric in densepose_metrics}
+        results_segm = {metric: -1 for metric in densepose_metrics}
+        return results_gps, results_gpsm, results_segm
+
+    coco_dt = coco_gt.loadRes(coco_results)
+
+    results = []
+    for eval_mode_name in ["GPS", "GPSM", "IOU"]:
+        eval_mode = getattr(DensePoseEvalMode, eval_mode_name)
+        coco_eval = DensePoseCocoEval(
+            coco_gt, coco_dt, "densepose", multi_storage, embedder, dpEvalMode=eval_mode
+        )
+        result = _derive_results_from_coco_eval(
+            coco_eval, eval_mode_name, densepose_metrics, class_names, min_threshold, img_ids
+        )
+        results.append(result)
+    return results
+
+
+def _get_densepose_metrics(min_threshold: float = 0.5):
+    metrics = ["AP"]
+    if min_threshold <= 0.201:
+        metrics += ["AP20"]
+    if min_threshold <= 0.301:
+        metrics += ["AP30"]
+    if min_threshold <= 0.401:
+        metrics += ["AP40"]
+    metrics.extend(["AP50", "AP75", "APm", "APl", "AR", "AR50", "AR75", "ARm", "ARl"])
+    return metrics
+
+
+def _derive_results_from_coco_eval(
+    coco_eval, eval_mode_name, metrics, class_names, min_threshold: float, img_ids
+):
+    if img_ids is not None:
+        coco_eval.params.imgIds = img_ids
+    coco_eval.params.iouThrs = np.linspace(
+        min_threshold, 0.95, int(np.round((0.95 - min_threshold) / 0.05)) + 1, endpoint=True
+    )
+    coco_eval.evaluate()
+    coco_eval.accumulate()
+    coco_eval.summarize()
+    results = {metric: float(coco_eval.stats[idx] * 100) for idx, metric in enumerate(metrics)}
+    logger = logging.getLogger(__name__)
+    logger.info(
+        f"Evaluation results for densepose, {eval_mode_name} metric: \n"
+        + create_small_table(results)
+    )
+    if class_names is None or len(class_names) <= 1:
+        return results
+
+    # Compute per-category AP, the same way as it is done in D2
+    # (see detectron2/evaluation/coco_evaluation.py):
+    precisions = coco_eval.eval["precision"]
+    # precision has dims (iou, recall, cls, area range, max dets)
+    assert len(class_names) == precisions.shape[2]
+
+    results_per_category = []
+    for idx, name in enumerate(class_names):
+        # area range index 0: all area ranges
+        # max dets index -1: typically 100 per image
+        precision = precisions[:, :, idx, 0, -1]
+        precision = precision[precision > -1]
+        ap = np.mean(precision) if precision.size else float("nan")
+        results_per_category.append((f"{name}", float(ap * 100)))
+
+    # tabulate it
+    n_cols = min(6, len(results_per_category) * 2)
+    results_flatten = list(itertools.chain(*results_per_category))
+    results_2d = itertools.zip_longest(*[results_flatten[i::n_cols] for i in range(n_cols)])
+    table = tabulate(
+        results_2d,
+        tablefmt="pipe",
+        floatfmt=".3f",
+        headers=["category", "AP"] * (n_cols // 2),
+        numalign="left",
+    )
+    logger.info(f"Per-category {eval_mode_name} AP: \n" + table)
+
+    results.update({"AP-" + name: ap for name, ap in results_per_category})
+    return results
+
+
+def build_densepose_evaluator_storage(cfg: CfgNode, output_folder: str):
+    storage_spec = cfg.DENSEPOSE_EVALUATION.STORAGE
+    if storage_spec == "none":
+        return None
+    evaluator_type = cfg.DENSEPOSE_EVALUATION.TYPE
+    # common output tensor sizes
+    hout = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE
+    wout = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE
+    n_csc = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS
+    # specific output tensors
+    if evaluator_type == "iuv":
+        n_fsc = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES + 1
+        schema = {
+            "coarse_segm": SizeData(dtype="float32", shape=(n_csc, hout, wout)),
+            "fine_segm": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
+            "u": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
+            "v": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
+        }
+    elif evaluator_type == "cse":
+        embed_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE
+        schema = {
+            "coarse_segm": SizeData(dtype="float32", shape=(n_csc, hout, wout)),
+            "embedding": SizeData(dtype="float32", shape=(embed_size, hout, wout)),
+        }
+    else:
+        raise ValueError(f"Unknown evaluator type: {evaluator_type}")
+    # storage types
+    if storage_spec == "ram":
+        storage = SingleProcessRamTensorStorage(schema, io.BytesIO())
+    elif storage_spec == "file":
+        fpath = os.path.join(output_folder, f"DensePoseEvaluatorStorage.{get_rank()}.bin")
+        PathManager.mkdirs(output_folder)
+        storage = SingleProcessFileTensorStorage(schema, fpath, "wb")
+    else:
+        raise ValueError(f"Unknown storage specification: {storage_spec}")
+    return storage

densepose/evaluation/mesh_alignment_evaluator.py

@@ -1,66 +1,66 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-
-import json
-import logging
-from typing import List, Optional
-import torch
-from torch import nn
-
-from detectron2.utils.file_io import PathManager
-
-from densepose.structures.mesh import create_mesh
-
-
-class MeshAlignmentEvaluator:
-    """
-    Class for evaluation of 3D mesh alignment based on the learned vertex embeddings
-    """
-
-    def __init__(self, embedder: nn.Module, mesh_names: Optional[List[str]]):
-        self.embedder = embedder
-        # use the provided mesh names if not None and not an empty list
-        self.mesh_names = mesh_names if mesh_names else embedder.mesh_names
-        self.logger = logging.getLogger(__name__)
-        with PathManager.open(
-            "https://dl.fbaipublicfiles.com/densepose/data/cse/mesh_keyvertices_v0.json", "r"
-        ) as f:
-            self.mesh_keyvertices = json.load(f)
-
-    def evaluate(self):
-        ge_per_mesh = {}
-        gps_per_mesh = {}
-        for mesh_name_1 in self.mesh_names:
-            avg_errors = []
-            avg_gps = []
-            embeddings_1 = self.embedder(mesh_name_1)
-            keyvertices_1 = self.mesh_keyvertices[mesh_name_1]
-            keyvertex_names_1 = list(keyvertices_1.keys())
-            keyvertex_indices_1 = [keyvertices_1[name] for name in keyvertex_names_1]
-            for mesh_name_2 in self.mesh_names:
-                if mesh_name_1 == mesh_name_2:
-                    continue
-                embeddings_2 = self.embedder(mesh_name_2)
-                keyvertices_2 = self.mesh_keyvertices[mesh_name_2]
-                sim_matrix_12 = embeddings_1[keyvertex_indices_1].mm(embeddings_2.T)
-                vertices_2_matching_keyvertices_1 = sim_matrix_12.argmax(axis=1)
-                mesh_2 = create_mesh(mesh_name_2, embeddings_2.device)
-                geodists = mesh_2.geodists[
-                    vertices_2_matching_keyvertices_1,
-                    [keyvertices_2[name] for name in keyvertex_names_1],
-                ]
-                Current_Mean_Distances = 0.255
-                gps = (-(geodists**2) / (2 * (Current_Mean_Distances**2))).exp()
-                avg_errors.append(geodists.mean().item())
-                avg_gps.append(gps.mean().item())
-
-            ge_mean = torch.as_tensor(avg_errors).mean().item()
-            gps_mean = torch.as_tensor(avg_gps).mean().item()
-            ge_per_mesh[mesh_name_1] = ge_mean
-            gps_per_mesh[mesh_name_1] = gps_mean
-        ge_mean_global = torch.as_tensor(list(ge_per_mesh.values())).mean().item()
-        gps_mean_global = torch.as_tensor(list(gps_per_mesh.values())).mean().item()
-        per_mesh_metrics = {
-            "GE": ge_per_mesh,
-            "GPS": gps_per_mesh,
-        }
-        return ge_mean_global, gps_mean_global, per_mesh_metrics
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import json
+import logging
+from typing import List, Optional
+import torch
+from torch import nn
+
+from detectron2.utils.file_io import PathManager
+
+from densepose.structures.mesh import create_mesh
+
+
+class MeshAlignmentEvaluator:
+    """
+    Class for evaluation of 3D mesh alignment based on the learned vertex embeddings
+    """
+
+    def __init__(self, embedder: nn.Module, mesh_names: Optional[List[str]]):
+        self.embedder = embedder
+        # use the provided mesh names if not None and not an empty list
+        self.mesh_names = mesh_names if mesh_names else embedder.mesh_names
+        self.logger = logging.getLogger(__name__)
+        with PathManager.open(
+            "https://dl.fbaipublicfiles.com/densepose/data/cse/mesh_keyvertices_v0.json", "r"
+        ) as f:
+            self.mesh_keyvertices = json.load(f)
+
+    def evaluate(self):
+        ge_per_mesh = {}
+        gps_per_mesh = {}
+        for mesh_name_1 in self.mesh_names:
+            avg_errors = []
+            avg_gps = []
+            embeddings_1 = self.embedder(mesh_name_1)
+            keyvertices_1 = self.mesh_keyvertices[mesh_name_1]
+            keyvertex_names_1 = list(keyvertices_1.keys())
+            keyvertex_indices_1 = [keyvertices_1[name] for name in keyvertex_names_1]
+            for mesh_name_2 in self.mesh_names:
+                if mesh_name_1 == mesh_name_2:
+                    continue
+                embeddings_2 = self.embedder(mesh_name_2)
+                keyvertices_2 = self.mesh_keyvertices[mesh_name_2]
+                sim_matrix_12 = embeddings_1[keyvertex_indices_1].mm(embeddings_2.T)
+                vertices_2_matching_keyvertices_1 = sim_matrix_12.argmax(axis=1)
+                mesh_2 = create_mesh(mesh_name_2, embeddings_2.device)
+                geodists = mesh_2.geodists[
+                    vertices_2_matching_keyvertices_1,
+                    [keyvertices_2[name] for name in keyvertex_names_1],
+                ]
+                Current_Mean_Distances = 0.255
+                gps = (-(geodists**2) / (2 * (Current_Mean_Distances**2))).exp()
+                avg_errors.append(geodists.mean().item())
+                avg_gps.append(gps.mean().item())
+
+            ge_mean = torch.as_tensor(avg_errors).mean().item()
+            gps_mean = torch.as_tensor(avg_gps).mean().item()
+            ge_per_mesh[mesh_name_1] = ge_mean
+            gps_per_mesh[mesh_name_1] = gps_mean
+        ge_mean_global = torch.as_tensor(list(ge_per_mesh.values())).mean().item()
+        gps_mean_global = torch.as_tensor(list(gps_per_mesh.values())).mean().item()
+        per_mesh_metrics = {
+            "GE": ge_per_mesh,
+            "GPS": gps_per_mesh,
+        }
+        return ge_mean_global, gps_mean_global, per_mesh_metrics

densepose/evaluation/tensor_storage.py

@@ -1,239 +1,239 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import io
-import numpy as np
-import os
-from dataclasses import dataclass
-from functools import reduce
-from operator import mul
-from typing import BinaryIO, Dict, Optional, Tuple
-import torch
-
-from detectron2.utils.comm import gather, get_rank
-from detectron2.utils.file_io import PathManager
-
-
-@dataclass
-class SizeData:
-    dtype: str
-    shape: Tuple[int]
-
-
-def _calculate_record_field_size_b(data_schema: Dict[str, SizeData], field_name: str) -> int:
-    schema = data_schema[field_name]
-    element_size_b = np.dtype(schema.dtype).itemsize
-    record_field_size_b = reduce(mul, schema.shape) * element_size_b
-    return record_field_size_b
-
-
-def _calculate_record_size_b(data_schema: Dict[str, SizeData]) -> int:
-    record_size_b = 0
-    for field_name in data_schema:
-        record_field_size_b = _calculate_record_field_size_b(data_schema, field_name)
-        record_size_b += record_field_size_b
-    return record_size_b
-
-
-def _calculate_record_field_sizes_b(data_schema: Dict[str, SizeData]) -> Dict[str, int]:
-    field_sizes_b = {}
-    for field_name in data_schema:
-        field_sizes_b[field_name] = _calculate_record_field_size_b(data_schema, field_name)
-    return field_sizes_b
-
-
-class SingleProcessTensorStorage:
-    """
-    Compact tensor storage to keep tensor data of predefined size and type.
-    """
-
-    def __init__(self, data_schema: Dict[str, SizeData], storage_impl: BinaryIO):
-        """
-        Construct tensor storage based on information on data shape and size.
-        Internally uses numpy to interpret the type specification.
-        The storage must support operations `seek(offset, whence=os.SEEK_SET)` and
-        `read(size)` to be able to perform the `get` operation.
-        The storage must support operation `write(bytes)` to be able to perform
-        the `put` operation.
-
-        Args:
-            data_schema (dict: str -> SizeData): dictionary which maps tensor name
-                to its size data (shape and data type), e.g.
-                ```
-                {
-                  "coarse_segm": SizeData(dtype="float32", shape=(112, 112)),
-                  "embedding": SizeData(dtype="float32", shape=(16, 112, 112)),
-                }
-                ```
-            storage_impl (BinaryIO): io instance that handles file-like seek, read
-                and write operations, e.g. a file handle or a memory buffer like io.BytesIO
-        """
-        self.data_schema = data_schema
-        self.record_size_b = _calculate_record_size_b(data_schema)
-        self.record_field_sizes_b = _calculate_record_field_sizes_b(data_schema)
-        self.storage_impl = storage_impl
-        self.next_record_id = 0
-
-    def get(self, record_id: int) -> Dict[str, torch.Tensor]:
-        """
-        Load tensors from the storage by record ID
-
-        Args:
-            record_id (int): Record ID, for which to load the data
-
-        Return:
-            dict: str -> tensor: tensor name mapped to tensor data, recorded under the provided ID
-        """
-        self.storage_impl.seek(record_id * self.record_size_b, os.SEEK_SET)
-        data_bytes = self.storage_impl.read(self.record_size_b)
-        assert len(data_bytes) == self.record_size_b, (
-            f"Expected data size {self.record_size_b} B could not be read: "
-            f"got {len(data_bytes)} B"
-        )
-        record = {}
-        cur_idx = 0
-        # it's important to read and write in the same order
-        for field_name in sorted(self.data_schema):
-            schema = self.data_schema[field_name]
-            field_size_b = self.record_field_sizes_b[field_name]
-            chunk = data_bytes[cur_idx : cur_idx + field_size_b]
-            data_np = np.frombuffer(
-                chunk, dtype=schema.dtype, count=reduce(mul, schema.shape)
-            ).reshape(schema.shape)
-            record[field_name] = torch.from_numpy(data_np)
-            cur_idx += field_size_b
-        return record
-
-    def put(self, data: Dict[str, torch.Tensor]) -> int:
-        """
-        Store tensors in the storage
-
-        Args:
-            data (dict: str -> tensor): data to store, a dictionary which maps
-                tensor names into tensors; tensor shapes must match those specified
-                in data schema.
-        Return:
-            int: record ID, under which the data is stored
-        """
-        # it's important to read and write in the same order
-        for field_name in sorted(self.data_schema):
-            assert (
-                field_name in data
-            ), f"Field '{field_name}' not present in data: data keys are {data.keys()}"
-            value = data[field_name]
-            assert value.shape == self.data_schema[field_name].shape, (
-                f"Mismatched tensor shapes for field '{field_name}': "
-                f"expected {self.data_schema[field_name].shape}, got {value.shape}"
-            )
-            data_bytes = value.cpu().numpy().tobytes()
-            assert len(data_bytes) == self.record_field_sizes_b[field_name], (
-                f"Expected field {field_name} to be of size "
-                f"{self.record_field_sizes_b[field_name]} B, got {len(data_bytes)} B"
-            )
-            self.storage_impl.write(data_bytes)
-        record_id = self.next_record_id
-        self.next_record_id += 1
-        return record_id
-
-
-class SingleProcessFileTensorStorage(SingleProcessTensorStorage):
-    """
-    Implementation of a single process tensor storage which stores data in a file
-    """
-
-    def __init__(self, data_schema: Dict[str, SizeData], fpath: str, mode: str):
-        self.fpath = fpath
-        assert "b" in mode, f"Tensor storage should be opened in binary mode, got '{mode}'"
-        if "w" in mode:
-            # pyre-fixme[6]: For 2nd argument expected `Union[typing_extensions.Liter...
-            file_h = PathManager.open(fpath, mode)
-        elif "r" in mode:
-            local_fpath = PathManager.get_local_path(fpath)
-            file_h = open(local_fpath, mode)
-        else:
-            raise ValueError(f"Unsupported file mode {mode}, supported modes: rb, wb")
-        super().__init__(data_schema, file_h)  # pyre-ignore[6]
-
-
-class SingleProcessRamTensorStorage(SingleProcessTensorStorage):
-    """
-    Implementation of a single process tensor storage which stores data in RAM
-    """
-
-    def __init__(self, data_schema: Dict[str, SizeData], buf: io.BytesIO):
-        super().__init__(data_schema, buf)
-
-
-class MultiProcessTensorStorage:
-    """
-    Representation of a set of tensor storages created by individual processes,
-    allows to access those storages from a single owner process. The storages
-    should either be shared or broadcasted to the owner process.
-    The processes are identified by their rank, data is uniquely defined by
-    the rank of the process and the record ID.
-    """
-
-    def __init__(self, rank_to_storage: Dict[int, SingleProcessTensorStorage]):
-        self.rank_to_storage = rank_to_storage
-
-    def get(self, rank: int, record_id: int) -> Dict[str, torch.Tensor]:
-        storage = self.rank_to_storage[rank]
-        return storage.get(record_id)
-
-    def put(self, rank: int, data: Dict[str, torch.Tensor]) -> int:
-        storage = self.rank_to_storage[rank]
-        return storage.put(data)
-
-
-class MultiProcessFileTensorStorage(MultiProcessTensorStorage):
-    def __init__(self, data_schema: Dict[str, SizeData], rank_to_fpath: Dict[int, str], mode: str):
-        rank_to_storage = {
-            rank: SingleProcessFileTensorStorage(data_schema, fpath, mode)
-            for rank, fpath in rank_to_fpath.items()
-        }
-        super().__init__(rank_to_storage)  # pyre-ignore[6]
-
-
-class MultiProcessRamTensorStorage(MultiProcessTensorStorage):
-    def __init__(self, data_schema: Dict[str, SizeData], rank_to_buffer: Dict[int, io.BytesIO]):
-        rank_to_storage = {
-            rank: SingleProcessRamTensorStorage(data_schema, buf)
-            for rank, buf in rank_to_buffer.items()
-        }
-        super().__init__(rank_to_storage)  # pyre-ignore[6]
-
-
-def _ram_storage_gather(
-    storage: SingleProcessRamTensorStorage, dst_rank: int = 0
-) -> Optional[MultiProcessRamTensorStorage]:
-    storage.storage_impl.seek(0, os.SEEK_SET)
-    # TODO: overhead, pickling a bytes object, can just pass bytes in a tensor directly
-    # see detectron2/utils.comm.py
-    data_list = gather(storage.storage_impl.read(), dst=dst_rank)
-    if get_rank() != dst_rank:
-        return None
-    rank_to_buffer = {i: io.BytesIO(data_list[i]) for i in range(len(data_list))}
-    multiprocess_storage = MultiProcessRamTensorStorage(storage.data_schema, rank_to_buffer)
-    return multiprocess_storage
-
-
-def _file_storage_gather(
-    storage: SingleProcessFileTensorStorage,
-    dst_rank: int = 0,
-    mode: str = "rb",
-) -> Optional[MultiProcessFileTensorStorage]:
-    storage.storage_impl.close()
-    fpath_list = gather(storage.fpath, dst=dst_rank)
-    if get_rank() != dst_rank:
-        return None
-    rank_to_fpath = {i: fpath_list[i] for i in range(len(fpath_list))}
-    return MultiProcessFileTensorStorage(storage.data_schema, rank_to_fpath, mode)
-
-
-def storage_gather(
-    storage: SingleProcessTensorStorage, dst_rank: int = 0
-) -> Optional[MultiProcessTensorStorage]:
-    if isinstance(storage, SingleProcessRamTensorStorage):
-        return _ram_storage_gather(storage, dst_rank)
-    elif isinstance(storage, SingleProcessFileTensorStorage):
-        return _file_storage_gather(storage, dst_rank)
-    raise Exception(f"Unsupported storage for gather operation: {storage}")
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import io
+import numpy as np
+import os
+from dataclasses import dataclass
+from functools import reduce
+from operator import mul
+from typing import BinaryIO, Dict, Optional, Tuple
+import torch
+
+from detectron2.utils.comm import gather, get_rank
+from detectron2.utils.file_io import PathManager
+
+
+@dataclass
+class SizeData:
+    dtype: str
+    shape: Tuple[int]
+
+
+def _calculate_record_field_size_b(data_schema: Dict[str, SizeData], field_name: str) -> int:
+    schema = data_schema[field_name]
+    element_size_b = np.dtype(schema.dtype).itemsize
+    record_field_size_b = reduce(mul, schema.shape) * element_size_b
+    return record_field_size_b
+
+
+def _calculate_record_size_b(data_schema: Dict[str, SizeData]) -> int:
+    record_size_b = 0
+    for field_name in data_schema:
+        record_field_size_b = _calculate_record_field_size_b(data_schema, field_name)
+        record_size_b += record_field_size_b
+    return record_size_b
+
+
+def _calculate_record_field_sizes_b(data_schema: Dict[str, SizeData]) -> Dict[str, int]:
+    field_sizes_b = {}
+    for field_name in data_schema:
+        field_sizes_b[field_name] = _calculate_record_field_size_b(data_schema, field_name)
+    return field_sizes_b
+
+
+class SingleProcessTensorStorage:
+    """
+    Compact tensor storage to keep tensor data of predefined size and type.
+    """
+
+    def __init__(self, data_schema: Dict[str, SizeData], storage_impl: BinaryIO):
+        """
+        Construct tensor storage based on information on data shape and size.
+        Internally uses numpy to interpret the type specification.
+        The storage must support operations `seek(offset, whence=os.SEEK_SET)` and
+        `read(size)` to be able to perform the `get` operation.
+        The storage must support operation `write(bytes)` to be able to perform
+        the `put` operation.
+
+        Args:
+            data_schema (dict: str -> SizeData): dictionary which maps tensor name
+                to its size data (shape and data type), e.g.
+                ```
+                {
+                  "coarse_segm": SizeData(dtype="float32", shape=(112, 112)),
+                  "embedding": SizeData(dtype="float32", shape=(16, 112, 112)),
+                }
+                ```
+            storage_impl (BinaryIO): io instance that handles file-like seek, read
+                and write operations, e.g. a file handle or a memory buffer like io.BytesIO
+        """
+        self.data_schema = data_schema
+        self.record_size_b = _calculate_record_size_b(data_schema)
+        self.record_field_sizes_b = _calculate_record_field_sizes_b(data_schema)
+        self.storage_impl = storage_impl
+        self.next_record_id = 0
+
+    def get(self, record_id: int) -> Dict[str, torch.Tensor]:
+        """
+        Load tensors from the storage by record ID
+
+        Args:
+            record_id (int): Record ID, for which to load the data
+
+        Return:
+            dict: str -> tensor: tensor name mapped to tensor data, recorded under the provided ID
+        """
+        self.storage_impl.seek(record_id * self.record_size_b, os.SEEK_SET)
+        data_bytes = self.storage_impl.read(self.record_size_b)
+        assert len(data_bytes) == self.record_size_b, (
+            f"Expected data size {self.record_size_b} B could not be read: "
+            f"got {len(data_bytes)} B"
+        )
+        record = {}
+        cur_idx = 0
+        # it's important to read and write in the same order
+        for field_name in sorted(self.data_schema):
+            schema = self.data_schema[field_name]
+            field_size_b = self.record_field_sizes_b[field_name]
+            chunk = data_bytes[cur_idx : cur_idx + field_size_b]
+            data_np = np.frombuffer(
+                chunk, dtype=schema.dtype, count=reduce(mul, schema.shape)
+            ).reshape(schema.shape)
+            record[field_name] = torch.from_numpy(data_np)
+            cur_idx += field_size_b
+        return record
+
+    def put(self, data: Dict[str, torch.Tensor]) -> int:
+        """
+        Store tensors in the storage
+
+        Args:
+            data (dict: str -> tensor): data to store, a dictionary which maps
+                tensor names into tensors; tensor shapes must match those specified
+                in data schema.
+        Return:
+            int: record ID, under which the data is stored
+        """
+        # it's important to read and write in the same order
+        for field_name in sorted(self.data_schema):
+            assert (
+                field_name in data
+            ), f"Field '{field_name}' not present in data: data keys are {data.keys()}"
+            value = data[field_name]
+            assert value.shape == self.data_schema[field_name].shape, (
+                f"Mismatched tensor shapes for field '{field_name}': "
+                f"expected {self.data_schema[field_name].shape}, got {value.shape}"
+            )
+            data_bytes = value.cpu().numpy().tobytes()
+            assert len(data_bytes) == self.record_field_sizes_b[field_name], (
+                f"Expected field {field_name} to be of size "
+                f"{self.record_field_sizes_b[field_name]} B, got {len(data_bytes)} B"
+            )
+            self.storage_impl.write(data_bytes)
+        record_id = self.next_record_id
+        self.next_record_id += 1
+        return record_id
+
+
+class SingleProcessFileTensorStorage(SingleProcessTensorStorage):
+    """
+    Implementation of a single process tensor storage which stores data in a file
+    """
+
+    def __init__(self, data_schema: Dict[str, SizeData], fpath: str, mode: str):
+        self.fpath = fpath
+        assert "b" in mode, f"Tensor storage should be opened in binary mode, got '{mode}'"
+        if "w" in mode:
+            # pyre-fixme[6]: For 2nd argument expected `Union[typing_extensions.Liter...
+            file_h = PathManager.open(fpath, mode)
+        elif "r" in mode:
+            local_fpath = PathManager.get_local_path(fpath)
+            file_h = open(local_fpath, mode)
+        else:
+            raise ValueError(f"Unsupported file mode {mode}, supported modes: rb, wb")
+        super().__init__(data_schema, file_h)  # pyre-ignore[6]
+
+
+class SingleProcessRamTensorStorage(SingleProcessTensorStorage):
+    """
+    Implementation of a single process tensor storage which stores data in RAM
+    """
+
+    def __init__(self, data_schema: Dict[str, SizeData], buf: io.BytesIO):
+        super().__init__(data_schema, buf)
+
+
+class MultiProcessTensorStorage:
+    """
+    Representation of a set of tensor storages created by individual processes,
+    allows to access those storages from a single owner process. The storages
+    should either be shared or broadcasted to the owner process.
+    The processes are identified by their rank, data is uniquely defined by
+    the rank of the process and the record ID.
+    """
+
+    def __init__(self, rank_to_storage: Dict[int, SingleProcessTensorStorage]):
+        self.rank_to_storage = rank_to_storage
+
+    def get(self, rank: int, record_id: int) -> Dict[str, torch.Tensor]:
+        storage = self.rank_to_storage[rank]
+        return storage.get(record_id)
+
+    def put(self, rank: int, data: Dict[str, torch.Tensor]) -> int:
+        storage = self.rank_to_storage[rank]
+        return storage.put(data)
+
+
+class MultiProcessFileTensorStorage(MultiProcessTensorStorage):
+    def __init__(self, data_schema: Dict[str, SizeData], rank_to_fpath: Dict[int, str], mode: str):
+        rank_to_storage = {
+            rank: SingleProcessFileTensorStorage(data_schema, fpath, mode)
+            for rank, fpath in rank_to_fpath.items()
+        }
+        super().__init__(rank_to_storage)  # pyre-ignore[6]
+
+
+class MultiProcessRamTensorStorage(MultiProcessTensorStorage):
+    def __init__(self, data_schema: Dict[str, SizeData], rank_to_buffer: Dict[int, io.BytesIO]):
+        rank_to_storage = {
+            rank: SingleProcessRamTensorStorage(data_schema, buf)
+            for rank, buf in rank_to_buffer.items()
+        }
+        super().__init__(rank_to_storage)  # pyre-ignore[6]
+
+
+def _ram_storage_gather(
+    storage: SingleProcessRamTensorStorage, dst_rank: int = 0
+) -> Optional[MultiProcessRamTensorStorage]:
+    storage.storage_impl.seek(0, os.SEEK_SET)
+    # TODO: overhead, pickling a bytes object, can just pass bytes in a tensor directly
+    # see detectron2/utils.comm.py
+    data_list = gather(storage.storage_impl.read(), dst=dst_rank)
+    if get_rank() != dst_rank:
+        return None
+    rank_to_buffer = {i: io.BytesIO(data_list[i]) for i in range(len(data_list))}
+    multiprocess_storage = MultiProcessRamTensorStorage(storage.data_schema, rank_to_buffer)
+    return multiprocess_storage
+
+
+def _file_storage_gather(
+    storage: SingleProcessFileTensorStorage,
+    dst_rank: int = 0,
+    mode: str = "rb",
+) -> Optional[MultiProcessFileTensorStorage]:
+    storage.storage_impl.close()
+    fpath_list = gather(storage.fpath, dst=dst_rank)
+    if get_rank() != dst_rank:
+        return None
+    rank_to_fpath = {i: fpath_list[i] for i in range(len(fpath_list))}
+    return MultiProcessFileTensorStorage(storage.data_schema, rank_to_fpath, mode)
+
+
+def storage_gather(
+    storage: SingleProcessTensorStorage, dst_rank: int = 0
+) -> Optional[MultiProcessTensorStorage]:
+    if isinstance(storage, SingleProcessRamTensorStorage):
+        return _ram_storage_gather(storage, dst_rank)
+    elif isinstance(storage, SingleProcessFileTensorStorage):
+        return _file_storage_gather(storage, dst_rank)
+    raise Exception(f"Unsupported storage for gather operation: {storage}")

densepose/modeling/__init__.py

@@ -1,13 +1,13 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-from .confidence import DensePoseConfidenceModelConfig, DensePoseUVConfidenceType
-from .filter import DensePoseDataFilter
-from .inference import densepose_inference
-from .utils import initialize_module_params
-from .build import (
-    build_densepose_data_filter,
-    build_densepose_embedder,
-    build_densepose_head,
-    build_densepose_losses,
-    build_densepose_predictor,
-)
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+from .confidence import DensePoseConfidenceModelConfig, DensePoseUVConfidenceType
+from .filter import DensePoseDataFilter
+from .inference import densepose_inference
+from .utils import initialize_module_params
+from .build import (
+    build_densepose_data_filter,
+    build_densepose_embedder,
+    build_densepose_head,
+    build_densepose_losses,
+    build_densepose_predictor,
+)