From f111cb002b9c6065fdf6bb274ce5857a9e875e8c Mon Sep 17 00:00:00 2001
From: chenshijun <csj_sky@126.com>
Date: 星期三, 05 六月 2019 15:38:49 +0800
Subject: [PATCH] face rectangle
---
code/datasets.py | 75 ++++++++++++++++++++++++++++++++-----
1 files changed, 64 insertions(+), 11 deletions(-)
diff --git a/code/datasets.py b/code/datasets.py
index a28c584..e8ab9f4 100644
--- a/code/datasets.py
+++ b/code/datasets.py
@@ -1,6 +1,7 @@
import os
import numpy as np
import cv2
+import pandas as pd
import torch
from torch.utils.data.dataset import Dataset
@@ -16,6 +17,65 @@
with open(file_path) as f:
lines = f.read().splitlines()
return lines
+
+class Synhead(Dataset):
+ def __init__(self, data_dir, csv_path, transform, test=False):
+ column_names = ['path', 'bbox_x_min', 'bbox_y_min', 'bbox_x_max', 'bbox_y_max', 'yaw', 'pitch', 'roll']
+ tmp_df = pd.read_csv(csv_path, sep=',', names=column_names, index_col=False, encoding="utf-8-sig")
+ self.data_dir = data_dir
+ self.transform = transform
+ self.X_train = tmp_df['path']
+ self.y_train = tmp_df[['bbox_x_min', 'bbox_y_min', 'bbox_x_max', 'bbox_y_max', 'yaw', 'pitch', 'roll']]
+ self.length = len(tmp_df)
+ self.test = test
+
+ def __getitem__(self, index):
+ path = os.path.join(self.data_dir, self.X_train.iloc[index]).strip('.jpg') + '.png'
+ img = Image.open(path)
+ img = img.convert('RGB')
+
+ x_min, y_min, x_max, y_max, yaw, pitch, roll = self.y_train.iloc[index]
+ x_min = float(x_min); x_max = float(x_max)
+ y_min = float(y_min); y_max = float(y_max)
+ yaw = -float(yaw); pitch = float(pitch); roll = float(roll)
+
+ # k = 0.2 to 0.40
+ k = np.random.random_sample() * 0.2 + 0.2
+ x_min -= 0.6 * k * abs(x_max - x_min)
+ y_min -= 2 * k * abs(y_max - y_min)
+ x_max += 0.6 * k * abs(x_max - x_min)
+ y_max += 0.6 * k * abs(y_max - y_min)
+
+ width, height = img.size
+ # Crop the face
+ img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max)))
+
+ # Flip?
+ rnd = np.random.random_sample()
+ if rnd < 0.5:
+ yaw = -yaw
+ roll = -roll
+ img = img.transpose(Image.FLIP_LEFT_RIGHT)
+
+ # Blur?
+ rnd = np.random.random_sample()
+ if rnd < 0.05:
+ img = img.filter(ImageFilter.BLUR)
+
+ # Bin values
+ bins = np.array(range(-99, 102, 3))
+ binned_pose = np.digitize([yaw, pitch, roll], bins) - 1
+
+ labels = torch.LongTensor(binned_pose)
+ cont_labels = torch.FloatTensor([yaw, pitch, roll])
+
+ if self.transform is not None:
+ img = self.transform(img)
+
+ return img, labels, cont_labels, self.X_train[index]
+
+ def __len__(self):
+ return self.length
class Pose_300W_LP(Dataset):
# Head pose from 300W-LP dataset
@@ -36,7 +96,6 @@
img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext))
img = img.convert(self.image_mode)
mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext)
- shape_path = os.path.join(self.data_dir, self.y_train[index] + '_shape.npy')
# Crop the face loosely
pt2d = utils.get_pt2d_from_mat(mat_path)
@@ -76,11 +135,8 @@
bins = np.array(range(-99, 102, 3))
binned_pose = np.digitize([yaw, pitch, roll], bins) - 1
- # Get shape
- shape = np.load(shape_path)
-
# Get target tensors
- labels = torch.LongTensor(np.concatenate((binned_pose, shape), axis = 0))
+ labels = binned_pose
cont_labels = torch.FloatTensor([yaw, pitch, roll])
if self.transform is not None:
@@ -111,7 +167,6 @@
img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext))
img = img.convert(self.image_mode)
mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext)
- shape_path = os.path.join(self.data_dir, self.y_train[index] + '_shape.npy')
# Crop the face loosely
pt2d = utils.get_pt2d_from_mat(mat_path)
@@ -129,7 +184,8 @@
img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max)))
# We get the pose in radians
- pose = utils.get_ypr_fro # Head pose from AFLW2000 datasetp.pi
+ pose = utils.get_ypr_from_mat(mat_path)
+ pitch = pose[0] * 180 / np.pi
yaw = pose[1] * 180 / np.pi
roll = pose[2] * 180 / np.pi
@@ -154,11 +210,8 @@
bins = np.array(range(-99, 102, 3))
binned_pose = np.digitize([yaw, pitch, roll], bins) - 1
- # Get shape
- shape = np.load(shape_path)
-
# Get target tensors
- labels = torch.LongTensor(np.concatenate((binned_pose, shape), axis = 0))
+ labels = binned_pose
cont_labels = torch.FloatTensor([yaw, pitch, roll])
if self.transform is not None:
--
Gitblit v1.8.0