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