def lpips(images1: np.ndarray,
images2: np.ndarray,
batch_size: int = 64,
metric_type: str = "net-lin",
reduce=True):
assert metric_type in [
"net-lin",
"l2",
"ssim",
]
check_shape(images1, images2)
n_batches = int(np.ceil(images1.shape[0] / batch_size))
model = PerceptualLoss(model='net-lin',
net='alex',
use_gpu=torch.cuda.is_available())
distances = np.zeros((images1.shape[0]), dtype=np.float32)
for i in range(n_batches):
start_idx = i * batch_size
end_idx = (i + 1) * batch_size
im1 = images1[start_idx:end_idx]
im2 = images2[start_idx:end_idx]
im1 = torch_utils.image_to_torch(im1, normalize_img=True)
im2 = torch_utils.image_to_torch(im2, normalize_img=True)
with torch.no_grad():
dists = model(im1, im2, normalize=False).cpu().numpy().squeeze()
distances[start_idx:end_idx] = dists
if reduce:
return distances.mean()
assert batch_size == 1
return distances
def inpaint_images(images: np.ndarray, masks: np.ndarray, generator):
z = None
fakes = torch.zeros(
(images.shape[0], images.shape[-1], images.shape[1], images.shape[2]),
dtype=torch.float32)
masks = pre_process_masks(masks)
inputs = [im * mask for im, mask in zip(images, masks)]
images = [
torch_utils.image_to_torch(im, cuda=False, normalize_img=True)
for im in images
]
masks = [torch_utils.mask_to_torch(mask, cuda=False) for mask in masks]
with torch.no_grad():
for idx, (im, mask) in enumerate(
tqdm.tqdm(zip(images, masks), total=len(images))):
im = torch_utils.to_cuda(im)
mask = torch_utils.to_cuda(mask)
assert im.shape[0] == mask.shape[0]
assert im.shape[2:] == mask.shape[2:],\
f"im shape: {im.shape}, mask shape: {mask.shape}"
z = truncated_z(im, generator.z_shape, 0)
condition = mask * im
fake = generator(condition, mask, z)
fakes[idx:(idx + 1)] = fake.cpu()
fakes = torch_utils.image_to_numpy(fakes, denormalize=True) * 255
return fakes, inputs
def detect_keypoints(img, keypoint_threshold=.3):
img = image_to_torch(img, cuda=True)[0]
with torch.no_grad():
outputs = model([img])
# Shape: [N persons, K keypoints, (x,y,visibility)]
keypoints = outputs[0]["keypoints"]
scores = outputs[0]["scores"]
assert list(scores) == sorted(list(scores))[::-1]
mask = scores > keypoint_threshold
keypoints = keypoints[mask, :, :2]
return keypoints.cpu().numpy()
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
def batch_detect_keypoints(images, keypoint_threshold=.3):
images = [image_to_torch(im, cuda=False)[0] for im in images]
batch_size = 16
keypoints = []
scores = []
if len(images) > 0:
num_batches = int(np.ceil(len(images) / batch_size))
with torch.no_grad():
for i in tqdm.trange(num_batches, desc="Keypoint inference"):
images_ = images[i * batch_size:(i + 1) * batch_size]
images_ = [to_cuda(_) for _ in images_]
outputs = model(images_)
images_ = [_.cpu() for _ in images_]
keypoints += [o["keypoints"].cpu() for o in outputs]
scores += [o["scores"].cpu() for o in outputs]
for i in range(len(scores)):
im_scores = scores[i]
im_keypoints = keypoints[i]
mask = im_scores > keypoint_threshold
keypoints[i] = im_keypoints[mask, :, :2].numpy()
return keypoints
def batch_detect_keypoints(self, images):
images = [image_to_torch(im, cuda=False)[0] for im in images]
keypoints = []
if len(images) > 0:
num_batches = int(np.ceil(len(images) / self.batch_size))
with torch.no_grad():
for i in range(num_batches):
images_ = images[i * self.batch_size:(i + 1) *
self.batch_size]
images_ = [to_cuda(_) for _ in images_]
outputs = self.model(images_)
images_ = [_.cpu() for _ in images_]
kps = [o["keypoints"].cpu() for o in outputs]
score = [o["scores"].cpu() for o in outputs]
masks = [
imscore >= self.keypoint_threshold for imscore in score
]
kps = [
kp[mask, :, :2].numpy()
for kp, mask in zip(kps, masks)
]
keypoints += kps
return keypoints
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)
print("Results saved to:", savedir)
def anonymize_images(self,
images: np.ndarray,
image_annotations: typing.List[ImageAnnotation]
) -> typing.List[np.ndarray]:
anonymized_images = []
for im_idx, image_annotation in enumerate(image_annotations):
# pre-process
imsize = self.inference_imsize
condition = torch.zeros(
(len(image_annotation), 3, imsize, imsize),
dtype=torch.float32)
mask = torch.zeros((len(image_annotation), 1, imsize, imsize))
landmarks = torch.empty(
(len(image_annotation), self.pose_size), dtype=torch.float32)
for face_idx in range(len(image_annotation)):
face, mask_ = image_annotation.get_face(face_idx, imsize)
condition[face_idx] = torch_utils.image_to_torch(
face, cuda=False, normalize_img=True
)
mask[face_idx, 0] = torch.from_numpy(mask_).float()
kp = image_annotation.aligned_keypoint(face_idx)
landmarks[face_idx] = kp[:, :self.pose_size]
img = condition
condition = condition * mask
z = infer.truncated_z(
condition, self.cfg.models.generator.z_shape,
self.truncation_level)
batches = dict(
condition=condition,
mask=mask,
landmarks=landmarks,
z=z,
img=img
)
# Inference
anonymized_faces = np.zeros((
len(image_annotation), imsize, imsize, 3), dtype=np.float32)
for idx, batch in enumerate(
batched_iterator(batches, self.batch_size)):
face = self._get_face(batch)
face = torch_utils.image_to_numpy(
face, to_uint8=False, denormalize=True)
start = idx * self.batch_size
anonymized_faces[start:start + self.batch_size] = face
anonymized_image = image_annotation.stitch_faces(anonymized_faces)
anonymized_images.append(anonymized_image)
if self.save_debug:
num_faces = len(batches["condition"])
for face_idx in range(num_faces):
orig_face = torch_utils.image_to_numpy(
batches["img"][face_idx], denormalize=True, to_uint8=True)
condition = torch_utils.image_to_numpy(
batches["condition"][face_idx],
denormalize=True, to_uint8=True)
fake_face = anonymized_faces[face_idx]
fake_face = (fake_face * 255).astype(np.uint8)
to_save = np.concatenate(
(orig_face, condition, fake_face), axis=1)
filepath = self.debug_directory.joinpath(
f"im{im_idx}_face{face_idx}.png")
cv2.imwrite(str(filepath), to_save[:, :, ::-1])
return anonymized_images
