Skip to content

Commit

Permalink
Added missing rgb2bgr conversion (#1609)
Browse files Browse the repository at this point in the history
Co-authored-by: Ofri Masad <ofrimasad@users.noreply.github.com>
  • Loading branch information
BloodAxe and ofrimasad committed Nov 5, 2023
1 parent 885e2f0 commit b90f4d8
Showing 1 changed file with 18 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from super_gradients.common.decorators.factory_decorator import resolve_param
from super_gradients.common.factories.transforms_factory import TransformsFactory
from super_gradients.common.factories.type_factory import TypeFactory
from super_gradients.common.object_names import Datasets
from super_gradients.common.object_names import Datasets, Processings
from super_gradients.common.registry.registry import register_dataset
from super_gradients.training.datasets.data_formats.bbox_formats.xywh import xywh_to_xyxy, xyxy_to_xywh
from super_gradients.training.datasets.pose_estimation_datasets.abstract_pose_estimation_dataset import AbstractPoseEstimationDataset
Expand Down Expand Up @@ -206,3 +206,20 @@ def _get_crowd_mask(self, anno, img_info) -> np.ndarray:
m += mask

return (m < 0.5).astype(np.float32)

def get_dataset_preprocessing_params(self) -> dict:
"""
This method returns a dictionary of parameters describing preprocessing steps to be applied to the dataset.
:return:
"""
rgb_to_bgr = {Processings.ReverseImageChannels: {}}
image_to_tensor = {Processings.ImagePermute: {"permutation": (2, 0, 1)}}
pipeline = [rgb_to_bgr] + self.transforms.get_equivalent_preprocessing() + [image_to_tensor]
params = dict(
conf=0.05,
image_processor={Processings.ComposeProcessing: {"processings": pipeline}},
edge_links=self.edge_links,
edge_colors=self.edge_colors,
keypoint_colors=self.keypoint_colors,
)
return params

0 comments on commit b90f4d8

Please sign in to comment.