|
- # coding=utf-8
- # Copyright 2021 The Google Research Authors.
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
-
- """Library for loading train and eval data.
- This libary contains two functions, make_train_iterator for generating a
- training data iterator from multiple sources of different formats and
- make_eval_function for creating an evaluation function that evaluates on
- data from multiple sources of different formats.
- """
-
- # pylint:disable=g-importing-member
- from functools import partial
- import glob
- import mindspore.ops as ops
- import numpy as np
- import random
- import cv2
- import mindspore.dataset.vision.py_transforms as py_vision
- import mindspore.dataset.vision.c_transforms as c_vision
- import mindspore.dataset as ds
- from mindspore.dataset.transforms.py_transforms import Compose
- # pylint:disable=g-long-lambda
-
-
- def create_dataset_train(
- train_on,
- height,
- width,
- shuffle_buffer_size,
- batch_size,
- seq_len,
- rank=0,
- group_size=1,
- ):
- """Build joint training iterator for all data in train_on.
- Args:
- train_on: string of the format 'format0:path0;format1:path1', e.g.
- 'kitti:/usr/local/home/...'.
- height: int, height to which the images will be resized or cropped.
- width: int, width to which the images will be resized or cropped.
- shuffle_buffer_size: int, size that will be used for the shuffle buffer.
- batch_size: int, batch size for the iterator.
- seq_len: int, number of frames per sequences (at the moment this should
- always be 2)
- seed: A seed for a random number generator, controls shuffling of data.
- mode: str, will be passed on to the data iterator class. Can be used to
- specify different settings within the data iterator.
- Returns:
- A tf.data.Iterator that produces batches of images of shape [batch
- size, sequence length=3, height, width, channels=3]
- """
- data_format, path = train_on.split(':')
-
- file = glob.glob(path + '/*mindrecord')
- if group_size == 1:
- mindrecord_dataset = ds.MindDataset(file,
- #num_parallel_workers=8,
- shuffle=True
- )
- else:
- mindrecord_dataset = ds.MindDataset(file,
- num_parallel_workers=8,
- num_shards=group_size,
- shard_id=rank,
- shuffle=True
- )
-
-
- def roll_op(data):
- r = random.randint(0,2)
- return np.roll(data, r, axis=-1)
-
- def reverse_color_op(data):
- r = random.randint(0,1)
- if(r == 1):
- return np.flip(data, -1)
- else:
- return data
-
- def reverse_height_op(data):
- r = random.randint(0,1)
- if(r == 1):
- return np.flip(data, -3)
- else:
- return data
-
- def mul_255(data):
- return data*255
-
- def stack_op(data1, data2):
- return np.stack((data1, data2), axis=0), np.stack((data1, data2), axis=0)
-
- if 'sintel' in data_format:
- transforms_list1 = Compose([py_vision.Decode(),
- py_vision.Resize([height,width]),
- reverse_height_op,
- py_vision.ToTensor()])
-
- transforms_list2 = [py_vision.Decode(),
- py_vision.Resize([height,width]),
- reverse_height_op,
- roll_op,
- reverse_color_op,
- py_vision.ToPIL(),
- py_vision.RandomColorAdjust(hue=0.5),
- py_vision.ToTensor()]
- elif 'chairs' in data_format:
- transforms_list1 = Compose([py_vision.Decode(),
- py_vision.Resize([height,width]),
- py_vision.ToTensor()])
-
- transforms_list2 = [py_vision.Decode(),
- py_vision.Resize([height,width]),
- roll_op,
- reverse_color_op,
- py_vision.ToPIL(),
- py_vision.RandomColorAdjust(hue=0.5),
- py_vision.ToTensor()]
- else:
- print('Unknown data format "{}"'.format(data_format))
- '''
- transforms_list3 = [py_vision.Decode(),
- py_vision.Resize([height,width]),
- reverse_height_op,
- py_vision.ToPIL(),
- py_vision.Grayscale(),
- py_vision.ToTensor(),
- mul_255]
- '''
- mindrecord_dataset = mindrecord_dataset.shuffle(1024)
- mindrecord_dataset = mindrecord_dataset.repeat()
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list1, input_columns=["images1_data"], output_columns=["image1_without_photo_aug"])
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list1, input_columns=["images2_data"], output_columns=["image2_without_photo_aug"])
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list2, input_columns=["images1_aug"], output_columns=["image1"])
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list2, input_columns=["images2_aug"], output_columns=["image2"])
- #mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list3, input_columns=["images1_gray"])
- #mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list3, input_columns=["images2_gray"])
-
-
- mindrecord_dataset = mindrecord_dataset.map(operations=stack_op,
- input_columns=["image1_without_photo_aug", "image2_without_photo_aug"],
- output_columns=["images_without_photo_aug", "images_without_photo_aug_copy"])
-
- mindrecord_dataset = mindrecord_dataset.map(operations=stack_op,
- input_columns=["image1", "image2"],
- output_columns=["images", "images_copy"])
-
-
- '''
- mindrecord_dataset = mindrecord_dataset.map(operations=stack_op,
- input_columns=["images1_gray", "images2_gray"],
- output_columns=["images_gray", "images_gray_copy"])
- '''
-
- '''
- if 'sintel' in data_format:
- mindrecord_dataset = mindrecord_dataset.repeat(2)
- mindrecord_dataset = mindrecord_dataset.shuffle(1024)
- #mindrecord_dataset = mindrecord_dataset.take(1000)
- '''
- mindrecord_dataset = mindrecord_dataset.batch(1)
-
- return mindrecord_dataset
-
- def create_dataset_eval(
- eval_on,
- height,
- width,
- shuffle_buffer_size,
- batch_size,
- seq_len
- ):
- """Build joint training iterator for all data in train_on.
- Args:
- eval_on: string of the format 'format0:path0;format1:path1', e.g.
- 'kitti:/usr/local/home/...'.
- height: int, height to which the images will be resized or cropped.
- width: int, width to which the images will be resized or cropped.
- shuffle_buffer_size: int, size that will be used for the shuffle buffer.
- batch_size: int, batch size for the iterator.
- seq_len: int, number of frames per sequences (at the moment this should
- always be 2)
- seed: A seed for a random number generator, controls shuffling of data.
- mode: str, will be passed on to the data iterator class. Can be used to
- specify different settings within the data iterator.
- Returns:
- A tf.data.Iterator that produces batches of images of shape [batch
- size, sequence length=3, height, width, channels=3]
- """
- data_format, path = eval_on.split(':')
-
- file = glob.glob(path + '/*mindrecord')
-
- mindrecord_dataset = ds.MindDataset(file)
-
- for data in mindrecord_dataset.create_dict_iterator(output_numpy=True):
- origin_height = data["height"].item()
- origin_width = data["width"].item()
- break
-
-
- def mul_255(data):
- return data*255
-
- def reshape_flow_op(data):
- return np.reshape(data, (origin_height, origin_width, 2))[Ellipsis, ::-1]
-
- def resize_op(data):
- return cv2.resize(data, (width, height))
-
- def expand_op(data):
- return np.expand_dims(data, -1)
-
- def scal_op(data):
- scaling = np.reshape([
- float(height) / float(origin_height),
- float(width) / float(origin_width)
- ], [1, 1, 2])
- return data * scaling
-
- def flow_valid_op(data1, data2):
- data2 = np.ones_like(data1[:1, Ellipsis])
- return data1, data2
-
- def reshape_occlusion_op(data):
- return np.reshape(data, (origin_height, origin_width, 1))
-
- def transpose_op(data):
- transpose_x = np.transpose(data, axes=[2, 0, 1])
- return transpose_x
-
- def stack_op(data1, data2):
- return np.stack((data1, data2), axis=0), np.stack((data1, data2), axis=0)
-
- transforms_list1 = Compose([py_vision.Decode(),
- py_vision.ToTensor()])
-
- transforms_flow = [reshape_flow_op,
- transpose_op]
-
- transforms_occlusion = [reshape_occlusion_op,
- transpose_op]
-
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list1, input_columns=["images1_data"])
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_list1, input_columns=["images2_data"])
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_flow, input_columns=["flow_uv"])
- mindrecord_dataset = mindrecord_dataset.map(operations=flow_valid_op, input_columns=["flow_uv", "flow_path"], output_columns=["flow_uv", "flow_valid"])
- if 'sintel' in data_format:
- mindrecord_dataset = mindrecord_dataset.map(operations=transforms_occlusion, input_columns=["occlusion_mask"])
-
- mindrecord_dataset = mindrecord_dataset.map(operations=stack_op,
- input_columns=["images1_data", "images2_data"],
- output_columns=["images", "images_copy"])
-
- mindrecord_dataset = mindrecord_dataset.batch(1)
- mindrecord_dataset = mindrecord_dataset.shuffle(1024)
- #mindrecord_dataset = mindrecord_dataset.repeat(21)
- #for data in mindrecord_dataset.creat_dict_iterator():
- # print(data["images"])
- return mindrecord_dataset
-
|