from torch.utils.data import Dataset, DataLoader

 # 這里添加：判斷image是否為None,如果為None，則在原來的batch中清除掉，這樣就可以在迭代中避免出錯了
 if isinstance(batch, list):
 batch = [(image, image_id) for (image, image_id) in batch if image is not None]
 if batch==[]:
 return (None,None)

# -*-coding: utf-8 -*-
"""
 @Project: pytorch-learning-tutorials
 @File : dataset_collate.py
 @Author : panjq
 @E-mail : pan_jinquan@163.com
 @Date : 2019-06-07 17:09:13
"""
 
r""""Contains definitions of the methods used by the _DataLoaderIter workers to
collate samples fetched from dataset into Tensor(s).
These **needs** to be in global scope since Py2 doesn't support serializing
static methods.
"""
import torch
import re
from torch._six import container_abcs, string_classes, int_classes 
_use_shared_memory = False
r"""Whether to use shared memory in default_collate"""
 
np_str_obj_array_pattern = re.compile(r'[SaUO]')
 
error_msg_fmt = "batch must contain tensors, numbers, dicts or lists; found {}"
 
numpy_type_map = {
 'float64': torch.DoubleTensor,
 'float32': torch.FloatTensor,
 'float16': torch.HalfTensor,
 'int64': torch.LongTensor,
 'int32': torch.IntTensor,
 'int16': torch.ShortTensor,
 'int8': torch.CharTensor,
 'uint8': torch.ByteTensor,
}
 
def collate_fn(batch):
 '''
 collate_fn (callable, optional): merges a list of samples to form a mini-batch.
 該函數參考touch的default_collate函數，也是DataLoader的默認的校對方法，當batch中含有None等數據時，
 默認的default_collate校隊方法會出現錯誤
 一種的解決方法是：
 判斷batch中image是否為None,如果為None，則在原來的batch中清除掉，這樣就可以在迭代中避免出錯了
 :param batch:
 :return:
 '''
 r"""Puts each data field into a tensor with outer dimension batch size"""
 # 這里添加：判斷image是否為None,如果為None，則在原來的batch中清除掉，這樣就可以在迭代中避免出錯了
 if isinstance(batch, list):
 batch = [(image, image_id) for (image, image_id) in batch if image is not None]
 if batch==[]:
 return (None,None)
 
 elem_type = type(batch[0])
 if isinstance(batch[0], torch.Tensor):
 out = None
 if _use_shared_memory:
  # If we're in a background process, concatenate directly into a
  # shared memory tensor to avoid an extra copy
  numel = sum([x.numel() for x in batch])
  storage = batch[0].storage()._new_shared(numel)
  out = batch[0].new(storage)
 return torch.stack(batch, 0, out=out)
 elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \

  and elem_type.__name__ != 'string_':
 elem = batch[0]
 if elem_type.__name__ == 'ndarray':
  # array of string classes and object
  if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
  raise TypeError(error_msg_fmt.format(elem.dtype))
 
  return collate_fn([torch.from_numpy(b) for b in batch])
 if elem.shape == (): # scalars
  py_type = float if elem.dtype.name.startswith('float') else int
  return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))
 elif isinstance(batch[0], float):
 return torch.tensor(batch, dtype=torch.float64)
 elif isinstance(batch[0], int_classes):
 return torch.tensor(batch)
 elif isinstance(batch[0], string_classes):
 return batch
 elif isinstance(batch[0], container_abcs.Mapping):
 return {key: collate_fn([d[key] for d in batch]) for key in batch[0]}
 elif isinstance(batch[0], tuple) and hasattr(batch[0], '_fields'): # namedtuple
 return type(batch[0])(*(collate_fn(samples) for samples in zip(*batch)))
 elif isinstance(batch[0], container_abcs.Sequence):
 transposed = zip(*batch)#ok
 return [collate_fn(samples) for samples in transposed]
 
 raise TypeError((error_msg_fmt.format(type(batch[0]))))

# -*-coding: utf-8 -*-
"""
 @Project: pytorch-learning-tutorials
 @File : dataset.py
 @Author : panjq
 @E-mail : pan_jinquan@163.com
 @Date : 2019-03-07 18:45:06
"""
import torch
from torch.autograd import Variable
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import numpy as np
from utils import dataset_collate
import os
import cv2
from PIL import Image
def read_image(path,mode='RGB'):
 '''
 :param path:
 :param mode: RGB or L
 :return:
 '''
 return Image.open(path).convert(mode)
 
class TorchDataset(Dataset):
 def __init__(self, image_id_list, image_dir, resize_height=256, resize_width=256, repeat=1, transform=None):
 '''
 :param filename: 數據文件TXT：格式：imge_name.jpg label1_id labe2_id
 :param image_dir: 圖片路徑：image_dir+imge_name.jpg構成圖片的完整路徑
 :param resize_height 為None時，不進行縮放
 :param resize_width 為None時，不進行縮放，
    PS：當參數resize_height或resize_width其中一個為None時，可實現等比例縮放
 :param repeat: 所有樣本數據重復次數，默認循環一次，當repeat為None時，表示無限循環sys.maxsize
 :param transform:預處理
 '''
 self.image_dir = image_dir
 self.image_id_list=image_id_list
 self.len = len(image_id_list)
 self.repeat = repeat
 self.resize_height = resize_height
 self.resize_width = resize_width
 self.transform= transform
 
 def __getitem__(self, i):
 index = i % self.len
 # print("i={},index={}".format(i, index))
 image_id = self.image_id_list[index]
 image_path = os.path.join(self.image_dir, image_id)
 img = self.load_data(image_path)
 
 if img is None:
  return None,image_id
 img = self.data_preproccess(img)
 return img,image_id
 
 def __len__(self):
 if self.repeat == None:
  data_len = 10000000
 else:
  data_len = len(self.image_id_list) * self.repeat
 return data_len
 
 def load_data(self, path):
 '''
 加載數據
 :param path:
 :param resize_height:
 :param resize_width:
 :param normalization: 是否歸一化
 :return:
 '''
 try:
  image = read_image(path)
 except Exception as e:
  image=None
  print(e)
 # image = image_processing.read_image(path)#用opencv讀取圖像
 return image
 
 def data_preproccess(self, data):
 '''
 數據預處理
 :param data:
 :return:
 '''
 if self.transform is not None:
  data = self.transform(data)
 return data
 
if __name__=='__main__':
 
 resize_height = 224
 resize_width = 224
 image_id_list=["1.jpg","ddd.jpg","111.jpg","3.jpg","4.jpg","5.jpg","6.jpg","7.jpg","8.jpg","9.jpg"]
 image_dir="../dataset/test_images/images"
 # 相關預處理的初始化
 '''class torchvision.transforms.ToTensor把shape=(H,W,C)的像素值范圍為[0, 255]的PIL.Image或者numpy.ndarray數據
 # 轉換成shape=(C,H,W)的像素數據，并且被歸一化到[0.0, 1.0]的torch.FloatTensor類型。
 '''
 train_transform = transforms.Compose([
 transforms.Resize(size=(resize_height, resize_width)),
 # transforms.RandomHorizontalFlip(),#隨機翻轉圖像
 transforms.RandomCrop(size=(resize_height, resize_width), padding=4), # 隨機裁剪
 transforms.ToTensor(), # 吧shape=(H,W,C)->換成shape=(C,H,W),并且歸一化到[0.0, 1.0]的torch.FloatTensor類型
 # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))#給定均值(R,G,B) 方差（R，G，B），將會把Tensor正則化
 ])
 
 epoch_num=2 #總樣本循環次數
 batch_size=5 #訓練時的一組數據的大小
 train_data_nums=10
 max_iterate=int((train_data_nums+batch_size-1)/batch_size*epoch_num) #總迭代次數
 
 train_data = TorchDataset(image_id_list=image_id_list,
    image_dir=image_dir,
    resize_height=resize_height,
    resize_width=resize_width,
    repeat=1,
    transform=train_transform)
 # 使用默認的default_collate會報錯
 # train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=False)
 # 使用自定義的collate_fn
 train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=False, collate_fn=dataset_collate.collate_fn)
 
 
 # [1]使用epoch方法迭代，TorchDataset的參數repeat=1
 for epoch in range(epoch_num):
 for step,(batch_image, batch_label) in enumerate(train_loader):
  if batch_image is None and batch_label is None:
  print("batch_image:{},batch_label:{}".format(batch_image, batch_label))
  continue
  image=batch_image[0,:]
  image=image.numpy()#image=np.array(image)
  image = image.transpose(1, 2, 0) # 通道由[c,h,w]->[h,w,c]
  cv2.imshow("image",image)
  cv2.waitKey(2000)
  print("batch_image.shape:{},batch_label:{}".format(batch_image.shape,batch_label))
  # batch_x, batch_y = Variable(batch_x), Variable(batch_y)