|
- # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
- #
- # 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.
-
- import os
- import glob
- import random
-
- import paddle
- import numpy as np
- from PIL import Image
-
- from paddleseg.cvlibs import manager
- from paddleseg.transforms import Compose
-
- # Random seed is set to ensure that after shuffling dataset per epoch during multi-gpu training, the data sequences of all gpus are consistent.
- random.seed(100)
-
-
- @manager.DATASETS.add_component
- class AutoNueCrop(paddle.io.Dataset):
- """
- You need to to download and convert the [India Driving Dataset](https://idd.insaan.iiit.ac.in/evaluation/autonue21/#bm5)
- following the instructions of Segmentation Track. IDD_Dectection dataset also need for pseudo-labeling.
-
- And then, you need to organize data following the below structure.
-
- IDD_Segmentation
- |
- |--leftImg8bit
- | |--train
- | |--val
- | |--test
- |
- |--gtFine
- | |--train
- | |--val
- | |--test
-
- Args:
- transforms (list): Transforms for image.
- dataset_root (str): Cityscapes dataset directory.
- mode (str, optional): Which part of dataset to use. it is one of ('train', 'val', 'test'). Default: 'train'.
- coarse_multiple (float|int, optional): Multiple of the amount of coarse data relative to fine data. Default: 1
- add_val (bool, optional): Whether to add val set in training. Default: False
- """
-
- def __init__(self,
- transforms,
- dataset_root,
- mode='train',
- coarse_multiple=1,
- add_val=False):
- self.dataset_root = dataset_root
- self.transforms = Compose(transforms)
- self.file_list = list()
- mode = mode.lower()
- self.mode = mode
- self.num_classes = 26
- self.ignore_index = 255
- self.coarse_multiple = coarse_multiple
-
- if mode not in ['train', 'val', 'test']:
- raise ValueError(
- "mode should be 'train', 'val' or 'test', but got {}.".format(
- mode))
-
- if self.transforms is None:
- raise ValueError("`transforms` is necessary, but it is None.")
-
- img_dir = os.path.join(self.dataset_root, 'leftImg8bit')
- label_dir = os.path.join(self.dataset_root, 'gtFine')
- if self.dataset_root is None or not os.path.isdir(
- self.dataset_root) or not os.path.isdir(
- img_dir) or not os.path.isdir(label_dir):
- raise ValueError(
- "The dataset is not Found or the folder structure is nonconfoumance."
- )
-
- label_files = sorted(
- glob.glob(
- os.path.join(label_dir, mode, '*',
- '*_gtFine_labellevel3Ids.png')))
- img_files = sorted(
- glob.glob(os.path.join(img_dir, mode, '*', '*_leftImg8bit.*')))
-
- self.file_list = [
- [img_path, label_path]
- for img_path, label_path in zip(img_files, label_files)
- ]
- # for ii in range(len(self.file_list)):
- # print(self.file_list[ii])
- # print(len(self.file_list))
- self.num_files = len(self.file_list)
- self.total_num_files = self.num_files
-
- if mode == 'train':
- # whether to add val set in training
- if add_val:
- label_files = sorted(
- glob.glob(
- os.path.join(label_dir, 'val', '*',
- '*_gtFine_labellevel3Ids.png')))
- img_files = sorted(
- glob.glob(
- os.path.join(img_dir, 'val', '*', '*_leftImg8bit.*')))
- val_file_list = [
- [img_path, label_path]
- for img_path, label_path in zip(img_files, label_files)
- ]
- self.file_list.extend(val_file_list)
- for ii in range(len(self.file_list)):
- print(self.file_list[ii])
- print(len(self.file_list))
- self.num_files = len(self.file_list)
- self.total_num_files = self.num_files
-
- def __getitem__(self, idx):
- if self.mode == 'test':
- image_path, label_path = self.file_list[idx]
- im, _ = self.transforms(im=image_path)
- (h, w) = im.shape[1:]
- im1 = im[:, 540:540 + 720, 320:320 + 1280]
- im2 = im[:, 540:540 + 720, 960:960 + 1280]
- im3 = im[:, 540:540 + 720, 1600:1600 + 1280]
- return im1, im2, im3, image_path
- elif self.mode == 'val':
- image_path, label_path = self.file_list[idx]
- im, label = self.transforms(im=image_path, label=label_path)
- (h, w) = im.shape[1:]
- im1 = im[:, 540:540 + 720, 320:320 + 1280]
- im2 = im[:, 540:540 + 720, 960:960 + 1280]
- im3 = im[:, 540:540 + 720, 1600:1600 + 1280]
- label = label[540:540 + 720, 320:1600 + 1280]
- return im1, im2, im3, label
- else:
- if idx >= self.num_files:
- image_path, label_path = self.coarse_file_list[idx -
- self.num_files]
- else:
- image_path, label_path = self.file_list[idx]
-
- im, label = self.transforms(im=image_path, label=label_path)
- (h, w) = im.shape[1:]
-
- start_w = np.linspace(320, 1600, 5).tolist()
- np.random.shuffle(start_w)
- start = int(start_w[0])
- crop_im = im[:, 540:540 + 720, start:(start + 1280)]
- crop_label = label[540:540 + 720, start:(start + 1280)]
- return crop_im, crop_label
-
- def __len__(self):
- return self.total_num_files
|