Hotfix/sg 000 fix sigmas (#1503)

* Fixed a bug where OKS sigmas had two incorrect values

* Fixed a bug where OKS sigmas had two incorrect values

* Allow user setting the SUPER_GRADIENTS_COCO_DATASET_DIR env variable to run test without need to edit path location by hand

* Remove excess stuff

(cherry picked from commit 307f7bd)

documentation/source/PoseEstimation.md

@@ -284,7 +284,7 @@ Then, the full configuration file should look like this:
 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]
+oks_sigmas: [0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089]
 
 train_dataset_params:
   image_paths: /my_new_dataset/train/images 

src/super_gradients/recipes/dataset_params/coco_pose_estimation_common_dataset_params.yaml

@@ -0,0 +1,72 @@
+# This is not "true" dataset params, one cannot use it to instantiate dataloaders
+# But it contains skeleton definitions for COCO2017 dataset and exists to avoid
+# duplication of those parameters in other dataset params
+
+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, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089]
+
+flip_indexes: [ 0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15,]
+
+edge_links:
+  - [0, 1]
+  - [0, 2]
+  - [1, 2]
+  - [1, 3]
+  - [2, 4]
+  - [3, 5]
+  - [4, 6]
+  - [5, 6]
+  - [5, 7]
+  - [5, 11]
+  - [6, 8]
+  - [6, 12]
+  - [7, 9]
+  - [8, 10]
+  - [11, 12]
+  - [11, 13]
+  - [12, 14]
+  - [13, 15]
+  - [14, 16]
+
+edge_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]

src/super_gradients/recipes/dataset_params/coco_pose_estimation_dataset_params.yaml

@@ -1,71 +1,6 @@
-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: [ 0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15,]
-
-edge_links:
-  - [0, 1]
-  - [0, 2]
-  - [1, 2]
-  - [1, 3]
-  - [2, 4]
-  - [3, 5]
-  - [4, 6]
-  - [5, 6]
-  - [5, 7]
-  - [5, 11]
-  - [6, 8]
-  - [6, 12]
-  - [7, 9]
-  - [8, 10]
-  - [11, 12]
-  - [11, 13]
-  - [12, 14]
-  - [13, 15]
-  - [14, 16]
-
-edge_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]
+defaults:
+  - coco_pose_estimation_common_dataset_params
+  - _self_
 
 
 train_dataset_params:

src/super_gradients/recipes/dataset_params/coco_pose_estimation_rescoring_dataset_params.yaml

@@ -1,29 +1,6 @@
-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]
-
-edge_links:
-  - [0, 1]
-  - [0, 2]
-  - [1, 2]
-  - [1, 3]
-  - [2, 4]
-  - [3, 5]
-  - [4, 6]
-  - [5, 6]
-  - [5, 7]
-  - [5, 11]
-  - [6, 8]
-  - [6, 12]
-  - [7, 9]
-  - [8, 10]
-  - [11, 12]
-  - [11, 13]
-  - [12, 14]
-  - [13, 15]
-  - [14, 16]
-
+defaults:
+  - coco_pose_estimation_common_dataset_params
+  - _self_
 
 train_dataset_params:
   pkl_file: REPLACE_ME_WITH_PATH_TO_TRAIN_DATASET.PKL

src/super_gradients/training/metrics/pose_estimation_metrics.py

@@ -104,7 +104,7 @@ def __init__(
 
         if oks_sigmas is None:
             if num_joints == 17:
-                oks_sigmas = np.array([0.26, 0.25, 0.25, 0.35, 0.35, 0.79, 0.79, 0.72, 0.72, 0.62, 0.62, 1.07, 1.07, 0.87, 0.87, 0.89, 0.89]) / 10.0
+                oks_sigmas = np.array([0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089])
             else:
                 oks_sigmas = np.array([0.1] * num_joints)
                 logger.warning(

tests/integration_tests/pose_estimation_models_test.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 import torch
@@ -14,14 +15,17 @@
 
 class PoseEstimationModelsIntegrationTest(unittest.TestCase):
     def setUp(self):
-        self.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]
+        self.oks_sigmas = [0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089]
         self.flip_indexes = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
+        # This is for easy testing on local machine - you can set this environment variable to your own COCO dataset location
+        self.data_dir = os.environ.get("SUPER_GRADIENTS_COCO_DATASET_DIR", "/data/coco")
 
     def test_dekr_model(self):
         val_loader = get_data_loader(
             "coco_pose_estimation_dekr_dataset_params",
             COCOKeypointsDataset,
             train=False,
+            dataset_params=dict(data_dir=self.data_dir),
             dataloader_params=dict(num_workers=0),
         )
 
@@ -54,6 +58,7 @@ def test_dekr_model_with_tta(self):
             "coco_pose_estimation_dekr_dataset_params",
             COCOKeypointsDataset,
             train=False,
+            dataset_params=dict(data_dir=self.data_dir),
             dataloader_params=dict(num_workers=0),
         )
 
@@ -84,6 +89,7 @@ def test_dekr_model_with_rescoring(self):
             "coco_pose_estimation_dekr_dataset_params",
             COCOKeypointsDataset,
             train=False,
+            dataset_params=dict(data_dir=self.data_dir),
             dataloader_params=dict(batch_size=1, num_workers=0),
         )