predict() support for pose estimation models

Makefile

@@ -6,3 +6,11 @@ integration_tests:
 
 yolo_nas_integration_tests:
 	python -m unittest tests/integration_tests/yolo_nas_integration_test.py
+
+recipe_accuracy_tests:
+	python3.8 src/super_gradients/examples/convert_recipe_example/convert_recipe_example.py --config-name=cifar10_conversion_params experiment_name=shortened_cifar10_resnet_accuracy_test
+	python3.8 src/super_gradients/train_from_recipe.py --config-name=coco2017_pose_dekr_w32_no_dc experiment_name=shortened_coco2017_pose_dekr_w32_ap_test epochs=1 batch_size=4 val_batch_size=8 training_hyperparams.lr_warmup_steps=0 training_hyperparams.average_best_models=False training_hyperparams.max_train_batches=1000 training_hyperparams.max_valid_batches=100 multi_gpu=DDP num_gpus=4
+	python3.8 src/super_gradients/train_from_recipe.py --config-name=cifar10_resnet               experiment_name=shortened_cifar10_resnet_accuracy_test   epochs=100 training_hyperparams.average_best_models=False multi_gpu=DDP num_gpus=4
+	python3.8 src/super_gradients/train_from_recipe.py --config-name=coco2017_yolox               experiment_name=shortened_coco2017_yolox_n_map_test      epochs=10  architecture=yolox_n training_hyperparams.loss=yolox_fast_loss training_hyperparams.average_best_models=False multi_gpu=DDP num_gpus=4
+	python3.8 src/super_gradients/train_from_recipe.py --config-name=cityscapes_regseg48          experiment_name=shortened_cityscapes_regseg48_iou_test   epochs=10 training_hyperparams.average_best_models=False multi_gpu=DDP num_gpus=4
+	coverage run --source=super_gradients -m unittest tests/deci_core_recipe_test_suite_runner.py

src/super_gradients/common/object_names.py

@@ -413,8 +413,10 @@ class Processings:
     DetectionCenterPadding = "DetectionCenterPadding"
     DetectionLongestMaxSizeRescale = "DetectionLongestMaxSizeRescale"
     DetectionBottomRightPadding = "DetectionBottomRightPadding"
-    ImagePermute = "ImagePermute"
     DetectionRescale = "DetectionRescale"
+    KeypointsLongestMaxSizeRescale = "KeypointsLongestMaxSizeRescale"
+    KeypointsBottomRightPadding = "KeypointsBottomRightPadding"
+    ImagePermute = "ImagePermute"
     ReverseImageChannels = "ReverseImageChannels"
     NormalizeImage = "NormalizeImage"
     ComposeProcessing = "ComposeProcessing"

src/super_gradients/recipes/dataset_params/coco_pose_estimation_dataset_params.yaml

@@ -3,8 +3,7 @@ num_joints: 17
 # OKs sigma values take from https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py#L523
 oks_sigmas: [0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 1.007, 1.007, 0.087, 0.087, 0.089, 0.089]
 
-flip_indexes_heatmap: [ 0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15, 17]
-flip_indexes_offset: [ 0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15,]
+flip_indexes: [ 0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15,]
 
 joint_links:
   - [0, 1]
@@ -27,6 +26,47 @@ joint_links:
   - [13, 15]
   - [14, 16]
 
+joint_colors:
+ - [214, 39, 40]  # Nose -> LeftEye
+ - [148, 103, 189]  # Nose -> RightEye
+ - [44, 160, 44]  # LeftEye -> RightEye
+ - [140, 86, 75]  # LeftEye -> LeftEar
+ - [227, 119, 194]  # RightEye -> RightEar
+ - [127, 127, 127]  # LeftEar -> LeftShoulder
+ - [188, 189, 34]  # RightEar -> RightShoulder
+ - [127, 127, 127]  # Shoulders
+ - [188, 189, 34]  # LeftShoulder -> LeftElbow
+ - [140, 86, 75]  # LeftTorso
+ - [23, 190, 207]  # RightShoulder -> RightElbow
+ - [227, 119, 194]  # RightTorso
+ - [31, 119, 180]  # LeftElbow -> LeftArm
+ - [255, 127, 14]  # RightElbow -> RightArm
+ - [148, 103, 189]  # Waist
+ - [255, 127, 14]  # Left Hip -> Left Knee
+ - [214, 39, 40]  # Right Hip -> Right Knee
+ - [31, 119, 180]  # Left Knee -> Left Ankle
+ - [44, 160, 44]  # Right Knee -> Right Ankle
+
+
+keypoint_colors:
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+  - [31, 119, 180]
+  - [148, 103, 189]
+
 
 train_dataset_params:
   data_dir: /data/coco # root path to coco data
@@ -36,6 +76,10 @@ train_dataset_params:
   include_empty_samples: False
   min_instance_area: 64
 
+  joint_links: ${dataset_params.joint_links}
+  joint_colors: ${dataset_params.joint_colors}
+  keypoint_colors: ${dataset_params.keypoint_colors}
+
   transforms:
     - KeypointsLongestMaxSize:
         max_height: 640
@@ -49,7 +93,7 @@ train_dataset_params:
 
     - KeypointsRandomHorizontalFlip:
         # Note these indexes are COCO-specific. If you're using a different dataset, you'll need to change these accordingly.
-        flip_index: [ 0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15 ]
+        flip_index: ${dataset_params.flip_indexes}
         prob: 0.5
 
     - KeypointsRandomAffineTransform:
@@ -75,6 +119,11 @@ val_dataset_params:
   json_file: annotations/person_keypoints_val2017.json
   include_empty_samples: True
   min_instance_area: 128
+
+  joint_links: ${dataset_params.joint_links}
+  joint_colors: ${dataset_params.joint_colors}
+  keypoint_colors: ${dataset_params.keypoint_colors}
+
   transforms:
     - KeypointsLongestMaxSize:
         max_height: 640

src/super_gradients/scripts/generate_rescoring_training_data/__main__.py

@@ -115,7 +115,7 @@ def main(cfg: DictConfig) -> None:
     )
 
     # model = DEKRWrapper(model, apply_sigmoid=True).cuda().eval()
-    model = DEKRHorisontalFlipWrapper(model, cfg.dataset_params.flip_indexes_heatmap, cfg.dataset_params.flip_indexes_offset, apply_sigmoid=True).cuda().eval()
+    model = DEKRHorisontalFlipWrapper(model, cfg.dataset_params.flip_indexes, apply_sigmoid=True).cuda().eval()
 
     post_prediction_callback = cfg.post_prediction_callback
 

src/super_gradients/training/datasets/pose_estimation_datasets/base_keypoints.py

@@ -1,11 +1,12 @@
 import abc
-from typing import Tuple, List, Mapping, Any, Dict
+from typing import Tuple, List, Mapping, Any, Dict, Union
 
 import numpy as np
 import torch
 from torch.utils.data.dataloader import default_collate, Dataset
 
 from super_gradients.common.abstractions.abstract_logger import get_logger
+from super_gradients.common.object_names import Processings
 from super_gradients.common.registry.registry import register_collate_function
 from super_gradients.training.datasets.pose_estimation_datasets.target_generators import KeypointsTargetsGenerator
 from super_gradients.training.transforms.keypoint_transforms import KeypointsCompose, KeypointTransform
@@ -24,6 +25,9 @@ def __init__(
         target_generator: KeypointsTargetsGenerator,
         transforms: List[KeypointTransform],
         min_instance_area: float,
+        joint_links: Union[List[Tuple[int, int]], np.ndarray],
+        joint_colors: Union[List[Tuple[int, int, int]], np.ndarray, None],
+        keypoint_colors: Union[List[Tuple[int, int, int]], np.ndarray, None],
     ):
         """
 
@@ -36,6 +40,9 @@ def __init__(
         self.target_generator = target_generator
         self.transforms = KeypointsCompose(transforms)
         self.min_instance_area = min_instance_area
+        self.joint_links = joint_links
+        self.joint_colors = joint_colors
+        self.keypoint_colors = keypoint_colors
 
     @abc.abstractmethod
     def __len__(self) -> int:
@@ -95,6 +102,21 @@ def filter_joints(self, joints: np.ndarray, image: np.ndarray) -> np.ndarray:
 
         return joints
 
+    def get_dataset_preprocessing_params(self):
+        """
+
+        :return:
+        """
+        pipeline = self.transforms.get_equivalent_preprocessing()
+        params = dict(
+            conf=0.25,
+            image_processor={Processings.ComposeProcessing: {"processings": pipeline}},
+            joint_links=self.joint_links,
+            joint_names=self.joint_colors,
+            keypoint_colors=self.keypoint_colors,
+        )
+        return params
+
 
 @register_collate_function()
 class KeypointsCollate:

src/super_gradients/training/datasets/pose_estimation_datasets/coco_keypoints.py

@@ -1,5 +1,5 @@
 import os
-from typing import Tuple, List, Mapping, Any
+from typing import Tuple, List, Mapping, Any, Union
 
 import cv2
 import numpy as np
@@ -8,7 +8,7 @@
 from torch import Tensor
 
 from super_gradients.common.abstractions.abstract_logger import get_logger
-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.common.decorators.factory_decorator import resolve_param
 from super_gradients.common.factories.target_generator_factory import TargetGeneratorsFactory
@@ -37,6 +37,9 @@ def __init__(
         target_generator,
         transforms: List[KeypointTransform],
         min_instance_area: float,
+        joint_links: Union[List[Tuple[int, int]], np.ndarray],
+        joint_colors: Union[List[Tuple[int, int, int]], np.ndarray, None],
+        keypoint_colors: Union[List[Tuple[int, int, int]], np.ndarray, None],
     ):
         """
 
@@ -49,8 +52,17 @@ def __init__(
             See DEKRTargetsGenerator for an example.
         :param transforms: Transforms to be applied to the image & keypoints
         :param min_instance_area: Minimum area of an instance to be included in the dataset
+        :param joint_links: List of joint links to be visualized on the image
+        :param joint_colors: List of colors for each joint link
         """
-        super().__init__(transforms=transforms, target_generator=target_generator, min_instance_area=min_instance_area)
+        super().__init__(
+            transforms=transforms,
+            target_generator=target_generator,
+            min_instance_area=min_instance_area,
+            joint_links=joint_links,
+            joint_colors=joint_colors,
+            keypoint_colors=keypoint_colors,
+        )
         self.root = data_dir
         self.images_dir = os.path.join(data_dir, images_dir)
         self.json_file = os.path.join(data_dir, json_file)
@@ -190,3 +202,17 @@ def get_mask(self, anno, img_info) -> np.ndarray:
                     m += mask
 
         return (m < 0.5).astype(np.float32)
+
+    def get_dataset_preprocessing_params(self):
+        """
+
+        :return:
+        """
+        pipeline = [Processings.ReverseImageChannels] + self.transforms.get_equivalent_preprocessing()
+        params = dict(
+            conf=0.25,
+            image_processor={Processings.ComposeProcessing: {"processings": pipeline}},
+            joint_links=self.joint_links,
+            joint_names=self.joint_colors,
+        )
+        return params