def __init__(self, base_dir=None, split='train', affine_augmenter=None, image_augmenter=None,
target_size=224, filename=None, use_bined=False, n_class=4, debug=False):
self.base_dir = base_dir
self.base_dir = Path(base_dir)
self.split = split
self.use_bined = use_bined
self.n_class = n_class
self.debug = debug
self.img_paths = []
self.bbox = []
self.labels = []
self.euler_binned = []
with open(self.base_dir / filename) as f:
for i, line in enumerate(f.readlines()):
ls = line.strip()
mat_path = self.base_dir / ls.replace('.jpg', '.mat')
bbox, pose = get_pt_ypr_from_mat(mat_path, pt3d=True)
if True and (abs(pose[0])>99 or abs(pose[1])>99 or abs(pose[2])>99):
continue
if use_bined:
yaw_pitch_bins = np.array([-60, -40, -20, 20, 40, 60])
roll_bins = np.array(range(-81, 82, 9))
self.euler_binned.append([np.digitize(pose[0], yaw_pitch_bins),np.digitize(pose[1], yaw_pitch_bins),np.digitize(pose[2], roll_bins)])
self.labels.append(np.array(pose))
self.bbox.append(bbox)
self.img_paths.append(ls)
self.labels_sort_idx = np.argsort(-np.mean(np.abs(self.labels), axis=1))
if 'train' in self.split:
self.resizer = albu.Compose([albu.SmallestMaxSize(target_size, p=1.),
albu.RandomScale(scale_limit=(-0.2, 0.2), p=0.1),
albu.PadIfNeeded(min_height=target_size, min_width=target_size, value=0, p=1),
albu.RandomCrop(target_size, target_size, p=1.)])
else:
# self.resizer = albu.Compose([albu.Resize(target_size[0], target_size[1], p=1.)])
self.resizer = albu.Compose([albu.SmallestMaxSize(target_size, p=1.),
albu.CenterCrop(target_size, target_size, p=1.)])
self.affine_augmenter = affine_augmenter
self.image_augmenter = image_augmenter
def __init__(self,
base_dir=None,
filename=None,
n_class=3,
target_size=224,
affine_augmenter=None,
image_augmenter=None,
debug=False):
print("[INFO] Initing ALFW2000Dataset.")
self.base_dir = Path(base_dir)
self.n_class = n_class
self.target_size = target_size
self.affine_augmenter = affine_augmenter
self.image_augmenter = image_augmenter
self.debug = debug
self.img_paths = []
self.bboxs = []
self.labels = []
# Read img, bbox, pose
with open(self.base_dir / filename) as f:
for i, line in enumerate(f.readlines()):
ls = line.strip()
mat_path = self.base_dir / ls.replace('.jpg', '.mat')
bbox, pose = get_pt_ypr_from_mat(mat_path, pt3d=True)
if True and (abs(pose[0]) > 99 or abs(pose[1]) > 99
or abs(pose[2]) > 99):
continue
self.labels.append(np.array(pose))
self.bboxs.append(bbox)
self.img_paths.append(ls)
# labels_sort_idx is created by the magnitude of 3 args (yall, pitch, row)
self.labels_sort_idx = np.argsort(
-np.mean(np.abs(self.labels), axis=1))
self.resizer = albu.Compose([
albu.SmallestMaxSize(target_size, p=1.),
albu.CenterCrop(target_size, target_size, p=1.)
])
def __init__(self,
base_dir=None,
split='train',
affine_augmenter=None,
image_augmenter=None,
target_size=224,
filename=None,
use_bined=False,
n_class=3,
debug=False):
self.base_dir = base_dir
self.base_dir = Path(base_dir)
self.split = split
self.use_bined = use_bined
self.debug = debug
self.n_class = n_class
self.ids = []
self.bboxs = []
self.labels = []
self.ids_index = []
self.euler_binned = []
with open(self.base_dir / filename) as f:
for i, line in enumerate(f.readlines()):
ls = line.strip()
id_index = []
bboxs = []
labels = []
euler_binned = []
for j in range(100):
img_path = self.base_dir / (ls + '_%d.jpg' % j)
if not os.path.exists(img_path):
break
mat_path = str(img_path).replace('.jpg', '.mat')
bbox, pose = get_pt_ypr_from_mat(mat_path)
if False and (abs(pose[1]) > 99 or abs(pose[0]) > 99
or abs(pose[2]) > 99):
continue
if use_bined:
yaw_pitch_bins = np.array([-60, -40, -20, 20, 40, 60])
roll_bins = np.array(range(-81, 82, 9))
euler_binned.append([
np.digitize(pose[0], yaw_pitch_bins),
np.digitize(pose[1], yaw_pitch_bins),
np.digitize(pose[2], roll_bins)
])
id_index.append(j)
bboxs.append(bbox)
labels.append(np.array(pose))
self.labels.append(labels)
self.bboxs.append(bboxs)
self.ids.append(ls)
self.ids_index.append(id_index)
if use_bined:
self.euler_binned.append(np.array(euler_binned))
if 'train' in self.split:
self.resizer = albu.Compose([
albu.RandomScale((-0.5, 0), p=0.01),
albu.SmallestMaxSize(int(target_size * 1.1), p=1.),
albu.RandomScale(scale_limit=(-0.1, 0.1), p=1),
albu.PadIfNeeded(min_height=target_size,
min_width=target_size,
value=0,
p=1),
albu.RandomCrop(target_size, target_size, p=1.)
])
else:
# self.resizer = albu.Compose([albu.Resize(target_size[0], target_size[1], p=1.)])
self.resizer = albu.Compose([
albu.SmallestMaxSize(target_size, p=1.),
albu.CenterCrop(target_size, target_size, p=1.)
])
self.affine_augmenter = affine_augmenter
self.image_augmenter = image_augmenter
def visualize_AFLW2000(model, base_dir, target_size, filename, save_imgs_path):
print("[INFO] Initing ALFW2000Dataset.")
base_dir = Path(base_dir)
target_size = target_size
img_paths = []
bboxs = []
labels = []
pred_poses = []
with open(base_dir / filename) as f:
for i, line in enumerate(tqdm.tqdm(f.readlines()[:])):
ls = line.strip()
mat_path = base_dir / ls.replace('.jpg', '.mat')
bbox, pose = get_pt_ypr_from_mat(mat_path, pt3d=True)
if True and (abs(pose[0])>99 or abs(pose[1])>99 or abs(pose[2])>99):
continue
labels.append(np.array(pose))
bboxs.append(bbox)
img_paths.append(ls)
resizer = albu.Compose([albu.SmallestMaxSize(target_size, p=1.),
albu.CenterCrop(target_size, target_size, p=1.)])
for index in tqdm.tqdm(range(0, len(img_paths[:]))):
img_path = base_dir / img_paths[index]
img_name = os.path.basename(img_paths[index])
# print(img_path)
bbox = change_bbox(bboxs[index], 2, use_forehead=False)
# bbox = bboxs[index]
raw_img = Image.open(img_path)
img = np.array(raw_img.crop(bbox))
img = img[:,:,::-1]
img = resizer(image=img)
img = preprocess(np.array(img['image']))
# print(img)
# Convert to tensor and transform to [-1, 1]
img = torch.FloatTensor(img).permute(2, 0, 1).unsqueeze_(0)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
img.to(device)
pred = model(img)
pred = pred.cpu().detach().numpy()
pred_poses.append(pred[0])
if save_imgs_path:
# print(bbox)
# print(pred[0])
b, g, r = raw_img.split()
raw_img = Image.merge("RGB", (r, g, b))
drawed_img = draw_annotates(np.array(raw_img), np.array([bbox]), np.array([pred[0]]))
cv2.imwrite(os.path.join(save_imgs_path, img_name), drawed_img)
pred_poses = np.array(pred_poses)
labels = np.array(labels)
delta = np.absolute(labels - pred_poses)
delta = np.sum(delta, axis=1) / 3
# print(delta)
sortted_delta = np.sort(delta)
# print(sortted_delta)
index_of_max_delta = [np.where(delta==value_delta) for value_delta in sortted_delta[-5:]]
index_of_min_delta = [np.where(delta==value_delta) for value_delta in sortted_delta[:5]]
max_delta_dir = os.path.join(save_imgs_path, "max_delta")
os.system(f"rm -rf \"{max_delta_dir}\"")
min_delta_dir = os.path.join(save_imgs_path, "min_delta")
os.system(f"rm -rf \"{min_delta_dir}\"")
# print(index_of_max_delta)
for index in index_of_max_delta:
# print(index[0][0])
Path(max_delta_dir).mkdir(parents=True, exist_ok=True)
img_name = os.path.basename(img_paths[index[0][0]])
path_drawed_img = os.path.join(save_imgs_path, img_name)
os.system(f"cp -i \"{path_drawed_img}\" \"{max_delta_dir}\"")
for index in index_of_min_delta:
# print(index[0][0])
Path(min_delta_dir).mkdir(parents=True, exist_ok=True)
img_name = os.path.basename(img_paths[index[0][0]])
path_drawed_img = os.path.join(save_imgs_path, img_name)
os.system(f"cp -i \"{path_drawed_img}\" \"{min_delta_dir}\"")