[Refactor] data flow (#1956)

* [WIP] Refactor data flow

* model return

* [WIP] Refactor data flow

* support data_samples is optional

* fix benchmark

* fix base

* minors

* rebase

* fix api

* ut

* fix api inference

* comments

* docstring

* docstring

* docstring

* fix bug of slide inference

* add assert c > 1

Author: MeowZheng

mmseg/apis/inference.py

@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import warnings
+from collections import defaultdict
 from pathlib import Path
 from typing import Optional, Sequence, Union
 
@@ -11,9 +12,9 @@
 from mmengine.runner import load_checkpoint
 from mmengine.utils import mkdir_or_exist
 
-from mmseg.data import SegDataSample
 from mmseg.models import BaseSegmentor
 from mmseg.registry import MODELS
+from mmseg.structures import SegDataSample
 from mmseg.utils import SampleList, dataset_aliases, get_classes, get_palette
 from mmseg.visualization import SegLocalVisualizer
 
@@ -50,7 +51,6 @@ def init_model(config: Union[str, Path, Config],
     model = MODELS.build(config.model)
     if checkpoint is not None:
         checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
-
         dataset_meta = checkpoint['meta'].get('dataset_meta', None)
         # save the dataset_meta in the model for convenience
         if 'dataset_meta' in checkpoint.get('meta', {}):
@@ -108,14 +108,15 @@ def _preprare_data(imgs: ImageType, model: BaseSegmentor):
     # a pipeline for each inference
     pipeline = Compose(cfg.test_pipeline)
 
-    data = []
+    data = defaultdict(list)
     for img in imgs:
         if isinstance(img, np.ndarray):
             data_ = dict(img=img)
         else:
             data_ = dict(img_path=img)
         data_ = pipeline(data_)
-        data.append(data_)
+        data['inputs'].append(data_['inputs'])
+        data['data_samples'].append(data_['data_samples'])
 
     return data, is_batch
 
@@ -187,11 +188,12 @@ def show_result_pyplot(model: BaseSegmentor,
         save_dir=save_dir,
         alpha=opacity)
     visualizer.dataset_meta = dict(
-        classes=model.CLASSES, palette=model.PALETTE)
+        classes=model.dataset_meta['classes'],
+        palette=model.dataset_meta['palette'])
     visualizer.add_datasample(
         name=title,
         image=image,
-        pred_sample=result[0],
+        data_sample=result[0],
         draw_gt=draw_gt,
         draw_pred=draw_pred,
         wait_time=wait_time,

mmseg/datasets/transforms/formatting.py

@@ -78,7 +78,7 @@ def transform(self, results: dict) -> dict:
             if key in results:
                 img_meta[key] = results[key]
         data_sample.set_metainfo(img_meta)
-        packed_results['data_sample'] = data_sample
+        packed_results['data_samples'] = data_sample
 
         return packed_results
 

mmseg/engine/hooks/visualization_hook.py

@@ -66,15 +66,15 @@ def __init__(self,
     def _after_iter(self,
                     runner: Runner,
                     batch_idx: int,
-                    data_batch: Sequence[dict],
+                    data_batch: dict,
                     outputs: Sequence[SegDataSample],
                     mode: str = 'val') -> None:
         """Run after every ``self.interval`` validation iterations.
 
         Args:
             runner (:obj:`Runner`): The runner of the validation process.
             batch_idx (int): The index of the current batch in the val loop.
-            data_batch (Sequence[dict]): Data from dataloader.
+            data_batch (dict): Data from dataloader.
             outputs (Sequence[:obj:`SegDataSample`]): Outputs from model.
             mode (str): mode (str): Current mode of runner. Defaults to 'val'.
         """
@@ -85,18 +85,16 @@ def _after_iter(self,
             self.file_client = FileClient(**self.file_client_args)
 
         if self.every_n_inner_iters(batch_idx, self.interval):
-            for input_data, output in zip(data_batch, outputs):
-                img_path = input_data['data_sample'].img_path
+            for output in outputs:
+                img_path = output.img_path
                 img_bytes = self.file_client.get(img_path)
                 img = mmcv.imfrombytes(img_bytes, channel_order='rgb')
                 window_name = f'{mode}_{osp.basename(img_path)}'
 
-                gt_sample = input_data['data_sample']
                 self._visualizer.add_datasample(
                     window_name,
                     img,
-                    gt_sample=gt_sample,
-                    pred_sample=output,
+                    data_sample=output,
                     show=self.show,
                     wait_time=self.wait_time,
                     step=runner.iter)

mmseg/evaluation/metrics/citys_metric.py

@@ -49,25 +49,24 @@ def __init__(self,
         self.to_label_id = to_label_id
         self.suffix = suffix
 
-    def process(self, data_batch: Sequence[dict],
-                predictions: Sequence[dict]) -> None:
-        """Process one batch of data and predictions.
+    def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
+        """Process one batch of data and data_samples.
 
         The processed results should be stored in ``self.results``, which will
         be used to computed the metrics when all batches have been processed.
 
         Args:
-            data_batch (Sequence[dict]): A batch of data from the dataloader.
-            predictions (Sequence[dict]): A batch of outputs from the model.
+            data_batch (dict): A batch of data from the dataloader.
+            data_samples (Sequence[dict]): A batch of outputs from the model.
         """
         mkdir_or_exist(self.suffix)
 
-        for pred in predictions:
-            pred_label = pred['pred_sem_seg']['data'][0].cpu().numpy()
+        for data_sample in data_samples:
+            pred_label = data_sample['pred_sem_seg']['data'][0].cpu().numpy()
             # results2img
             if self.to_label_id:
                 pred_label = self._convert_to_label_id(pred_label)
-            basename = osp.splitext(osp.basename(pred['img_path']))[0]
+            basename = osp.splitext(osp.basename(data_sample['img_path']))[0]
             png_filename = osp.join(self.suffix, f'{basename}.png')
             output = Image.fromarray(pred_label.astype(np.uint8)).convert('P')
             import cityscapesscripts.helpers.labels as CSLabels

mmseg/evaluation/metrics/iou_metric.py

@@ -47,22 +47,20 @@ def __init__(self,
         self.nan_to_num = nan_to_num
         self.beta = beta
 
-    def process(self, data_batch: Sequence[dict],
-                predictions: Sequence[dict]) -> None:
-        """Process one batch of data and predictions.
+    def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
+        """Process one batch of data and data_samples.
 
         The processed results should be stored in ``self.results``, which will
         be used to computed the metrics when all batches have been processed.
 
         Args:
-            data_batch (Sequence[dict]): A batch of data from the dataloader.
-            predictions (Sequence[dict]): A batch of outputs from the model.
+            data_batch (dict): A batch of data from the dataloader.
+            data_samples (Sequence[dict]): A batch of outputs from the model.
         """
         num_classes = len(self.dataset_meta['classes'])
-        for data, pred in zip(data_batch, predictions):
-            pred_label = pred['pred_sem_seg']['data'].squeeze()
-            label = data['data_sample']['gt_sem_seg']['data'].squeeze().to(
-                pred_label)
+        for data_sample in data_samples:
+            pred_label = data_sample['pred_sem_seg']['data'].squeeze()
+            label = data_sample['gt_sem_seg']['data'].squeeze().to(pred_label)
             self.results.append(
                 self.intersect_and_union(pred_label, label, num_classes,
                                          self.ignore_index))

mmseg/models/data_preprocessor.py

@@ -1,13 +1,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from numbers import Number
-from typing import List, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Optional, Sequence
 
 import torch
 from mmengine.model import BaseDataPreprocessor
-from torch import Tensor
 
 from mmseg.registry import MODELS
-from mmseg.utils import OptSampleList, stack_batch
+from mmseg.utils import stack_batch
 
 
 @MODELS.register_module()
@@ -87,22 +86,20 @@ def __init__(self,
         # TODO: support batch augmentations.
         self.batch_augments = batch_augments
 
-    def forward(self,
-                data: Sequence[dict],
-                training: bool = False) -> Tuple[Tensor, OptSampleList]:
+    def forward(self, data: dict, training: bool = False) -> Dict[str, Any]:
         """Perform normalization、padding and bgr2rgb conversion based on
         ``BaseDataPreprocessor``.
 
         Args:
-            data (Sequence[dict]): data sampled from dataloader.
+            data (dict): data sampled from dataloader.
             training (bool): Whether to enable training time augmentation.
 
         Returns:
-            Tuple[torch.Tensor, Optional[list]]: Data in the same format as the
-            model input.
+            Dict: Data in the same format as the model input.
         """
-        inputs, batch_data_samples = self.collate_data(data)
-
+        data = self.cast_data(data)  # type: ignore
+        inputs = data['inputs']
+        data_samples = data.get('data_samples', None)
         # TODO: whether normalize should be after stack_batch
         if self.channel_conversion and inputs[0].size(0) == 3:
             inputs = [_input[[2, 1, 0], ...] for _input in inputs]
@@ -113,20 +110,23 @@ def forward(self,
             inputs = [_input.float() for _input in inputs]
 
         if training:
-            batch_inputs, batch_data_samples = stack_batch(
+            assert data_samples is not None, ('During training, ',
+                                              '`data_samples` must be define.')
+            inputs, data_samples = stack_batch(
                 inputs=inputs,
-                batch_data_samples=batch_data_samples,
+                data_samples=data_samples,
                 size=self.size,
                 size_divisor=self.size_divisor,
                 pad_val=self.pad_val,
                 seg_pad_val=self.seg_pad_val)
 
             if self.batch_augments is not None:
-                inputs, batch_data_samples = self.batch_augments(
-                    inputs, batch_data_samples)
-            return batch_inputs, batch_data_samples
+                inputs, data_samples = self.batch_augments(
+                    inputs, data_samples)
+            return dict(inputs=inputs, data_samples=data_samples)
         else:
             assert len(inputs) == 1, (
                 'Batch inference is not support currently, '
                 'as the image size might be different in a batch')
-            return torch.stack(inputs, dim=0), batch_data_samples
+            return dict(
+                inputs=torch.stack(inputs, dim=0), data_samples=data_samples)