[Feature] Add results2img, format_results for ade dataset

mmseg/datasets/ade.py

@@ -1,3 +1,10 @@
+import os.path as osp
+import tempfile
+
+import mmcv
+import numpy as np
+from PIL import Image
+
 from .builder import DATASETS
 from .custom import CustomDataset
 
@@ -82,3 +89,81 @@ def __init__(self, **kwargs):
             seg_map_suffix='.png',
             reduce_zero_label=True,
             **kwargs)
+
+    def results2img(self, results, imgfile_prefix, to_label_id):
+        """Write the segmentation results to images.
+
+        Args:
+            results (list[list | tuple | ndarray]): Testing results of the
+                dataset.
+            imgfile_prefix (str): The filename prefix of the png files.
+                If the prefix is "somepath/xxx",
+                the png files will be named "somepath/xxx.png".
+            to_label_id (bool): whether convert output to label_id for
+                submission
+
+        Returns:
+            list[str: str]: result txt files which contains corresponding
+            semantic segmentation images.
+        """
+        mmcv.mkdir_or_exist(imgfile_prefix)
+        result_files = []
+        prog_bar = mmcv.ProgressBar(len(self))
+        for idx in range(len(self)):
+            result = results[idx]
+
+            filename = self.img_infos[idx]['filename']
+            basename = osp.splitext(osp.basename(filename))[0]
+
+            # save seg_logit
+            if len(result.shape) == 3:
+                # print(result.shape)
+                npy_filename = osp.join(imgfile_prefix, f'{basename}.npy')
+                np.save(npy_filename, result)
+                result_files.append(npy_filename)
+
+            # save seg_pred
+            if len(result.shape) == 2:
+                # print(result.shape)
+                png_filename = osp.join(imgfile_prefix, f'{basename}.png')
+                result = result + 1
+
+                output = Image.fromarray(result.astype(np.uint8))
+                output.save(png_filename)
+                result_files.append(png_filename)
+
+            prog_bar.update()
+
+        return result_files
+
+    def format_results(self, results, imgfile_prefix=None, to_label_id=True):
+        """Format the results into dir (standard format for Cityscapes
+        evaluation).
+
+        Args:
+            results (list): Testing results of the dataset.
+            imgfile_prefix (str | None): The prefix of images files. It
+                includes the file path and the prefix of filename, e.g.,
+                "a/b/prefix". If not specified, a temp file will be created.
+                Default: None.
+            to_label_id (bool): whether convert output to label_id for
+                submission. Default: False
+
+        Returns:
+            tuple: (result_files, tmp_dir), result_files is a list containing
+               the image paths, tmp_dir is the temporal directory created
+                for saving json/png files when img_prefix is not specified.
+        """
+
+        assert isinstance(results, list), 'results must be a list'
+        assert len(results) == len(self), (
+            'The length of results is not equal to the dataset len: '
+            f'{len(results)} != {len(self)}')
+
+        if imgfile_prefix is None:
+            tmp_dir = tempfile.TemporaryDirectory()
+            imgfile_prefix = tmp_dir.name
+        else:
+            tmp_dir = None
+        result_files = self.results2img(results, imgfile_prefix, to_label_id)
+        return result_files, tmp_dir