data.py

import cv2
import numpy as np

import mindspore.dataset as de

cv2.setNumThreads(0)


class SegDataset:
    def __init__(
        self,
        image_mean,
        image_std,
        data_file="",
        batch_size=32,
        crop_size=512,
        max_scale=2.0,
        min_scale=0.5,
        ignore_label=255,
        num_classes=21,
        num_readers=2,
        num_parallel_calls=4,
        shard_id=None,
        shard_num=None,
        shuffle=True,
    ):
        self.data_file = data_file
        self.batch_size = batch_size
        self.crop_size = crop_size
        self.image_mean = np.array(image_mean, dtype=np.float32)
        self.image_std = np.array(image_std, dtype=np.float32)
        self.max_scale = max_scale
        self.min_scale = min_scale
        self.ignore_label = ignore_label
        self.num_classes = num_classes
        self.num_readers = num_readers
        self.num_parallel_calls = num_parallel_calls
        self.shard_id = shard_id
        self.shard_num = shard_num
        self.shuffle = shuffle
        assert max_scale > min_scale

    def train_preprocess_(self, image, label):
        # bgr image
        image_out = cv2.imdecode(np.frombuffer(image, dtype=np.uint8), cv2.IMREAD_COLOR)
        label_out = cv2.imdecode(np.frombuffer(label, dtype=np.uint8), cv2.IMREAD_GRAYSCALE)

        # random scaling
        sc = np.random.uniform(self.min_scale, self.max_scale)
        new_h, new_w = int(sc * image_out.shape[0]), int(sc * image_out.shape[1])
        image_out = cv2.resize(image_out, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
        label_out = cv2.resize(label_out, (new_w, new_h), interpolation=cv2.INTER_NEAREST)

        # mean std
        image_out = (image_out - self.image_mean) / self.image_std

        # pad or random crop
        h_, w_ = max(new_h, self.crop_size), max(new_w, self.crop_size)
        pad_h, pad_w = h_ - new_h, w_ - new_w
        if pad_h > 0 or pad_w > 0:
            image_out = cv2.copyMakeBorder(image_out, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0)
            label_out = cv2.copyMakeBorder(label_out, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=self.ignore_label)
        offset_h = np.random.randint(0, h_ - self.crop_size + 1)
        offset_w = np.random.randint(0, w_ - self.crop_size + 1)
        image_out = image_out[offset_h : offset_h + self.crop_size, offset_w : offset_w + self.crop_size, :]
        label_out = label_out[offset_h : offset_h + self.crop_size, offset_w : offset_w + self.crop_size]

        # random flip
        if np.random.uniform(0.0, 1.0) > 0.5:
            image_out = image_out[:, ::-1, :]
            label_out = label_out[:, ::-1]

        image_out = image_out.transpose((2, 0, 1))
        image_out = image_out.copy()
        label_out = label_out.copy()

        return image_out, label_out

    def get_dataset(self):
        data_set = de.MindDataset(
            dataset_files=self.data_file,
            columns_list=["data", "label"],
            shuffle=True,
            num_parallel_workers=self.num_readers,
            num_shards=self.shard_num,
            shard_id=self.shard_id,
        )

        input_columns = ["data", "label"]
        output_columns = ["data", "label"]
        transforms_list = self.train_preprocess_

        data_set = data_set.map(
            operations=transforms_list,
            input_columns=input_columns,
            output_columns=output_columns,
            num_parallel_workers=self.num_parallel_calls,
        )

        if self.shuffle:
            data_set = data_set.shuffle(buffer_size=self.batch_size * 10)

        data_set = data_set.batch(self.batch_size, drop_remainder=True)

        return data_set


def create_segment_dataset(
    name,
    data_dir,
    is_training=False,
    args=None,
    shard_id=None,
    shard_num=None,
):
    if name == "voc" or name == "vocaug":
        if is_training:
            dataset = SegDataset(
                image_mean=args.image_mean,
                image_std=args.image_std,
                data_file=data_dir,
                batch_size=args.batch_size,
                crop_size=args.crop_size,
                max_scale=args.max_scale,
                min_scale=args.min_scale,
                ignore_label=args.ignore_label,
                num_classes=args.num_classes,
                num_readers=args.num_parallel_workers,
                num_parallel_calls=args.num_parallel_workers,
                shard_id=shard_id,
                shard_num=shard_num,
                shuffle=args.shuffle,
            )
            return dataset.get_dataset()

        else:
            with open(data_dir) as f:
                eval_data_list = f.readlines()
            return eval_data_list

    else:
        raise NotImplementedError