fix

app.py

@@ -9,15 +9,15 @@ os.system("pip install ftfy regex tqdm")
 os.system("pip install --no-deps git+https://github.com/openai/CLIP.git")
 os.system("pip install git+https://github.com/facebookresearch/detectron2@96c752ce821a3340e27edd51c28a00665dd32a30#subdirectory=projects/DensePose")
 os.system("pip install --no-deps git+https://github.com/hukkelas/DSFD-Pytorch-Inference")
-sys.path.insert(0, Path(os.getcwd(), "deep_privacy2"))
+sys.path.insert(0, Path(os.getcwd(), "deep_privacy"))
 os.environ["TORCH_HOME"] = "torch_home"
 from dp2 import utils
 from gradio_demos.modules import ExampleDemo, WebcamDemo
 
-cfg_face = utils.load_config("deep_privacy2/configs/anonymizers/face.py")
+cfg_face = utils.load_config("deep_privacy/configs/anonymizers/face.py")
 for key in ["person_G_cfg", "cse_person_G_cfg", "face_G_cfg", "car_G_cfg"]:
     if key in cfg_face.anonymizer:
-        cfg_face.anonymizer[key] = Path("deep_privacy2", cfg_face.anonymizer[key])
+        cfg_face.anonymizer[key] = Path("deep_privacy", cfg_face.anonymizer[key])
 
 
 anonymizer_face = instantiate(cfg_face.anonymizer, load_cache=False)

deep_privacy/.gitignore

@@ -0,0 +1,54 @@
+# FILES
+*.yaml
+*.pkl
+*.flist
+*.zip
+*.out
+*.npy
+*.gz
+*.ckpt
+*.pth
+*.log
+*.pyc
+*.csv
+*.yml
+*.ods
+*.ods#
+*.json
+build_docker.sh
+
+# Images / Videos
+#*.png
+#*.jpg
+*.jpeg
+*.m4a
+*.mkv
+*.mp4
+
+# Directories created by inpaintron
+.cache/
+test_examples/
+.vscode
+__pycache__
+.debug/
+**/.ipynb_checkpoints/**
+outputs/
+
+
+# From pip setup
+build/
+*.egg-info
+*.egg
+.npm/
+
+# From dockerfile
+.bash_history
+.viminfo
+.local/
+*.pickle
+*.onnx
+
+
+sbatch_files/
+figures/
+image_dump/

deep_privacy/CHANGELOG.md

@@ -0,0 +1,13 @@
+# Changelog
+
+## 23.03.2023
+- Quality of life improvements
+- Add support for refined keypoints for the FDH dataset. 
+- Add FDF128 dataset loader with webdataset.
+- Support for using detector and anonymizer from DeepPrivacy1.
+- Update visualization of keypoints
+- Fix bug for upsampling/downsampling in the anonymization pipeline.
+- Support for keypoint-guided face anonymization.
+- Add ViTPose + Mask-RCNN detection model for keypoint-guided full-body anonymization.
+- Set caching of detections to False as default, as it can produce unexpected behaviour. For example, using a different score threshold requires re-run of detector.
+- Add Gradio Demos for face and body anonymization

deep_privacy/Dockerfile

@@ -0,0 +1,47 @@
+FROM nvcr.io/nvidia/pytorch:22.08-py3
+ARG UID=1000
+ARG UNAME=testuser
+ARG WANDB_API_KEY
+RUN useradd -ms /bin/bash -u $UID $UNAME && \
+    mkdir -p /home/${UNAME} &&\
+    chown -R $UID /home/${UNAME}
+WORKDIR /home/${UNAME}
+ENV DEBIAN_FRONTEND="noninteractive"
+ENV WANDB_API_KEY=$WANDB_API_KEY
+ENV TORCH_HOME=/home/${UNAME}/.cache
+
+# OPTIONAL - DeepPrivacy2 uses these environment variables to set directories outside the current working directory
+#ENV BASE_DATASET_DIR=/work/haakohu/datasets
+#ENV BASE_OUTPUT_DIR=/work/haakohu/outputs
+#ENV FBA_METRICS_CACHE=/work/haakohu/metrics_cache
+
+RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  qt5-default -y
+RUN pip install git+https://github.com/facebookresearch/detectron2@96c752ce821a3340e27edd51c28a00665dd32a30#subdirectory=projects/DensePose
+COPY setup.py setup.py
+RUN pip install \
+    numpy>=1.20 \
+    matplotlib \
+    cython \
+    tensorboard \
+    tqdm \
+    ninja==1.10.2 \
+    opencv-python==4.5.5.64 \
+    moviepy \
+    pyspng \
+    git+https://github.com/hukkelas/DSFD-Pytorch-Inference \
+    wandb \ 
+    termcolor \
+    git+https://github.com/hukkelas/torch_ops.git \
+    git+https://github.com/wmuron/motpy@c77f85d27e371c0a298e9a88ca99292d9b9cbe6b \
+    fast_pytorch_kmeans \
+    einops_exts  \ 
+    einops \ 
+    regex \
+    setuptools==59.5.0 \
+    resize_right==0.0.2 \
+    pillow \
+    scipy==1.7.1 \
+    webdataset==0.2.26 \
+    scikit-image \
+    timm==0.6.7
+RUN pip install --no-deps torch_fidelity==0.3.0 clip@git+https://github.com/openai/CLIP.git@b46f5ac7587d2e1862f8b7b1573179d80dcdd620

deep_privacy/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

deep_privacy/anonymize.py

@@ -0,0 +1,255 @@
+import hashlib
+from typing import Optional
+import click
+import tops
+import numpy as np
+import tqdm
+import moviepy.editor as mp
+import cv2
+from tops.config import instantiate
+from pathlib import Path
+from PIL import Image
+from dp2 import utils
+from detectron2.data.detection_utils import _apply_exif_orientation
+from tops import logger
+from dp2.utils.bufferless_video_capture import BufferlessVideoCapture
+
+
+def show_video(video_path):
+    video_cap = cv2.VideoCapture(str(video_path))
+    while video_cap.isOpened():
+        ret,  frame = video_cap.read()
+        cv2.imshow("Frame", frame)
+        key = cv2.waitKey(25)
+        if key == ord("q"):
+            break
+    video_cap.release()
+    cv2.destroyAllWindows()
+
+
+class ImageIndexTracker:
+
+    def __init__(self, fn) -> None:
+        self.fn = fn
+        self.idx = 0
+    
+    def fl_image(self, frame):
+        self.idx += 1
+        return self.fn(frame, self.idx-1)
+
+
+def anonymize_video(
+        video_path, output_path: Path,
+        anonymizer, visualize: bool, max_res: int,
+        start_time: int, fps: int,
+        end_time: int,
+        visualize_detection: bool,
+        track: bool, 
+        synthesis_kwargs,
+        **kwargs):
+    video = mp.VideoFileClip(str(video_path))
+    if track:
+        anonymizer.initialize_tracker(video.fps)
+
+    def process_frame(frame, idx):
+        frame = np.array(resize(Image.fromarray(frame), max_res))
+        cache_id = hashlib.md5(frame).hexdigest()
+        frame = utils.im2torch(frame, to_float=False, normalize=False)[0]
+        cache_id_ = cache_id +  str(idx)
+        synthesis_kwargs["cache_id"] = cache_id_
+        if visualize_detection:
+            anonymized = anonymizer.visualize_detection(frame, cache_id=cache_id_)
+        else:
+            anonymized = anonymizer(frame, **synthesis_kwargs)
+        anonymized = utils.im2numpy(anonymized)
+        if visualize:
+            cv2.imshow("frame", anonymized[:, :, ::-1])
+            key = cv2.waitKey(1)
+            if key == ord("q"):
+                exit()
+        return anonymized
+    video: mp.VideoClip = video.subclip(start_time, end_time)
+
+    if fps is not None:
+        video = video.set_fps(fps)
+
+    video = video.fl_image(ImageIndexTracker(process_frame).fl_image)
+    if str(output_path).endswith(".avi"):
+        output_path = str(output_path).replace(".avi", ".mp4")
+    if not output_path.parent.exists():
+        output_path.parent.mkdir(parents=True)
+    video.write_videofile(str(output_path))
+
+
+def resize(frame: Image.Image, max_res):
+    if max_res is None:
+        return frame
+    f = max(*[x/max_res for x in frame.size], 1)
+    if f  == 1:
+        return frame
+    new_shape = [int(x/f) for x in frame.size]
+    return frame.resize(new_shape,  resample=Image.BILINEAR)
+
+
+def anonymize_image(
+        image_path, output_path: Path, visualize: bool,
+        anonymizer, max_res: int,
+        visualize_detection: bool,
+        synthesis_kwargs,
+        **kwargs):
+    with Image.open(image_path) as im:
+        im = _apply_exif_orientation(im)
+        orig_im_mode = im.mode
+
+        im = im.convert("RGB")
+        im = resize(im, max_res)
+    im = np.array(im)
+    md5_ = hashlib.md5(im).hexdigest()
+    im = utils.im2torch(np.array(im), to_float=False, normalize=False)[0]
+    synthesis_kwargs["cache_id"] = md5_
+    if visualize_detection:
+        im_ = anonymizer.visualize_detection(tops.to_cuda(im), cache_id=md5_)
+    else:
+        im_ = anonymizer(im, **synthesis_kwargs)
+    im_ = utils.im2numpy(im_)
+    if visualize:
+        while True:
+            cv2.imshow("frame", im_[:, :, ::-1])
+            key = cv2.waitKey(0)
+            if key == ord("q"):
+                break
+            elif key == ord("u"):
+                im_ = utils.im2numpy(anonymizer(im, **synthesis_kwargs))
+    im = Image.fromarray(im_).convert(orig_im_mode)
+    if output_path is not None:
+        output_path.parent.mkdir(exist_ok=True, parents=True)
+        im.save(output_path, optimize=False, quality=100)
+        print(f"Saved to: {output_path}")
+
+
+def anonymize_file(input_path: Path, output_path: Optional[Path], **kwargs):
+    if output_path is not None and output_path.is_file():
+        logger.warn(f"Overwriting previous file: {output_path}")
+    if tops.is_image(input_path):
+        anonymize_image(input_path, output_path, **kwargs)
+    elif tops.is_video(input_path):
+        anonymize_video(input_path, output_path, **kwargs)
+    else:
+        logger.log(f"Filepath not a video or image file: {input_path}")
+
+
+def anonymize_directory(input_dir: Path, output_dir: Path, **kwargs):
+    for childname in tqdm.tqdm(input_dir.iterdir()):
+        childpath = input_dir.joinpath(childname.name)
+        output_path = output_dir.joinpath(childname.name)
+        if not childpath.is_file():
+            anonymize_directory(childpath, output_path, **kwargs)
+        else:
+            assert childpath.is_file()
+            anonymize_file(childpath, output_path, **kwargs)
+
+def anonymize_webcam(
+        anonymizer, max_res: int,
+        synthesis_kwargs,
+        visualize_detection,
+        track: bool, 
+        **kwargs):
+    import time
+    cap = BufferlessVideoCapture(0, width=1920, height=1080)
+    t = time.time()
+    frames = 0
+    if track:
+        anonymizer.initialize_tracker(fps=5) # FPS used for tracking objects
+    while True:
+        # Capture frame-by-frame
+        ret, frame = cap.read()
+        frame = Image.fromarray(frame[:, :, ::-1])
+        frame = resize(frame, max_res)
+        frame = np.array(frame)
+        im = utils.im2torch(np.array(frame), to_float=False, normalize=False)[0]
+        if visualize_detection:
+            im_ = anonymizer.visualize_detection(tops.to_cuda(im))
+        else:
+            im_ = anonymizer(im, **synthesis_kwargs)
+        im_ = utils.im2numpy(im_)
+
+        frames += 1
+        delta = time.time() - t
+        fps = "?"
+        if delta > 1e-6:
+            fps = frames / delta
+        print(f"FPS: {fps:.3f}", end="\r")
+        cv2.imshow('frame', im_[:, :, ::-1])
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            break
+
+
+@click.command()
+@click.argument("config_path", type=click.Path(exists=True))
+@click.option("-i", "--input_path", help="Input path. Accepted inputs: images, videos, directories.")
+@click.option("-o", "--output_path", default=None, type=click.Path(), help="Output path to save. Can be directory or file.")
+@click.option("-v","--visualize", default=False, is_flag=True, help="Visualize the result")
+@click.option("--max-res", default=None, type=int, help="Maximum resolution  of height/wideo")
+@click.option("--start-time", "--st", default=0, type=int, help="Start time (second) for vide anonymization")
+@click.option("--end-time", "--et", default=None, type=int, help="End time (second) for vide anonymization")
+@click.option("--fps", default=None, type=int, help="FPS for anonymization")
+@click.option("--detection-score-threshold", "--dst", default=.3, type=click.FloatRange(0, 1), help="Detection threshold, threshold applied for all detection models.")
+@click.option("--visualize-detection", "--vd",default=False, is_flag=True, help="Visualize only detections without running anonymization.")
+@click.option("--multi-modal-truncation", "--mt", default=False, is_flag=True, help="Enable multi-modal truncation proposed by: https://arxiv.org/pdf/2202.12211.pdf")
+@click.option("--cache", default=False, is_flag=True, help="Enable detection caching. Will save and load detections from cache.")
+@click.option("--amp", default=True, is_flag=True, help="Use automatic mixed precision for generator forward pass")
+@click.option("-t", "--truncation_value", default=0, type=click.FloatRange(0, 1), help="Latent interpolation truncation value.")
+@click.option("--track", default=False, is_flag=True, help="Track detections over frames. Will use the same latent variable (z) for tracked identities.")
+@click.option("--seed", default=0, type=int, help="Set random seed for generating images.")
+@click.option("--person-generator", default=None, help="Config path to unconditional person generator", type=click.Path())
+@click.option("--cse-person-generator", default=None, help="Config path to CSE-guided person generator", type=click.Path())
+@click.option("--webcam", default=False, is_flag=True, help="Read image from webcam feed.")
+def anonymize_path(
+        config_path,
+        input_path,
+        output_path,
+        detection_score_threshold: float,
+        visualize_detection: bool,
+        cache: bool,
+        seed: int,
+        person_generator: str,
+        cse_person_generator: str,
+        webcam: bool,
+        **kwargs):
+    """
+    config_path: Specify the path to the anonymization model to use.
+    """
+    tops.set_seed(seed)
+    cfg = utils.load_config(config_path)
+    if person_generator is not None:
+        cfg.anonymizer.person_G_cfg = person_generator
+    if cse_person_generator is not None:
+        cfg.anonymizer.cse_person_G_cfg = cse_person_generator
+    cfg.detector.score_threshold = detection_score_threshold
+    utils.print_config(cfg)
+
+    anonymizer = instantiate(cfg.anonymizer, load_cache=cache)
+    synthesis_kwargs = ["amp", "multi_modal_truncation", "truncation_value"]
+    synthesis_kwargs = {k: kwargs.pop(k) for k in synthesis_kwargs}
+
+    kwargs["anonymizer"] = anonymizer
+    kwargs["visualize_detection"] = visualize_detection
+    kwargs["synthesis_kwargs"] = synthesis_kwargs
+    if webcam:
+        anonymize_webcam(**kwargs)
+        return
+    input_path = Path(input_path)
+    output_path = Path(output_path) if output_path is not None else None
+    if output_path is None and not kwargs["visualize"]:
+        logger.log("Output path not set. Setting visualize to True")
+        kwargs["visualize"] = True
+    if input_path.is_dir():
+        assert output_path is None or not output_path.is_file()
+        anonymize_directory(input_path, output_path, **kwargs)
+    else:
+        anonymize_file(input_path, output_path, **kwargs)
+    
+    
+if __name__ == "__main__":
+
+    anonymize_path()

deep_privacy/attribute_guided_demo.py

@@ -0,0 +1,144 @@
+from collections import defaultdict
+import gradio
+import numpy as np
+import torch
+import cv2
+from PIL import Image
+from dp2 import utils
+from tops.config import instantiate
+import tops
+import gradio.inputs
+from stylemc import get_and_cache_direction, get_styles
+
+
+class GuidedDemo:
+    def __init__(self, face_anonymizer, cfg_face) -> None:
+        self.anonymizer = face_anonymizer
+        assert sum([x is not None for x in list(face_anonymizer.generators.values())]) == 1
+        self.generator = [x for x in list(face_anonymizer.generators.values()) if x is not None][0]
+        face_G_cfg = utils.load_config(cfg_face.anonymizer.face_G_cfg)
+        face_G_cfg.train.batch_size = 1
+        self.dl = instantiate(face_G_cfg.data.val.loader)
+        self.cache_dir = face_G_cfg.output_dir
+        self.precompute_edits()
+        
+    def precompute_edits(self):
+        self.precomputed_edits = set()
+        for edit in self.precomputed_edits:
+            get_and_cache_direction(self.cache_dir, self.dl, self.generator, edit)
+        if self.cache_dir.joinpath("stylemc_cache").is_dir():
+            for path in self.cache_dir.joinpath("stylemc_cache").iterdir():
+                text_prompt = path.stem.replace("_", " ")
+                self.precomputed_edits.add(text_prompt)
+                print(text_prompt)
+        self.edits = defaultdict(defaultdict)
+
+    def anonymize(self, img, show_boxes: bool, current_box_idx: int, current_styles, current_boxes, update_identity, edits, cache_id=None):
+        if not isinstance(img, torch.Tensor):
+            img, cache_id = pil2torch(img)
+            img = tops.to_cuda(img)
+
+        current_box_idx = current_box_idx % len(current_boxes)
+        edited_styles = [s.clone() for s in current_styles]
+        for face_idx, face_edits in edits.items():
+            for prompt, strength in face_edits.items():
+                direction = get_and_cache_direction(self.cache_dir, self.dl, self.generator, prompt)
+                edited_styles[int(face_idx)] += direction * strength
+            update_identity[int(face_idx)] = True
+        assert img.dtype == torch.uint8
+        img = self.anonymizer(
+            img, truncation_value=0,
+            multi_modal_truncation=True, amp=True,
+            cache_id=cache_id,
+            all_styles=edited_styles,
+            update_identity=update_identity)
+        update_identity = [True for i in range(len(update_identity))]
+        img = utils.im2numpy(img)
+        if show_boxes:
+            x0, y0, x1, y1 = [int(_) for _ in current_boxes[int(current_box_idx)]]
+            img = cv2.rectangle(img, (x0, y0), (x1, y1), (255, 0, 0), 1)
+        return img, update_identity
+
+    def update_image(self, img, show_boxes):
+        img, cache_id = pil2torch(img)
+        img = tops.to_cuda(img)
+        det = self.anonymizer.detector.forward_and_cache(img, cache_id, load_cache=True)[0]
+        current_styles = []
+        for i in range(len(det)):
+            # Need to do forward pass to register all affine modules.
+            batch = det.get_crop(i, img)
+            batch["condition"] = batch["img"].float()
+
+            s = get_styles(
+                np.random.randint(0, 999999),self.generator,
+            batch, truncation_value=0)
+            current_styles.append(s)
+        update_identity = [True for i in range(len(det))]
+        current_boxes = np.array(det.boxes)
+        edits = defaultdict(defaultdict)
+        cur_face_idx = -1 % len(current_boxes)
+        img, update_identity = self.anonymize(img, show_boxes, cur_face_idx, current_styles, current_boxes, update_identity, edits, cache_id=cache_id)
+        return img, current_styles, current_boxes, update_identity, edits, cur_face_idx
+
+    def change_face(self, change, cur_face_idx, current_boxes, input_image, show_boxes, current_styles, update_identity, edits):
+        cur_face_idx = (cur_face_idx+change) % len(current_boxes)
+        img, update_identity = self.anonymize(input_image, show_boxes, cur_face_idx, current_styles, current_boxes, update_identity, edits)
+        return img, update_identity, cur_face_idx
+
+    def add_style(self, face_idx: int, prompt: str, strength: float, input_image, show_boxes, current_styles, current_boxes, update_identity, edits):
+        face_idx = face_idx % len(current_boxes)
+        edits[face_idx][prompt] = strength
+        img, update_identity = self.anonymize(input_image, show_boxes, face_idx, current_styles, current_boxes, update_identity, edits)
+        return img, update_identity, edits
+    
+    def setup_interface(self):
+        current_styles = gradio.State()
+        current_boxes = gradio.State(None)
+        update_identity = gradio.State([])
+        edits = gradio.State([])
+        with gradio.Row():
+            input_image = gradio.Image(
+                type="pil", label="Upload your image or try the example below!",source="webcam")
+            output_image = gradio.Image(type="numpy", label="Output")
+        with gradio.Row():
+            update_btn = gradio.Button("Update Anonymization").style(full_width=True)
+        with gradio.Row():
+            show_boxes = gradio.Checkbox(value=True, label="Show Selected")
+            cur_face_idx = gradio.Number(value=-1,label="Current", interactive=False)
+            previous = gradio.Button("Previous Person")
+            next_ = gradio.Button("Next Person")
+        with gradio.Row():
+            text_prompt = gradio.Textbox(
+                placeholder=" | ".join(list(self.precomputed_edits)),
+                label="Text Prompt for Edit")
+            edit_strength = gradio.Slider(0, 5, step=.01)
+            add_btn = gradio.Button("Add Edit")
+            add_btn.click(self.add_style, inputs=[cur_face_idx, text_prompt, edit_strength, input_image, show_boxes, current_styles, current_boxes, update_identity, edits], outputs=[output_image, update_identity, edits])
+        update_btn.click(self.update_image, inputs=[input_image, show_boxes], outputs=[output_image, current_styles, current_boxes, update_identity, edits, cur_face_idx])
+        input_image.change(self.update_image, inputs=[input_image, show_boxes], outputs=[output_image, current_styles, current_boxes, update_identity, edits, cur_face_idx])
+        previous.click(self.change_face, inputs=[gradio.State(-1), cur_face_idx, current_boxes, input_image, show_boxes, current_styles, update_identity, edits], outputs=[output_image, update_identity, cur_face_idx])
+        next_.click(self.change_face, inputs=[gradio.State(1), cur_face_idx, current_boxes, input_image, show_boxes, current_styles, update_identity, edits], outputs=[output_image, update_identity, cur_face_idx])
+
+        show_boxes.change(self.anonymize, inputs=[input_image, show_boxes, cur_face_idx, current_styles, current_boxes, update_identity, edits], outputs=[output_image, update_identity])
+
+
+def pil2torch(img: Image.Image):
+    img = img.convert("RGB")
+    img = np.array(img)
+    img = np.rollaxis(img, 2)
+    return torch.from_numpy(img), None
+
+
+cfg_face = utils.load_config("configs/anonymizers/face.py")
+anonymizer_face = instantiate(cfg_face.anonymizer, load_cache=False)
+anonymizer_face.initialize_tracker(fps=1)
+
+
+with gradio.Blocks() as demo:
+    gradio.Markdown("# <center> DeepPrivacy2 - Realistic Image Anonymization </center>")
+    gradio.Markdown("### <center> Håkon Hukkelås, Rudolf Mester, Frank Lindseth </center>")
+    with gradio.Tab("Text-Guided Anonymization"):
+        GuidedDemo(anonymizer_face, cfg_face).setup_interface()
+
+
+demo.launch()

deep_privacy/configs/anonymizers/FB_cse.py

@@ -0,0 +1,28 @@
+from dp2.anonymizer import Anonymizer
+from dp2.detection.person_detector import CSEPersonDetector
+from ..defaults import common
+from tops.config import LazyCall as L
+from dp2.generator.dummy_generators import MaskOutGenerator
+
+
+maskout_G = L(MaskOutGenerator)(noise="constant")
+
+detector = L(CSEPersonDetector)(
+    mask_rcnn_cfg=dict(),
+    cse_cfg=dict(),
+    cse_post_process_cfg=dict(
+        target_imsize=(288, 160),
+        exp_bbox_cfg=dict(percentage_background=0.3, axis_minimum_expansion=.1),
+        exp_bbox_filter=dict(minimum_area=32*32, min_bbox_ratio_inside=0, aspect_ratio_range=[0, 99999]),
+        iou_combine_threshold=0.4,
+        dilation_percentage=0.02,
+        normalize_embedding=False
+    ),
+    score_threshold=0.3,
+    cache_directory=common.output_dir.joinpath("cse_person_detection_cache")
+)
+
+anonymizer = L(Anonymizer)(
+    detector="${detector}",
+    cse_person_G_cfg="configs/fdh/styleganL.py",    
+)

deep_privacy/configs/anonymizers/FB_cse_mask.py

@@ -0,0 +1,29 @@
+from dp2.anonymizer import Anonymizer
+from dp2.detection.person_detector import CSEPersonDetector
+from ..defaults import common
+from tops.config import LazyCall as L
+from dp2.generator.dummy_generators import MaskOutGenerator
+
+
+maskout_G = L(MaskOutGenerator)(noise="constant")
+
+detector = L(CSEPersonDetector)(
+    mask_rcnn_cfg=dict(),
+    cse_cfg=dict(),
+    cse_post_process_cfg=dict(
+        target_imsize=(288, 160),
+        exp_bbox_cfg=dict(percentage_background=0.3, axis_minimum_expansion=.1),
+        exp_bbox_filter=dict(minimum_area=32*32, min_bbox_ratio_inside=0, aspect_ratio_range=[0, 99999]),
+        iou_combine_threshold=0.4,
+        dilation_percentage=0.02,
+        normalize_embedding=False
+    ),
+    score_threshold=0.3,
+    cache_directory=common.output_dir.joinpath("cse_person_detection_cache")
+)
+
+anonymizer = L(Anonymizer)(
+    detector="${detector}",
+    person_G_cfg="configs/fdh/styleganL_nocse.py",
+    cse_person_G_cfg="configs/fdh/styleganL.py",    
+)

deep_privacy/configs/anonymizers/FB_cse_mask_face.py

@@ -0,0 +1,29 @@
+from dp2.anonymizer import Anonymizer
+from dp2.detection.cse_mask_face_detector import CSeMaskFaceDetector
+from ..defaults import common
+from tops.config import LazyCall as L
+
+detector = L(CSeMaskFaceDetector)(
+    mask_rcnn_cfg=dict(),
+    face_detector_cfg=dict(),
+    face_post_process_cfg=dict(target_imsize=(256, 256), fdf128_expand=False),
+    cse_cfg=dict(),
+    cse_post_process_cfg=dict(
+        target_imsize=(288, 160),
+        exp_bbox_cfg=dict(percentage_background=0.3, axis_minimum_expansion=.1),
+        exp_bbox_filter=dict(minimum_area=32*32, min_bbox_ratio_inside=0, aspect_ratio_range=[0, 99999]),
+        iou_combine_threshold=0.4,
+        dilation_percentage=0.02,
+        normalize_embedding=False
+    ),
+    score_threshold=0.3,
+    cache_directory=common.output_dir.joinpath("cse_mask_face_detection_cache")
+)
+
+anonymizer = L(Anonymizer)(
+    detector="${detector}",
+    face_G_cfg="configs/fdf/stylegan.py",
+    person_G_cfg="configs/fdh/styleganL_nocse.py",
+    cse_person_G_cfg="configs/fdh/styleganL.py",
+    car_G_cfg="configs/generators/dummy/pixelation8.py"
+)

deep_privacy/configs/anonymizers/deep_privacy1.py

@@ -0,0 +1,15 @@
+from .face_fdf128 import anonymizer, common, detector
+from dp2.detection.deep_privacy1_detector import DeepPrivacy1Detector
+from tops.config import LazyCall as L
+
+anonymizer.update(
+    face_G_cfg="configs/fdf/deep_privacy1.py",
+)
+
+anonymizer.detector = L(DeepPrivacy1Detector)(
+    face_detector_cfg=dict(name="DSFDDetector", clip_boxes=True),
+    face_post_process_cfg=dict(target_imsize=(128, 128), fdf128_expand=True),
+    score_threshold=0.3,
+    keypoint_threshold=0.3,
+    cache_directory=common.output_dir.joinpath("deep_privacy1_cache")
+)

deep_privacy/configs/anonymizers/face.py

@@ -0,0 +1,17 @@
+from dp2.anonymizer import Anonymizer
+from dp2.detection.face_detector import FaceDetector
+from ..defaults import common
+from tops.config import LazyCall as L
+
+
+detector = L(FaceDetector)(
+    face_detector_cfg=dict(name="DSFDDetector", clip_boxes=True),
+    face_post_process_cfg=dict(target_imsize=(256, 256), fdf128_expand=False),
+    score_threshold=0.3,
+    cache_directory=common.output_dir.joinpath("face_detection_cache"),
+)
+
+anonymizer = L(Anonymizer)(
+    detector="${detector}",
+    face_G_cfg="configs/fdf/stylegan.py",
+)

deep_privacy/configs/anonymizers/face_fdf128.py

@@ -0,0 +1,18 @@
+from dp2.anonymizer import Anonymizer
+from dp2.detection.face_detector import FaceDetector
+from ..defaults import common
+from tops.config import LazyCall as L
+
+
+detector = L(FaceDetector)(
+    face_detector_cfg=dict(name="DSFDDetector", clip_boxes=True),
+    face_post_process_cfg=dict(target_imsize=(128, 128), fdf128_expand=True),
+    score_threshold=0.3,
+    cache_directory=common.output_dir.joinpath("face_detection_cache")
+)
+
+
+anonymizer = L(Anonymizer)(
+    detector="${detector}",
+    face_G_cfg="configs/fdf/stylegan_fdf128.py",
+)

deep_privacy/configs/anonymizers/market1501/blackout.py

@@ -0,0 +1,8 @@
+from ..FB_cse_mask_face import anonymizer, detector, common
+
+detector.score_threshold = .1
+detector.face_detector_cfg.confidence_threshold = .5
+detector.cse_cfg.score_thres = 0.3
+anonymizer.generators.face_G_cfg = None
+anonymizer.generators.person_G_cfg = "configs/generators/dummy/maskout.py"
+anonymizer.generators.cse_person_G_cfg = "configs/generators/dummy/maskout.py"

deep_privacy/configs/anonymizers/market1501/person.py

@@ -0,0 +1,6 @@
+from ..FB_cse_mask_face import anonymizer, detector, common
+
+detector.score_threshold = .1
+detector.face_detector_cfg.confidence_threshold = .5
+detector.cse_cfg.score_thres = 0.3
+anonymizer.generators.face_G_cfg = None

deep_privacy/configs/anonymizers/market1501/pixelation16.py

@@ -0,0 +1,8 @@
+from ..FB_cse_mask_face import anonymizer, detector, common
+
+detector.score_threshold = .1
+detector.face_detector_cfg.confidence_threshold = .5
+detector.cse_cfg.score_thres = 0.3
+anonymizer.generators.face_G_cfg = None
+anonymizer.generators.person_G_cfg = "configs/generators/dummy/pixelation16.py"
+anonymizer.generators.cse_person_G_cfg = "configs/generators/dummy/pixelation16.py"

deep_privacy/configs/anonymizers/market1501/pixelation8.py

@@ -0,0 +1,8 @@
+from ..FB_cse_mask_face import anonymizer, detector, common
+
+detector.score_threshold = .1
+detector.face_detector_cfg.confidence_threshold = .5
+detector.cse_cfg.score_thres = 0.3
+anonymizer.generators.face_G_cfg = None
+anonymizer.generators.person_G_cfg = "configs/generators/dummy/pixelation8.py"
+anonymizer.generators.cse_person_G_cfg = "configs/generators/dummy/pixelation8.py"

deep_privacy/configs/datasets/coco_cse.py

@@ -0,0 +1,69 @@
+import os
+from pathlib import Path
+from tops.config import LazyCall as L
+import torch
+import functools
+from dp2.data.datasets.coco_cse import CocoCSE
+from dp2.data.build import get_dataloader
+from dp2.data.transforms.transforms import CreateEmbedding, Normalize, Resize, ToFloat, CreateCondition, RandomHorizontalFlip
+from dp2.data.transforms.stylegan2_transform import StyleGANAugmentPipe
+from dp2.metrics.torch_metrics import compute_metrics_iteratively
+from .utils import final_eval_fn
+
+
+dataset_base_dir = os.environ["BASE_DATASET_DIR"] if "BASE_DATASET_DIR" in os.environ else "data"
+metrics_cache = os.environ["FBA_METRICS_CACHE"] if "FBA_METRICS_CACHE" in os.environ else ".cache"
+data_dir = Path(dataset_base_dir, "coco_cse")
+data = dict(
+    imsize=(288, 160),
+    im_channels=3,
+    semantic_nc=26,
+    cse_nc=16,
+    train=dict(
+        dataset=L(CocoCSE)(data_dir.joinpath("train"), transform=None, normalize_E=False),
+        loader=L(get_dataloader)(
+            shuffle=True, num_workers=6, drop_last=True, prefetch_factor=2,
+            batch_size="${train.batch_size}",
+            dataset="${..dataset}",
+            infinite=True,
+            gpu_transform=L(torch.nn.Sequential)(*[
+                L(ToFloat)(),
+                L(StyleGANAugmentPipe)(
+                    rotate=0.5, rotate_max=.05,
+                    xint=.5, xint_max=0.05,
+                    scale=.5, scale_std=.05,
+                    aniso=0.5, aniso_std=.05,
+                    xfrac=.5, xfrac_std=.05,
+                    brightness=.5, brightness_std=.05,
+                    contrast=.5, contrast_std=.1,
+                    hue=.5, hue_max=.05,
+                    saturation=.5, saturation_std=.5,
+                    imgfilter=.5, imgfilter_std=.1),
+                L(RandomHorizontalFlip)(p=0.5),
+                L(CreateEmbedding)(),
+                L(Resize)(size="${data.imsize}"),
+                L(Normalize)(mean=[.5, .5, .5], std=[.5, .5, .5], inplace=True),
+                L(CreateCondition)(),
+            ])
+        )
+    ),
+    val=dict(
+        dataset=L(CocoCSE)(data_dir.joinpath("val"), transform=None, normalize_E=False),
+        loader=L(get_dataloader)(
+            shuffle=False, num_workers=6, drop_last=True, prefetch_factor=2,
+            batch_size="${train.batch_size}",
+            dataset="${..dataset}",
+            infinite=False,
+            gpu_transform=L(torch.nn.Sequential)(*[
+                L(ToFloat)(),
+                L(CreateEmbedding)(),
+                L(Resize)(size="${data.imsize}"),
+                L(Normalize)(mean=[.5, .5, .5], std=[.5, .5, .5], inplace=True),
+                L(CreateCondition)(),
+            ])
+        )
+    ),
+    # Training evaluation might do optimizations to reduce compute overhead. E.g. compute with AMP.
+    train_evaluation_fn=functools.partial(compute_metrics_iteratively, cache_directory=Path(metrics_cache, "coco_cse_val"), include_two_fake=False),
+    evaluation_fn=functools.partial(final_eval_fn, cache_directory=Path(metrics_cache, "coco_cse_val_final"), include_two_fake=True)
+)

deep_privacy/configs/datasets/fdf128.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+from functools import partial
+from dp2.data.datasets.fdf import FDFDataset
+from .fdf256 import data, dataset_base_dir, metrics_cache, final_eval_fn, train_eval_fn
+
+data_dir = Path(dataset_base_dir, "fdf")
+data.train.dataset.dirpath = data_dir.joinpath("train")
+data.val.dataset.dirpath = data_dir.joinpath("val")
+data.imsize = (128, 128)
+        
+
+data.train_evaluation_fn = partial(
+    train_eval_fn, cache_directory=Path(metrics_cache, "fdf128_val_train"))
+data.evaluation_fn = partial(
+    final_eval_fn, cache_directory=Path(metrics_cache, "fdf128_val_final"))
+
+data.train.dataset.update(
+    _target_ = FDFDataset,
+    imsize="${data.imsize}"
+)
+data.val.dataset.update(
+    _target_ = FDFDataset,
+    imsize="${data.imsize}"
+)

deep_privacy/configs/datasets/fdf256.py

@@ -0,0 +1,55 @@
+import os
+from pathlib import Path
+from tops.config import LazyCall as L
+import torch
+import functools
+from dp2.data.datasets.fdf import FDF256Dataset
+from dp2.data.build import get_dataloader
+from dp2.data.transforms.transforms import Normalize, Resize, ToFloat, CreateCondition, RandomHorizontalFlip
+from .utils import final_eval_fn, train_eval_fn
+
+
+dataset_base_dir = os.environ["BASE_DATASET_DIR"] if "BASE_DATASET_DIR" in os.environ else "data"
+metrics_cache = os.environ["FBA_METRICS_CACHE"] if "FBA_METRICS_CACHE" in os.environ else ".cache"
+data_dir = Path(dataset_base_dir, "fdf256")
+data = dict(
+    imsize=(256, 256),
+    im_channels=3,
+    semantic_nc=None,
+    cse_nc=None,
+    n_keypoints=None,
+    train=dict(
+        dataset=L(FDF256Dataset)(dirpath=data_dir.joinpath("train"), transform=None, load_keypoints=False),
+        loader=L(get_dataloader)(
+            shuffle=True, num_workers=3, drop_last=True, prefetch_factor=2,
+            batch_size="${train.batch_size}",
+            dataset="${..dataset}",
+            infinite=True,
+            gpu_transform=L(torch.nn.Sequential)(*[
+                L(ToFloat)(),
+                L(RandomHorizontalFlip)(p=0.5),
+                L(Resize)(size="${data.imsize}"),
+                L(Normalize)(mean=[.5, .5, .5], std=[.5, .5, .5], inplace=True),
+                L(CreateCondition)(),
+            ])
+        )
+    ),
+    val=dict(
+        dataset=L(FDF256Dataset)(dirpath=data_dir.joinpath("val"), transform=None, load_keypoints=False),
+        loader=L(get_dataloader)(
+            shuffle=False, num_workers=3, drop_last=False, prefetch_factor=2,
+            batch_size="${train.batch_size}",
+            dataset="${..dataset}",
+            infinite=False,
+            gpu_transform=L(torch.nn.Sequential)(*[
+                L(ToFloat)(),
+                L(Resize)(size="${data.imsize}"),
+                L(Normalize)(mean=[.5, .5, .5], std=[.5, .5, .5], inplace=True),
+                L(CreateCondition)(),
+            ])
+        )
+    ),
+    # Training evaluation might do optimizations to reduce compute overhead. E.g. compute with AMP.
+    train_evaluation_fn=functools.partial(train_eval_fn, cache_directory=Path(metrics_cache, "fdf_val_train")),
+    evaluation_fn=functools.partial(final_eval_fn, cache_directory=Path(metrics_cache, "fdf_val"))
+)

deep_privacy/configs/datasets/fdh.py

@@ -0,0 +1,90 @@
+import os
+from pathlib import Path
+from tops.config import LazyCall as L
+import torch
+import functools
+from dp2.data.datasets.fdh import get_dataloader_fdh_wds
+from dp2.data.utils import get_coco_flipmap
+from dp2.data.transforms.transforms import (
+    Normalize,
+    ToFloat,
+    CreateCondition,
+    RandomHorizontalFlip,
+    CreateEmbedding,
+)
+from dp2.metrics.torch_metrics import compute_metrics_iteratively
+from dp2.metrics.fid_clip import compute_fid_clip
+from dp2.metrics.ppl import calculate_ppl
+from .utils import train_eval_fn
+
+
+def final_eval_fn(*args, **kwargs):
+    result = compute_metrics_iteratively(*args, **kwargs)
+    result2 = calculate_ppl(*args, **kwargs, upsample_size=(288, 160))
+    result3 = compute_fid_clip(*args, **kwargs)
+    assert all(key not in result for key in result2)
+    result.update(result2)
+    result.update(result3)
+    return result
+
+
+def get_cache_directory(imsize, subset):
+    return Path(metrics_cache, f"{subset}{imsize[0]}")
+
+dataset_base_dir = (
+    os.environ["BASE_DATASET_DIR"] if "BASE_DATASET_DIR" in os.environ else "data"
+)
+metrics_cache = (
+    os.environ["FBA_METRICS_CACHE"] if "FBA_METRICS_CACHE" in os.environ else ".cache"
+)
+data_dir = Path(dataset_base_dir, "fdh")
+data = dict(
+    imsize=(288, 160),
+    im_channels=3,
+    cse_nc=16,
+    n_keypoints=17,
+    train=dict(
+        loader=L(get_dataloader_fdh_wds)(
+            path=data_dir.joinpath("train", "out-{000000..001423}.tar"),
+            batch_size="${train.batch_size}",
+            num_workers=6,
+            transform=L(torch.nn.Sequential)(
+                L(RandomHorizontalFlip)(p=0.5, flip_map=get_coco_flipmap()),
+            ),
+            gpu_transform=L(torch.nn.Sequential)(
+                L(ToFloat)(norm=False, keys=["img", "mask", "E_mask", "maskrcnn_mask"]),
+                L(CreateEmbedding)(embed_path=data_dir.joinpath("embed_map.torch")),
+                L(Normalize)(mean=[0.5*255, 0.5*255, 0.5*255], std=[0.5*255, 0.5*255, 0.5*255], inplace=True),
+                L(CreateCondition)(),
+            ),
+            infinite=True,
+            shuffle=True,
+            partial_batches=False,
+            load_embedding=True,
+            keypoints_split="train",
+            load_new_keypoints=False
+        )
+    ),
+    val=dict(
+        loader=L(get_dataloader_fdh_wds)(
+            path=data_dir.joinpath("val", "out-{000000..000023}.tar"),
+            batch_size="${train.batch_size}",
+            num_workers=6,
+            transform=None,
+            gpu_transform="${data.train.loader.gpu_transform}",
+            infinite=False,
+            shuffle=False,
+            partial_batches=True,
+            load_embedding=True,
+            keypoints_split="val",
+            load_new_keypoints="${data.train.loader.load_new_keypoints}"
+        )
+    ),
+    # Training evaluation might do optimizations to reduce compute overhead. E.g. compute with AMP.
+    train_evaluation_fn=L(functools.partial)(
+        train_eval_fn, cache_directory=L(get_cache_directory)(imsize="${data.imsize}", subset="fdh"),
+        data_len=30_000),
+    evaluation_fn=L(functools.partial)(
+        final_eval_fn, cache_directory=L(get_cache_directory)(imsize="${data.imsize}", subset="fdh_eval"), 
+        data_len=30_000)
+)

deep_privacy/configs/datasets/utils.py

@@ -0,0 +1,21 @@
+from dp2.metrics.ppl import calculate_ppl
+from dp2.metrics.torch_metrics import compute_metrics_iteratively
+from dp2.metrics.fid_clip import compute_fid_clip
+
+
+def final_eval_fn(*args, **kwargs):
+    result = compute_metrics_iteratively(*args, **kwargs)
+    result2 = calculate_ppl(*args, **kwargs,)
+    result3 = compute_fid_clip(*args, **kwargs)
+    assert all(key not in result for key in result2)
+    result.update(result2)
+    result.update(result3)
+    return result
+
+
+def train_eval_fn(*args, **kwargs):
+    result = compute_metrics_iteratively(*args, **kwargs)
+    result2 = compute_fid_clip(*args, **kwargs)
+    assert all(key not in result for key in result2)
+    result.update(result2)
+    return result

deep_privacy/configs/defaults.py

@@ -0,0 +1,53 @@
+import pathlib
+import os
+import torch
+from tops.config import LazyCall as L
+
+if "PRETRAINED_CHECKPOINTS_PATH" in os.environ:
+    PRETRAINED_CHECKPOINTS_PATH = pathlib.Path(os.environ["PRETRAINED_CHECKPOINTS_PATH"])
+else:
+    PRETRAINED_CHECKPOINTS_PATH = pathlib.Path("pretrained_checkpoints")
+if "BASE_OUTPUT_DIR" in os.environ:
+    BASE_OUTPUT_DIR = pathlib.Path(os.environ["BASE_OUTPUT_DIR"])
+else:
+    BASE_OUTPUT_DIR = pathlib.Path("outputs")
+
+
+
+common = dict(
+    logger_backend=["wandb", "stdout", "json", "image_dumper"],
+    wandb_project="deep_privacy2",
+    output_dir=BASE_OUTPUT_DIR,
+    experiment_name=None, # Optional experiment name to show on wandb
+)
+
+train = dict(
+    batch_size=32,
+    seed=0,
+    ims_per_log=1024,
+    ims_per_val=int(200e3),
+    max_images_to_train=int(12e6),
+    amp=dict(
+        enabled=True,
+        scaler_D=L(torch.cuda.amp.GradScaler)(init_scale=2**16, growth_factor=4, growth_interval=100, enabled="${..enabled}"),
+        scaler_G=L(torch.cuda.amp.GradScaler)(init_scale=2**16, growth_factor=4, growth_interval=100, enabled="${..enabled}"),
+    ),
+    fp16_ddp_accumulate=False, # All gather gradients in fp16?
+    broadcast_buffers=False,
+    bias_act_plugin_enabled=True,
+    grid_sample_gradfix_enabled=True,
+    conv2d_gradfix_enabled=False,
+    channels_last=False,
+    compile_G=dict(
+        enabled=False,
+        mode="default" # default, reduce-overhead or max-autotune
+    ),
+    compile_D=dict(
+        enabled=False,
+        mode="default" # default, reduce-overhead or max-autotune
+    )
+)
+
+# exponential moving average
+EMA = dict(rampup=0.05)
+

deep_privacy/configs/discriminators/sg2_discriminator.py

@@ -0,0 +1,43 @@
+from tops.config import LazyCall as L
+from dp2.discriminator import SG2Discriminator
+import torch
+from dp2.loss import StyleGAN2Loss
+
+
+discriminator = L(SG2Discriminator)(
+    imsize="${data.imsize}",
+    im_channels="${data.im_channels}",
+    min_fmap_resolution=4,
+    max_cnum_mul=8,
+    cnum=80,
+    input_condition=True,
+    conv_clamp=256,
+    input_cse=False,
+    cse_nc="${data.cse_nc}",
+    fix_residual=False,
+)
+
+
+loss_fnc = L(StyleGAN2Loss)(
+    lazy_regularization=True,
+    lazy_reg_interval=16,
+    r1_opts=dict(lambd=5, mask_out=False, mask_out_scale=False),
+    EP_lambd=0.001,
+    pl_reg_opts=dict(weight=0, batch_shrink=2,start_nimg=int(1e6), pl_decay=0.01)
+)
+
+def build_D_optim(type, lr, betas, lazy_regularization, lazy_reg_interval, **kwargs):
+    if lazy_regularization:
+        # From Analyzing and improving the image quality of stylegan, CVPR 2020
+        c = lazy_reg_interval / (lazy_reg_interval + 1)
+        betas = [beta ** c for beta in betas]
+        lr *= c
+        print(f"Lazy regularization on. Setting lr to: {lr}, betas to: {betas}")
+    return type(lr=lr, betas=betas, **kwargs)
+
+
+D_optim = L(build_D_optim)(
+    type=torch.optim.Adam, lr=0.001, betas=(0.0, 0.99),
+    lazy_regularization="${loss_fnc.lazy_regularization}",
+    lazy_reg_interval="${loss_fnc.lazy_reg_interval}")
+G_optim = L(torch.optim.Adam)(lr=0.001, betas=(0.0, 0.99))

deep_privacy/configs/fdf/deep_privacy1.py

@@ -0,0 +1,9 @@
+from tops.config import LazyCall as L 
+from dp2.generator.deep_privacy1 import MSGGenerator
+from ..datasets.fdf128 import data
+from ..defaults import common, train
+
+generator = L(MSGGenerator)()
+
+common.model_url = "https://folk.ntnu.no/haakohu/checkpoints/fdf128_model512.ckpt"
+common.model_md5sum = "6cc8b285bdc1fcdfc64f5db7c521d0a6"

deep_privacy/configs/fdf/stylegan.py

@@ -0,0 +1,14 @@
+from ..generators.stylegan_unet import generator
+from ..datasets.fdf256 import data
+from ..discriminators.sg2_discriminator import discriminator, G_optim, D_optim, loss_fnc
+from ..defaults import train, common, EMA
+
+train.max_images_to_train = int(35e6)
+G_optim.lr = 0.002
+D_optim.lr = 0.002
+generator.input_cse = False
+loss_fnc.r1_opts.lambd = 1
+train.ims_per_val = int(2e6)
+
+common.model_url = "https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/89660f04-5c11-4dbf-adac-cbe2f11b0aeea25cbf78-7558-475a-b3c7-03f5c10b7934646b0720-ca0a-4d53-aded-daddbfa45c9e"
+common.model_md5sum = "e8e32190528af2ed75f0cb792b7f2b07"

deep_privacy/configs/fdf/stylegan_fdf128.py

@@ -0,0 +1,17 @@
+from ..discriminators.sg2_discriminator import discriminator, G_optim, D_optim, loss_fnc
+from ..datasets.fdf128 import data
+from ..generators.stylegan_unet import generator
+from ..defaults import train, common, EMA
+from tops.config import LazyCall as L
+
+G_optim.lr = 0.002
+D_optim.lr = 0.002
+generator.update(cnum=128, max_cnum_mul=4, input_cse=False)
+loss_fnc.r1_opts.lambd = 0.1
+
+train.update(ims_per_val=int(2e6), batch_size=64, max_images_to_train=int(35e6))
+
+common.update(
+    model_url="https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/66d803c0-55ce-44c0-9d53-815c2c0e6ba4eb458409-9e91-45d1-bce0-95c8a47a57218b102fdf-bea3-44dc-aac4-0fb1d370ef1c",
+    model_md5sum="bccd4403e7c9bca682566ff3319e8176"
+)

deep_privacy/configs/fdh/styleganL.py

@@ -0,0 +1,16 @@
+from tops.config import LazyCall as L
+from ..generators.stylegan_unet import generator
+from ..datasets.fdh import data
+from ..discriminators.sg2_discriminator import discriminator, G_optim, D_optim, loss_fnc
+from ..defaults import train, common, EMA
+
+train.max_images_to_train = int(50e6)
+train.batch_size = 64
+G_optim.lr = 0.002
+D_optim.lr = 0.002
+data.train.loader.num_workers = 4
+train.ims_per_val = int(1e6)
+loss_fnc.r1_opts.lambd = .1
+
+common.model_url = "https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/21841da7-2546-4ce3-8460-909b3a63c58b13aac1a1-c778-4c8d-9b69-3e5ed2cde9de1524e76e-7aa6-4dd8-b643-52abc9f0792c"
+common.model_md5sum = "3411478b5ec600a4219cccf4499732bd"

deep_privacy/configs/fdh/styleganL_nocse.py

@@ -0,0 +1,14 @@
+from tops.config import LazyCall as L
+from ..generators.stylegan_unet import generator
+from ..datasets.fdh import data
+from ..discriminators.sg2_discriminator import discriminator, G_optim, D_optim, loss_fnc
+from ..defaults import train, common, EMA
+
+train.max_images_to_train = int(50e6)
+G_optim.lr = 0.002
+D_optim.lr = 0.002
+generator.input_cse = False
+data.load_embeddings = False
+common.model_url = "https://folk.ntnu.no/haakohu/checkpoints/deep_privacy2/fdh_styleganL_nocse.ckpt"
+common.model_md5sum = "fda0d809741bc67487abada793975c37"
+generator.fix_errors = False

deep_privacy/configs/generators/stylegan_unet.py

@@ -0,0 +1,22 @@
+from dp2.generator.stylegan_unet import StyleGANUnet
+from tops.config import LazyCall as L
+
+generator = L(StyleGANUnet)(
+    imsize="${data.imsize}",
+    im_channels="${data.im_channels}",
+    min_fmap_resolution=8,
+    cnum=64,
+    max_cnum_mul=8,
+    n_middle_blocks=0,
+    z_channels=512,
+    mask_output=True,
+    conv_clamp=256,
+    input_cse=True,
+    scale_grad=True,
+    cse_nc="${data.cse_nc}",
+    w_dim=512,
+    n_keypoints="${data.n_keypoints}",
+    input_keypoints=False,
+    input_keypoint_indices=[],
+    fix_errors=True
+)

deep_privacy/dp2/anonymizer/__init__.py

@@ -0,0 +1 @@
+from .anonymizer import Anonymizer

deep_privacy/dp2/anonymizer/anonymizer.py

@@ -0,0 +1,163 @@
+from pathlib import Path
+from typing import Union, Optional
+import numpy as np
+import torch
+import tops
+import torchvision.transforms.functional as F
+from motpy import Detection, MultiObjectTracker
+from dp2.utils import load_config
+from dp2.infer import build_trained_generator
+from dp2.detection.structures import CSEPersonDetection, FaceDetection, PersonDetection, VehicleDetection
+
+
+def load_generator_from_cfg_path(cfg_path: Union[str, Path]):
+    cfg = load_config(cfg_path)
+    G = build_trained_generator(cfg)
+    tops.logger.log(f"Loaded generator from: {cfg_path}")
+    return G
+
+
+class Anonymizer:
+
+    def __init__(
+            self,
+            detector,
+            load_cache: bool = False,
+            person_G_cfg: Optional[Union[str, Path]] = None,
+            cse_person_G_cfg: Optional[Union[str, Path]] = None,
+            face_G_cfg: Optional[Union[str, Path]] = None,
+            car_G_cfg: Optional[Union[str, Path]] = None,
+    ) -> None:
+        self.detector = detector
+        self.generators = {k: None for k in [CSEPersonDetection, PersonDetection, FaceDetection, VehicleDetection]}
+        self.load_cache = load_cache
+        if cse_person_G_cfg is not None:
+            self.generators[CSEPersonDetection] = load_generator_from_cfg_path(cse_person_G_cfg)
+        if person_G_cfg is not None:
+            self.generators[PersonDetection] = load_generator_from_cfg_path(person_G_cfg)
+        if face_G_cfg is not None:
+            self.generators[FaceDetection] = load_generator_from_cfg_path(face_G_cfg)
+        if car_G_cfg is not None:
+            self.generators[VehicleDetection] = load_generator_from_cfg_path(car_G_cfg)
+
+    def initialize_tracker(self, fps: float):
+        self.tracker = MultiObjectTracker(dt=1/fps)
+        self.track_to_z_idx = dict()
+
+    def reset_tracker(self):
+        self.track_to_z_idx = dict()
+
+    def forward_G(self,
+                  G,
+                  batch,
+                  multi_modal_truncation: bool,
+                  amp: bool,
+                  z_idx: int,
+                  truncation_value: float,
+                  idx: int,
+                  all_styles=None):
+        batch["img"] = F.normalize(batch["img"].float(), [0.5*255, 0.5*255, 0.5*255], [0.5*255, 0.5*255, 0.5*255])
+        batch["img"] = batch["img"].float()
+        batch["condition"] = batch["mask"].float() * batch["img"]
+
+        with torch.cuda.amp.autocast(amp):
+            z = None
+            if z_idx is not None:
+                state = np.random.RandomState(seed=z_idx[idx])
+                z = state.normal(size=(1, G.z_channels)).astype(np.float32)
+                z = tops.to_cuda(torch.from_numpy(z))
+
+            if all_styles is not None:
+                anonymized_im = G(**batch, s=iter(all_styles[idx]))["img"]
+            elif multi_modal_truncation:
+                w_indices = None
+                if z_idx is not None:
+                    w_indices = [z_idx[idx] % len(G.style_net.w_centers)]
+                anonymized_im = G.multi_modal_truncate(
+                    **batch, truncation_value=truncation_value,
+                    w_indices=w_indices,
+                    z=z
+                )["img"]
+            else:
+                anonymized_im = G.sample(**batch, truncation_value=truncation_value, z=z)["img"]
+        anonymized_im = (anonymized_im+1).div(2).clamp(0, 1).mul(255)
+        return anonymized_im
+
+    @torch.no_grad()
+    def anonymize_detections(self,
+                             im, detection,
+                             update_identity=None,
+                             **synthesis_kwargs
+                             ):
+        G = self.generators[type(detection)]
+        if G is None:
+            return im
+        C, H, W = im.shape
+        if update_identity is None:
+            update_identity = [True for i in range(len(detection))]
+        for idx in range(len(detection)):
+            if not update_identity[idx]:
+                continue
+            batch = detection.get_crop(idx, im)
+            x0, y0, x1, y1 = batch.pop("boxes")[0]
+            batch = {k: tops.to_cuda(v) for k, v in batch.items()}
+            anonymized_im = self.forward_G(G, batch, **synthesis_kwargs, idx=idx)
+
+            gim = F.resize(anonymized_im[0], (y1-y0, x1-x0), interpolation=F.InterpolationMode.BICUBIC, antialias=True)
+            mask = F.resize(batch["mask"][0], (y1-y0, x1-x0), interpolation=F.InterpolationMode.NEAREST).squeeze(0)
+            # Remove padding
+            pad = [max(-x0, 0), max(-y0, 0)]
+            pad = [*pad, max(x1-W, 0), max(y1-H, 0)]
+            def remove_pad(x): return x[..., pad[1]:x.shape[-2]-pad[3], pad[0]:x.shape[-1]-pad[2]]
+
+            gim = remove_pad(gim)
+            mask = remove_pad(mask) > 0.5
+            x0, y0 = max(x0, 0), max(y0, 0)
+            x1, y1 = min(x1, W), min(y1, H)
+            mask = mask.logical_not()[None].repeat(3, 1, 1)
+
+            im[:, y0:y1, x0:x1][mask] = gim[mask].round().clamp(0, 255).byte()
+        return im
+
+    def visualize_detection(self, im: torch.Tensor, cache_id: str = None) -> torch.Tensor:
+        all_detections = self.detector.forward_and_cache(im, cache_id, load_cache=self.load_cache)
+        im = im.cpu()
+        for det in all_detections:
+            im = det.visualize(im)
+        return im
+
+    @torch.no_grad()
+    def forward(self, im: torch.Tensor, cache_id: str = None, track=True, detections=None, **synthesis_kwargs) -> torch.Tensor:
+        assert im.dtype == torch.uint8
+        im = tops.to_cuda(im)
+        all_detections = detections
+        if detections is None:
+            if self.load_cache:
+                all_detections = self.detector.forward_and_cache(im, cache_id)
+            else:
+                all_detections = self.detector(im)
+        if hasattr(self, "tracker") and track:
+            [_.pre_process() for _ in all_detections]
+            boxes = np.concatenate([_.boxes for _ in all_detections])
+            boxes = [Detection(box) for box in boxes]
+            self.tracker.step(boxes)
+            track_ids = self.tracker.detections_matched_ids
+            z_idx = []
+            for track_id in track_ids:
+                if track_id not in self.track_to_z_idx:
+                    self.track_to_z_idx[track_id] = np.random.randint(0, 2**32-1)
+                z_idx.append(self.track_to_z_idx[track_id])
+            z_idx = np.array(z_idx)
+            idx_offset = 0
+
+        for detection in all_detections:
+            zs = None
+            if hasattr(self, "tracker") and track:
+                zs = z_idx[idx_offset:idx_offset+len(detection)]
+                idx_offset += len(detection)
+            im = self.anonymize_detections(im, detection, z_idx=zs, **synthesis_kwargs)
+
+        return im.cpu()
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)

deep_privacy/dp2/anonymizer/histogram_match_anonymizers.py

@@ -0,0 +1,93 @@
+
+import torch
+import tops
+import numpy as np
+from kornia.color import rgb_to_hsv
+from dp2 import utils
+from kornia.enhance import histogram
+from .anonymizer import Anonymizer
+import torchvision.transforms.functional as F
+from skimage.exposure import match_histograms
+from kornia.filters import gaussian_blur2d
+
+
+class LatentHistogramMatchAnonymizer(Anonymizer):
+
+    def forward_G(
+        self,
+        G,
+        batch,
+        multi_modal_truncation: bool,
+        amp: bool,
+        z_idx: int,
+        truncation_value: float,
+        idx: int,
+        n_sampling_steps: int = 1,
+        all_styles=None,
+    ):
+        batch["img"] = F.normalize(batch["img"].float(), [0.5*255, 0.5*255, 0.5*255], [0.5*255, 0.5*255, 0.5*255])
+        batch["img"] = batch["img"].float()
+        batch["condition"] = batch["mask"].float() * batch["img"]
+
+        assert z_idx is None and all_styles is None, "Arguments not supported with n_sampling_steps > 1."
+        real_hls = rgb_to_hsv(utils.denormalize_img(batch["img"]))
+        real_hls[:, 0] /= 2 * torch.pi
+        indices = [1, 2]
+        hist_kwargs = dict(
+            bins=torch.linspace(0, 1, 256, dtype=torch.float32, device=tops.get_device()),
+            bandwidth=torch.tensor(1., device=tops.get_device()))
+        real_hist = [histogram(real_hls[:, i].flatten(start_dim=1), **hist_kwargs) for i in indices]
+        for j in range(n_sampling_steps):
+            if j == 0:
+                if multi_modal_truncation:
+                    w = G.style_net.multi_modal_truncate(
+                        truncation_value=truncation_value, **batch, w_indices=None).detach()
+                else:
+                    w = G.style_net.get_truncated(truncation_value, **batch).detach()
+                assert z_idx is None and all_styles is None, "Arguments not supported with n_sampling_steps > 1."
+                w.requires_grad = True
+                optim = torch.optim.Adam([w])
+            with torch.set_grad_enabled(True):
+                with torch.cuda.amp.autocast(amp):
+                    anonymized_im = G(**batch, truncation_value=None, w=w)["img"]
+                fake_hls = rgb_to_hsv(anonymized_im*0.5 + 0.5)
+                fake_hls[:, 0] /= 2 * torch.pi
+                fake_hist = [histogram(fake_hls[:, i].flatten(start_dim=1), **hist_kwargs) for i in indices]
+                dist = sum([utils.torch_wasserstein_loss(r, f) for r, f in zip(real_hist, fake_hist)])
+                dist.backward()
+                if w.grad.sum() == 0:
+                    break
+                assert w.grad.sum() != 0
+                optim.step()
+                optim.zero_grad()
+                if dist < 0.02:
+                    break
+        anonymized_im = (anonymized_im+1).div(2).clamp(0, 1).mul(255)
+        return anonymized_im
+
+
+class HistogramMatchAnonymizer(Anonymizer):
+
+    def forward_G(self, batch, *args, **kwargs):
+        rimg = batch["img"]
+        batch["img"] = F.normalize(batch["img"].float(), [0.5*255, 0.5*255, 0.5*255], [0.5*255, 0.5*255, 0.5*255])
+        batch["img"] = batch["img"].float()
+        batch["condition"] = batch["mask"].float() * batch["img"]
+
+        anonymized_im = super().forward_G(batch, *args, **kwargs)
+
+        equalized_gim = match_histograms(tops.im2numpy(anonymized_im.round().clamp(0, 255).byte()), tops.im2numpy(rimg))
+        if equalized_gim.dtype != np.uint8:
+            equalized_gim = equalized_gim.astype(np.float32)
+            assert equalized_gim.dtype == np.float32, equalized_gim.dtype
+            equalized_gim = tops.im2torch(equalized_gim, to_float=False)[0]
+        else:
+            equalized_gim = tops.im2torch(equalized_gim, to_float=False).float()[0]
+        equalized_gim = equalized_gim.to(device=rimg.device)
+        assert equalized_gim.dtype == torch.float32
+        gaussian_mask = 1 - (batch["maskrcnn_mask"][0].repeat(3, 1, 1) > 0.5).float()
+
+        gaussian_mask = gaussian_blur2d(gaussian_mask[None], kernel_size=[19, 19], sigma=[10, 10])[0]
+        gaussian_mask = gaussian_mask / gaussian_mask.max()
+        anonymized_im = gaussian_mask * equalized_gim + (1-gaussian_mask) * anonymized_im
+        return anonymized_im

deep_privacy/dp2/data/build.py

@@ -0,0 +1,40 @@
+import torch
+import tops
+from .utils import collate_fn
+
+
+def get_dataloader(
+        dataset, gpu_transform: torch.nn.Module,
+        num_workers,
+        batch_size,
+        infinite: bool,
+        drop_last: bool,
+        prefetch_factor: int,
+        shuffle,
+        channels_last=False
+    ):
+    sampler = None
+    dl_kwargs = dict(
+        pin_memory=True,
+    )
+    if infinite:
+        sampler = tops.InfiniteSampler(
+            dataset, rank=tops.rank(),
+            num_replicas=tops.world_size(),
+            shuffle=shuffle
+        )
+    elif tops.world_size() > 1:
+        sampler = torch.utils.data.DistributedSampler(
+            dataset, shuffle=shuffle, num_replicas=tops.world_size(), rank=tops.rank())
+        dl_kwargs["drop_last"] = drop_last
+    else:
+        dl_kwargs["shuffle"] = shuffle
+        dl_kwargs["drop_last"] = drop_last
+    dataloader = torch.utils.data.DataLoader(
+        dataset, sampler=sampler, collate_fn=collate_fn,
+        batch_size=batch_size,
+        num_workers=num_workers, prefetch_factor=prefetch_factor,
+        **dl_kwargs
+    )
+    dataloader = tops.DataPrefetcher(dataloader, gpu_transform, channels_last=channels_last)
+    return dataloader

deep_privacy/dp2/data/datasets/coco_cse.py

@@ -0,0 +1,68 @@
+import pickle
+import torchvision
+import torch
+import pathlib
+import numpy as np
+from typing import Callable, Optional, Union
+from torch.hub import get_dir as get_hub_dir
+
+
+def cache_embed_stats(embed_map: torch.Tensor):
+    mean = embed_map.mean(dim=0, keepdim=True)
+    rstd = ((embed_map - mean).square().mean(dim=0, keepdim=True)+1e-8).rsqrt()
+
+    cache = dict(mean=mean, rstd=rstd, embed_map=embed_map)
+    path = pathlib.Path(get_hub_dir(), f"embed_map_stats.torch")
+    path.parent.mkdir(exist_ok=True, parents=True)
+    torch.save(cache, path)
+
+
+class CocoCSE(torch.utils.data.Dataset):
+
+    def __init__(self,
+                 dirpath: Union[str, pathlib.Path],
+                 transform: Optional[Callable],
+                 normalize_E: bool,):
+        dirpath = pathlib.Path(dirpath)
+        self.dirpath = dirpath
+
+        self.transform = transform
+        assert self.dirpath.is_dir(),\
+            f"Did not find dataset at: {dirpath}"
+        self.image_paths, self.embedding_paths = self._load_impaths()
+        self.embed_map = torch.from_numpy(np.load(self.dirpath.joinpath("embed_map.npy")))
+        mean = self.embed_map.mean(dim=0, keepdim=True)
+        rstd = ((self.embed_map - mean).square().mean(dim=0, keepdim=True)+1e-8).rsqrt()
+        self.embed_map = (self.embed_map - mean) * rstd
+        cache_embed_stats(self.embed_map)
+
+    def _load_impaths(self):
+        image_dir = self.dirpath.joinpath("images")
+        image_paths = list(image_dir.glob("*.png"))
+        image_paths.sort()
+        embedding_paths = [
+            self.dirpath.joinpath("embedding", x.stem + ".npy") for x in image_paths
+        ]
+        return image_paths, embedding_paths
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        im = torchvision.io.read_image(str(self.image_paths[idx]))
+        vertices, mask, border = np.split(np.load(self.embedding_paths[idx]), 3, axis=-1)
+        vertices = torch.from_numpy(vertices.squeeze()).long()
+        mask = torch.from_numpy(mask.squeeze()).float()
+        border = torch.from_numpy(border.squeeze()).float()
+        E_mask = 1 - mask - border
+        batch = {
+            "img": im,
+            "vertices": vertices[None],
+            "mask": mask[None],
+            "embed_map": self.embed_map,
+            "border": border[None],
+            "E_mask": E_mask[None]
+        }
+        if self.transform is None:
+            return batch
+        return self.transform(batch)

deep_privacy/dp2/data/datasets/fdf.py

@@ -0,0 +1,128 @@
+import pathlib
+from typing import Tuple
+import numpy as np
+import torch
+import pathlib
+try:
+    import pyspng
+    PYSPNG_IMPORTED = True
+except ImportError:
+    PYSPNG_IMPORTED = False
+    print("Could not load pyspng. Defaulting to pillow image backend.")
+    from PIL import Image
+from tops import logger
+
+
+class FDFDataset:
+
+    def __init__(self,
+                 dirpath,
+                 imsize: Tuple[int],
+                 load_keypoints: bool,
+                 transform):
+        dirpath = pathlib.Path(dirpath)
+        self.dirpath = dirpath
+        self.transform = transform
+        self.imsize = imsize[0]
+        self.load_keypoints = load_keypoints
+        assert self.dirpath.is_dir(),\
+            f"Did not find dataset at: {dirpath}"
+        image_dir = self.dirpath.joinpath("images", str(self.imsize))
+        self.image_paths = list(image_dir.glob("*.png"))
+        assert len(self.image_paths) > 0,\
+            f"Did not find images in: {image_dir}"
+        self.image_paths.sort(key=lambda x: int(x.stem))
+        self.landmarks = np.load(self.dirpath.joinpath("landmarks.npy")).reshape(-1, 7, 2).astype(np.float32)
+
+        self.bounding_boxes = torch.load(self.dirpath.joinpath("bounding_box", f"{self.imsize}.torch"))
+        assert len(self.image_paths) == len(self.bounding_boxes)
+        assert len(self.image_paths) == len(self.landmarks)
+        logger.log(
+            f"Dataset loaded from: {dirpath}. Number of samples:{len(self)}, imsize={imsize}")
+
+    def get_mask(self, idx):
+        mask = torch.ones((1, self.imsize, self.imsize), dtype=torch.bool)
+        bounding_box = self.bounding_boxes[idx]
+        x0, y0, x1, y1 = bounding_box
+        mask[:, y0:y1, x0:x1] = 0
+        return mask
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, index):
+        impath = self.image_paths[index]
+        if PYSPNG_IMPORTED:
+            with open(impath, "rb") as fp:
+                im = pyspng.load(fp.read())
+        else:
+            with Image.open(impath) as fp:
+                im = np.array(fp)
+        im = torch.from_numpy(np.rollaxis(im, -1, 0))
+        masks = self.get_mask(index)
+        landmark = self.landmarks[index]
+        batch = {
+            "img": im,
+            "mask": masks,
+        }
+        if self.load_keypoints:
+            batch["keypoints"] = landmark
+        if self.transform is None:
+            return batch
+        return self.transform(batch)
+
+
+class FDF256Dataset:
+
+    def __init__(self,
+                 dirpath,
+                 load_keypoints: bool,
+                 transform):
+        dirpath = pathlib.Path(dirpath)
+        self.dirpath = dirpath
+        self.transform = transform
+        self.load_keypoints = load_keypoints
+        assert self.dirpath.is_dir(),\
+            f"Did not find dataset at: {dirpath}"
+        image_dir = self.dirpath.joinpath("images")
+        self.image_paths = list(image_dir.glob("*.png"))
+        assert len(self.image_paths) > 0,\
+            f"Did not find images in: {image_dir}"
+        self.image_paths.sort(key=lambda x: int(x.stem))
+        self.landmarks = np.load(self.dirpath.joinpath("landmarks.npy")).reshape(-1, 7, 2).astype(np.float32)
+        self.bounding_boxes = torch.from_numpy(np.load(self.dirpath.joinpath("bounding_box.npy")))
+        assert len(self.image_paths) == len(self.bounding_boxes)
+        assert len(self.image_paths) == len(self.landmarks)
+        logger.log(
+            f"Dataset loaded from: {dirpath}. Number of samples:{len(self)}")
+
+    def get_mask(self, idx):
+        mask = torch.ones((1, 256, 256), dtype=torch.bool)
+        bounding_box = self.bounding_boxes[idx]
+        x0, y0, x1, y1 = bounding_box
+        mask[:, y0:y1, x0:x1] = 0
+        return mask
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, index):
+        impath = self.image_paths[index]
+        if PYSPNG_IMPORTED:
+            with open(impath, "rb") as fp:
+                im = pyspng.load(fp.read())
+        else:
+            with Image.open(impath) as fp:
+                im = np.array(fp)
+        im = torch.from_numpy(np.rollaxis(im, -1, 0))
+        masks = self.get_mask(index)
+        landmark = self.landmarks[index]
+        batch = {
+            "img": im,
+            "mask": masks,
+        }
+        if self.load_keypoints:
+            batch["keypoints"] = landmark
+        if self.transform is None:
+            return batch
+        return self.transform(batch)

deep_privacy/dp2/data/datasets/fdf128_wds.py

@@ -0,0 +1,96 @@
+import torch
+import tops
+import numpy as np
+import io
+import webdataset as wds
+import os
+from ..utils import png_decoder, get_num_workers, collate_fn
+
+
+def kp_decoder(x):
+    # Keypoints are between [0, 1] for webdataset
+    keypoints = torch.from_numpy(np.load(io.BytesIO(x))).float().view(7, 2).clamp(0, 1)
+    keypoints = torch.cat((keypoints, torch.ones((7, 1))), dim=-1)
+    return keypoints
+
+
+def bbox_decoder(x):
+    return torch.from_numpy(np.load(io.BytesIO(x))).float().view(4)
+
+
+class BBoxToMask:
+
+    def __call__(self, sample):
+        imsize = sample["image.png"].shape[-1]
+        bbox = sample["bounding_box.npy"] * imsize
+        x0, y0, x1, y1 = np.round(bbox).astype(np.int64)
+        mask = torch.ones((1, imsize, imsize), dtype=torch.bool)
+        mask[:, y0:y1, x0:x1] = 0
+        sample["mask"] = mask
+        return sample
+
+
+def get_dataloader_fdf_wds(
+        path,
+        batch_size: int,
+        num_workers: int,
+        transform: torch.nn.Module,
+        gpu_transform: torch.nn.Module,
+        infinite: bool,
+        shuffle: bool,
+        partial_batches: bool,
+        sample_shuffle=10_000,
+        tar_shuffle=100,
+        channels_last=False,
+    ):
+    # Need to set this for split_by_node to work.
+    os.environ["RANK"] = str(tops.rank())
+    os.environ["WORLD_SIZE"] = str(tops.world_size())
+    if infinite:
+        pipeline = [wds.ResampledShards(str(path))]
+    else:
+        pipeline = [wds.SimpleShardList(str(path))]
+    if shuffle:
+        pipeline.append(wds.shuffle(tar_shuffle))
+    pipeline.extend([
+        wds.split_by_node,
+        wds.split_by_worker,
+    ])
+    if shuffle:
+        pipeline.append(wds.shuffle(sample_shuffle))
+
+    decoder = [
+        wds.handle_extension("image.png", png_decoder),
+        wds.handle_extension("keypoints.npy", kp_decoder),
+    ]
+
+    rename_keys = [
+        ["img", "image.png"],
+        ["keypoints", "keypoints.npy"],
+        ["__key__", "__key__"],
+        ["mask", "mask"]
+    ]
+
+    pipeline.extend([
+        wds.tarfile_to_samples(),
+        wds.decode(*decoder),
+    ])
+    pipeline.append(wds.map(BBoxToMask()))
+    pipeline.extend([
+        wds.batched(batch_size, collation_fn=collate_fn, partial=partial_batches),
+        wds.rename_keys(*rename_keys),
+    ])
+
+    if transform is not None:
+        pipeline.append(wds.map(transform))
+    pipeline = wds.DataPipeline(*pipeline)
+    if infinite:
+        pipeline = pipeline.repeat(nepochs=1000000)
+
+    loader = wds.WebLoader(
+        pipeline, batch_size=None, shuffle=False,
+        num_workers=get_num_workers(num_workers),
+        persistent_workers=True,
+    )
+    loader = tops.DataPrefetcher(loader, gpu_transform, channels_last=channels_last, to_float=False)
+    return loader

deep_privacy/dp2/data/datasets/fdh.py

@@ -0,0 +1,142 @@
+import torch
+import tops
+import numpy as np
+import io
+import webdataset as wds
+import os
+import json
+from pathlib import Path
+from ..utils import png_decoder, mask_decoder, get_num_workers, collate_fn
+
+
+def kp_decoder(x):
+    # Keypoints are between [0, 1] for webdataset
+    keypoints = torch.from_numpy(np.load(io.BytesIO(x))).float()
+    def check_outside(x): return (x < 0).logical_or(x > 1)
+    is_outside = check_outside(keypoints[:, 0]).logical_or(
+        check_outside(keypoints[:, 1])
+    )
+    keypoints[:, 2] = (keypoints[:, 2] > 0).logical_and(is_outside.logical_not())
+    return keypoints
+
+
+def vertices_decoder(x):
+    vertices = torch.from_numpy(np.load(io.BytesIO(x)).astype(np.int32))
+    return vertices.squeeze()[None]
+
+
+class InsertNewKeypoints:
+
+    def __init__(self, keypoints_path: Path) -> None:
+        with open(keypoints_path, "r") as fp:
+            self.keypoints = json.load(fp)
+
+    def __call__(self, sample):
+        key = sample["__key__"]
+        keypoints = torch.tensor(self.keypoints[key], dtype=torch.float32)
+        def check_outside(x): return (x < 0).logical_or(x > 1)
+        is_outside = check_outside(keypoints[:, 0]).logical_or(
+            check_outside(keypoints[:, 1])
+        )
+        keypoints[:, 2] = (keypoints[:, 2] > 0).logical_and(is_outside.logical_not())
+
+        sample["keypoints.npy"] = keypoints
+        return sample
+
+
+def get_dataloader_fdh_wds(
+        path,
+        batch_size: int,
+        num_workers: int,
+        transform: torch.nn.Module,
+        gpu_transform: torch.nn.Module,
+        infinite: bool,
+        shuffle: bool,
+        partial_batches: bool,
+        load_embedding: bool,
+        sample_shuffle=10_000,
+        tar_shuffle=100,
+        read_condition=False,
+        channels_last=False,
+        load_new_keypoints=False,
+        keypoints_split=None,
+    ):
+    # Need to set this for split_by_node to work.
+    os.environ["RANK"] = str(tops.rank())
+    os.environ["WORLD_SIZE"] = str(tops.world_size())
+    if infinite:
+        pipeline = [wds.ResampledShards(str(path))]
+    else:
+        pipeline = [wds.SimpleShardList(str(path))]
+    if shuffle:
+        pipeline.append(wds.shuffle(tar_shuffle))
+    pipeline.extend([
+        wds.split_by_node,
+        wds.split_by_worker,
+    ])
+    if shuffle:
+        pipeline.append(wds.shuffle(sample_shuffle))
+
+    decoder = [
+        wds.handle_extension("image.png", png_decoder),
+        wds.handle_extension("mask.png", mask_decoder),
+        wds.handle_extension("maskrcnn_mask.png", mask_decoder),
+        wds.handle_extension("keypoints.npy", kp_decoder),
+    ]
+
+    rename_keys = [
+        ["img", "image.png"], ["mask", "mask.png"],
+        ["keypoints", "keypoints.npy"], ["maskrcnn_mask", "maskrcnn_mask.png"],
+        ["__key__", "__key__"]
+    ]
+    if load_embedding:
+        decoder.extend([
+            wds.handle_extension("vertices.npy", vertices_decoder),
+            wds.handle_extension("E_mask.png", mask_decoder)
+        ])
+        rename_keys.extend([
+            ["vertices", "vertices.npy"],
+            ["E_mask", "e_mask.png"]
+        ])
+
+    if read_condition:
+        decoder.append(
+            wds.handle_extension("condition.png", png_decoder)
+        )
+        rename_keys.append(["condition", "condition.png"])
+
+    pipeline.extend([
+        wds.tarfile_to_samples(),
+        wds.decode(*decoder),
+
+    ])
+    if load_new_keypoints:
+        assert keypoints_split in ["train", "val"]
+        keypoint_url = "https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/1eb88522-8b91-49c7-b56a-ed98a9c7888cef9c0429-a385-4248-abe3-8682de26d041f268aed1-7c88-4677-baad-7623c2ee330f"
+        file_name = "fdh_keypoints_val-050133b34d.json"
+        if keypoints_split == "train":
+            keypoint_url = "https://api.loke.aws.unit.no/dlr-gui-backend-resources-content/v2/contents/links/3e828b1c-d6c0-4622-90bc-1b2cce48ccfff14ab45d-0a5c-431d-be13-7e60580765bd7938601c-e72e-41d9-8836-fffc49e76f58"
+            file_name = "fdh_keypoints_train-2cff11f69a.json"
+        # Set check_hash=True if you suspect download is incorrect.
+        filepath = tops.download_file(keypoint_url, file_name=file_name, check_hash=False)
+        pipeline.append(
+            wds.map(InsertNewKeypoints(filepath))
+        )
+    pipeline.extend([
+        wds.batched(batch_size, collation_fn=collate_fn, partial=partial_batches),
+        wds.rename_keys(*rename_keys),
+    ])
+
+    if transform is not None:
+        pipeline.append(wds.map(transform))
+    pipeline = wds.DataPipeline(*pipeline)
+    if infinite:
+        pipeline = pipeline.repeat(nepochs=1000000)
+
+    loader = wds.WebLoader(
+        pipeline, batch_size=None, shuffle=False,
+        num_workers=get_num_workers(num_workers),
+        persistent_workers=True,
+    )
+    loader = tops.DataPrefetcher(loader, gpu_transform, channels_last=channels_last, to_float=False)
+    return loader

deep_privacy/dp2/data/transforms/__init__.py

@@ -0,0 +1,2 @@
+from .transforms import RandomCrop, CreateCondition, CreateEmbedding, Resize, ToFloat, Normalize
+from .stylegan2_transform import StyleGANAugmentPipe

deep_privacy/dp2/data/transforms/functional.py

@@ -0,0 +1,57 @@
+import torchvision.transforms.functional as F
+import torch
+import pickle
+from tops import download_file, assert_shape
+from typing import Dict
+from functools import lru_cache
+
+global symmetry_transform
+
+
+@lru_cache(maxsize=1)
+def get_symmetry_transform(symmetry_url):
+    file_name = download_file(symmetry_url)
+    with open(file_name, "rb") as fp:
+        symmetry = pickle.load(fp)
+    return torch.from_numpy(symmetry["vertex_transforms"]).long()
+
+
+hflip_handled_cases = set([
+    "keypoints", "img", "mask", "border", "semantic_mask", "vertices", "E_mask", "embed_map", "condition",
+    "embedding", "vertx2cat", "maskrcnn_mask", "__key__"])
+
+
+def hflip(container: Dict[str, torch.Tensor], flip_map=None) -> Dict[str, torch.Tensor]:
+    container["img"] = F.hflip(container["img"])
+    if "condition" in container:
+        container["condition"] = F.hflip(container["condition"])
+    if "embedding" in container:
+        container["embedding"] = F.hflip(container["embedding"])
+    assert all([key in hflip_handled_cases for key in container]), container.keys()
+    if "keypoints" in container:
+        assert flip_map is not None
+        if container["keypoints"].ndim == 3:
+            keypoints = container["keypoints"][:, flip_map, :]
+            keypoints[:, :,  0] = 1 - keypoints[:, :,  0]
+        else:
+            assert_shape(container["keypoints"], (None, 3))
+            keypoints = container["keypoints"][flip_map, :]
+            keypoints[:, 0] = 1 - keypoints[:, 0]
+        container["keypoints"] = keypoints
+    if "mask" in container:
+        container["mask"] = F.hflip(container["mask"])
+    if "border" in container:
+        container["border"] = F.hflip(container["border"])
+    if "semantic_mask" in container:
+        container["semantic_mask"] = F.hflip(container["semantic_mask"])
+    if "vertices" in container:
+        symmetry_transform = get_symmetry_transform(
+            "https://dl.fbaipublicfiles.com/densepose/meshes/symmetry/symmetry_smpl_27554.pkl")
+        container["vertices"] = F.hflip(container["vertices"])
+        symmetry_transform_ = symmetry_transform.to(container["vertices"].device)
+        container["vertices"] = symmetry_transform_[container["vertices"].long()]
+    if "E_mask" in container:
+        container["E_mask"] = F.hflip(container["E_mask"])
+    if "maskrcnn_mask" in container:
+        container["maskrcnn_mask"] = F.hflip(container["maskrcnn_mask"])
+    return container

deep_privacy/dp2/data/transforms/stylegan2_transform.py

@@ -0,0 +1,394 @@
+import numpy as np
+import scipy.signal
+import torch
+try:
+    from sg3_torch_utils import misc
+    from sg3_torch_utils.ops import upfirdn2d
+    from sg3_torch_utils.ops import grid_sample_gradfix
+    from sg3_torch_utils.ops import conv2d_gradfix
+except:
+    pass
+#----------------------------------------------------------------------------
+# Coefficients of various wavelet decomposition low-pass filters.
+
+wavelets = {
+    'haar': [0.7071067811865476, 0.7071067811865476],
+    'db1':  [0.7071067811865476, 0.7071067811865476],
+    'db2':  [-0.12940952255092145, 0.22414386804185735, 0.836516303737469, 0.48296291314469025],
+    'db3':  [0.035226291882100656, -0.08544127388224149, -0.13501102001039084, 0.4598775021193313, 0.8068915093133388, 0.3326705529509569],
+    'db4':  [-0.010597401784997278, 0.032883011666982945, 0.030841381835986965, -0.18703481171888114, -0.02798376941698385, 0.6308807679295904, 0.7148465705525415, 0.23037781330885523],
+    'db5':  [0.003335725285001549, -0.012580751999015526, -0.006241490213011705, 0.07757149384006515, -0.03224486958502952, -0.24229488706619015, 0.13842814590110342, 0.7243085284385744, 0.6038292697974729, 0.160102397974125],
+    'db6':  [-0.00107730108499558, 0.004777257511010651, 0.0005538422009938016, -0.031582039318031156, 0.02752286553001629, 0.09750160558707936, -0.12976686756709563, -0.22626469396516913, 0.3152503517092432, 0.7511339080215775, 0.4946238903983854, 0.11154074335008017],
+    'db7':  [0.0003537138000010399, -0.0018016407039998328, 0.00042957797300470274, 0.012550998556013784, -0.01657454163101562, -0.03802993693503463, 0.0806126091510659, 0.07130921926705004, -0.22403618499416572, -0.14390600392910627, 0.4697822874053586, 0.7291320908465551, 0.39653931948230575, 0.07785205408506236],
+    'db8':  [-0.00011747678400228192, 0.0006754494059985568, -0.0003917403729959771, -0.00487035299301066, 0.008746094047015655, 0.013981027917015516, -0.04408825393106472, -0.01736930100202211, 0.128747426620186, 0.00047248457399797254, -0.2840155429624281, -0.015829105256023893, 0.5853546836548691, 0.6756307362980128, 0.3128715909144659, 0.05441584224308161],
+    'sym2': [-0.12940952255092145, 0.22414386804185735, 0.836516303737469, 0.48296291314469025],
+    'sym3': [0.035226291882100656, -0.08544127388224149, -0.13501102001039084, 0.4598775021193313, 0.8068915093133388, 0.3326705529509569],
+    'sym4': [-0.07576571478927333, -0.02963552764599851, 0.49761866763201545, 0.8037387518059161, 0.29785779560527736, -0.09921954357684722, -0.012603967262037833, 0.0322231006040427],
+    'sym5': [0.027333068345077982, 0.029519490925774643, -0.039134249302383094, 0.1993975339773936, 0.7234076904024206, 0.6339789634582119, 0.01660210576452232, -0.17532808990845047, -0.021101834024758855, 0.019538882735286728],
+    'sym6': [0.015404109327027373, 0.0034907120842174702, -0.11799011114819057, -0.048311742585633, 0.4910559419267466, 0.787641141030194, 0.3379294217276218, -0.07263752278646252, -0.021060292512300564, 0.04472490177066578, 0.0017677118642428036, -0.007800708325034148],
+    'sym7': [0.002681814568257878, -0.0010473848886829163, -0.01263630340325193, 0.03051551316596357, 0.0678926935013727, -0.049552834937127255, 0.017441255086855827, 0.5361019170917628, 0.767764317003164, 0.2886296317515146, -0.14004724044296152, -0.10780823770381774, 0.004010244871533663, 0.010268176708511255],
+    'sym8': [-0.0033824159510061256, -0.0005421323317911481, 0.03169508781149298, 0.007607487324917605, -0.1432942383508097, -0.061273359067658524, 0.4813596512583722, 0.7771857517005235, 0.3644418948353314, -0.05194583810770904, -0.027219029917056003, 0.049137179673607506, 0.003808752013890615, -0.01495225833704823, -0.0003029205147213668, 0.0018899503327594609],
+}
+
+#----------------------------------------------------------------------------
+# Helpers for constructing transformation matrices.
+
+
+def matrix(*rows, device=None):
+    assert all(len(row) == len(rows[0]) for row in rows)
+    elems = [x for row in rows for x in row]
+    ref = [x for x in elems if isinstance(x, torch.Tensor)]
+    if len(ref) == 0:
+        return misc.constant(np.asarray(rows), device=device)
+    assert device is None or device == ref[0].device
+    elems = [x if isinstance(x, torch.Tensor) else misc.constant(x, shape=ref[0].shape, device=ref[0].device) for x in elems]
+    return torch.stack(elems, dim=-1).reshape(ref[0].shape + (len(rows), -1))
+
+
+def translate2d(tx, ty, **kwargs):
+    return matrix(
+        [1, 0, tx],
+        [0, 1, ty],
+        [0, 0, 1],
+        **kwargs)
+
+
+def translate3d(tx, ty, tz, **kwargs):
+    return matrix(
+        [1, 0, 0, tx],
+        [0, 1, 0, ty],
+        [0, 0, 1, tz],
+        [0, 0, 0, 1],
+        **kwargs)
+
+
+def scale2d(sx, sy, **kwargs):
+    return matrix(
+        [sx, 0,  0],
+        [0,  sy, 0],
+        [0,  0,  1],
+        **kwargs)
+
+
+def scale3d(sx, sy, sz, **kwargs):
+    return matrix(
+        [sx, 0,  0,  0],
+        [0,  sy, 0,  0],
+        [0,  0,  sz, 0],
+        [0,  0,  0,  1],
+        **kwargs)
+
+
+def rotate2d(theta, **kwargs):
+    return matrix(
+        [torch.cos(theta), torch.sin(-theta), 0],
+        [torch.sin(theta), torch.cos(theta),  0],
+        [0,                0,                 1],
+        **kwargs)
+
+
+def rotate3d(v, theta, **kwargs):
+    vx = v[..., 0]; vy = v[..., 1]; vz = v[..., 2]
+    s = torch.sin(theta); c = torch.cos(theta); cc = 1 - c
+    return matrix(
+        [vx*vx*cc+c,    vx*vy*cc-vz*s, vx*vz*cc+vy*s, 0],
+        [vy*vx*cc+vz*s, vy*vy*cc+c,    vy*vz*cc-vx*s, 0],
+        [vz*vx*cc-vy*s, vz*vy*cc+vx*s, vz*vz*cc+c,    0],
+        [0,             0,             0,             1],
+        **kwargs)
+
+
+def translate2d_inv(tx, ty, **kwargs):
+    return translate2d(-tx, -ty, **kwargs)
+
+
+def scale2d_inv(sx, sy, **kwargs):
+    return scale2d(1 / sx, 1 / sy, **kwargs)
+
+
+def rotate2d_inv(theta, **kwargs):
+    return rotate2d(-theta, **kwargs)
+
+
+class StyleGANAugmentPipe(torch.nn.Module):
+    def __init__(self,
+        rotate90=0, xint=0, xint_max=0.125,
+        scale=0, rotate=0, aniso=0, xfrac=0, scale_std=0.2, rotate_max=1, aniso_std=0.2, xfrac_std=0.125,
+        brightness=0, contrast=0, lumaflip=0, hue=0, saturation=0, brightness_std=0.2, contrast_std=0.5,
+        hue_max=1, saturation_std=1,
+        imgfilter=0, imgfilter_bands=[1,1,1,1], imgfilter_std=1,
+        ):
+        super().__init__()
+        self.register_buffer('p', torch.ones([]))       # Overall multiplier for augmentation probability.
+
+        # Pixel blitting.
+        self.rotate90         = float(rotate90)         # Probability multiplier for 90 degree rotations.
+        self.xint             = float(xint)             # Probability multiplier for integer translation.
+        self.xint_max         = float(xint_max)         # Range of integer translation, relative to image dimensions.
+
+        # General geometric transformations.
+        self.scale            = float(scale)            # Probability multiplier for isotropic scaling.
+        self.rotate           = float(rotate)           # Probability multiplier for arbitrary rotation.
+        self.aniso            = float(aniso)            # Probability multiplier for anisotropic scaling.
+        self.xfrac            = float(xfrac)            # Probability multiplier for fractional translation.
+        self.scale_std        = float(scale_std)        # Log2 standard deviation of isotropic scaling.
+        self.rotate_max       = float(rotate_max)       # Range of arbitrary rotation, 1 = full circle.
+        self.aniso_std        = float(aniso_std)        # Log2 standard deviation of anisotropic scaling.
+        self.xfrac_std        = float(xfrac_std)        # Standard deviation of frational translation, relative to image dimensions.
+
+        # Color transformations.
+        self.brightness       = float(brightness)       # Probability multiplier for brightness.
+        self.contrast         = float(contrast)         # Probability multiplier for contrast.
+        self.lumaflip         = float(lumaflip)         # Probability multiplier for luma flip.
+        self.hue              = float(hue)              # Probability multiplier for hue rotation.
+        self.saturation       = float(saturation)       # Probability multiplier for saturation.
+        self.brightness_std   = float(brightness_std)   # Standard deviation of brightness.
+        self.contrast_std     = float(contrast_std)     # Log2 standard deviation of contrast.
+        self.hue_max          = float(hue_max)          # Range of hue rotation, 1 = full circle.
+        self.saturation_std   = float(saturation_std)   # Log2 standard deviation of saturation.
+
+        # Image-space filtering.
+        self.imgfilter        = float(imgfilter)        # Probability multiplier for image-space filtering.
+        self.imgfilter_bands  = list(imgfilter_bands)   # Probability multipliers for individual frequency bands.
+        self.imgfilter_std    = float(imgfilter_std)    # Log2 standard deviation of image-space filter amplification.
+
+        # Setup orthogonal lowpass filter for geometric augmentations.
+        self.register_buffer('Hz_geom', upfirdn2d.setup_filter(wavelets['sym6']))
+
+        # Construct filter bank for image-space filtering.
+        Hz_lo = np.asarray(wavelets['sym2'])            # H(z)
+        Hz_hi = Hz_lo * ((-1) ** np.arange(Hz_lo.size)) # H(-z)
+        Hz_lo2 = np.convolve(Hz_lo, Hz_lo[::-1]) / 2    # H(z) * H(z^-1) / 2
+        Hz_hi2 = np.convolve(Hz_hi, Hz_hi[::-1]) / 2    # H(-z) * H(-z^-1) / 2
+        Hz_fbank = np.eye(4, 1)                         # Bandpass(H(z), b_i)
+        for i in range(1, Hz_fbank.shape[0]):
+            Hz_fbank = np.dstack([Hz_fbank, np.zeros_like(Hz_fbank)]).reshape(Hz_fbank.shape[0], -1)[:, :-1]
+            Hz_fbank = scipy.signal.convolve(Hz_fbank, [Hz_lo2])
+            Hz_fbank[i, (Hz_fbank.shape[1] - Hz_hi2.size) // 2 : (Hz_fbank.shape[1] + Hz_hi2.size) // 2] += Hz_hi2
+        self.register_buffer('Hz_fbank', torch.as_tensor(Hz_fbank, dtype=torch.float32))
+
+    def forward(self, batch, debug_percentile=None):
+        images = batch["img"]
+        batch["vertices"] = batch["vertices"].float()
+        assert isinstance(images, torch.Tensor) and images.ndim == 4
+        batch_size, num_channels, height, width = images.shape
+        device = images.device
+        self.Hz_fbank = self.Hz_fbank.to(device)
+        self.Hz_geom = self.Hz_geom.to(device)
+        if debug_percentile is not None:
+            debug_percentile = torch.as_tensor(debug_percentile, dtype=torch.float32, device=device)
+
+        # -------------------------------------
+        # Select parameters for pixel blitting.
+        # -------------------------------------
+
+        # Initialize inverse homogeneous 2D transform: G_inv @ pixel_out ==> pixel_in
+        I_3 = torch.eye(3, device=device)
+        G_inv = I_3
+
+        # Apply integer translation with probability (xint * strength).
+        if self.xint > 0:
+            t = (torch.rand([batch_size, 2], device=device) * 2 - 1) * self.xint_max
+            t = torch.where(torch.rand([batch_size, 1], device=device) < self.xint * self.p, t, torch.zeros_like(t))
+            if debug_percentile is not None:
+                t = torch.full_like(t, (debug_percentile * 2 - 1) * self.xint_max)
+            G_inv = G_inv @ translate2d_inv(torch.round(t[:,0] * width), torch.round(t[:,1] * height))
+
+        # --------------------------------------------------------
+        # Select parameters for general geometric transformations.
+        # --------------------------------------------------------
+
+        # Apply isotropic scaling with probability (scale * strength).
+        if self.scale > 0:
+            s = torch.exp2(torch.randn([batch_size], device=device) * self.scale_std)
+            s = torch.where(torch.rand([batch_size], device=device) < self.scale * self.p, s, torch.ones_like(s))
+            if debug_percentile is not None:
+                s = torch.full_like(s, torch.exp2(torch.erfinv(debug_percentile * 2 - 1) * self.scale_std))
+            G_inv = G_inv @ scale2d_inv(s, s)
+
+        # Apply pre-rotation with probability p_rot.
+        p_rot = 1 - torch.sqrt((1 - self.rotate * self.p).clamp(0, 1)) # P(pre OR post) = p
+        if self.rotate > 0:
+            theta = (torch.rand([batch_size], device=device) * 2 - 1) * np.pi * self.rotate_max
+            theta = torch.where(torch.rand([batch_size], device=device) < p_rot, theta, torch.zeros_like(theta))
+            if debug_percentile is not None:
+                theta = torch.full_like(theta, (debug_percentile * 2 - 1) * np.pi * self.rotate_max)
+            G_inv = G_inv @ rotate2d_inv(-theta) # Before anisotropic scaling.
+
+        # Apply anisotropic scaling with probability (aniso * strength).
+        if self.aniso > 0:
+            s = torch.exp2(torch.randn([batch_size], device=device) * self.aniso_std)
+            s = torch.where(torch.rand([batch_size], device=device) < self.aniso * self.p, s, torch.ones_like(s))
+            if debug_percentile is not None:
+                s = torch.full_like(s, torch.exp2(torch.erfinv(debug_percentile * 2 - 1) * self.aniso_std))
+            G_inv = G_inv @ scale2d_inv(s, 1 / s)
+
+        # Apply post-rotation with probability p_rot.
+        if self.rotate > 0:
+            theta = (torch.rand([batch_size], device=device) * 2 - 1) * np.pi * self.rotate_max
+            theta = torch.where(torch.rand([batch_size], device=device) < p_rot, theta, torch.zeros_like(theta))
+            if debug_percentile is not None:
+                theta = torch.zeros_like(theta)
+            G_inv = G_inv @ rotate2d_inv(-theta) # After anisotropic scaling.
+
+        # Apply fractional translation with probability (xfrac * strength).
+        if self.xfrac > 0:
+            t = torch.randn([batch_size, 2], device=device) * self.xfrac_std
+            t = torch.where(torch.rand([batch_size, 1], device=device) < self.xfrac * self.p, t, torch.zeros_like(t))
+            if debug_percentile is not None:
+                t = torch.full_like(t, torch.erfinv(debug_percentile * 2 - 1) * self.xfrac_std)
+            G_inv = G_inv @ translate2d_inv(t[:,0] * width, t[:,1] * height)
+
+        # ----------------------------------
+        # Execute geometric transformations.
+        # ----------------------------------
+
+        # Execute if the transform is not identity.
+        if G_inv is not I_3:
+            # Calculate padding.
+            cx = (width - 1) / 2
+            cy = (height - 1) / 2
+            cp = matrix([-cx, -cy, 1], [cx, -cy, 1], [cx, cy, 1], [-cx, cy, 1], device=device) # [idx, xyz]
+            cp = G_inv @ cp.t() # [batch, xyz, idx]
+            Hz_pad = self.Hz_geom.shape[0] // 4
+            margin = cp[:, :2, :].permute(1, 0, 2).flatten(1) # [xy, batch * idx]
+            margin = torch.cat([-margin, margin]).max(dim=1).values # [x0, y0, x1, y1]
+            margin = margin + misc.constant([Hz_pad * 2 - cx, Hz_pad * 2 - cy] * 2, device=device)
+            margin = margin.max(misc.constant([0, 0] * 2, device=device))
+            margin = margin.min(misc.constant([width-1, height-1] * 2, device=device))
+            mx0, my0, mx1, my1 = margin.ceil().to(torch.int32)
+
+            # Pad image and adjust origin.
+            images = torch.nn.functional.pad(input=images, pad=[mx0,mx1,my0,my1], mode='reflect')
+            batch["mask"]  = torch.nn.functional.pad(input=batch["mask"], pad=[mx0,mx1,my0,my1], mode='constant', value=1.0)
+            batch["E_mask"]  = torch.nn.functional.pad(input=batch["E_mask"], pad=[mx0,mx1,my0,my1], mode='constant', value=0.0)
+            batch["vertices"]  = torch.nn.functional.pad(input=batch["vertices"], pad=[mx0,mx1,my0,my1], mode='constant', value=0.0)
+            G_inv = translate2d((mx0 - mx1) / 2, (my0 - my1) / 2) @ G_inv
+
+            # Upsample.
+            images = upfirdn2d.upsample2d(x=images, f=self.Hz_geom, up=2)
+            batch["mask"] = torch.nn.functional.interpolate(batch["mask"], scale_factor=2, mode="nearest")
+            batch["E_mask"] = torch.nn.functional.interpolate(batch["E_mask"], scale_factor=2, mode="nearest")
+            batch["vertices"] = torch.nn.functional.interpolate(batch["vertices"], scale_factor=2, mode="nearest")
+            G_inv = scale2d(2, 2, device=device) @ G_inv @ scale2d_inv(2, 2, device=device)
+            G_inv = translate2d(-0.5, -0.5, device=device) @ G_inv @ translate2d_inv(-0.5, -0.5, device=device)
+
+            # Execute transformation.
+            shape = [batch_size, num_channels, (height + Hz_pad * 2) * 2, (width + Hz_pad * 2) * 2]
+            G_inv = scale2d(2 / images.shape[3], 2 / images.shape[2], device=device) @ G_inv @ scale2d_inv(2 / shape[3], 2 / shape[2], device=device)
+            grid = torch.nn.functional.affine_grid(theta=G_inv[:,:2,:], size=shape, align_corners=False)
+            images = grid_sample_gradfix.grid_sample(images, grid)
+            
+            batch["mask"] = torch.nn.functional.grid_sample(
+                input=batch["mask"], grid=grid, mode='nearest', padding_mode="border", align_corners=False)
+            batch["E_mask"] = torch.nn.functional.grid_sample(
+                input=batch["E_mask"], grid=grid, mode='nearest', padding_mode="border", align_corners=False)
+            batch["vertices"] = torch.nn.functional.grid_sample(
+                input=batch["vertices"], grid=grid, mode='nearest', padding_mode="border", align_corners=False)
+            
+
+            # Downsample and crop.
+            images = upfirdn2d.downsample2d(x=images, f=self.Hz_geom, down=2, padding=-Hz_pad*2, flip_filter=True)
+            batch["mask"] = torch.nn.functional.interpolate(batch["mask"][:, :, Hz_pad*2:-Hz_pad*2, Hz_pad*2:-Hz_pad*2], scale_factor=.5, mode="nearest", recompute_scale_factor=False)
+            batch["E_mask"] = torch.nn.functional.interpolate(batch["E_mask"][:, :, Hz_pad*2:-Hz_pad*2, Hz_pad*2:-Hz_pad*2], scale_factor=.5, mode="nearest", recompute_scale_factor=False)
+            batch["vertices"] = torch.nn.functional.interpolate(batch["vertices"][:, :, Hz_pad*2:-Hz_pad*2, Hz_pad*2:-Hz_pad*2], scale_factor=.5, mode="nearest", recompute_scale_factor=False)
+        # --------------------------------------------
+        # Select parameters for color transformations.
+        # --------------------------------------------
+
+        # Initialize homogeneous 3D transformation matrix: C @ color_in ==> color_out
+        I_4 = torch.eye(4, device=device)
+        C = I_4
+
+        # Apply brightness with probability (brightness * strength).
+        if self.brightness > 0:
+            b = torch.randn([batch_size], device=device) * self.brightness_std
+            b = torch.where(torch.rand([batch_size], device=device) < self.brightness * self.p, b, torch.zeros_like(b))
+            if debug_percentile is not None:
+                b = torch.full_like(b, torch.erfinv(debug_percentile * 2 - 1) * self.brightness_std)
+            C = translate3d(b, b, b) @ C
+
+        # Apply contrast with probability (contrast * strength).
+        if self.contrast > 0:
+            c = torch.exp2(torch.randn([batch_size], device=device) * self.contrast_std)
+            c = torch.where(torch.rand([batch_size], device=device) < self.contrast * self.p, c, torch.ones_like(c))
+            if debug_percentile is not None:
+                c = torch.full_like(c, torch.exp2(torch.erfinv(debug_percentile * 2 - 1) * self.contrast_std))
+            C = scale3d(c, c, c) @ C
+
+        # Apply luma flip with probability (lumaflip * strength).
+        v = misc.constant(np.asarray([1, 1, 1, 0]) / np.sqrt(3), device=device) # Luma axis.
+
+        # Apply hue rotation with probability (hue * strength).
+        if self.hue > 0 and num_channels > 1:
+            theta = (torch.rand([batch_size], device=device) * 2 - 1) * np.pi * self.hue_max
+            theta = torch.where(torch.rand([batch_size], device=device) < self.hue * self.p, theta, torch.zeros_like(theta))
+            if debug_percentile is not None:
+                theta = torch.full_like(theta, (debug_percentile * 2 - 1) * np.pi * self.hue_max)
+            C = rotate3d(v, theta) @ C # Rotate around v.
+
+        # Apply saturation with probability (saturation * strength).
+        if self.saturation > 0 and num_channels > 1:
+            s = torch.exp2(torch.randn([batch_size, 1, 1], device=device) * self.saturation_std)
+            s = torch.where(torch.rand([batch_size, 1, 1], device=device) < self.saturation * self.p, s, torch.ones_like(s))
+            if debug_percentile is not None:
+                s = torch.full_like(s, torch.exp2(torch.erfinv(debug_percentile * 2 - 1) * self.saturation_std))
+            C = (v.ger(v) + (I_4 - v.ger(v)) * s) @ C
+
+        # ------------------------------
+        # Execute color transformations.
+        # ------------------------------
+
+        # Execute if the transform is not identity.
+        if C is not I_4:
+            images = images.reshape([batch_size, num_channels, height * width])
+            if num_channels == 3:
+                images = C[:, :3, :3] @ images + C[:, :3, 3:]
+            elif num_channels == 1:
+                C = C[:, :3, :].mean(dim=1, keepdims=True)
+                images = images * C[:, :, :3].sum(dim=2, keepdims=True) + C[:, :, 3:]
+            else:
+                raise ValueError('Image must be RGB (3 channels) or L (1 channel)')
+            images = images.reshape([batch_size, num_channels, height, width])
+
+        # ----------------------
+        # Image-space filtering.
+        # ----------------------
+
+        if self.imgfilter > 0:
+            num_bands = self.Hz_fbank.shape[0]
+            assert len(self.imgfilter_bands) == num_bands
+            expected_power = misc.constant(np.array([10, 1, 1, 1]) / 13, device=device) # Expected power spectrum (1/f).
+
+            # Apply amplification for each band with probability (imgfilter * strength * band_strength).
+            g = torch.ones([batch_size, num_bands], device=device) # Global gain vector (identity).
+            for i, band_strength in enumerate(self.imgfilter_bands):
+                t_i = torch.exp2(torch.randn([batch_size], device=device) * self.imgfilter_std)
+                t_i = torch.where(torch.rand([batch_size], device=device) < self.imgfilter * self.p * band_strength, t_i, torch.ones_like(t_i))
+                if debug_percentile is not None:
+                    t_i = torch.full_like(t_i, torch.exp2(torch.erfinv(debug_percentile * 2 - 1) * self.imgfilter_std)) if band_strength > 0 else torch.ones_like(t_i)
+                t = torch.ones([batch_size, num_bands], device=device)                  # Temporary gain vector.
+                t[:, i] = t_i                                                           # Replace i'th element.
+                t = t / (expected_power * t.square()).sum(dim=-1, keepdims=True).sqrt() # Normalize power.
+                g = g * t                                                               # Accumulate into global gain.
+
+            # Construct combined amplification filter.
+            Hz_prime = g @ self.Hz_fbank                                    # [batch, tap]
+            Hz_prime = Hz_prime.unsqueeze(1).repeat([1, num_channels, 1])   # [batch, channels, tap]
+            Hz_prime = Hz_prime.reshape([batch_size * num_channels, 1, -1]) # [batch * channels, 1, tap]
+
+            # Apply filter.
+            p = self.Hz_fbank.shape[1] // 2
+            images = images.reshape([1, batch_size * num_channels, height, width])
+            images = torch.nn.functional.pad(input=images, pad=[p,p,p,p], mode='reflect')
+            images = conv2d_gradfix.conv2d(input=images, weight=Hz_prime.unsqueeze(2), groups=batch_size*num_channels)
+            images = conv2d_gradfix.conv2d(input=images, weight=Hz_prime.unsqueeze(3), groups=batch_size*num_channels)
+            images = images.reshape([batch_size, num_channels, height, width])
+
+        # ------------------------
+        # Image-space corruptions.
+        # ------------------------
+        batch["img"] = images
+        batch["vertices"] = batch["vertices"].long()
+        batch["border"] = 1 - batch["E_mask"] - batch["mask"]
+        return batch

deep_privacy/dp2/data/transforms/transforms.py

@@ -0,0 +1,277 @@
+from pathlib import Path
+from typing import Dict, List
+import torchvision
+import torch
+import tops
+import torchvision.transforms.functional as F
+from .functional import hflip
+import numpy as np
+from dp2.utils.vis_utils import get_coco_keypoints
+from PIL import Image, ImageDraw
+from typing import Tuple
+
+
+class RandomHorizontalFlip(torch.nn.Module):
+
+    def __init__(self, p: float,  flip_map=None, **kwargs):
+        super().__init__()
+        self.flip_ratio = p
+        self.flip_map = flip_map
+        if self.flip_ratio is None:
+            self.flip_ratio = 0.5
+        assert 0 <= self.flip_ratio <= 1
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        if torch.rand(1) > self.flip_ratio:
+            return container
+        return hflip(container, self.flip_map)
+
+
+class CenterCrop(torch.nn.Module):
+    """
+    Performs the transform on the image.
+    NOTE: Does not transform the mask to improve runtime.
+    """
+
+    def __init__(self, size: List[int]):
+        super().__init__()
+        self.size = tuple(size)
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        min_size = min(container["img"].shape[1], container["img"].shape[2])
+        if min_size < self.size[0]:
+            container["img"] = F.center_crop(container["img"], min_size)
+            container["img"] = F.resize(container["img"], self.size)
+            return container
+        container["img"] = F.center_crop(container["img"], self.size)
+        return container
+
+
+class Resize(torch.nn.Module):
+    """
+    Performs the transform on the image.
+    NOTE: Does not transform the mask to improve runtime.
+    """
+
+    def __init__(self, size, interpolation=F.InterpolationMode.BILINEAR):
+        super().__init__()
+        self.size = tuple(size)
+        self.interpolation = interpolation
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        container["img"] = F.resize(container["img"], self.size, self.interpolation, antialias=True)
+        if "semantic_mask" in container:
+            container["semantic_mask"] = F.resize(
+                container["semantic_mask"], self.size, F.InterpolationMode.NEAREST)
+        if "embedding" in container:
+            container["embedding"] = F.resize(
+                container["embedding"], self.size, self.interpolation)
+        if "mask" in container:
+            container["mask"] = F.resize(
+                container["mask"], self.size, F.InterpolationMode.NEAREST)
+        if "E_mask" in container:
+            container["E_mask"] = F.resize(
+                container["E_mask"], self.size, F.InterpolationMode.NEAREST)
+        if "maskrcnn_mask" in container:
+            container["maskrcnn_mask"] = F.resize(
+                container["maskrcnn_mask"], self.size, F.InterpolationMode.NEAREST)
+        if "vertices" in container:
+            container["vertices"] = F.resize(
+                container["vertices"], self.size, F.InterpolationMode.NEAREST)
+        return container
+
+    def __repr__(self):
+        repr = super().__repr__()
+        vars_ = dict(size=self.size, interpolation=self.interpolation)
+        return repr + " " + " ".join([f"{k}: {v}" for k, v in vars_.items()])
+
+
+class Normalize(torch.nn.Module):
+    """
+    Performs the transform on the image.
+    NOTE: Does not transform the mask to improve runtime.
+    """
+
+    def __init__(self, mean, std, inplace, keys=["img"]):
+        super().__init__()
+        self.mean = mean
+        self.std = std
+        self.inplace = inplace
+        self.keys = keys
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        for key in self.keys:
+            container[key] = F.normalize(container[key], self.mean, self.std, self.inplace)
+        return container
+
+    def __repr__(self):
+        repr = super().__repr__()
+        vars_ = dict(mean=self.mean, std=self.std, inplace=self.inplace)
+        return repr + " " + " ".join([f"{k}: {v}" for k, v in vars_.items()])
+
+
+class ToFloat(torch.nn.Module):
+
+    def __init__(self, keys=["img"], norm=True) -> None:
+        super().__init__()
+        self.keys = keys
+        self.gain = 255 if norm else 1
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        for key in self.keys:
+            container[key] = container[key].float() / self.gain
+        return container
+
+
+class RandomCrop(torchvision.transforms.RandomCrop):
+    """
+    Performs the transform on the image.
+    NOTE: Does not transform the mask to improve runtime.
+    """
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        container["img"] = super().forward(container["img"])
+        return container
+
+
+class CreateCondition(torch.nn.Module):
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        if container["img"].dtype == torch.uint8:
+            container["condition"] = container["img"] * container["mask"].byte() + (1-container["mask"].byte()) * 127
+            return container
+        container["condition"] = container["img"] * container["mask"]
+        return container
+
+
+class CreateEmbedding(torch.nn.Module):
+
+    def __init__(self, embed_path: Path, cuda=True) -> None:
+        super().__init__()
+        self.embed_map = torch.load(embed_path, map_location=torch.device("cpu"))
+        if cuda:
+            self.embed_map = tops.to_cuda(self.embed_map)
+
+    def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        vertices = container["vertices"]
+        if vertices.ndim == 3:
+            embedding = self.embed_map[vertices.long()].squeeze(dim=0)
+            embedding = embedding.permute(2, 0, 1) * container["E_mask"]
+            pass
+        else:
+            assert vertices.ndim == 4
+            embedding = self.embed_map[vertices.long()].squeeze(dim=1)
+            embedding = embedding.permute(0, 3, 1, 2) * container["E_mask"]
+        container["embedding"] = embedding
+        container["embed_map"] = self.embed_map.clone()
+        return container
+
+
+class InsertJointMap(torch.nn.Module):
+
+    def __init__(self, imsize: Tuple) -> None:
+        super().__init__()
+        self.imsize = imsize
+        knames = get_coco_keypoints()[0]
+        knames = knames + ["neck", "mid_hip"]
+        connectivity = {
+            "nose": ["left_eye", "right_eye", "neck"],
+            "left_eye": ["right_eye", "left_ear"],
+            "right_eye": ["right_ear"],
+            "left_shoulder": ["right_shoulder", "left_elbow", "left_hip"],
+            "right_shoulder": ["right_elbow", "right_hip"],
+            "left_elbow": ["left_wrist"],
+            "right_elbow": ["right_wrist"],
+            "left_hip": ["right_hip", "left_knee"],
+            "right_hip": ["right_knee"],
+            "left_knee": ["left_ankle"],
+            "right_knee": ["right_ankle"],
+            "neck": ["mid_hip", "nose"],
+        }
+        category = {
+            ("nose", "left_eye"): 0,  # head
+            ("nose", "right_eye"): 0,  # head
+            ("nose", "neck"): 0,  # head
+            ("left_eye", "right_eye"): 0,  # head
+            ("left_eye", "left_ear"): 0,  # head
+            ("right_eye", "right_ear"): 0,  # head
+            ("left_shoulder", "left_elbow"): 1,  # left arm
+            ("left_elbow", "left_wrist"): 1,  # left arm
+            ("right_shoulder", "right_elbow"): 2,  # right arm
+            ("right_elbow", "right_wrist"): 2,  # right arm
+            ("left_shoulder", "right_shoulder"): 3,  # body
+            ("left_shoulder", "left_hip"): 3,  # body
+            ("right_shoulder", "right_hip"): 3,  # body
+            ("left_hip", "right_hip"): 3,  # body
+            ("left_hip", "left_knee"): 4,  # left leg
+            ("left_knee", "left_ankle"): 4,  # left leg
+            ("right_hip", "right_knee"): 5,  # right leg
+            ("right_knee", "right_ankle"): 5,  # right leg
+            ("neck", "mid_hip"): 3,  # body
+            ("neck", "nose"): 0,  # head
+        }
+        self.indices2category = {
+            tuple([knames.index(n) for n in k]): v for k, v in category.items()
+        }
+        self.connectivity_indices = {
+            knames.index(k): [knames.index(v_) for v_ in v]
+            for k, v in connectivity.items()
+        }
+        self.l_shoulder = knames.index("left_shoulder")
+        self.r_shoulder = knames.index("right_shoulder")
+        self.l_hip = knames.index("left_hip")
+        self.r_hip = knames.index("right_hip")
+        self.l_eye = knames.index("left_eye")
+        self.r_eye = knames.index("right_eye")
+        self.nose = knames.index("nose")
+        self.neck = knames.index("neck")
+
+    def create_joint_map(self, N, H, W, keypoints):
+        joint_maps = np.zeros((N, H, W), dtype=np.uint8)
+        for bidx, keypoints in enumerate(keypoints):
+            assert keypoints.shape == (17, 3), keypoints.shape
+            keypoints = torch.cat((keypoints, torch.zeros(2, 3)))
+            visible = keypoints[:, -1] > 0
+
+            if visible[self.l_shoulder] and visible[self.r_shoulder]:
+                neck = (keypoints[self.l_shoulder]
+                        + (keypoints[self.r_shoulder] - keypoints[self.l_shoulder]) / 2)
+                keypoints[-2] = neck
+                visible[-2] = 1
+            if visible[self.l_hip] and visible[self.r_hip]:
+                mhip = (keypoints[self.l_hip]
+                        + (keypoints[self.r_hip] - keypoints[self.l_hip]) / 2
+                        )
+                keypoints[-1] = mhip
+                visible[-1] = 1
+
+            keypoints[:, 0] *= W
+            keypoints[:, 1] *= H
+            joint_map = Image.fromarray(np.zeros((H, W), dtype=np.uint8))
+            draw = ImageDraw.Draw(joint_map)
+            for fidx in self.connectivity_indices.keys():
+                for tidx in self.connectivity_indices[fidx]:
+                    if visible[fidx] == 0 or visible[tidx] == 0:
+                        continue
+                    c = self.indices2category[(fidx, tidx)]
+                    s = tuple(keypoints[fidx, :2].round().long().numpy().tolist())
+                    e = tuple(keypoints[tidx, :2].round().long().numpy().tolist())
+                    draw.line((s, e), width=1, fill=c + 1)
+            if visible[self.nose] == 0 and visible[self.neck] == 1:
+                m_eye = (
+                    keypoints[self.l_eye]
+                    + (keypoints[self.r_eye] - keypoints[self.l_eye]) / 2
+                )
+                s = tuple(m_eye[:2].round().long().numpy().tolist())
+                e = tuple(keypoints[self.neck, :2].round().long().numpy().tolist())
+                c = self.indices2category[(self.nose, self.neck)]
+                draw.line((s, e), width=1, fill=c + 1)
+            joint_map = np.array(joint_map)
+
+            joint_maps[bidx] = np.array(joint_map)
+        return joint_maps[:, None]
+
+    def forward(self, batch):
+        batch["joint_map"] = torch.from_numpy(self.create_joint_map(
+            batch["img"].shape[0], *self.imsize, batch["keypoints"]))
+        return batch

deep_privacy/dp2/data/utils.py

@@ -0,0 +1,122 @@
+import torch
+from PIL import Image
+import numpy as np
+import multiprocessing
+import io
+from tops import logger
+from torch.utils.data._utils.collate import default_collate
+
+try:
+    import pyspng
+
+    PYSPNG_IMPORTED = True
+except ImportError:
+    PYSPNG_IMPORTED = False
+    print("Could not load pyspng. Defaulting to pillow image backend.")
+    from PIL import Image
+
+
+def get_fdf_keypoints():
+    return get_coco_keypoints()[:7]
+
+
+def get_fdf_flipmap():
+    keypoints = get_fdf_keypoints()
+    keypoint_flip_map = {
+        "left_eye": "right_eye",
+        "left_ear": "right_ear",
+        "left_shoulder": "right_shoulder",
+    }
+    for key, value in list(keypoint_flip_map.items()):
+        keypoint_flip_map[value] = key
+    keypoint_flip_map["nose"] = "nose"
+    keypoint_flip_map_idx = []
+    for source in keypoints:
+        keypoint_flip_map_idx.append(keypoints.index(keypoint_flip_map[source]))
+    return keypoint_flip_map_idx
+
+
+def get_coco_keypoints():
+    return [
+        "nose",
+        "left_eye",
+        "right_eye",  # 2
+        "left_ear",
+        "right_ear",  # 4
+        "left_shoulder",
+        "right_shoulder",  # 6
+        "left_elbow",
+        "right_elbow",  # 8
+        "left_wrist",
+        "right_wrist",  # 10
+        "left_hip",
+        "right_hip",  # 12
+        "left_knee",
+        "right_knee",  # 14
+        "left_ankle",
+        "right_ankle",  # 16
+    ]
+
+
+def get_coco_flipmap():
+    keypoints = get_coco_keypoints()
+    keypoint_flip_map = {
+        "left_eye": "right_eye",
+        "left_ear": "right_ear",
+        "left_shoulder": "right_shoulder",
+        "left_elbow": "right_elbow",
+        "left_wrist": "right_wrist",
+        "left_hip": "right_hip",
+        "left_knee": "right_knee",
+        "left_ankle": "right_ankle",
+    }
+    for key, value in list(keypoint_flip_map.items()):
+        keypoint_flip_map[value] = key
+    keypoint_flip_map["nose"] = "nose"
+    keypoint_flip_map_idx = []
+    for source in keypoints:
+        keypoint_flip_map_idx.append(keypoints.index(keypoint_flip_map[source]))
+    return keypoint_flip_map_idx
+
+
+def mask_decoder(x):
+    mask = torch.from_numpy(np.array(Image.open(io.BytesIO(x)))).squeeze()[None]
+    mask = mask > 0  # This fixes bug causing  maskf.loat().max() == 255.
+    return mask
+
+
+def png_decoder(x):
+    if PYSPNG_IMPORTED:
+        return torch.from_numpy(np.rollaxis(pyspng.load(x), 2))
+    with Image.open(io.BytesIO(x)) as im:
+        im = torch.from_numpy(np.rollaxis(np.array(im.convert("RGB")), 2))
+    return im
+
+
+def jpg_decoder(x):
+    with Image.open(io.BytesIO(x)) as im:
+        im = torch.from_numpy(np.rollaxis(np.array(im.convert("RGB")), 2))
+    return im
+
+
+def get_num_workers(num_workers: int):
+    n_cpus = multiprocessing.cpu_count()
+    if num_workers > n_cpus:
+        logger.warn(f"Setting the number of workers to match cpu count: {n_cpus}")
+        return n_cpus
+    return num_workers
+
+
+def collate_fn(batch):
+    elem = batch[0]
+    ignore_keys = set(["embed_map", "vertx2cat"])
+    batch_ = {
+        key: default_collate([d[key] for d in batch])
+        for key in elem
+        if key not in ignore_keys
+    }
+    if "embed_map" in elem:
+        batch_["embed_map"] = elem["embed_map"]
+    if "vertx2cat" in elem:
+        batch_["vertx2cat"] = elem["vertx2cat"]
+    return batch_

deep_privacy/dp2/detection/__init__.py

@@ -0,0 +1,3 @@
+from .cse_mask_face_detector import CSeMaskFaceDetector
+from .person_detector import CSEPersonDetector
+from .structures import PersonDetection, VehicleDetection, FaceDetection

deep_privacy/dp2/detection/base.py

@@ -0,0 +1,42 @@
+import pickle
+import torch
+import lzma
+from pathlib import Path
+from tops import logger
+
+
+class BaseDetector:
+
+    def __init__(self, cache_directory: str) -> None:
+        if cache_directory is not None:
+            self.cache_directory = Path(cache_directory, str(self.__class__.__name__))
+            self.cache_directory.mkdir(exist_ok=True, parents=True)
+
+    def save_to_cache(self, detection, cache_path: Path, after_preprocess=True):
+        logger.log(f"Caching detection to: {cache_path}")
+        with lzma.open(cache_path, "wb") as fp:
+            torch.save(
+                [det.state_dict(after_preprocess=after_preprocess) for det in detection], fp,
+                pickle_protocol=pickle.HIGHEST_PROTOCOL)
+
+    def load_from_cache(self, cache_path: Path):
+        logger.log(f"Loading detection from cache path: {cache_path}")
+        with lzma.open(cache_path, "rb") as fp:
+            state_dict = torch.load(fp)
+        return [
+            state["cls"].from_state_dict(state_dict=state) for state in state_dict
+        ]
+
+    def forward_and_cache(self, im: torch.Tensor, cache_id: str, load_cache: bool):
+        if cache_id is None:
+            return self.forward(im)
+        cache_path = self.cache_directory.joinpath(cache_id + ".torch")
+        if cache_path.is_file() and load_cache:
+            try:
+                return self.load_from_cache(cache_path)
+            except Exception as e:
+                logger.warn(f"The cache file was corrupted: {cache_path}")
+                exit()
+        detections = self.forward(im)
+        self.save_to_cache(detections, cache_path)
+        return detections

deep_privacy/dp2/detection/box_utils.py

@@ -0,0 +1,104 @@
+import numpy as np
+
+
+def expand_bbox_to_ratio(bbox, imshape, target_aspect_ratio):
+    x0, y0, x1, y1 = [int(_) for _ in bbox]
+    h, w = y1 - y0, x1 - x0
+    cur_ratio = h / w
+
+    if cur_ratio == target_aspect_ratio:
+        return [x0, y0, x1, y1]
+    if cur_ratio < target_aspect_ratio:
+        target_height = int(w*target_aspect_ratio)
+        y0, y1 = expand_axis(y0, y1, target_height, imshape[0])
+    else:
+        target_width = int(h/target_aspect_ratio)
+        x0, x1 = expand_axis(x0, x1, target_width, imshape[1])
+    return x0, y0, x1, y1
+
+
+def expand_axis(start, end, target_width, limit):
+    # Can return a bbox outside of limit
+    cur_width = end - start
+    start = start - (target_width-cur_width)//2
+    end = end + (target_width-cur_width)//2
+    if end - start != target_width:
+        end += 1
+    assert end - start == target_width
+    if start < 0 and end > limit:
+        return start, end
+    if start < 0 and end < limit:
+        to_shift = min(0 - start, limit - end)
+        start += to_shift
+        end += to_shift
+    if end > limit and start > 0:
+        to_shift = min(end - limit, start)
+        end -= to_shift
+        start -= to_shift
+    assert end - start == target_width
+    return start, end
+
+
+def expand_box(bbox, imshape, mask, percentage_background: float):
+    assert isinstance(bbox[0], int)
+    assert 0 < percentage_background < 1
+    # Percentage in S
+    mask_pixels = mask.long().sum().cpu()
+    total_pixels = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+    percentage_mask = mask_pixels / total_pixels
+    if (1 - percentage_mask) > percentage_background:
+        return bbox
+    target_pixels = mask_pixels / (1 - percentage_background)
+    x0, y0, x1, y1 = bbox
+    H = y1 - y0
+    W = x1 - x0
+    p = np.sqrt(target_pixels/(H*W))
+    target_width = int(np.ceil(p * W))
+    target_height = int(np.ceil(p * H))
+    x0, x1 = expand_axis(x0, x1, target_width, imshape[1])
+    y0, y1 = expand_axis(y0, y1, target_height, imshape[0])
+    return [x0, y0, x1, y1]
+
+
+def expand_axises_by_percentage(bbox_XYXY, imshape, percentage):
+    x0, y0, x1, y1 = bbox_XYXY
+    H = y1 - y0
+    W = x1 - x0
+    expansion = int(((H*W)**0.5) * percentage)
+    new_width = W + expansion
+    new_height = H + expansion
+    x0, x1 = expand_axis(x0, x1, min(new_width, imshape[1]), imshape[1])
+    y0, y1 = expand_axis(y0, y1, min(new_height, imshape[0]), imshape[0])
+    return [x0, y0, x1, y1]
+
+
+def get_expanded_bbox(
+        bbox_XYXY,
+        imshape,
+        mask,
+        percentage_background: float,
+        axis_minimum_expansion: float,
+        target_aspect_ratio: float):
+    bbox_XYXY = bbox_XYXY.long().cpu().numpy().tolist()
+    # Expand each axis of the bounding box by a minimum percentage
+    bbox_XYXY = expand_axises_by_percentage(bbox_XYXY, imshape, axis_minimum_expansion)
+    # Find the minimum bbox with the aspect ratio. Can be outside of imshape
+    bbox_XYXY = expand_bbox_to_ratio(bbox_XYXY, imshape, target_aspect_ratio)
+    # Expands square box such that X% of the bbox is background
+    bbox_XYXY = expand_box(bbox_XYXY, imshape, mask, percentage_background)
+    assert isinstance(bbox_XYXY[0], (int, np.int64))
+    return bbox_XYXY
+
+
+def include_box(bbox, minimum_area, aspect_ratio_range, min_bbox_ratio_inside, imshape):
+    def area_inside_ratio(bbox, imshape):
+        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+        area_inside = (min(bbox[2], imshape[1]) - max(0, bbox[0])) * (min(imshape[0], bbox[3]) - max(0, bbox[1]))
+        return area_inside / area
+    ratio = (bbox[3] - bbox[1]) / (bbox[2] - bbox[0])
+    area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
+    if area_inside_ratio(bbox, imshape) < min_bbox_ratio_inside:
+        return False
+    if ratio <= aspect_ratio_range[0] or ratio >= aspect_ratio_range[1] or area < minimum_area:
+        return False
+    return True

deep_privacy/dp2/detection/box_utils_fdf.py

@@ -0,0 +1,202 @@
+"""
+The FDF dataset expands bound boxes differently from what is used for CSE.
+"""
+
+import numpy as np
+
+
+def quadratic_bounding_box(x0, y0, width, height, imshape):
+    # We assume that we can create a image that is quadratic without
+    # minimizing any of the sides
+    assert width <= min(imshape[:2])
+    assert height <= min(imshape[:2])
+    min_side = min(height, width)
+    if height != width:
+        side_diff = abs(height - width)
+        # Want to extend the shortest side
+        if min_side == height:
+            # Vertical side
+            height += side_diff
+            if height > imshape[0]:
+                # Take full frame, and shrink width
+                y0 = 0
+                height = imshape[0]
+
+                side_diff = abs(height - width)
+                width -= side_diff
+                x0 += side_diff // 2
+            else:
+                y0 -= side_diff // 2
+                y0 = max(0, y0)
+        else:
+            # Horizontal side
+            width += side_diff
+            if width > imshape[1]:
+                # Take full frame width, and shrink height
+                x0 = 0
+                width = imshape[1]
+
+                side_diff = abs(height - width)
+                height -= side_diff
+                y0 += side_diff // 2
+            else:
+                x0 -= side_diff // 2
+                x0 = max(0, x0)
+        # Check that bbox goes outside image
+        x1 = x0 + width
+        y1 = y0 + height
+        if imshape[1] < x1:
+            diff = x1 - imshape[1]
+            x0 -= diff
+        if imshape[0] < y1:
+            diff = y1 - imshape[0]
+            y0 -= diff
+    assert x0 >= 0, "Bounding box outside image."
+    assert y0 >= 0, "Bounding box outside image."
+    assert x0 + width <= imshape[1], "Bounding box outside image."
+    assert y0 + height <= imshape[0], "Bounding box outside image."
+    return x0, y0, width, height
+
+
+def expand_bounding_box(bbox, percentage, imshape):
+    orig_bbox = bbox.copy()
+    x0, y0, x1, y1 = bbox
+    width = x1 - x0
+    height = y1 - y0
+    x0, y0, width, height = quadratic_bounding_box(
+        x0, y0, width, height, imshape)
+    expanding_factor = int(max(height, width) * percentage)
+
+    possible_max_expansion = [(imshape[0] - width) // 2,
+                              (imshape[1] - height) // 2,
+                              expanding_factor]
+
+    expanding_factor = min(possible_max_expansion)
+    # Expand height
+
+    if expanding_factor > 0:
+
+        y0 = y0 - expanding_factor
+        y0 = max(0, y0)
+
+        height += expanding_factor * 2
+        if height > imshape[0]:
+            y0 -= (imshape[0] - height)
+            height = imshape[0]
+
+        if height + y0 > imshape[0]:
+            y0 -= (height + y0 - imshape[0])
+
+        # Expand width
+        x0 = x0 - expanding_factor
+        x0 = max(0, x0)
+
+        width += expanding_factor * 2
+        if width > imshape[1]:
+            x0 -= (imshape[1] - width)
+            width = imshape[1]
+
+        if width + x0 > imshape[1]:
+            x0 -= (width + x0 - imshape[1])
+    y1 = y0 + height
+    x1 = x0 + width
+    assert y0 >= 0, "Y0 is minus"
+    assert height <= imshape[0], "Height is larger than image."
+    assert x0 + width <= imshape[1]
+    assert y0 + height <= imshape[0]
+    assert width == height, "HEIGHT IS NOT EQUAL WIDTH!!"
+    assert x0 >= 0, "Y0 is minus"
+    assert width <= imshape[1], "Height is larger than image."
+    # Check that original bbox is within new
+    x0_o, y0_o, x1_o, y1_o = orig_bbox
+    assert x0 <= x0_o, f"New bbox is outisde of original. O:{x0_o}, N: {x0}"
+    assert x1 >= x1_o, f"New bbox is outisde of original. O:{x1_o}, N: {x1}"
+    assert y0 <= y0_o, f"New bbox is outisde of original. O:{y0_o}, N: {y0}"
+    assert y1 >= y1_o, f"New bbox is outisde of original. O:{y1_o}, N: {y1}"
+
+    x0, y0, width, height = [int(_) for _ in [x0, y0, width, height]]
+    x1 = x0 + width
+    y1 = y0 + height
+    return np.array([x0, y0, x1, y1])
+
+
+def is_keypoint_within_bbox(x0, y0, x1, y1, keypoint):
+    keypoint = keypoint[:, :3]  # only nose + eyes are relevant
+    kp_X = keypoint[0, :]
+    kp_Y = keypoint[1, :]
+    within_X = np.all(kp_X >= x0) and np.all(kp_X <= x1)
+    within_Y = np.all(kp_Y >= y0) and np.all(kp_Y <= y1)
+    return within_X and within_Y
+
+
+def expand_bbox_simple(bbox, percentage):
+    x0, y0, x1, y1 = bbox.astype(float)
+    width = x1 - x0
+    height = y1 - y0
+    x_c = int(x0) + width // 2
+    y_c = int(y0) + height // 2
+    avg_size = max(width, height)
+    new_width = avg_size * (1 + percentage)
+    x0 = x_c - new_width // 2
+    y0 = y_c - new_width // 2
+    x1 = x_c + new_width // 2
+    y1 = y_c + new_width // 2
+    return np.array([x0, y0, x1, y1]).astype(int)
+
+
+def pad_image(im, bbox, pad_value):
+    x0, y0, x1, y1 = bbox
+    if x0 < 0:
+        pad_im = np.zeros((im.shape[0], abs(x0), im.shape[2]),
+                          dtype=np.uint8) + pad_value
+        im = np.concatenate((pad_im, im), axis=1)
+        x1 += abs(x0)
+        x0 = 0
+    if y0 < 0:
+        pad_im = np.zeros((abs(y0), im.shape[1], im.shape[2]),
+                          dtype=np.uint8) + pad_value
+        im = np.concatenate((pad_im, im), axis=0)
+        y1 += abs(y0)
+        y0 = 0
+    if x1 >= im.shape[1]:
+        pad_im = np.zeros(
+            (im.shape[0], x1 - im.shape[1] + 1, im.shape[2]),
+            dtype=np.uint8) + pad_value
+        im = np.concatenate((im, pad_im), axis=1)
+    if y1 >= im.shape[0]:
+        pad_im = np.zeros(
+            (y1 - im.shape[0] + 1, im.shape[1], im.shape[2]),
+            dtype=np.uint8) + pad_value
+        im = np.concatenate((im, pad_im), axis=0)
+    return im[y0:y1, x0:x1]
+
+
+def clip_box(bbox, im):
+    bbox[0] = max(0, bbox[0])
+    bbox[1] = max(0, bbox[1])
+    bbox[2] = min(im.shape[1] - 1, bbox[2])
+    bbox[3] = min(im.shape[0] - 1, bbox[3])
+    return bbox
+
+
+def cut_face(im, bbox, simple_expand=False, pad_value=0, pad_im=True):
+    outside_im = (bbox < 0).any() or bbox[2] > im.shape[1] or bbox[3] > im.shape[0]
+    if simple_expand or (outside_im and pad_im):
+        return pad_image(im, bbox, pad_value)
+    bbox = clip_box(bbox, im)
+    x0, y0, x1, y1 = bbox
+    return im[y0:y1, x0:x1]
+
+
+def expand_bbox(
+        bbox_ltrb, imshape, simple_expand, default_to_simple=False,
+        expansion_factor=0.35):
+    assert bbox_ltrb.shape == (4,), f"BBox shape was: {bbox_ltrb.shape}"
+    bbox = bbox_ltrb.astype(float)
+    # FDF256 uses simple expand with ratio 0.4
+    if simple_expand:
+        return expand_bbox_simple(bbox, 0.4)
+    try:
+        return expand_bounding_box(bbox, expansion_factor, imshape)
+    except AssertionError:
+        return expand_bbox_simple(bbox, expansion_factor * 2)