import numpy as np
from skimage import io
from skimage import transform
import matplotlib.pyplot as plt
import os
import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms
from torchvision.utils import make_grid

# step1: 定義MyDataset類， 繼承Dataset, 重寫抽象方法：__len()__, __getitem()__
class MyDataset(Dataset):

 def __init__(self, root_dir, names_file, transform=None):
 self.root_dir = root_dir
 self.names_file = names_file
 self.transform = transform
 self.size = 0
 self.names_list = []

 if not os.path.isfile(self.names_file):
  print(self.names_file + 'does not exist!')
 file = open(self.names_file)
 for f in file:
  self.names_list.append(f)
  self.size += 1

 def __len__(self):
 return self.size

 def __getitem__(self, idx):
 image_path = self.root_dir + self.names_list[idx].split(' ')[0]
 if not os.path.isfile(image_path):
  print(image_path + 'does not exist!')
  return None
 image = io.imread(image_path) # use skitimage
 label = int(self.names_list[idx].split(' ')[1])

 sample = {'image': image, 'label': label}
 if self.transform:
  sample = self.transform(sample)

 return sample

train_dataset = MyDataset(root_dir='./data/train',
    names_file='./data/train/train.txt',
    transform=None)

plt.figure()
for (cnt,i) in enumerate(train_dataset):
 image = i['image']
 label = i['label']

 ax = plt.subplot(4, 4, cnt+1)
 ax.axis('off')
 ax.imshow(image)
 ax.set_title('label {}'.format(label))
 plt.pause(0.001)

 if cnt == 15:
 break

# # 變換Resize
class Resize(object):

 def __init__(self, output_size: tuple):
 self.output_size = output_size

 def __call__(self, sample):
 # 圖像
 image = sample['image']
 # 使用skitimage.transform對圖像進行縮放
 image_new = transform.resize(image, self.output_size)
 return {'image': image_new, 'label': sample['label']}

# # 變換ToTensor
class ToTensor(object):

 def __call__(self, sample):
 image = sample['image']
 image_new = np.transpose(image, (2, 0, 1))
 return {'image': torch.from_numpy(image_new),
  'label': sample['label']}

# 對原始的訓練數據集進行變換
transformed_trainset = MyDataset(root_dir='./data/train',
    names_file='./data/train/train.txt',
    transform=transforms.Compose(
    [Resize((224,224)),
    ToTensor()]
    ))

# 使用DataLoader可以利用多線程，batch,shuffle等
trainset_dataloader = DataLoader(dataset=transformed_trainset,
     batch_size=4,
     shuffle=True,
     num_workers=4)

def show_images_batch(sample_batched):
 images_batch, labels_batch = \
 sample_batched['image'], sample_batched['label']
 grid = make_grid(images_batch)
 plt.imshow(grid.numpy().transpose(1, 2, 0))


# sample_batch: Tensor , NxCxHxW
plt.figure()
for i_batch, sample_batch in enumerate(trainset_dataloader):
 show_images_batch(sample_batch)
 plt.axis('off')
 plt.ioff()
 plt.show()


plt.show()

import torch
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import matplotlib.pyplot as plt
import numpy as np


# https://pytorch.org/tutorials/beginner/data_loading_tutorial.html

# data_transform = transforms.Compose([
#  transforms.RandomResizedCrop(224),
#  transforms.RandomHorizontalFlip(),
#  transforms.ToTensor(),
#  transforms.Normalize(mean=[0.485, 0.456, 0.406],
#       std=[0.229, 0.224, 0.225])
# ])

data_transform = transforms.Compose([
 transforms.Resize((224,224)),
 transforms.RandomHorizontalFlip(),
 transforms.ToTensor(),

])

train_dataset = datasets.ImageFolder(root='./data/train',transform=data_transform)
train_dataloader = DataLoader(dataset=train_dataset,
        batch_size=4,
        shuffle=True,
        num_workers=4)


def show_batch_images(sample_batch):
 labels_batch = sample_batch[1]
 images_batch = sample_batch[0]

 for i in range(4):
  label_ = labels_batch[i].item()
  image_ = np.transpose(images_batch[i], (1, 2, 0))
  ax = plt.subplot(1, 4, i + 1)
  ax.imshow(image_)
  ax.set_title(str(label_))
  ax.axis('off')
  plt.pause(0.01)


plt.figure()
for i_batch, sample_batch in enumerate(train_dataloader):
 show_batch_images(sample_batch)

 plt.show()