def process_pose(self, input_path, output_dir, seq_file_path):
if not self.cache_pose:
return
input_path_no_ext, input_ext = os.path.splitext(input_path)
# Load sequences from file
with open(seq_file_path, "rb") as fp: # Unpickling
seq_list = pickle.load(fp)
# Initialize transforms
img_transforms = img_landmarks_transforms.Compose([
Resize(224),
ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# For each sequence
for seq in seq_list:
curr_vid_name = os.path.basename(
input_path_no_ext) + '_seq%02d%s' % (seq.id, input_ext)
curr_vid_path = os.path.join(output_dir, curr_vid_name)
curr_pose_path = os.path.splitext(
curr_vid_path)[0] + self.pose_postfix
if os.path.isfile(curr_pose_path):
continue
print('=> Computing face poses for video: "%s"...' % curr_vid_name)
# Initialize input video
in_vid = VideoInferenceDataset(curr_vid_path,
transform=img_transforms)
in_vid_loader = DataLoader(in_vid,
batch_size=self.pose_batch_size,
num_workers=1,
pin_memory=True,
drop_last=False,
shuffle=False)
# For each batch of frames in the input video
seq_poses = []
for i, frame in enumerate(
tqdm(in_vid_loader, unit='batches', file=sys.stdout)):
frame = frame.to(self.device)
poses = self.face_pose(frame).div_(99.) # Yaw, Pitch, Roll
seq_poses.append(poses.cpu().numpy())
seq_poses = np.concatenate(seq_poses)
# Save poses to file
seq_poses_smoothed = smooth_poses(seq_poses, self.smooth_poses)
np.savez_compressed(curr_pose_path,
poses=seq_poses,
poses_smoothed=seq_poses_smoothed)
exp_name = os.path.splitext(os.path.basename(__file__))[0]
exp_dir = os.path.join('../results/segmentation', exp_name)
train_dataset = partial(ImageSegDataset,
'/data/datasets/celeba_mask_hq',
'img_list_train.txt',
'bboxes_train.npy',
seg_classes=3)
val_dataset = partial(ImageSegDataset,
'/data/datasets/celeba_mask_hq',
'img_list_val.txt',
'bboxes_val.npy',
seg_classes=3)
numpy_transforms = [
RandomRotation(30.0, ('cubic', 'nearest')),
Crop(),
Resize(256, ('cubic', 'nearest')),
RandomHorizontalFlip(),
ColorJitter(0.5, 0.5, 0.5, 0.5, filter=(True, False)),
RandomGaussianBlur(filter=(True, False))
]
resolutions = [256]
learning_rate = [1e-4]
epochs = [60]
iterations = ['40k']
batch_size = [48]
workers = 12
pretrained = False
optimizer = partial(optim.Adam, betas=(0.5, 0.999))
scheduler = partial(lr_scheduler.StepLR, step_size=10, gamma=0.5)
criterion = nn.CrossEntropyLoss()
model = partial(UNet)