import zipfile # 解壓
import pandas as pd # 操作數(shù)據(jù)
import os # 操作文件或文件夾
import cv2 # 圖像操作庫
import matplotlib.pyplot as plt # 圖像展示庫
from torch.utils.data import Dataset # PyTorch內(nèi)置對象
from torchvision import transforms # 圖像增廣轉(zhuǎn)換庫 PyTorch內(nèi)置
import torch 

# 解壓文件到指定目錄
def unzip_file(root_path, filename):
  full_path = os.path.join(root_path, filename)
  file = zipfile.ZipFile(full_path)
  file.extractall(root_path)
unzip_file(root_path, zip_filename)

# 讀入csv文件
face_landmarks = pd.read_csv(os.path.join(extract_path, csv_filename))

# pandas讀出的數(shù)據(jù)如想要操作索引 使用iloc
image_name = face_landmarks.iloc[:,0]
landmarks = face_landmarks.iloc[:,1:]

# 展示
def show_face(extract_path, image_file, face_landmark):
  plt.imshow(plt.imread(os.path.join(extract_path, image_file)), cmap='gray')
  point_x = face_landmark.to_numpy()[0::2]
  point_y = face_landmark.to_numpy()[1::2]
  plt.scatter(point_x, point_y, c='r', s=6)
  
show_face(extract_path, image_name.iloc[1], landmarks.iloc[1])

class FaceDataset(Dataset):
  def __init__(self, extract_path, csv_filename, transform=None):
    super(FaceDataset, self).__init__()
    self.extract_path = extract_path
    self.csv_filename = csv_filename
    self.transform = transform
    self.face_landmarks = pd.read_csv(os.path.join(extract_path, csv_filename))
  def __len__(self):
    return len(self.face_landmarks)
  def __getitem__(self, idx):
    image_name = self.face_landmarks.iloc[idx,0]
    landmarks = self.face_landmarks.iloc[idx,1:].astype('float32')
    point_x = landmarks.to_numpy()[0::2]
    point_y = landmarks.to_numpy()[1::2]
    image = plt.imread(os.path.join(self.extract_path, image_name))
    sample = {'image':image, 'point_x':point_x, 'point_y':point_y}
    if self.transform is not None:
      sample = self.transform(sample)
    return sample

face_dataset = FaceDataset(extract_path, csv_filename)
sample = face_dataset[0]
plt.imshow(sample['image'], cmap='gray')
plt.scatter(sample['point_x'], sample['point_y'], c='r', s=2)
plt.title('face')

class Rescale(object):
  def __init__(self, out_size):
    assert isinstance(out_size,tuple) or isinstance(out_size,int), 'out size isinstance int or tuple'
    self.out_size = out_size
  def __call__(self, sample):
    image, point_x, point_y = sample['image'], sample['point_x'], sample['point_y']
    new_h, new_w = self.out_size if isinstance(self.out_size,tuple) else (self.out_size, self.out_size)
    new_image = cv2.resize(image,(new_w, new_h))
    h, w = image.shape[0:2]
    new_y = new_h / h * point_y
    new_x = new_w / w * point_x
    return {'image':new_image, 'point_x':new_x, 'point_y':new_y}

class ToTensor(object):
  def __call__(self, sample):
    image, point_x, point_y = sample['image'], sample['point_x'], sample['point_y']
    new_image = image.transpose((2,0,1))
    return {'image':torch.from_numpy(new_image), 'point_x':torch.from_numpy(point_x), 'point_y':torch.from_numpy(point_y)}

transform = transforms.Compose([Rescale((1024, 512)), ToTensor()])
face_dataset = FaceDataset(extract_path, csv_filename, transform=transform)
sample = face_dataset[0]
plt.imshow(sample['image'].permute((1,2,0)), cmap='gray')
plt.scatter(sample['point_x'], sample['point_y'], c='r', s=2)
plt.title('face')

data_loader = DataLoader(face_dataset, batch_size=4, shuffle=True, num_workers=0)
for i in data_loader:
  print(i['image'].shape)
  break

torch.Size([4, 3, 1024, 512])