diff --git a/main.py b/main.py index b416079..b637ba2 100644 --- a/main.py +++ b/main.py @@ -170,15 +170,17 @@ def anonymize_image(model_file: str, input_file: str, output_file: str, img_size """ img = cv2.imread(input_file) if img is not None: - img = FaceCrop(align)(img)[0] - img = ZeroPaddingResize(img_size)(img) - img = FacialLandmarks478()(img) - img = Pix2PixTransformer(model_file, img_size, device)(img) - img = torch.squeeze(img) - img = (img + 1) / 2 - pil_img = transforms.ToPILImage()(img) - os.makedirs(os.path.dirname(output_file), exist_ok=True) - pil_img.save(output_file) + faces = FaceCrop(align)(img) + for idx, f in enumerate(faces): + img = f + img = ZeroPaddingResize(img_size)(img) + img = FacialLandmarks478()(img) + img = Pix2PixTransformer(model_file, img_size, device)(img) + img = torch.squeeze(img) + img = (img + 1) / 2 + pil_img = transforms.ToPILImage()(img) + os.makedirs(os.path.dirname(output_file), exist_ok=True) + pil_img.save(os.path.join(os.path.dirname(output_file), f'{idx}_{os.path.basename(output_file)}')) def anonymize_directory(model_file: str, input_directory: str, output_directory: str,