text: # -*- encoding: utf-8 -*-
  #File    :   cityscapes_dataset.py
  #Time    :   2021/06/01 17:37:10
  #Author  :   Chen 
  #Version :   1.0
  
  import os
  import os.path as osp
  import numpy as np
  import random
  import matplotlib.pyplot as plt
  import collections
  import torch
  import torchvision
  from torch.utils import data
  from PIL import Image
  
  class cityscapesDataSet(data.Dataset):
      def __init__(self, root, list_path, max_iters=None, crop_size=(321, 321), mean=(128, 128, 128), scale=True, mirror=True, ignore_label=255, set='val'):
          self.root = root
          self.list_path = list_path
          self.crop_size = crop_size
          self.scale = scale
          self.ignore_label = ignore_label
          self.mean = mean
          self.is_mirror = mirror
          # self.mean_bgr = np.array([104.00698793, 116.66876762, 122.67891434])
          self.img_ids = [i_id.strip() for i_id in open(list_path)]
          if not max_iters==None:
              self.img_ids = self.img_ids * int(np.ceil(float(max_iters) / len(self.img_ids)))
          self.files = []
          self.set = set
          # for split in ["train", "trainval", "val"]:
          for name in self.img_ids:
              img_file = osp.join(self.root, "%s/%s" % (self.set, name))
              self.files.append({
                  "img": img_file,
                  "name": name
              })
  
      def __len__(self):
          return len(self.files)
  
      def __getitem__(self, index):
          datafiles = self.files[index]
  
          image = Image.open(datafiles["img"]).convert('RGB')
          name = datafiles["name"]
  
          # resize
          image = image.resize(self.crop_size, Image.BICUBIC)
  
          image = np.asarray(image, np.float32)
  
          size = image.shape
          image = image[:, :, ::-1]  # change to BGR
          image -= self.mean
          image = image.transpose((2, 0, 1))
  
          return image.copy(), np.array(size), name
  
  class cityscapesPseudo(data.Dataset):
      def __init__(self, root, list_path, max_iters=None, crop_size=(321, 321), mean=(128, 128, 128), scale=True, mirror=True, ignore_label=255):
          self.root = root
          self.list_path = list_path
          self.crop_size = crop_size
          self.scale = scale
          self.ignore_label = ignore_label
          self.mean = mean
          self.is_mirror = mirror
          # self.mean_bgr = np.array([104.00698793, 116.66876762, 122.67891434])
          self.img_ids = [i_id.strip().split() for i_id in open(list_path)]
          if not max_iters==None:
              self.img_ids = self.img_ids * int(np.ceil(float(max_iters) / len(self.img_ids)))
          self.files = []
  
          for item in self.img_ids:
              image_path, label_path = item
              name = osp.splitext(osp.basename(label_path))[0]
              img_file = osp.join(self.root, image_path)
              label_file = osp.join(self.root, label_path)
              self.files.append({
                  "img": img_file,
                  "label": label_file,
                  "name": name
              })
          #print(self.files)
  
      def __len__(self):
          return len(self.files)
  
      def __getitem__(self, index):
          datafiles = self.files[index]
  
          image = Image.open(datafiles["img"]).convert('RGB')
          label = Image.open(datafiles["label"])
          name = datafiles["name"]
  
          # resize
          image = image.resize(self.crop_size, Image.BICUBIC)
          label = label.resize(self.crop_size, Image.NEAREST)
  
          #image, label = randomCrop(image, label, self.crop_size[0], self.crop_size[1])
  
          image = np.asarray(image, np.float32)
          label = np.asarray(label, np.float32)
  
          if self.is_mirror:
              flip = np.random.choice(2) * 2 - 1
              image = image[:, :, ::flip]
              label = label[:, ::flip]
  
          size = image.shape
          image = image[:, :, ::-1]  # change to BGR
          image -= self.mean
          image = image.transpose((2, 0, 1))
  
          return image.copy(), label.copy(), np.array(size), name
  
  
  
  
  if __name__ == '__main__':
      dst = GTA5DataSet("./data", is_transform=True)
      trainloader = data.DataLoader(dst, batch_size=4)
      for i, data in enumerate(trainloader):
          imgs, labels = data
          if i == 0:
              img = torchvision.utils.make_grid(imgs).numpy()
              img = np.transpose(img, (1, 2, 0))
              img = img[:, :, ::-1]
              plt.imshow(img)
              plt.show()
  