Add files using upload-large-folder tool

7dce762 verified over 1 year ago

12.9 kB

# Copyright (c) OpenMMLab. All rights reserved.
import mmcv
import numpy as np
from mmseg.datasets.builder import PIPELINES


@PIPELINES.register_module()
class SETR_Resize(object):
    """Resize images & seg.
    This transform resizes the input image to some scale. If the input dict
    contains the key "scale", then the scale in the input dict is used,
    otherwise the specified scale in the init method is used.
    ``img_scale`` can either be a tuple (single-scale) or a list of tuple
    (multi-scale). There are 3 multiscale modes:
    - ``ratio_range is not None``: randomly sample a ratio from the ratio range
    and multiply it with the image scale.
    - ``ratio_range is None and multiscale_mode == "range"``: randomly sample a
    scale from the a range.
    - ``ratio_range is None and multiscale_mode == "value"``: randomly sample a
    scale from multiple scales.
    Args:
        img_scale (tuple or list[tuple]): Images scales for resizing.
        multiscale_mode (str): Either "range" or "value".
        ratio_range (tuple[float]): (min_ratio, max_ratio)
        keep_ratio (bool): Whether to keep the aspect ratio when resizing the
            image.
    """
    def __init__(self,
                 img_scale=None,
                 multiscale_mode='range',
                 ratio_range=None,
                 keep_ratio=True,
                 crop_size=None,
                 setr_multi_scale=False):

        if img_scale is None:
            self.img_scale = None
        else:
            if isinstance(img_scale, list):
                self.img_scale = img_scale
            else:
                self.img_scale = [img_scale]
            # assert mmcv.is_list_of(self.img_scale, tuple)

        if ratio_range is not None:
            # mode 1: given a scale and a range of image ratio
            assert len(self.img_scale) == 1
        else:
            # mode 2: given multiple scales or a range of scales
            assert multiscale_mode in ['value', 'range']

        self.multiscale_mode = multiscale_mode
        self.ratio_range = ratio_range
        self.keep_ratio = keep_ratio
        self.crop_size = crop_size
        self.setr_multi_scale = setr_multi_scale

    @staticmethod
    def random_select(img_scales):
        """Randomly select an img_scale from given candidates.
        Args:
            img_scales (list[tuple]): Images scales for selection.
        Returns:
            (tuple, int): Returns a tuple ``(img_scale, scale_dix)``,
                where ``img_scale`` is the selected image scale and
                ``scale_idx`` is the selected index in the given candidates.
        """

        assert mmcv.is_list_of(img_scales, tuple)
        scale_idx = np.random.randint(len(img_scales))
        img_scale = img_scales[scale_idx]
        return img_scale, scale_idx

    @staticmethod
    def random_sample(img_scales):
        """Randomly sample an img_scale when ``multiscale_mode=='range'``.
        Args:
            img_scales (list[tuple]): Images scale range for sampling.
                There must be two tuples in img_scales, which specify the lower
                and uper bound of image scales.
        Returns:
            (tuple, None): Returns a tuple ``(img_scale, None)``, where
                ``img_scale`` is sampled scale and None is just a placeholder
                to be consistent with :func:`random_select`.
        """

        assert mmcv.is_list_of(img_scales, tuple) and len(img_scales) == 2
        img_scale_long = [max(s) for s in img_scales]
        img_scale_short = [min(s) for s in img_scales]
        long_edge = np.random.randint(
            min(img_scale_long),
            max(img_scale_long) + 1)
        short_edge = np.random.randint(
            min(img_scale_short),
            max(img_scale_short) + 1)
        img_scale = (long_edge, short_edge)
        return img_scale, None

    @staticmethod
    def random_sample_ratio(img_scale, ratio_range):
        """Randomly sample an img_scale when ``ratio_range`` is specified.
        A ratio will be randomly sampled from the range specified by
        ``ratio_range``. Then it would be multiplied with ``img_scale`` to
        generate sampled scale.
        Args:
            img_scale (tuple): Images scale base to multiply with ratio.
            ratio_range (tuple[float]): The minimum and maximum ratio to scale
                the ``img_scale``.
        Returns:
            (tuple, None): Returns a tuple ``(scale, None)``, where
                ``scale`` is sampled ratio multiplied with ``img_scale`` and
                None is just a placeholder to be consistent with
                :func:`random_select`.
        """

        assert isinstance(img_scale, tuple) and len(img_scale) == 2
        min_ratio, max_ratio = ratio_range
        assert min_ratio <= max_ratio
        ratio = np.random.random_sample() * (max_ratio - min_ratio) + min_ratio
        scale = int(img_scale[0] * ratio), int(img_scale[1] * ratio)
        return scale, None

    def _random_scale(self, results):
        """Randomly sample an img_scale according to ``ratio_range`` and
        ``multiscale_mode``.
        If ``ratio_range`` is specified, a ratio will be sampled and be
        multiplied with ``img_scale``.
        If multiple scales are specified by ``img_scale``, a scale will be
        sampled according to ``multiscale_mode``.
        Otherwise, single scale will be used.
        Args:
            results (dict): Result dict from :obj:`dataset`.
        Returns:
            dict: Two new keys 'scale` and 'scale_idx` are added into
                ``results``, which would be used by subsequent pipelines.
        """

        if self.ratio_range is not None:
            scale, scale_idx = self.random_sample_ratio(
                self.img_scale[0], self.ratio_range)
        elif len(self.img_scale) == 1:
            scale, scale_idx = self.img_scale[0], 0
        elif self.multiscale_mode == 'range':
            scale, scale_idx = self.random_sample(self.img_scale)
        elif self.multiscale_mode == 'value':
            scale, scale_idx = self.random_select(self.img_scale)
        else:
            raise NotImplementedError

        results['scale'] = scale
        results['scale_idx'] = scale_idx

    def _resize_img(self, results):
        """Resize images with ``results['scale']``."""

        if self.keep_ratio:
            if self.setr_multi_scale:
                if min(results['scale']) < self.crop_size[0]:
                    new_short = self.crop_size[0]
                else:
                    new_short = min(results['scale'])

                h, w = results['img'].shape[:2]
                if h > w:
                    new_h, new_w = new_short * h / w, new_short
                else:
                    new_h, new_w = new_short, new_short * w / h
                results['scale'] = (new_h, new_w)

            img, scale_factor = mmcv.imrescale(results['img'],
                                               results['scale'],
                                               return_scale=True)
            # the w_scale and h_scale has minor difference
            # a real fix should be done in the mmcv.imrescale in the future
            new_h, new_w = img.shape[:2]
            h, w = results['img'].shape[:2]
            w_scale = new_w / w
            h_scale = new_h / h
        else:
            img, w_scale, h_scale = mmcv.imresize(results['img'],
                                                  results['scale'],
                                                  return_scale=True)
        scale_factor = np.array([w_scale, h_scale, w_scale, h_scale],
                                dtype=np.float32)
        results['img'] = img
        results['img_shape'] = img.shape
        results['pad_shape'] = img.shape  # in case that there is no padding
        results['scale_factor'] = scale_factor
        results['keep_ratio'] = self.keep_ratio

    def _resize_seg(self, results):
        """Resize semantic segmentation map with ``results['scale']``."""
        for key in results.get('seg_fields', []):
            if self.keep_ratio:
                gt_seg = mmcv.imrescale(results[key],
                                        results['scale'],
                                        interpolation='nearest')
            else:
                gt_seg = mmcv.imresize(results[key],
                                       results['scale'],
                                       interpolation='nearest')
            results['gt_semantic_seg'] = gt_seg

    def __call__(self, results):
        """Call function to resize images, bounding boxes, masks, semantic
        segmentation map.
        Args:
            results (dict): Result dict from loading pipeline.
        Returns:
            dict: Resized results, 'img_shape', 'pad_shape', 'scale_factor',
                'keep_ratio' keys are added into result dict.
        """

        if 'scale' not in results:
            self._random_scale(results)
        self._resize_img(results)
        self._resize_seg(results)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += (f'(img_scale={self.img_scale}, '
                     f'multiscale_mode={self.multiscale_mode}, '
                     f'ratio_range={self.ratio_range}, '
                     f'keep_ratio={self.keep_ratio})')
        return repr_str


@PIPELINES.register_module()
class PadShortSide(object):
    """Pad the image & mask.
    Pad to the minimum size that is equal or larger than a number.
    Added keys are "pad_shape", "pad_fixed_size",
    Args:
        size (int, optional): Fixed padding size.
        pad_val (float, optional): Padding value. Default: 0.
        seg_pad_val (float, optional): Padding value of segmentation map.
            Default: 255.
    """
    def __init__(self, size=None, pad_val=0, seg_pad_val=255):
        self.size = size
        self.pad_val = pad_val
        self.seg_pad_val = seg_pad_val
        # only one of size and size_divisor should be valid
        assert size is not None

    def _pad_img(self, results):
        """Pad images according to ``self.size``."""
        h, w = results['img'].shape[:2]
        new_h = max(h, self.size)
        new_w = max(w, self.size)
        padded_img = mmcv.impad(results['img'],
                                shape=(new_h, new_w),
                                pad_val=self.pad_val)

        results['img'] = padded_img
        results['pad_shape'] = padded_img.shape
        # results['unpad_shape'] = (h, w)

    def _pad_seg(self, results):
        """Pad masks according to ``results['pad_shape']``."""
        for key in results.get('seg_fields', []):
            results[key] = mmcv.impad(results[key],
                                      shape=results['pad_shape'][:2],
                                      pad_val=self.seg_pad_val)

    def __call__(self, results):
        """Call function to pad images, masks, semantic segmentation maps.
        Args:
            results (dict): Result dict from loading pipeline.
        Returns:
            dict: Updated result dict.
        """
        h, w = results['img'].shape[:2]
        if h >= self.size and w >= self.size:  # 短边比窗口大，跳过
            pass
        else:
            self._pad_img(results)
            self._pad_seg(results)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(size={self.size}, pad_val={self.pad_val})'
        return repr_str


@PIPELINES.register_module()
class MapillaryHack(object):
    """map MV 65 class to 19 class like Cityscapes."""
    def __init__(self):
        self.map = [[13, 24, 41], [2, 15], [17], [6], [3],
                    [45, 47], [48], [50], [30], [29], [27], [19], [20, 21, 22],
                    [55], [61], [54], [58], [57], [52]]

        self.others = [i for i in range(66)]
        for i in self.map:
            for j in i:
                if j in self.others:
                    self.others.remove(j)

    def __call__(self, results):
        """Call function to process the image with gamma correction.
        Args:
            results (dict): Result dict from loading pipeline.
        Returns:
            dict: Processed results.
        """
        gt_map = results['gt_semantic_seg']
        # others -> 255
        new_gt_map = np.zeros_like(gt_map)

        for value in self.others:
            new_gt_map[gt_map == value] = 255

        for index, map in enumerate(self.map):
            for value in map:
                new_gt_map[gt_map == value] = index

        results['gt_semantic_seg'] = new_gt_map

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        return repr_str

	# Copyright (c) OpenMMLab. All rights reserved.
	import mmcv
	import numpy as np
	from mmseg.datasets.builder import PIPELINES


	@PIPELINES.register_module()
	class SETR_Resize(object):
	"""Resize images & seg.

	This transform resizes the input image to some scale. If the input dict
	contains the key "scale", then the scale in the input dict is used,
	otherwise the specified scale in the init method is used.

	``img_scale`` can either be a tuple (single-scale) or a list of tuple
	(multi-scale). There are 3 multiscale modes:

	- ``ratio_range is not None``: randomly sample a ratio from the ratio range
	and multiply it with the image scale.

	- ``ratio_range is None and multiscale_mode == "range"``: randomly sample a
	scale from the a range.

	- ``ratio_range is None and multiscale_mode == "value"``: randomly sample a
	scale from multiple scales.

	Args:
	img_scale (tuple or list[tuple]): Images scales for resizing.
	multiscale_mode (str): Either "range" or "value".
	ratio_range (tuple[float]): (min_ratio, max_ratio)
	keep_ratio (bool): Whether to keep the aspect ratio when resizing the
	image.
	"""
	def __init__(self,
	img_scale=None,
	multiscale_mode='range',
	ratio_range=None,
	keep_ratio=True,
	crop_size=None,
	setr_multi_scale=False):

	if img_scale is None:
	self.img_scale = None
	else:
	if isinstance(img_scale, list):
	self.img_scale = img_scale
	else:
	self.img_scale = [img_scale]
	# assert mmcv.is_list_of(self.img_scale, tuple)

	if ratio_range is not None:
	# mode 1: given a scale and a range of image ratio
	assert len(self.img_scale) == 1
	else:
	# mode 2: given multiple scales or a range of scales
	assert multiscale_mode in ['value', 'range']

	self.multiscale_mode = multiscale_mode
	self.ratio_range = ratio_range
	self.keep_ratio = keep_ratio
	self.crop_size = crop_size
	self.setr_multi_scale = setr_multi_scale

	@staticmethod
	def random_select(img_scales):
	"""Randomly select an img_scale from given candidates.

	Args:
	img_scales (list[tuple]): Images scales for selection.

	Returns:
	(tuple, int): Returns a tuple ``(img_scale, scale_dix)``,
	where ``img_scale`` is the selected image scale and
	``scale_idx`` is the selected index in the given candidates.
	"""

	assert mmcv.is_list_of(img_scales, tuple)
	scale_idx = np.random.randint(len(img_scales))
	img_scale = img_scales[scale_idx]
	return img_scale, scale_idx

	@staticmethod
	def random_sample(img_scales):
	"""Randomly sample an img_scale when ``multiscale_mode=='range'``.

	Args:
	img_scales (list[tuple]): Images scale range for sampling.
	There must be two tuples in img_scales, which specify the lower
	and uper bound of image scales.

	Returns:
	(tuple, None): Returns a tuple ``(img_scale, None)``, where
	``img_scale`` is sampled scale and None is just a placeholder
	to be consistent with :func:`random_select`.
	"""

	assert mmcv.is_list_of(img_scales, tuple) and len(img_scales) == 2
	img_scale_long = [max(s) for s in img_scales]
	img_scale_short = [min(s) for s in img_scales]
	long_edge = np.random.randint(
	min(img_scale_long),
	max(img_scale_long) + 1)
	short_edge = np.random.randint(
	min(img_scale_short),
	max(img_scale_short) + 1)
	img_scale = (long_edge, short_edge)
	return img_scale, None

	@staticmethod
	def random_sample_ratio(img_scale, ratio_range):
	"""Randomly sample an img_scale when ``ratio_range`` is specified.

	A ratio will be randomly sampled from the range specified by
	``ratio_range``. Then it would be multiplied with ``img_scale`` to
	generate sampled scale.

	Args:
	img_scale (tuple): Images scale base to multiply with ratio.
	ratio_range (tuple[float]): The minimum and maximum ratio to scale
	the ``img_scale``.

	Returns:
	(tuple, None): Returns a tuple ``(scale, None)``, where
	``scale`` is sampled ratio multiplied with ``img_scale`` and
	None is just a placeholder to be consistent with
	:func:`random_select`.
	"""

	assert isinstance(img_scale, tuple) and len(img_scale) == 2
	min_ratio, max_ratio = ratio_range
	assert min_ratio <= max_ratio
	ratio = np.random.random_sample() * (max_ratio - min_ratio) + min_ratio
	scale = int(img_scale[0] * ratio), int(img_scale[1] * ratio)
	return scale, None

	def _random_scale(self, results):
	"""Randomly sample an img_scale according to ``ratio_range`` and
	``multiscale_mode``.

	If ``ratio_range`` is specified, a ratio will be sampled and be
	multiplied with ``img_scale``.
	If multiple scales are specified by ``img_scale``, a scale will be
	sampled according to ``multiscale_mode``.
	Otherwise, single scale will be used.

	Args:
	results (dict): Result dict from :obj:`dataset`.

	Returns:
	dict: Two new keys 'scale` and 'scale_idx` are added into
	``results``, which would be used by subsequent pipelines.
	"""

	if self.ratio_range is not None:
	scale, scale_idx = self.random_sample_ratio(
	self.img_scale[0], self.ratio_range)
	elif len(self.img_scale) == 1:
	scale, scale_idx = self.img_scale[0], 0
	elif self.multiscale_mode == 'range':
	scale, scale_idx = self.random_sample(self.img_scale)
	elif self.multiscale_mode == 'value':
	scale, scale_idx = self.random_select(self.img_scale)
	else:
	raise NotImplementedError

	results['scale'] = scale
	results['scale_idx'] = scale_idx

	def _resize_img(self, results):
	"""Resize images with ``results['scale']``."""

	if self.keep_ratio:
	if self.setr_multi_scale:
	if min(results['scale']) < self.crop_size[0]:
	new_short = self.crop_size[0]
	else:
	new_short = min(results['scale'])

	h, w = results['img'].shape[:2]
	if h > w:
	new_h, new_w = new_short * h / w, new_short
	else:
	new_h, new_w = new_short, new_short * w / h
	results['scale'] = (new_h, new_w)

	img, scale_factor = mmcv.imrescale(results['img'],
	results['scale'],
	return_scale=True)
	# the w_scale and h_scale has minor difference
	# a real fix should be done in the mmcv.imrescale in the future
	new_h, new_w = img.shape[:2]
	h, w = results['img'].shape[:2]
	w_scale = new_w / w
	h_scale = new_h / h
	else:
	img, w_scale, h_scale = mmcv.imresize(results['img'],
	results['scale'],
	return_scale=True)
	scale_factor = np.array([w_scale, h_scale, w_scale, h_scale],
	dtype=np.float32)
	results['img'] = img
	results['img_shape'] = img.shape
	results['pad_shape'] = img.shape # in case that there is no padding
	results['scale_factor'] = scale_factor
	results['keep_ratio'] = self.keep_ratio

	def _resize_seg(self, results):
	"""Resize semantic segmentation map with ``results['scale']``."""
	for key in results.get('seg_fields', []):
	if self.keep_ratio:
	gt_seg = mmcv.imrescale(results[key],
	results['scale'],
	interpolation='nearest')
	else:
	gt_seg = mmcv.imresize(results[key],
	results['scale'],
	interpolation='nearest')
	results['gt_semantic_seg'] = gt_seg

	def __call__(self, results):
	"""Call function to resize images, bounding boxes, masks, semantic
	segmentation map.

	Args:
	results (dict): Result dict from loading pipeline.

	Returns:
	dict: Resized results, 'img_shape', 'pad_shape', 'scale_factor',
	'keep_ratio' keys are added into result dict.
	"""

	if 'scale' not in results:
	self._random_scale(results)
	self._resize_img(results)
	self._resize_seg(results)
	return results

	def __repr__(self):
	repr_str = self.__class__.__name__
	repr_str += (f'(img_scale={self.img_scale}, '
	f'multiscale_mode={self.multiscale_mode}, '
	f'ratio_range={self.ratio_range}, '
	f'keep_ratio={self.keep_ratio})')
	return repr_str


	@PIPELINES.register_module()
	class PadShortSide(object):
	"""Pad the image & mask.

	Pad to the minimum size that is equal or larger than a number.
	Added keys are "pad_shape", "pad_fixed_size",

	Args:
	size (int, optional): Fixed padding size.
	pad_val (float, optional): Padding value. Default: 0.
	seg_pad_val (float, optional): Padding value of segmentation map.
	Default: 255.
	"""
	def __init__(self, size=None, pad_val=0, seg_pad_val=255):
	self.size = size
	self.pad_val = pad_val
	self.seg_pad_val = seg_pad_val
	# only one of size and size_divisor should be valid
	assert size is not None

	def _pad_img(self, results):
	"""Pad images according to ``self.size``."""
	h, w = results['img'].shape[:2]
	new_h = max(h, self.size)
	new_w = max(w, self.size)
	padded_img = mmcv.impad(results['img'],
	shape=(new_h, new_w),
	pad_val=self.pad_val)

	results['img'] = padded_img
	results['pad_shape'] = padded_img.shape
	# results['unpad_shape'] = (h, w)

	def _pad_seg(self, results):
	"""Pad masks according to ``results['pad_shape']``."""
	for key in results.get('seg_fields', []):
	results[key] = mmcv.impad(results[key],
	shape=results['pad_shape'][:2],
	pad_val=self.seg_pad_val)

	def __call__(self, results):
	"""Call function to pad images, masks, semantic segmentation maps.

	Args:
	results (dict): Result dict from loading pipeline.

	Returns:
	dict: Updated result dict.
	"""
	h, w = results['img'].shape[:2]
	if h >= self.size and w >= self.size: # 短边比窗口大，跳过
	pass
	else:
	self._pad_img(results)
	self._pad_seg(results)
	return results

	def __repr__(self):
	repr_str = self.__class__.__name__
	repr_str += f'(size={self.size}, pad_val={self.pad_val})'
	return repr_str


	@PIPELINES.register_module()
	class MapillaryHack(object):
	"""map MV 65 class to 19 class like Cityscapes."""
	def __init__(self):
	self.map = [[13, 24, 41], [2, 15], [17], [6], [3],
	[45, 47], [48], [50], [30], [29], [27], [19], [20, 21, 22],
	[55], [61], [54], [58], [57], [52]]

	self.others = [i for i in range(66)]
	for i in self.map:
	for j in i:
	if j in self.others:
	self.others.remove(j)

	def __call__(self, results):
	"""Call function to process the image with gamma correction.

	Args:
	results (dict): Result dict from loading pipeline.

	Returns:
	dict: Processed results.
	"""
	gt_map = results['gt_semantic_seg']
	# others -> 255
	new_gt_map = np.zeros_like(gt_map)

	for value in self.others:
	new_gt_map[gt_map == value] = 255

	for index, map in enumerate(self.map):
	for value in map:
	new_gt_map[gt_map == value] = index

	results['gt_semantic_seg'] = new_gt_map

	return results

	def __repr__(self):
	repr_str = self.__class__.__name__
	return repr_str