def pre_process(im, keypoint, bbox, imsize, cuda=True):
bbox = to_numpy(bbox)
expanded_bbox = dataset_utils.expand_bbox(bbox, im.shape, SIMPLE_EXPAND,
default_to_simple=True,
expansion_factor1=0.35)
to_replace = dataset_utils.cut_face(im, expanded_bbox, SIMPLE_EXPAND)
new_bbox = shift_bbox(bbox, expanded_bbox, imsize)
new_keypoint = shift_and_scale_keypoint(keypoint, expanded_bbox)
to_replace = cv2.resize(to_replace, (imsize, imsize))
to_replace = cut_bounding_box(to_replace.copy(), torch.tensor(new_bbox),
1.0)
to_replace = torch_utils.image_to_torch(to_replace, cuda=cuda)
keypoint = keypoint_to_torch(new_keypoint)
torch_input = to_replace * 2 - 1
return torch_input, keypoint, expanded_bbox, new_bbox
truncation_levels = np.linspace(0.00, 3, num_iterations)
for i in range(num_iterations):
im = orig.copy()
p = percentages[i]
bbox = bounding_boxes[idx].clone().float()
width = bbox[2] - bbox[0]
height = bbox[3] - bbox[1]
bbox[0] = bbox[0] - p * width
bbox[2] = bbox[2] + p * width
bbox[1] = bbox[1] - p * height
bbox[3] = bbox[3] + p * height
bbox = bbox.long()
im = cut_bounding_box(im, bbox, generator.transition_value)
orig_to_save = im.copy()
im = torch_utils.image_to_torch(im, cuda=True, normalize_img=True)
im = generator(im, pose, z.clone())
im = torch_utils.image_to_numpy(im.squeeze(),
to_uint8=True,
denormalize=True)
im = np.concatenate((orig_to_save.astype(np.uint8), im), axis=0)
to_save = np.concatenate((to_save, im), axis=1)
ims_to_save.append(to_save)
savepath = os.path.join(savedir, f"result_image.jpg")
ims_to_save = np.concatenate(ims_to_save, axis=0)
plt.imsave(savepath, ims_to_save)