def __init__(self, transforms, bbox_params=None, keypoint_params=None, additional_targets=None, p=1.0): super(Compose, self).__init__([t for t in transforms if t is not None], p) self.processors = {} if bbox_params: if isinstance(bbox_params, dict): params = BboxParams(**bbox_params) elif isinstance(bbox_params, BboxParams): params = bbox_params else: raise ValueError("unknown format of bbox_params, please use `dict` or `BboxParams`") self.processors["bboxes"] = BboxProcessor(params, additional_targets) if keypoint_params: if isinstance(keypoint_params, dict): params = KeypointParams(**keypoint_params) elif isinstance(keypoint_params, KeypointParams): params = keypoint_params else: raise ValueError("unknown format of keypoint_params, please use `dict` or `KeypointParams`") self.processors["keypoints"] = KeypointsProcessor(params, additional_targets) if additional_targets is None: additional_targets = {} self.additional_targets = additional_targets for proc in self.processors.values(): proc.ensure_transforms_valid(self.transforms) self.add_targets(additional_targets)
def __init__( self, transforms: TransformsSeqType, bbox_params: typing.Optional[typing.Union[dict, "BboxParams"]] = None, keypoint_params: typing.Optional[typing.Union[ dict, "KeypointParams"]] = None, additional_targets: typing.Optional[typing.Dict[str, str]] = None, p: float = 1.0, ): super(Compose, self).__init__(transforms, p) self.processors: typing.Dict[str, typing.Union[BboxProcessor, KeypointsProcessor]] = {} if bbox_params: if isinstance(bbox_params, dict): b_params = BboxParams(**bbox_params) elif isinstance(bbox_params, BboxParams): b_params = bbox_params else: raise ValueError( "unknown format of bbox_params, please use `dict` or `BboxParams`" ) self.processors["bboxes"] = BboxProcessor(b_params, additional_targets) if keypoint_params: if isinstance(keypoint_params, dict): k_params = KeypointParams(**keypoint_params) elif isinstance(keypoint_params, KeypointParams): k_params = keypoint_params else: raise ValueError( "unknown format of keypoint_params, please use `dict` or `KeypointParams`" ) self.processors["keypoints"] = KeypointsProcessor( k_params, additional_targets) if additional_targets is None: additional_targets = {} self.additional_targets = additional_targets for proc in self.processors.values(): proc.ensure_transforms_valid(self.transforms) self.add_targets(additional_targets) self.is_check_args = True self._disable_check_args_for_transforms(self.transforms)