Enable hierarchical classification in MPA

data/car_tree_bug/annotations/multilabel_car_tree.json

@@ -1 +1,6 @@
-{"images": [["Slide18.PNG", ["car"]], ["Slide19.PNG", ["tree"]], ["Slide16.PNG", ["car", "tree"]]], "classes": ["car", "tree"]}
+{"images": [["Slide18.PNG", ["car"]],
+            ["Slide19.PNG", ["tree"]], 
+            ["Slide16.PNG", ["car", "tree"]]
+            ], 
+ "classes": ["car", "tree"]
+}

data/car_tree_bug/annotations/multilabel_default.json

@@ -1 +1,20 @@
-{"images": [["Slide10.PNG", ["tree", "car", "bug"]], ["Slide12.PNG", ["bug", "car", "tree"]], ["Slide14.PNG", ["tree", "bug", "car"]], ["Slide9.PNG", ["bug", "car", "tree"]], ["Slide11.PNG", ["car", "tree", "bug"]], ["Slide17.PNG", ["bug", "tree"]], ["Slide7.PNG", ["tree", "car", "bug"]], ["Slide6.PNG", ["car", "tree", "bug"]], ["Slide4.PNG", ["bug", "tree", "car"]], ["Slide20.PNG", ["bug"]], ["Slide5.PNG", ["bug", "tree", "car"]], ["Slide1.PNG", ["car", "bug", "tree"]], ["Slide3.PNG", ["car", "tree", "bug"]], ["Slide13.PNG", ["bug", "car", "tree"]], ["Slide2.PNG", ["bug", "car", "tree"]], ["Slide8.PNG", ["car", "tree", "bug"]], ["Slide15.PNG", ["bug", "car"]]], "classes": ["car", "tree", "bug"]}
+{"images": [["Slide10.PNG", ["tree", "car", "bug"]], 
+            ["Slide12.PNG", ["bug", "car", "tree"]], 
+            ["Slide14.PNG", ["tree", "bug", "car"]], 
+            ["Slide9.PNG", ["bug", "car", "tree"]], 
+            ["Slide11.PNG", ["car", "tree", "bug"]], 
+            ["Slide17.PNG", ["bug", "tree"]], 
+            ["Slide7.PNG", ["tree", "car", "bug"]], 
+            ["Slide6.PNG", ["car", "tree", "bug"]], 
+            ["Slide4.PNG", ["bug", "tree", "car"]], 
+            ["Slide20.PNG", ["bug"]], 
+            ["Slide5.PNG", ["bug", "tree", "car"]], 
+            ["Slide1.PNG", ["car", "bug", "tree"]], 
+            ["Slide3.PNG", ["car", "tree", "bug"]], 
+            ["Slide13.PNG", ["bug", "car", "tree"]], 
+            ["Slide2.PNG", ["bug", "car", "tree"]], 
+            ["Slide8.PNG", ["car", "tree", "bug"]], 
+            ["Slide15.PNG", ["bug", "car"]]
+            ], 
+"classes": ["car", "tree", "bug"]
+}

external/deep-object-reid/configs/ote_custom_classification/mobilenet_v3_large_1/main_model_multihead.yaml

@@ -11,7 +11,7 @@ lr_finder:
   n_trials: 15
 
 model:
-  name: 'mobilenetv3_large_21k'
+  name: 'mobilenetv3_large'
   type: 'multihead'
   pretrained: True
   save_all_chkpts: False
@@ -26,8 +26,8 @@ data:
   root: './'
   height: 224
   width: 224
-  norm_mean: [0, 0, 0]
-  norm_std: [1, 1, 1]
+  norm_mean: [0.485, 0.456, 0.406]
+  norm_std: [0.229, 0.224, 0.225]
   save_dir: 'output/mulitihead/mobilenetv3_large_1'
   workers: 6
   transforms:

external/deep-object-reid/configs/ote_custom_classification/mobilenet_v3_large_1/main_model_multilabel.yaml

@@ -26,8 +26,8 @@ data:
   root: './'
   height: 224
   width: 224
-  norm_mean: [0, 0, 0]
-  norm_std: [1, 1, 1]
+  norm_mean: [0.485, 0.456, 0.406]
+  norm_std: [0.229, 0.224, 0.225]
   save_dir: 'output/mulitilabel/mobilenetv3_large'
   workers: 6
   transforms:

external/deep-object-reid/torchreid_tasks/utils.py

@@ -54,6 +54,7 @@
 
 from torchreid.utils import set_model_attr, get_model_attr
 
+
 class ClassificationType(Enum):
     MULTICLASS = auto()
     MULTILABEL = auto()
@@ -121,16 +122,32 @@ def _load_text_annotation(annot_path, data_dir):
         out_data = []
         with open(annot_path) as f:
             annotation = json.load(f)
-            if not 'label_groups' in annotation:
+            if 'hierarchy' not in annotation:
                 all_classes = sorted(annotation['classes'])
                 annotation_type = ClassificationType.MULTILABEL
                 groups = [[c] for c in all_classes]
-            else: # load multihead
-                groups = annotation['label_groups']
+            else:  # load multihead
                 all_classes = []
-                for g in groups:
-                    for c in g:
-                        all_classes.append(c)
+                groups = annotation['hierarchy']
+
+                def add_subtask_labels(group):
+                    if isinstance(group, dict) and 'subtask' in group:
+                        subtask = group['subtask']
+                        if isinstance(subtask, list):
+                            for task in subtask:
+                                for task_label in task['labels']:
+                                    all_classes.append(task_label)
+                        elif isinstance(subtask, dict):
+                            for task_label in subtask['labels']:
+                                all_classes.append(task_label)
+                        add_subtask_labels(subtask)
+                    elif isinstance(group, list):
+                        for task in group:
+                            add_subtask_labels(task)
+                for group in groups:
+                    for label in group['labels']:
+                        all_classes.append(label)
+                    add_subtask_labels(group)
                 annotation_type = ClassificationType.MULTIHEAD
 
             images_info = annotation['images']
@@ -258,7 +275,7 @@ def get_multihead_class_info(label_schema: LabelSchemaEntity):
         head_idx_to_logits_range[i] = (last_logits_pos, last_logits_pos + len(g))
         last_logits_pos += len(g)
         for j, c in enumerate(g):
-            class_to_idx[c] = (i, j) # group idx and idx inside group
+            class_to_idx[c] = (i, j)  # group idx and idx inside group
             num_single_label_classes += 1
 
     # other labels are in multilabel group
@@ -302,7 +319,7 @@ def __init__(self, ote_dataset: DatasetEntity, labels: List[LabelEntity], multil
             if item_labels:
                 if not self.hierarchical:
                     for ote_lbl in item_labels:
-                        if not ote_lbl in ignored_labels:
+                        if ote_lbl not in ignored_labels:
                             class_indices.append(self.label_names.index(ote_lbl.name))
                         else:
                             class_indices.append(-1)
@@ -318,12 +335,12 @@ def __init__(self, ote_dataset: DatasetEntity, labels: List[LabelEntity], multil
                         if group_idx < num_cls_heads:
                             class_indices[group_idx] = in_group_idx
                         else:
-                            if not ote_lbl in ignored_labels:
+                            if ote_lbl not in ignored_labels:
                                 class_indices[num_cls_heads + in_group_idx] = 1
                             else:
                                 class_indices[num_cls_heads + in_group_idx] = -1
 
-            else: # this supposed to happen only on inference stage or if we have a negative in multilabel data
+            else:  # this supposed to happen only on inference stage or if we have a negative in multilabel data
                 if self.mixed_cls_heads_info:
                     class_indices = [-1]*(self.mixed_cls_heads_info['num_multiclass_heads'] + \
                                           self.mixed_cls_heads_info['num_multilabel_classes'])

external/model-preparation-algorithm/configs/classification/efficientnet_b0_cls_incr/model_hierarchical.py

@@ -0,0 +1,3 @@
+_base_ = [
+  '../../../submodule/models/classification/ote_efficientnet_b0_hierarchical.yaml',
+]

external/model-preparation-algorithm/configs/classification/efficientnet_v2_s_cls_incr/model_hierarchical.py

@@ -0,0 +1,3 @@
+_base_ = [
+  '../../../submodule/models/classification/ote_efficientnet_v2_s_hierarchical.yaml',
+]

external/model-preparation-algorithm/configs/classification/mobilenet_v3_large_075_cls_incr/model_hierarchical.py

@@ -0,0 +1,3 @@
+_base_ = [
+  '../../../submodule/models/classification/ote_mobilenet_v3_large_075_hierarchical.yaml',
+]

external/model-preparation-algorithm/configs/classification/mobilenet_v3_large_1_cls_incr/model_hierarchical.py

@@ -0,0 +1,3 @@
+_base_ = [
+  '../../../submodule/models/classification/ote_mobilenet_v3_large_hierarchical.yaml',
+]

external/model-preparation-algorithm/configs/classification/mobilenet_v3_small_cls_incr/model_hierarchical.py

@@ -0,0 +1,3 @@
+_base_ = [
+  '../../../submodule/models/classification/ote_mobilenet_v3_small_hierarchical.yaml',
+]

external/model-preparation-algorithm/mpa_tasks/apis/classification/task.py

@@ -8,6 +8,7 @@
 from typing import List, Optional
 
 import torch
+import numpy as np
 from mpa import MPAConstants
 
 from ote_sdk.configuration import cfg_helper
@@ -34,12 +35,14 @@
 from ote_sdk.entities.model import (ModelFormat, ModelOptimizationType)
 from ote_sdk.serialization.label_mapper import label_schema_to_bytes
 from ote_sdk.entities.scored_label import ScoredLabel
+from ote_sdk.utils.labels_utils import get_empty_label
 from torchreid_tasks.utils import TrainingProgressCallback
 from torchreid_tasks.utils import OTELoggerHook
-from torchreid_tasks.train_task import OTEClassificationTrainingTask
+from torchreid_tasks.utils import get_multihead_class_info as get_hierarchical_info
 from mpa_tasks.apis import BaseTask, TrainType
 from mpa_tasks.apis.classification import ClassificationConfig
 from mpa.utils.config_utils import MPAConfig
+from mpa.stage import Stage
 from mpa.utils.logger import get_logger
 from ote_sdk.entities.label import Domain
 
@@ -63,38 +66,24 @@ def __init__(self, task_environment: TaskEnvironment):
             self._labels = task_environment.get_labels(include_empty=True)
         else:
             self._labels = task_environment.get_labels(include_empty=False)
+        self._empty_label = get_empty_label(task_environment.label_schema)
+        self._multilabel = False
+        self._hierarchical = False
+
         self._multilabel = len(task_environment.label_schema.get_groups(False)) > 1 and \
                            len(task_environment.label_schema.get_groups(False)) == \
                            len(task_environment.get_labels(include_empty=False))  # noqa:E127
 
-        self._hierarchical = False
+        self._hierarchical_info = None
         if not self._multilabel and len(task_environment.label_schema.get_groups(False)) > 1:
             self._hierarchical = True
-            torchreid_env = self.convert_to_torchreid_env(task_environment)
-            self.torchreid_train_task = OTEClassificationTrainingTask(torchreid_env)
-
-    @staticmethod
-    def convert_to_torchreid_env(task_env):
-        model_template = task_env.model_template
-        rename_dict = {'model-preparation-algorithm': 'deep-object-reid',
-                       'classification' : 'ote_custom_classification',
-                       '_cls_incr' : ''
-                       }
-
-        for key, val in rename_dict.items():
-            model_template.model_template_path = model_template.model_template_path.replace(key, val)
-        model_template.entrypoints.base = 'torchreid_tasks.train_task.OTEClassificationTrainingTask'
-        model_template.framework = 'OTEClassification v1.2.3'
-        return task_env
+            self._hierarchical_info = get_hierarchical_info(task_environment.label_schema)
 
     def infer(self,
               dataset: DatasetEntity,
               inference_parameters: Optional[InferenceParameters] = None
               ) -> DatasetEntity:
         logger.info('called infer()')
-        if self._hierarchical:
-            self = self.torchreid_train_task
-            return self.infer(dataset, inference_parameters)
         stage_module = 'ClsInferrer'
         self._data_cfg = self._init_test_data_cfg(dataset)
         dataset = dataset.with_empty_annotations()
@@ -109,19 +98,52 @@ def infer(self,
 
         dataset_size = len(dataset)
         for i, (dataset_item, prediction_item) in enumerate(zip(dataset, predictions)):
-            label = []
+            item_labels = []
+            pos_thr = 0.5
+
+            if any(np.isnan(prediction_item)):
+                logger.info('Nan in prediction_item.')
+
             if self._multilabel:
-                pos_thr = 0.5
+                if max(prediction_item) < pos_thr:
+                    logger.info('Confidence is smaller than pos_thr, empty_label will be appended to item_labels.')
+                    item_labels.append(ScoredLabel(self._empty_label, probability=1.))
+                    continue
                 for cls_idx, pred_item in enumerate(prediction_item):
                     if pred_item > pos_thr:
                         cls_label = ScoredLabel(self.labels[cls_idx], probability=float(pred_item))
-                        label.append(cls_label)
+                        item_labels.append(cls_label)
+
+            elif self._hierarchical:
+                for head_idx in range(self._hierarchical_info['num_multiclass_heads']):
+                    logits_begin, logits_end = self._hierarchical_info['head_idx_to_logits_range'][head_idx]
+                    head_logits = prediction_item[logits_begin : logits_end]
+                    head_pred = np.argmax(head_logits)  # Assume logits already passed softmax
+                    label_str = self._hierarchical_info['all_groups'][head_idx][head_pred]
+                    ote_label = next(x for x in self._labels if x.name == label_str)
+                    item_labels.append(ScoredLabel(label=ote_label, probability=float(head_logits[head_pred])))
+
+                if self._hierarchical_info['num_multilabel_classes']:
+                    logits_begin, logits_end = self._hierarchical_info['num_single_label_classes'], -1
+                    head_logits = prediction_item[logits_begin : logits_end]
+                    for logit_idx, logit in enumerate(head_logits):
+                        if logit > pos_thr:  # Assume logits already passed sigmoid
+                            label_str = self._hierarchical_info['all_groups'][self._hierarchical_info['num_multiclass_heads']+logit_idx][0]
+                            ote_label = next(x for x in self._labels if x.name == label_str)
+                            item_labels.append(ScoredLabel(label=ote_label, probability=float(logit)))
+                item_labels = self._task_environment.label_schema.resolve_labels_probabilistic(item_labels)
+                if not item_labels:
+                    logger.info('item_labels is empty.')
+                    item_labels.append(ScoredLabel(self._empty_label, probability=1.))
+
             else:
                 label_idx = prediction_item.argmax()
                 cls_label = ScoredLabel(self._labels[label_idx], probability=float(prediction_item[label_idx]))
-                label.append(cls_label)
-            dataset_item.append_labels(label)
+                item_labels.append(cls_label)
+
+            dataset_item.append_labels(item_labels)
             update_progress_callback(int(i / dataset_size * 100))
+
         return dataset
 
     def evaluate(self,
@@ -141,10 +163,6 @@ def export(self,
                export_type: ExportType,
                output_model: ModelEntity):
         logger.info('Exporting the model')
-        if self._hierarchical:
-            self = self.torchreid_train_task
-            self.export(export_type, output_model)
-            return
         if export_type != ExportType.OPENVINO:
             raise RuntimeError(f'not supported export type {export_type}')
         output_model.model_format = ModelFormat.OPENVINO
@@ -206,10 +224,15 @@ def _init_model_cfg(self):
         base_dir = os.path.abspath(os.path.dirname(self.template_file_path))
         if self._multilabel:
             cfg_path = os.path.join(base_dir, 'model_multilabel.py')
+        elif self._hierarchical:
+            cfg_path = os.path.join(base_dir, 'model_hierarchical.py')
         else:
             cfg_path = os.path.join(base_dir, 'model.py')
         cfg = MPAConfig.fromfile(cfg_path)
         cfg.model.multilabel = self._multilabel
+        cfg.model.hierarchical = self._hierarchical
+        if self._hierarchical:
+            cfg.model.head.hierarchical_info = self._hierarchical_info
         return cfg
 
     def _init_test_data_cfg(self, dataset: DatasetEntity):
@@ -250,11 +273,24 @@ def patch_color_conversion(pipeline):
 
             if self._multilabel:
                 cfg.type = 'MPAMultilabelClsDataset'
+            elif self._hierarchical:
+                cfg.type = 'MPAHierarchicalClsDataset'
+                cfg.hierarchical_info = self._hierarchical_info
+                if subset == 'train':
+                    cfg.drop_last = True  # For stable hierarchical information indexing
             else:
                 cfg.type = 'MPAClsDataset'
+
+            # In train dataset, when sample size is smaller than batch size
+            if subset == 'train' and self._data_cfg:
+                train_data_cfg = Stage.get_train_data_cfg(self._data_cfg)
+                if (len(train_data_cfg.get('ote_dataset', [])) < self._recipe_cfg.data.get('samples_per_gpu', 2)):
+                    cfg.drop_last = False
+
             cfg.domain = domain
             cfg.ote_dataset = None
             cfg.labels = None
+            cfg.empty_label = self._empty_label
             for pipeline_step in cfg.pipeline:
                 if subset == 'train' and pipeline_step.type == 'Collect':
                     pipeline_step = BaseTask._get_meta_keys(pipeline_step)
@@ -264,6 +300,8 @@ def _patch_evaluation(self, config: MPAConfig):
         cfg = config.evaluation
         if self._multilabel:
             cfg.metric = ['accuracy-mlc', 'mAP', 'CP', 'OP', 'CR', 'OR', 'CF1', 'OF1']
+        elif self._hierarchical:
+            cfg.metric = ['MHAcc', 'avgClsAcc', 'mAP']
         else:
             cfg.metric = ['accuracy', 'class_accuracy']
 
@@ -302,10 +340,6 @@ def train(self,
               output_model: ModelEntity,
               train_parameters: Optional[TrainParameters] = None):
         logger.info('train()')
-        if self._hierarchical:
-            self = self.torchreid_train_task
-            self.train(dataset, output_model, train_parameters)
-            return
         # Check for stop signal between pre-eval and training.
         # If training is cancelled at this point,
         if self._should_stop:
@@ -383,6 +417,8 @@ def _generate_training_metrics_group(self, learning_curves) -> Optional[List[Met
 
         if self._multilabel:
             metric_key = 'val/accuracy-mlc'
+        elif self._hierarchical:
+            metric_key = 'val/MHAcc'
         else:
             metric_key = 'val/accuracy_top-1'