Upload 52 files

.gitattributes

@@ -41,3 +41,8 @@ example_videos/ex3.mp4 filter=lfs diff=lfs merge=lfs -text
 example_videos/ex4.mp4 filter=lfs diff=lfs merge=lfs -text
 example_videos/ex5.mp4 filter=lfs diff=lfs merge=lfs -text
 media/banner.gif filter=lfs diff=lfs merge=lfs -text
+src/pixel3dmm/preprocessing/facer/samples/data/ffhq_15723.jpg filter=lfs diff=lfs merge=lfs -text
+src/pixel3dmm/preprocessing/facer/samples/data/weirdface.jpg filter=lfs diff=lfs merge=lfs -text
+src/pixel3dmm/preprocessing/facer/samples/example_output/alignment.png filter=lfs diff=lfs merge=lfs -text
+src/pixel3dmm/preprocessing/facer/samples/example_output/detect.png filter=lfs diff=lfs merge=lfs -text
+src/pixel3dmm/preprocessing/facer/samples/example_output/parsing.png filter=lfs diff=lfs merge=lfs -text

src/pixel3dmm/preprocessing/facer/.gitignore

@@ -0,0 +1,134 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+samples/output
+.local/
+
+.token.txt
+.downloaded/

src/pixel3dmm/preprocessing/facer/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 FacePerceiver
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/pixel3dmm/preprocessing/facer/README.md

@@ -0,0 +1,187 @@
+# FACER
+
+Face related toolkit. This repo is still under construction to include more models.
+
+## Updates
+- [01/17/2025] FaRL Face Parsing is now available in [elliottzheng/batch-face](https://github.com/elliottzheng/batch-face?tab=readme-ov-file#face-parsing), which provides faster batch processing ability.
+- [14/05/2023] Face attribute recognition model trained on CelebA is available, check it out [here](./samples/face_attribute.ipynb).
+- [04/05/2023] Face alignment model trained on IBUG300W, AFLW19, WFLW dataset is available, check it out [here](./samples/face_alignment.ipynb).
+- [27/04/2023] Face parsing model trained on CelebM dataset is available, check it out [here](./samples/face_parsing.ipynb).
+
+## Install
+
+The easiest way to install it is using pip:
+
+```bash
+pip install git+https://github.com/FacePerceiver/facer.git@main
+```
+No extra setup needs, pretrained weights will be downloaded automatically.
+
+If you have trouble install from source, you can try install from PyPI:
+```bash
+pip install pyfacer
+```
+the PyPI version is not guaranteed to be the latest version, but we will try to keep it up to date.
+
+
+## Face Detection
+
+We simply wrap a retinaface detector for easy usage.
+```python
+import facer
+
+image = facer.hwc2bchw(facer.read_hwc('data/twogirls.jpg')).to(device=device)  # image: 1 x 3 x h x w
+
+face_detector = facer.face_detector('retinaface/mobilenet', device=device)
+with torch.inference_mode():
+    faces = face_detector(image)
+
+facer.show_bchw(facer.draw_bchw(image, faces))
+```
+![](./samples/example_output/detect.png)
+
+Check [this notebook](./samples/face_detect.ipynb) for full example.
+
+Please consider citing
+```
+@inproceedings{deng2020retinaface,
+  title={Retinaface: Single-shot multi-level face localisation in the wild},
+  author={Deng, Jiankang and Guo, Jia and Ververas, Evangelos and Kotsia, Irene and Zafeiriou, Stefanos},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={5203--5212},
+  year={2020}
+}
+```
+
+## Face Parsing
+
+We wrap the [FaRL](https://github.com/faceperceiver/farl) models for face parsing.
+```python
+import torch
+import facer
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+image = facer.hwc2bchw(facer.read_hwc('data/twogirls.jpg')).to(device=device)  # image: 1 x 3 x h x w
+
+face_detector = facer.face_detector('retinaface/mobilenet', device=device)
+with torch.inference_mode():
+    faces = face_detector(image)
+
+face_parser = facer.face_parser('farl/lapa/448', device=device) # optional "farl/celebm/448"
+
+with torch.inference_mode():
+    faces = face_parser(image, faces)
+
+seg_logits = faces['seg']['logits']
+seg_probs = seg_logits.softmax(dim=1)  # nfaces x nclasses x h x w
+n_classes = seg_probs.size(1)
+vis_seg_probs = seg_probs.argmax(dim=1).float()/n_classes*255
+vis_img = vis_seg_probs.sum(0, keepdim=True)
+facer.show_bhw(vis_img)
+facer.show_bchw(facer.draw_bchw(image, faces))
+```
+![](./samples/example_output/parsing.png)
+
+Check [this notebook](./samples/face_parsing.ipynb) for full example.
+
+Please consider citing
+```
+@inproceedings{zheng2022farl,
+  title={General facial representation learning in a visual-linguistic manner},
+  author={Zheng, Yinglin and Yang, Hao and Zhang, Ting and Bao, Jianmin and Chen, Dongdong and Huang, Yangyu and Yuan, Lu and Chen, Dong and Zeng, Ming and Wen, Fang},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={18697--18709},
+  year={2022}
+}
+``` 
+
+
+## Face Alignment
+
+We wrap the [FaRL](https://github.com/faceperceiver/farl) models for face alignment.
+```python
+import torch
+import cv2
+from matplotlib import pyplot as plt
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+import facer
+img_file = 'data/twogirls.jpg'
+# image: 1 x 3 x h x w
+image = facer.hwc2bchw(facer.read_hwc(img_file)).to(device=device)  
+
+face_detector = facer.face_detector('retinaface/mobilenet', device=device)
+with torch.inference_mode():
+    faces = face_detector(image)
+
+face_aligner = facer.face_aligner('farl/ibug300w/448', device=device) # optional: "farl/wflw/448", "farl/aflw19/448"
+
+with torch.inference_mode():
+    faces = face_aligner(image, faces)
+
+img = cv2.imread(img_file)[..., ::-1]
+vis_img = img.copy()
+for pts in faces['alignment']:
+    vis_img = facer.draw_landmarks(vis_img, None, pts.cpu().numpy())
+plt.imshow(vis_img)
+```
+![](./samples/example_output/alignment.png)
+
+Check [this notebook](./samples/face_alignment.ipynb) for full example.
+
+Please consider citing
+```
+@inproceedings{zheng2022farl,
+  title={General facial representation learning in a visual-linguistic manner},
+  author={Zheng, Yinglin and Yang, Hao and Zhang, Ting and Bao, Jianmin and Chen, Dongdong and Huang, Yangyu and Yuan, Lu and Chen, Dong and Zeng, Ming and Wen, Fang},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={18697--18709},
+  year={2022}
+}
+``` 
+
+## Face Attribute Recognition
+We wrap the [FaRL](https://github.com/faceperceiver/farl) models for face attribute recognition, the model achieves 92.06% accuracy on [CelebA](https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) dataset.
+
+```python
+import sys
+import torch
+import facer
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# image: 1 x 3 x h x w
+image = facer.hwc2bchw(facer.read_hwc("data/girl.jpg")).to(device=device)
+
+face_detector = facer.face_detector("retinaface/mobilenet", device=device)
+with torch.inference_mode():
+    faces = face_detector(image)
+
+face_attr = facer.face_attr("farl/celeba/224", device=device)
+with torch.inference_mode():
+    faces = face_attr(image, faces)
+
+labels = face_attr.labels
+face1_attrs = faces["attrs"][0] # get the first face's attributes
+
+print(labels)
+
+for prob, label in zip(face1_attrs, labels):
+    if prob > 0.5:
+        print(label, prob.item())
+```
+
+Check [this notebook](./samples/face_attribute.ipynb) for full example.
+
+Please consider citing
+```
+@inproceedings{zheng2022farl,
+  title={General facial representation learning in a visual-linguistic manner},
+  author={Zheng, Yinglin and Yang, Hao and Zhang, Ting and Bao, Jianmin and Chen, Dongdong and Huang, Yangyu and Yuan, Lu and Chen, Dong and Zeng, Ming and Wen, Fang},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={18697--18709},
+  year={2022}
+}
+``` 

src/pixel3dmm/preprocessing/facer/facer/__init__.py

@@ -0,0 +1,55 @@
+from typing import Optional, Tuple
+import torch
+
+from .io import read_hwc, write_hwc
+from .util import hwc2bchw, bchw2hwc, bchw2bhwc, bhwc2bchw, bhwc2hwc
+from .draw import draw_bchw, draw_landmarks
+from .show import show_bchw, show_bhw
+
+from .face_detection import FaceDetector
+from .face_parsing import FaceParser
+from .face_alignment import FaceAlignment
+from .face_attribute import FaceAttribute
+
+
+def _split_name(name: str) -> Tuple[str, Optional[str]]:
+    if '/' in name:
+        detector_type, conf_name = name.split('/', 1)
+    else:
+        detector_type, conf_name = name, None
+    return detector_type, conf_name
+
+
+def face_detector(name: str, device: torch.device, **kwargs) -> FaceDetector:
+    detector_type, conf_name = _split_name(name)
+    if detector_type == 'retinaface':
+        from .face_detection import RetinaFaceDetector
+        return RetinaFaceDetector(conf_name, **kwargs).to(device)
+    else:
+        raise RuntimeError(f'Unknown detector type: {detector_type}')
+
+
+def face_parser(name: str, device: torch.device, **kwargs) -> FaceParser:
+    parser_type, conf_name = _split_name(name)
+    if parser_type == 'farl':
+        from .face_parsing import FaRLFaceParser
+        return FaRLFaceParser(conf_name, device=device, **kwargs).to(device)
+    else:
+        raise RuntimeError(f'Unknown parser type: {parser_type}')
+
+
+def face_aligner(name: str, device: torch.device, **kwargs) -> FaceAlignment:
+    aligner_type, conf_name = _split_name(name)
+    if aligner_type == 'farl':
+        from .face_alignment import FaRLFaceAlignment
+        return FaRLFaceAlignment(conf_name, device=device, **kwargs).to(device)
+    else:
+        raise RuntimeError(f'Unknown aligner type: {aligner_type}')
+
+def face_attr(name: str, device: torch.device, **kwargs) -> FaceAttribute:
+    attr_type, conf_name = _split_name(name)
+    if attr_type == 'farl':
+        from .face_attribute import FaRLFaceAttribute
+        return FaRLFaceAttribute(conf_name, device=device, **kwargs).to(device)
+    else:
+        raise RuntimeError(f'Unknown attribute type: {attr_type}')

src/pixel3dmm/preprocessing/facer/facer/draw.py

@@ -0,0 +1,186 @@
+from typing import Dict, List
+import torch
+import colorsys
+import random
+import numpy as np
+from skimage.draw import line_aa, circle_perimeter_aa
+import cv2
+from .util import select_data
+
+
+def _gen_random_colors(N, bright=True):
+    brightness = 1.0 if bright else 0.7
+    hsv = [(i / N, 1, brightness) for i in range(N)]
+    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
+    random.shuffle(colors)
+    return colors
+
+
+_static_label_colors = [
+    np.array((1.0, 1.0, 1.0), np.float32),
+    np.array((255, 250, 79), np.float32) / 255.0,  # face
+    np.array([255, 125, 138], np.float32) / 255.0,  # lb
+    np.array([213, 32, 29], np.float32) / 255.0,  # rb
+    np.array([0, 144, 187], np.float32) / 255.0,  # le
+    np.array([0, 196, 253], np.float32) / 255.0,  # re
+    np.array([255, 129, 54], np.float32) / 255.0,  # nose
+    np.array([88, 233, 135], np.float32) / 255.0,  # ulip
+    np.array([0, 117, 27], np.float32) / 255.0,  # llip
+    np.array([255, 76, 249], np.float32) / 255.0,  # imouth
+    np.array((1.0, 0.0, 0.0), np.float32),  # hair
+    np.array((255, 250, 100), np.float32) / 255.0,  # lr
+    np.array((255, 250, 100), np.float32) / 255.0,  # rr
+    np.array((250, 245, 50), np.float32) / 255.0,  # neck
+    np.array((0.0, 1.0, 0.5), np.float32),  # cloth
+    np.array((1.0, 0.0, 0.5), np.float32),
+] + _gen_random_colors(256)
+
+_names_in_static_label_colors = [
+    'background', 'face', 'lb', 'rb', 'le', 're', 'nose',
+    'ulip', 'llip', 'imouth', 'hair', 'lr', 'rr', 'neck',
+    'cloth', 'eyeg', 'hat', 'earr'
+]
+
+
+def _blend_labels(image, labels, label_names_dict=None,
+                  default_alpha=0.6, color_offset=None):
+    assert labels.ndim == 2
+    bg_mask = labels == 0
+    if label_names_dict is None:
+        colors = _static_label_colors
+    else:
+        colors = [np.array((1.0, 1.0, 1.0), np.float32)]
+        for i in range(1, labels.max() + 1):
+            if isinstance(label_names_dict, dict) and i not in label_names_dict:
+                bg_mask = np.logical_or(bg_mask, labels == i)
+                colors.append(np.zeros((3)))
+                continue
+            label_name = label_names_dict[i]
+            if label_name in _names_in_static_label_colors:
+                color = _static_label_colors[
+                    _names_in_static_label_colors.index(
+                        label_name)]
+            else:
+                color = np.array((1.0, 1.0, 1.0), np.float32)
+            colors.append(color)
+
+    if color_offset is not None:
+        ncolors = []
+        for c in colors:
+            nc = np.array(c)
+            if (nc != np.zeros(3)).any():
+                nc += color_offset
+            ncolors.append(nc)
+        colors = ncolors
+
+    if image is None:
+        image = orig_image = np.zeros(
+            [labels.shape[0], labels.shape[1], 3], np.float32)
+        alpha = 1.0
+    else:
+        orig_image = image / np.max(image)
+        image = orig_image * (1.0 - default_alpha)
+        alpha = default_alpha
+    for i in range(1, np.max(labels) + 1):
+        image += alpha * \
+            np.tile(
+                np.expand_dims(
+                    (labels == i).astype(np.float32), -1),
+                [1, 1, 3]) * colors[(i) % len(colors)]
+    image[np.where(image > 1.0)] = 1.0
+    image[np.where(image < 0)] = 0.0
+    image[np.where(bg_mask)] = orig_image[np.where(bg_mask)]
+    return image
+
+
+def _draw_hwc(image: torch.Tensor, data: Dict[str, torch.Tensor]):
+    device = image.device
+    image = np.array(image.cpu().numpy(), copy=True)
+    dtype = image.dtype
+    h, w, _ = image.shape
+
+    draw_score_error = False
+    for tag, batch_content in data.items():
+        if tag == 'rects':
+            for cid, content in enumerate(batch_content):
+                x1, y1, x2, y2 = [int(v) for v in content]
+                y1, y2 = [max(min(v, h-1), 0) for v in [y1, y2]]
+                x1, x2 = [max(min(v, w-1), 0) for v in [x1, x2]]
+                for xx1, yy1, xx2, yy2 in [
+                    [x1, y1, x2, y1],
+                    [x1, y2, x2, y2],
+                    [x1, y1, x1, y2],
+                    [x2, y1, x2, y2]
+                ]:
+                    rr, cc, val = line_aa(yy1, xx1, yy2, xx2)
+                    val = val[:, None][:, [0, 0, 0]]
+                    image[rr, cc] = image[rr, cc] * (1.0-val) + val * 255
+
+                if 'scores' in data:
+                    try:
+                        import cv2
+                        score = data['scores'][cid].item()
+                        score_str = f'{score:0.3f}'
+                        image_c = np.array(image).copy()
+                        cv2.putText(image_c, score_str, org=(int(x1), int(y2)),
+                                    fontFace=cv2.FONT_HERSHEY_TRIPLEX,
+                                    fontScale=0.6, color=(255, 255, 255), thickness=1)
+                        image[:, :, :] = image_c
+                    except Exception as e:
+                        if not draw_score_error:
+                            print(f'Failed to draw scores on image.')
+                            print(e)
+                        draw_score_error = True
+
+        if tag == 'points':
+            for content in batch_content:
+                # content: npoints x 2
+                for x, y in content:
+                    x = max(min(int(x), w-1), 0)
+                    y = max(min(int(y), h-1), 0)
+                    rr, cc, val = circle_perimeter_aa(y, x, 1)
+                    valid = np.all([rr >= 0, rr < h, cc >= 0, cc < w], axis=0)
+                    rr = rr[valid]
+                    cc = cc[valid]
+                    val = val[valid]
+                    val = val[:, None][:, [0, 0, 0]]
+                    image[rr, cc] = image[rr, cc] * (1.0-val) + val * 255
+
+        if tag == 'seg':
+            label_names = batch_content['label_names']
+            for seg_logits in batch_content['logits']:
+                # content: nclasses x h x w
+                seg_probs = seg_logits.softmax(dim=0)
+                seg_labels = seg_probs.argmax(dim=0).cpu().numpy()
+                image = (_blend_labels(image.astype(np.float32) /
+                         255, seg_labels,
+                         label_names_dict=label_names) * 255).astype(dtype)
+
+    return torch.from_numpy(image).to(device=device)
+
+
+def draw_bchw(images: torch.Tensor, data: Dict[str, torch.Tensor]) -> torch.Tensor:
+    images2 = []
+    for image_id, image_chw in enumerate(images):
+        selected_data = select_data(image_id == data['image_ids'], data)
+        images2.append(
+            _draw_hwc(image_chw.permute(1, 2, 0), selected_data).permute(2, 0, 1))
+    return torch.stack(images2, dim=0)
+
+def draw_landmarks(img, bbox=None, landmark=None, color=(0, 255, 0)):
+    """
+    Input:
+    - img: gray or RGB
+    - bbox: type of BBox
+    - landmark: reproject landmark of (5L, 2L)
+    Output:
+    - img marked with landmark and bbox
+    """
+    img = cv2.UMat(img).get()
+    if bbox is not None:
+        x1, y1, x2, y2 = np.array(bbox)[:4].astype(np.int32)
+        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
+    if landmark is not None:
+        for x, y in np.array(landmark).astype(np.int32):
+            cv2.circle(img, (int(x), int(y)), 2, color, -1)
+    return img

src/pixel3dmm/preprocessing/facer/facer/face_alignment/__init__.py

@@ -0,0 +1,2 @@
+from .base import FaceAlignment
+from .farl import FaRLFaceAlignment

src/pixel3dmm/preprocessing/facer/facer/face_alignment/base.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+
+class FaceAlignment(nn.Module):
+    """ face alignment
+
+    Args:
+        images (torch.Tensor): b x c x h x w
+
+        data (Dict[str, Any]):
+
+            * image_ids (torch.Tensor): nfaces
+            * rects (torch.Tensor): nfaces x 4 (x1, y1, x2, y2)
+            * points (torch.Tensor): nfaces x 5 x 2 (x, y)
+
+    Returns:
+        data (Dict[str, Any]):
+
+            * image_ids (torch.Tensor): nfaces
+            * rects (torch.Tensor): nfaces x 4 (x1, y1, x2, y2)
+            * points (torch.Tensor): nfaces x 5 x 2 (x, y)
+            * alignment 
+    """
+    pass

src/pixel3dmm/preprocessing/facer/facer/face_alignment/farl.py

@@ -0,0 +1,180 @@
+from typing import Optional, Dict, Any
+import functools
+import torch
+import torch.nn.functional as F
+from .network import FaRLVisualFeatures, MMSEG_UPerHead, FaceAlignmentTransformer, denormalize_points, heatmap2points
+from ..transform import (get_face_align_matrix,
+                         make_inverted_tanh_warp_grid, make_tanh_warp_grid)
+from .base import FaceAlignment
+from ..util import download_jit
+import io
+
+pretrain_settings = {
+    'ibug300w/448': {
+        # inter_ocular 0.028835 epoch 60
+        'num_classes': 68,
+        'url': "https://github.com/FacePerceiver/facer/releases/download/models-v1/face_alignment.farl.ibug300w.main_ema_jit.pt",
+        'matrix_src_tag': 'points',
+        'get_matrix_fn': functools.partial(get_face_align_matrix,
+                                           target_shape=(448, 448), target_face_scale=0.8),
+        'get_grid_fn': functools.partial(make_tanh_warp_grid,
+                                         warp_factor=0.0, warped_shape=(448, 448)),
+        'get_inv_grid_fn': functools.partial(make_inverted_tanh_warp_grid,
+                                             warp_factor=0.0, warped_shape=(448, 448)),
+        
+    },
+    'aflw19/448': {
+        # diag 0.009329 epoch 15
+        'num_classes': 19,
+        'url': "https://github.com/FacePerceiver/facer/releases/download/models-v1/face_alignment.farl.aflw19.main_ema_jit.pt",
+        'matrix_src_tag': 'points',
+        'get_matrix_fn': functools.partial(get_face_align_matrix,
+                                           target_shape=(448, 448), target_face_scale=0.8),
+        'get_grid_fn': functools.partial(make_tanh_warp_grid,
+                                         warp_factor=0.0, warped_shape=(448, 448)),
+        'get_inv_grid_fn': functools.partial(make_inverted_tanh_warp_grid,
+                                             warp_factor=0.0, warped_shape=(448, 448)),
+    },
+    'wflw/448': {
+        # inter_ocular 0.038933 epoch 20
+        'num_classes': 98,
+        'url': "https://github.com/FacePerceiver/facer/releases/download/models-v1/face_alignment.farl.wflw.main_ema_jit.pt",
+        'matrix_src_tag': 'points',
+        'get_matrix_fn': functools.partial(get_face_align_matrix,
+                                           target_shape=(448, 448), target_face_scale=0.8),
+        'get_grid_fn': functools.partial(make_tanh_warp_grid,
+                                         warp_factor=0.0, warped_shape=(448, 448)),
+        'get_inv_grid_fn': functools.partial(make_inverted_tanh_warp_grid,
+                                             warp_factor=0.0, warped_shape=(448, 448)),
+    },
+
+}
+
+
+def load_face_alignment_model(model_path: str, num_classes=68):
+    backbone = FaRLVisualFeatures("base", None, forced_input_resolution=448, output_indices=None).cpu()
+    if "jit" in model_path:
+        extra_files = {"backbone": None}
+        heatmap_head = download_jit(model_path, map_location="cpu", _extra_files=extra_files)
+        backbone_weight_io = io.BytesIO(extra_files["backbone"])
+        backbone.load_state_dict(torch.load(backbone_weight_io))
+        # print("load from jit")
+    else:
+        channels = backbone.get_output_channel("base")
+        in_channels = [channels] * 4
+        num_classes = num_classes
+        heatmap_head = MMSEG_UPerHead(in_channels=in_channels, channels=channels, num_classes=num_classes) # this requires mmseg as a dependency
+        state = torch.load(model_path,map_location="cpu")["networks"]["main_ema"]
+        # print("load from checkpoint")
+
+    main_network = FaceAlignmentTransformer(backbone, heatmap_head, heatmap_act="sigmoid").cpu()
+
+    if "jit" not in model_path:
+        main_network.load_state_dict(state, strict=True)
+
+    return main_network
+
+
+
+class FaRLFaceAlignment(FaceAlignment):
+    """ The face alignment models from [FaRL](https://github.com/FacePerceiver/FaRL).
+
+    Please consider citing 
+    ```bibtex
+        @article{zheng2021farl,
+            title={General Facial Representation Learning in a Visual-Linguistic Manner},
+            author={Zheng, Yinglin and Yang, Hao and Zhang, Ting and Bao, Jianmin and Chen, 
+                Dongdong and Huang, Yangyu and Yuan, Lu and Chen, 
+                Dong and Zeng, Ming and Wen, Fang},
+            journal={arXiv preprint arXiv:2112.03109},
+            year={2021}
+        }
+    ```
+    """
+
+    def __init__(self, conf_name: Optional[str] = None,
+                 model_path: Optional[str] = None, device=None) -> None:
+        super().__init__()
+        if conf_name is None:
+            conf_name = 'ibug300w/448'
+        if model_path is None:
+            model_path = pretrain_settings[conf_name]['url']
+        self.conf_name = conf_name
+
+        setting  = pretrain_settings[self.conf_name]
+        self.net = load_face_alignment_model(model_path, num_classes = setting["num_classes"])
+        if device is not None:
+            self.net = self.net.to(device)
+
+        self.heatmap_interpolate_mode = 'bilinear'
+        self.eval()
+
+    def forward(self, images: torch.Tensor, data: Dict[str, Any]):
+        setting = pretrain_settings[self.conf_name]
+        images = images.float() / 255.0 # backbone 自带 normalize
+        _, _, h, w = images.shape
+
+        simages = images[data['image_ids']]
+        matrix = setting['get_matrix_fn'](data[setting['matrix_src_tag']])
+        grid = setting['get_grid_fn'](matrix=matrix, orig_shape=(h, w))
+        inv_grid = setting['get_inv_grid_fn'](matrix=matrix, orig_shape=(h, w))
+
+        w_images = F.grid_sample(
+            simages, grid, mode='bilinear', align_corners=False)
+
+        _, _, warp_h, warp_w = w_images.shape
+
+        heatmap_acted = self.net(w_images)
+
+        warpped_heatmap = F.interpolate(
+                            heatmap_acted, size=(warp_h, warp_w),
+                            mode=self.heatmap_interpolate_mode, align_corners=False)
+
+        pred_heatmap = F.grid_sample(
+            warpped_heatmap, inv_grid, mode='bilinear', align_corners=False)
+
+        landmark = heatmap2points(pred_heatmap)
+
+        landmark = denormalize_points(landmark, h, w)
+
+        data['alignment'] = landmark
+
+        return data
+
+
+if __name__=="__main__":
+    image = torch.randn(1, 3, 448, 448)
+
+    aligner1 = FaRLFaceAlignment("wflw/448")
+    
+    x1 =  aligner1.net(image)
+
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--jit_path", type=str, default=None)
+    args = parser.parse_args()
+
+    if args.jit_path is None:
+        exit(0)
+        
+    net  = aligner1.net.cpu()
+
+    features, _ = net.backbone(image)
+
+    # Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
+    traced_script_module = torch.jit.trace(net.heatmap_head, example_inputs=[features])
+
+    buffer = io.BytesIO()
+
+    torch.save(net.backbone.state_dict(), buffer)
+    
+    # Save to file
+    torch.jit.save(traced_script_module, args.jit_path,
+                   _extra_files={"backbone": buffer.getvalue()})
+
+    aligner2 = FaRLFaceAlignment(model_path=args.jit_path)
+
+    # compare the output
+    x2 = aligner2.net(image)
+    print(torch.allclose(x1, x2))

src/pixel3dmm/preprocessing/facer/facer/face_alignment/network/__init__.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from .common import (load_checkpoint, Activation, MLP, Residual)
+from .geometry import (normalize_points, denormalize_points,
+                       heatmap2points)
+from .mmseg import MMSEG_UPerHead
+from .transformers import FaRLVisualFeatures
+from torch import nn
+from typing import Optional, List, Tuple
+
+
+
+class FaceAlignmentTransformer(nn.Module):
+    """Face alignment transformer.
+
+    Args:
+        image (torch.Tensor): Float32 tensor with shape [b, 3, h, w], normalized to [0, 1].
+
+    Returns:
+        landmark (torch.Tensor): Float32 tensor with shape [b, npoints, 2], coordinates normalized to [0, 1].
+        aux_outputs:
+            heatmap (torch.Tensor): Float32 tensor with shape [b, npoints, S, S]
+    """
+
+    def __init__(self, backbone: nn.Module, heatmap_head: nn.Module,
+                 heatmap_act: Optional[str] = 'relu'):
+        super().__init__()
+        self.backbone = backbone
+        self.heatmap_head = heatmap_head
+        self.heatmap_act = Activation(heatmap_act)
+        self.float()
+
+    def forward(self, image):
+        features, _ = self.backbone(image)
+        heatmap = self.heatmap_head(features)  # b x npoints x s x s
+        heatmap_acted = self.heatmap_act(heatmap)
+        # landmark = heatmap2points(heatmap_acted)  # b x npoints x 2
+        # return landmark, {'heatmap': heatmap, 'heatmap_acted': heatmap_acted}
+        return heatmap_acted
+
+

src/pixel3dmm/preprocessing/facer/facer/face_alignment/network/common.py

@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import List, Optional, Tuple, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def load_checkpoint(net: nn.Module, checkpoint_path: str, network_name: str):
+    states = torch.load(open(checkpoint_path, 'rb'), map_location={
+                        'cuda:0': f'cuda:{torch.cuda.current_device()}'})
+    network_states = states['networks']
+    net.load_state_dict(network_states[network_name])
+    return net
+
+
+class Activation(nn.Module):
+    def __init__(self, name: Optional[str], **kwargs):
+        super().__init__()
+        if name == 'relu':
+            self.fn = F.relu
+        elif name == 'softplus':
+            self.fn = F.softplus
+        elif name == 'gelu':
+            self.fn = F.gelu
+        elif name == 'sigmoid':
+            self.fn = torch.sigmoid
+        elif name == 'sigmoid_x':
+            self.epsilon = kwargs.get('epsilon', 1e-3)
+            self.fn = lambda x: torch.clamp(
+                x.sigmoid() * (1.0 + self.epsilon*2.0) - self.epsilon,
+                min=0.0, max=1.0)
+        elif name == None:
+            self.fn = lambda x: x
+        else:
+            raise RuntimeError(f'Unknown activation name: {name}')
+
+    def forward(self, x):
+        return self.fn(x)
+
+
+class MLP(nn.Module):
+    def __init__(self, channels: List[int], act: Optional[str]):
+        super().__init__()
+        assert len(channels) > 1
+        layers = []
+        for i in range(len(channels)-1):
+            layers.append(nn.Linear(channels[i], channels[i+1]))
+            if i+1 < len(channels):
+                layers.append(Activation(act))
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class Residual(nn.Module):
+    def __init__(self, net: nn.Module, res_weight_init: Optional[float] = 0.0):
+        super().__init__()
+        self.net = net
+        if res_weight_init is not None:
+            self.res_weight = nn.Parameter(torch.tensor(res_weight_init))
+        else:
+            self.res_weight = None
+
+    def forward(self, x):
+        if self.res_weight is not None:
+            return self.res_weight * self.net(x) + x
+        else:
+            return self.net(x) + x
+
+
+class SE(nn.Module):
+    def __init__(self, channel: int, r: int = 1):
+        super().__init__()
+        self.branch = nn.Sequential(
+            nn.Conv2d(channel, channel//r, (1, 1)),
+            nn.ReLU(),
+            nn.Conv2d(channel//r, channel, (1, 1)),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        # x: b x channel x h x w
+        v = x.mean([2, 3], keepdim=True)  # b x channel x 1 x 1
+        v = self.branch(v)  # b x channel x 1 x 1
+        return x * v
+
+

src/pixel3dmm/preprocessing/facer/facer/face_alignment/network/geometry.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Tuple, Union
+
+import torch
+
+
+def normalize_points(points: torch.Tensor, h: int, w: int) -> torch.Tensor:
+    """ Normalize coordinates to [0, 1].
+    """
+    return (points + 0.5) / torch.tensor([[[w, h]]]).to(points)
+
+
+def denormalize_points(normalized_points: torch.Tensor, h: int, w: int) -> torch.Tensor:
+    """ Reverse normalize_points.
+    """
+    return normalized_points * torch.tensor([[[w, h]]]).to(normalized_points) - 0.5
+
+
+def heatmap2points(heatmap, t_scale: Union[None, float, torch.Tensor] = None):
+    """ Heatmaps -> normalized points [b x npoints x 2(XY)].
+    """
+    dtype = heatmap.dtype
+    _, _, h, w = heatmap.shape
+
+    # 0 ~ h-1, 0 ~ w-1
+    yy, xx = torch.meshgrid(
+        torch.arange(h).float(),
+        torch.arange(w).float())
+
+    yy = yy.view(1, 1, h, w).to(heatmap)
+    xx = xx.view(1, 1, h, w).to(heatmap)
+
+    if t_scale is not None:
+        heatmap = (heatmap * t_scale).exp()
+    heatmap_sum = torch.clamp(heatmap.sum([2, 3]), min=1e-6)
+
+    yy_coord = (yy * heatmap).sum([2, 3]) / heatmap_sum  # b x npoints
+    xx_coord = (xx * heatmap).sum([2, 3]) / heatmap_sum  # b x npoints
+
+    points = torch.stack([xx_coord, yy_coord], dim=-1)  # b x npoints x 2
+
+    normalized_points = normalize_points(points, h, w)
+    return normalized_points

src/pixel3dmm/preprocessing/facer/facer/face_alignment/network/mmseg.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch.nn as nn
+
+
+class MMSEG_UPerHead(nn.Module):
+    """Wraps the UPerHead from mmseg for segmentation.
+    """
+
+    def __init__(self, num_classes: int,
+                 in_channels: list = [384, 384, 384, 384], channels: int = 512):
+        super().__init__()
+
+        from mmseg.models.decode_heads import UPerHead
+        self.head = UPerHead(
+            in_channels=in_channels,
+            in_index=[0, 1, 2, 3],
+            pool_scales=(1, 2, 3, 6),
+            channels=channels,
+            dropout_ratio=0.1,
+            num_classes=num_classes,
+            norm_cfg=dict(type='SyncBN', requires_grad=True),
+            align_corners=False,
+            loss_decode=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0))
+
+    def forward(self, inputs):
+        return self.head(inputs)

src/pixel3dmm/preprocessing/facer/facer/face_alignment/network/transformers.py

@@ -0,0 +1,173 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import math
+
+from typing import Optional, List, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+from ... import farl
+
+
+def _make_fpns(vision_patch_size: int, output_channels: int):
+    if vision_patch_size in {16, 14}:
+        fpn1 = nn.Sequential(
+            nn.ConvTranspose2d(output_channels, output_channels,
+                               kernel_size=2, stride=2),
+            nn.SyncBatchNorm(output_channels),
+            nn.GELU(),
+            nn.ConvTranspose2d(output_channels, output_channels, kernel_size=2, stride=2))
+
+        fpn2 = nn.ConvTranspose2d(
+            output_channels, output_channels, kernel_size=2, stride=2)
+        fpn3 = nn.Identity()
+        fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
+        return nn.ModuleList([fpn1, fpn2, fpn3, fpn4])
+    elif vision_patch_size == 8:
+        fpn1 = nn.Sequential(nn.ConvTranspose2d(
+            output_channels, output_channels, kernel_size=2, stride=2))
+        fpn2 = nn.Identity()
+        fpn3 = nn.MaxPool2d(kernel_size=2, stride=2)
+        fpn4 = nn.MaxPool2d(kernel_size=4, stride=4)
+        return nn.ModuleList([fpn1, fpn2, fpn3, fpn4])
+    else:
+        raise NotImplementedError()
+
+
+def _resize_pe(pe: torch.Tensor, new_size: int, mode: str = 'bicubic', num_tokens: int = 1) -> torch.Tensor:
+    """Resize positional embeddings.
+
+    Args: 
+        pe (torch.Tensor): A tensor with shape (num_tokens + old_size ** 2, width). pe[0, :] is the CLS token.
+
+    Returns:
+        torch.Tensor: A tensor with shape (num_tokens + new_size **2, width).
+    """
+    l, w = pe.shape
+    old_size = int(math.sqrt(l-num_tokens))
+    assert old_size ** 2 + num_tokens == l
+    return torch.cat([
+        pe[:num_tokens, :],
+        F.interpolate(pe[num_tokens:, :].reshape(1, old_size, old_size, w).permute(0, 3, 1, 2),
+                      (new_size, new_size), mode=mode, align_corners=False).view(w, -1).t()], dim=0)
+
+
+class FaRLVisualFeatures(nn.Module):
+    """Extract features from FaRL visual encoder.
+
+    Args:
+        image (torch.Tensor): Float32 tensor with shape [b, 3, h, w], 
+            normalized to [0, 1].
+
+    Returns:
+        List[torch.Tensor]: A list of features.
+    """
+    image_mean: torch.Tensor
+    image_std: torch.Tensor
+    output_channels: int
+    num_outputs: int
+
+    def __init__(self, model_type: str,
+                 model_path: Optional[str] = None, output_indices: Optional[List[int]] = None,
+                 forced_input_resolution: Optional[int] = None,
+                 apply_fpn: bool = True):
+        super().__init__()
+        self.visual = farl.load_farl(model_type, model_path)
+
+        vision_patch_size = self.visual.conv1.weight.shape[-1]
+
+        self.input_resolution = self.visual.input_resolution
+        if forced_input_resolution is not None and \
+                self.input_resolution != forced_input_resolution:
+            # resizing the positonal embeddings
+            self.visual.positional_embedding = nn.Parameter(
+                _resize_pe(self.visual.positional_embedding,
+                           forced_input_resolution//vision_patch_size))
+            self.input_resolution = forced_input_resolution
+
+        self.output_channels = self.visual.transformer.width
+
+        if output_indices is None:
+            output_indices = self.__class__.get_default_output_indices(
+                model_type)
+        self.output_indices = output_indices
+        self.num_outputs = len(output_indices)
+
+        self.register_buffer('image_mean', torch.tensor(
+            [0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1))
+        self.register_buffer('image_std', torch.tensor(
+            [0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1))
+
+        if apply_fpn:
+            self.fpns = _make_fpns(vision_patch_size, self.output_channels)
+        else:
+            self.fpns = None
+
+    @staticmethod
+    def get_output_channel(model_type):
+        if model_type == 'base':
+            return 768
+        if model_type == 'large':
+            return 1024
+        if model_type == 'huge':
+            return 1280
+
+    @staticmethod
+    def get_default_output_indices(model_type):
+        if model_type == 'base':
+            return [3, 5, 7, 11]
+        if model_type == 'large':
+            return [7, 11, 15, 23]
+        if model_type == 'huge':
+            return [8, 14, 20, 31]
+
+    def forward(self, image: torch.Tensor) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+        # b x 3 x res x res
+        _, _, input_h, input_w = image.shape
+        if input_h != self.input_resolution or input_w != self.input_resolution:
+            image = F.interpolate(image, self.input_resolution,
+                                  mode='bilinear', align_corners=False)
+
+        image = (image - self.image_mean.to(image.device)) / self.image_std.to(image.device)
+
+        x = image.to(self.visual.conv1.weight.data)
+
+        x = self.visual.conv1(x)  # shape = [*, width, grid, grid]
+        N, _, S, S = x.shape
+
+        # shape = [*, width, grid ** 2]
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.visual.class_embedding.to(x.dtype) +
+                       torch.zeros(x.shape[0], 1, x.shape[-1],
+                                   dtype=x.dtype, device=x.device),
+                       x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+
+        x = x + self.visual.positional_embedding.to(x.dtype)
+
+        x = self.visual.ln_pre(x)
+
+        x = x.permute(1, 0, 2).contiguous()  # NLD -> LND
+
+        features = []
+        cls_tokens = []
+        for blk in self.visual.transformer.resblocks:
+            x = blk(x)  # [S ** 2 + 1, N, D]
+            # if idx in self.output_indices:
+            feature = x[1:, :, :].permute(
+                1, 2, 0).view(N, -1, S, S).contiguous().float()
+            features.append(feature)
+            cls_tokens.append(x[0, :, :])
+
+        features = [features[ind] for ind in self.output_indices]
+        cls_tokens = [cls_tokens[ind] for ind in self.output_indices]
+
+        if self.fpns is not None:
+            for i, fpn in enumerate(self.fpns):
+                features[i] = fpn(features[i])
+
+        return features, cls_tokens

src/pixel3dmm/preprocessing/facer/facer/face_attribute/__init__.py

@@ -0,0 +1,2 @@
+from .base import FaceAttribute
+from .farl import FaRLFaceAttribute

src/pixel3dmm/preprocessing/facer/facer/face_attribute/base.py

@@ -0,0 +1,24 @@
+import torch.nn as nn
+
+class FaceAttribute(nn.Module):
+    """ face attribute base class
+
+    Args:
+        images (torch.Tensor): b x c x h x w
+
+        data (Dict[str, Any]):
+
+            * image_ids (torch.Tensor): nfaces
+            * rects (torch.Tensor): nfaces x 4 (x1, y1, x2, y2)
+            * points (torch.Tensor): nfaces x 5 x 2 (x, y)
+
+    Returns:
+        data (Dict[str, Any]):
+
+            * image_ids (torch.Tensor): nfaces
+            * rects (torch.Tensor): nfaces x 4 (x1, y1, x2, y2)
+            * points (torch.Tensor): nfaces x 5 x 2 (x, y)
+            * attrs (Dict[str, Any]):
+                * logits (torch.Tensor): nfaces x nclasses  
+    """
+    pass

src/pixel3dmm/preprocessing/facer/facer/face_attribute/farl.py

@@ -0,0 +1,158 @@
+from typing import Optional, Dict, Any
+import functools
+import torch
+import torch.nn.functional as F
+from ..transform import get_face_align_matrix, make_tanh_warp_grid
+from .base import FaceAttribute
+from ..farl import farl_classification
+from ..util import download_jit
+import numpy as np
+from torchvision.transforms import Normalize
+
+
+def get_std_points_xray(out_size=256, mid_size=500):
+    std_points_256 = np.array(
+        [
+            [85.82991, 85.7792],
+            [169.0532, 84.3381],
+            [127.574, 137.0006],
+            [90.6964, 174.7014],
+            [167.3069, 173.3733],
+        ]
+    )
+    std_points_256[:, 1] += 30
+    old_size = 256
+    mid = mid_size / 2
+    new_std_points = std_points_256 - old_size / 2 + mid
+    target_pts = new_std_points * out_size / mid_size
+    target_pts = torch.from_numpy(target_pts).float()
+    return target_pts
+
+
+pretrain_settings = {
+    "celeba/224": {
+        # acc 92.06617474555969
+        "num_classes": 40,
+        "layers": [11],
+        "url": "https://github.com/FacePerceiver/facer/releases/download/models-v1/face_attribute.farl.celeba.pt",
+        "matrix_src_tag": "points",
+        "get_matrix_fn": functools.partial(
+            get_face_align_matrix,
+            target_shape=(224, 224),
+            target_pts=get_std_points_xray(out_size=224, mid_size=500),
+        ),
+        "get_grid_fn": functools.partial(
+            make_tanh_warp_grid, warp_factor=0.0, warped_shape=(224, 224)
+        ),
+        "classes": [
+            "5_o_Clock_Shadow",
+            "Arched_Eyebrows",
+            "Attractive",
+            "Bags_Under_Eyes",
+            "Bald",
+            "Bangs",
+            "Big_Lips",
+            "Big_Nose",
+            "Black_Hair",
+            "Blond_Hair",
+            "Blurry",
+            "Brown_Hair",
+            "Bushy_Eyebrows",
+            "Chubby",
+            "Double_Chin",
+            "Eyeglasses",
+            "Goatee",
+            "Gray_Hair",
+            "Heavy_Makeup",
+            "High_Cheekbones",
+            "Male",
+            "Mouth_Slightly_Open",
+            "Mustache",
+            "Narrow_Eyes",
+            "No_Beard",
+            "Oval_Face",
+            "Pale_Skin",
+            "Pointy_Nose",
+            "Receding_Hairline",
+            "Rosy_Cheeks",
+            "Sideburns",
+            "Smiling",
+            "Straight_Hair",
+            "Wavy_Hair",
+            "Wearing_Earrings",
+            "Wearing_Hat",
+            "Wearing_Lipstick",
+            "Wearing_Necklace",
+            "Wearing_Necktie",
+            "Young",
+        ],
+    }
+}
+
+
+def load_face_attr(model_path, num_classes=40, layers=[11]):
+    model = farl_classification(num_classes=num_classes, layers=layers)
+    state_dict = download_jit(model_path, jit=False)
+    model.load_state_dict(state_dict)
+    return model
+
+
+class FaRLFaceAttribute(FaceAttribute):
+    """The face attribute recognition models from [FaRL](https://github.com/FacePerceiver/FaRL).
+
+    Please consider citing
+    ```bibtex
+        @article{zheng2021farl,
+            title={General Facial Representation Learning in a Visual-Linguistic Manner},
+            author={Zheng, Yinglin and Yang, Hao and Zhang, Ting and Bao, Jianmin and Chen,
+                Dongdong and Huang, Yangyu and Yuan, Lu and Chen,
+                Dong and Zeng, Ming and Wen, Fang},
+            journal={arXiv preprint arXiv:2112.03109},
+            year={2021}
+        }
+    ```
+    """
+
+    def __init__(
+        self,
+        conf_name: Optional[str] = None,
+        model_path: Optional[str] = None,
+        device=None,
+    ) -> None:
+        super().__init__()
+        if conf_name is None:
+            conf_name = "celeba/224"
+        if model_path is None:
+            model_path = pretrain_settings[conf_name]["url"]
+        self.conf_name = conf_name
+
+        setting = pretrain_settings[self.conf_name]
+        self.labels = setting["classes"]
+        self.net = load_face_attr(model_path, num_classes=setting["num_classes"], layers = setting["layers"])
+        if device is not None:
+            self.net = self.net.to(device)
+        self.normalize = Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
+        self.eval()
+
+    def forward(self, images: torch.Tensor, data: Dict[str, Any]):
+        setting = pretrain_settings[self.conf_name]
+        images = images.float() / 255.0  
+        _, _, h, w = images.shape
+
+        simages = images[data["image_ids"]]
+        matrix = setting["get_matrix_fn"](data[setting["matrix_src_tag"]])
+        grid = setting["get_grid_fn"](matrix=matrix, orig_shape=(h, w))
+
+        w_images = F.grid_sample(simages, grid, mode="bilinear", align_corners=False)
+        w_images = self.normalize(w_images)
+
+        outputs = self.net(w_images)
+        probs = torch.sigmoid(outputs)
+
+        data["attrs"] = probs
+
+        return data
+
+
+if __name__ == "__main__":
+    model = FaRLFaceAttribute()

src/pixel3dmm/preprocessing/facer/facer/face_detection/__init__.py

@@ -0,0 +1,2 @@
+from .base import FaceDetector
+from .retinaface import RetinaFaceDetector

src/pixel3dmm/preprocessing/facer/facer/face_detection/base.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+
+
+class FaceDetector(nn.Module):
+    """ face detector
+
+    Args:
+        images (torch.Tensor): b x c x h x w
+
+    Returns:
+        data (Dict[str, torch.Tensor]):
+
+            * rects: nfaces x 4 (x1, y1, x2, y2)
+            * points: nfaces x 5 x 2 (x, y)
+            * scores: nfaces
+            * image_ids: nfaces
+    """
+    pass

src/pixel3dmm/preprocessing/facer/facer/face_detection/retinaface.py

@@ -0,0 +1,677 @@
+# largely borrowed from https://github.dev/elliottzheng/batch-face/face_detection/alignment.py
+
+from typing import Dict, List, Optional, Tuple
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models._utils as _utils
+from .base import FaceDetector
+
+
+from itertools import product as product
+from math import ceil
+
+
+pretrained_urls = {
+    "mobilenet": "https://github.com/elliottzheng/face-detection/releases/download/0.0.1/mobilenet0.25_Final.pth",
+    "resnet50": "https://github.com/elliottzheng/face-detection/releases/download/0.0.1/Resnet50_Final.pth"
+}
+
+
+def conv_bn(inp, oup, stride=1, leaky=0):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True),
+    )
+
+
+def conv_bn_no_relu(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+    )
+
+
+def conv_bn1X1(inp, oup, stride, leaky=0):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True),
+    )
+
+
+def conv_dw(inp, oup, stride, leaky=0.1):
+    return nn.Sequential(
+        nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
+        nn.BatchNorm2d(inp),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True),
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True),
+    )
+
+
+class SSH(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(SSH, self).__init__()
+        assert out_channel % 4 == 0
+        leaky = 0
+        if out_channel <= 64:
+            leaky = 0.1
+        self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)
+
+        self.conv5X5_1 = conv_bn(
+            in_channel, out_channel // 4, stride=1, leaky=leaky)
+        self.conv5X5_2 = conv_bn_no_relu(
+            out_channel // 4, out_channel // 4, stride=1)
+
+        self.conv7X7_2 = conv_bn(
+            out_channel // 4, out_channel // 4, stride=1, leaky=leaky
+        )
+        self.conv7x7_3 = conv_bn_no_relu(
+            out_channel // 4, out_channel // 4, stride=1)
+
+    def forward(self, input):
+        conv3X3 = self.conv3X3(input)
+
+        conv5X5_1 = self.conv5X5_1(input)
+        conv5X5 = self.conv5X5_2(conv5X5_1)
+
+        conv7X7_2 = self.conv7X7_2(conv5X5_1)
+        conv7X7 = self.conv7x7_3(conv7X7_2)
+
+        out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
+        out = F.relu(out)
+        return out
+
+
+class FPN(nn.Module):
+    def __init__(self, in_channels_list, out_channels):
+        super(FPN, self).__init__()
+        leaky = 0
+        if out_channels <= 64:
+            leaky = 0.1
+        self.output1 = conv_bn1X1(
+            in_channels_list[0], out_channels, stride=1, leaky=leaky
+        )
+        self.output2 = conv_bn1X1(
+            in_channels_list[1], out_channels, stride=1, leaky=leaky
+        )
+        self.output3 = conv_bn1X1(
+            in_channels_list[2], out_channels, stride=1, leaky=leaky
+        )
+
+        self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
+        self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
+
+    def forward(self, input):
+        # names = list(input.keys())
+        input = list(input.values())
+
+        output1 = self.output1(input[0])
+        output2 = self.output2(input[1])
+        output3 = self.output3(input[2])
+
+        up3 = F.interpolate(
+            output3, size=[output2.size(2), output2.size(3)], mode="nearest"
+        )
+        output2 = output2 + up3
+        output2 = self.merge2(output2)
+
+        up2 = F.interpolate(
+            output2, size=[output1.size(2), output1.size(3)], mode="nearest"
+        )
+        output1 = output1 + up2
+        output1 = self.merge1(output1)
+
+        out = [output1, output2, output3]
+        return out
+
+
+class MobileNetV1(nn.Module):
+    def __init__(self):
+        super(MobileNetV1, self).__init__()
+        self.stage1 = nn.Sequential(
+            conv_bn(3, 8, 2, leaky=0.1),  # 3
+            conv_dw(8, 16, 1),  # 7
+            conv_dw(16, 32, 2),  # 11
+            conv_dw(32, 32, 1),  # 19
+            conv_dw(32, 64, 2),  # 27
+            conv_dw(64, 64, 1),  # 43
+        )
+        self.stage2 = nn.Sequential(
+            conv_dw(64, 128, 2),  # 43 + 16 = 59
+            conv_dw(128, 128, 1),  # 59 + 32 = 91
+            conv_dw(128, 128, 1),  # 91 + 32 = 123
+            conv_dw(128, 128, 1),  # 123 + 32 = 155
+            conv_dw(128, 128, 1),  # 155 + 32 = 187
+            conv_dw(128, 128, 1),  # 187 + 32 = 219
+        )
+        self.stage3 = nn.Sequential(
+            conv_dw(128, 256, 2),  # 219 +3 2 = 241
+            conv_dw(256, 256, 1),  # 241 + 64 = 301
+        )
+        self.avg = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(256, 1000)
+
+    def forward(self, x):
+        x = self.stage1(x)
+        x = self.stage2(x)
+        x = self.stage3(x)
+        x = self.avg(x)
+        # x = self.model(x)
+        x = x.view(-1, 256)
+        x = self.fc(x)
+        return x
+
+
+class ClassHead(nn.Module):
+    def __init__(self, inchannels=512, num_anchors=3):
+        super(ClassHead, self).__init__()
+        self.num_anchors = num_anchors
+        self.conv1x1 = nn.Conv2d(
+            inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0
+        )
+
+    def forward(self, x):
+        out = self.conv1x1(x)
+        out = out.permute(0, 2, 3, 1).contiguous()
+        return out.view(out.shape[0], -1, 2)
+
+
+class BboxHead(nn.Module):
+    def __init__(self, inchannels=512, num_anchors=3):
+        super(BboxHead, self).__init__()
+        self.conv1x1 = nn.Conv2d(
+            inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0
+        )
+
+    def forward(self, x):
+        out = self.conv1x1(x)
+        out = out.permute(0, 2, 3, 1).contiguous()
+        return out.view(out.shape[0], -1, 4)
+
+
+class LandmarkHead(nn.Module):
+    def __init__(self, inchannels=512, num_anchors=3):
+        super(LandmarkHead, self).__init__()
+        self.conv1x1 = nn.Conv2d(
+            inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0
+        )
+
+    def forward(self, x):
+        out = self.conv1x1(x)
+        out = out.permute(0, 2, 3, 1).contiguous()
+        return out.view(out.shape[0], -1, 10)
+
+
+class RetinaFace(nn.Module):
+    def __init__(self, cfg=None, phase="train"):
+        """
+        :param cfg:  Network related settings.
+        :param phase: train or test.
+        """
+        super(RetinaFace, self).__init__()
+        self.phase = phase
+        backbone = None
+        if cfg["name"] == "mobilenet0.25":
+            backbone = MobileNetV1()
+        elif cfg["name"] == "Resnet50":
+            import torchvision.models as models
+            backbone = models.resnet50(pretrained=cfg["pretrain"])
+
+        self.body = _utils.IntermediateLayerGetter(
+            backbone, cfg["return_layers"])
+        in_channels_stage2 = cfg["in_channel"]
+        in_channels_list = [
+            in_channels_stage2 * 2,
+            in_channels_stage2 * 4,
+            in_channels_stage2 * 8,
+        ]
+        out_channels = cfg["out_channel"]
+        self.fpn = FPN(in_channels_list, out_channels)
+        self.ssh1 = SSH(out_channels, out_channels)
+        self.ssh2 = SSH(out_channels, out_channels)
+        self.ssh3 = SSH(out_channels, out_channels)
+
+        self.ClassHead = self._make_class_head(
+            fpn_num=3, inchannels=cfg["out_channel"])
+        self.BboxHead = self._make_bbox_head(
+            fpn_num=3, inchannels=cfg["out_channel"])
+        self.LandmarkHead = self._make_landmark_head(
+            fpn_num=3, inchannels=cfg["out_channel"]
+        )
+
+    def _make_class_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+        classhead = nn.ModuleList()
+        for i in range(fpn_num):
+            classhead.append(ClassHead(inchannels, anchor_num))
+        return classhead
+
+    def _make_bbox_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+        bboxhead = nn.ModuleList()
+        for i in range(fpn_num):
+            bboxhead.append(BboxHead(inchannels, anchor_num))
+        return bboxhead
+
+    def _make_landmark_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+        landmarkhead = nn.ModuleList()
+        for i in range(fpn_num):
+            landmarkhead.append(LandmarkHead(inchannels, anchor_num))
+        return landmarkhead
+
+    def forward(self, inputs):
+        out = self.body(inputs)
+
+        # FPN
+        fpn = self.fpn(out)
+
+        # SSH
+        feature1 = self.ssh1(fpn[0])
+        feature2 = self.ssh2(fpn[1])
+        feature3 = self.ssh3(fpn[2])
+        features = [feature1, feature2, feature3]
+
+        bbox_regressions = torch.cat(
+            [self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1
+        )
+        classifications = torch.cat(
+            [self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1
+        )
+        ldm_regressions = torch.cat(
+            [self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1
+        )
+
+        if self.phase == "train":
+            output = (bbox_regressions, classifications, ldm_regressions)
+        else:
+            output = (
+                bbox_regressions,
+                F.softmax(classifications, dim=-1),
+                ldm_regressions,
+            )
+        return output
+
+
+# Adapted from https://github.com/Hakuyume/chainer-ssd
+def decode(loc: torch.Tensor, priors: torch.Tensor, variances: Tuple[float, float]) -> torch.Tensor:
+    boxes = torch.cat(
+        (
+            priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+            priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1]),
+        ),
+        1,
+    )
+    boxes[:, :2] -= boxes[:, 2:] / 2
+    boxes[:, 2:] += boxes[:, :2]
+    return boxes
+
+
+def decode_landm(pre: torch.Tensor, priors: torch.Tensor, variances: Tuple[float, float]) -> torch.Tensor:
+    landms = torch.cat(
+        (
+            priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
+        ),
+        dim=1,
+    )
+    return landms
+
+
+def nms(dets: torch.Tensor, thresh: float) -> List[int]:
+    """Pure Python NMS baseline."""
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    x2 = dets[:, 2]
+    y2 = dets[:, 3]
+    scores = dets[:, 4]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = torch.flip(scores.argsort(), [0])
+
+    keep = []
+    while order.numel() > 0:
+        i = order[0].item()
+        keep.append(i)
+        xx1 = torch.maximum(x1[i], x1[order[1:]])
+        yy1 = torch.maximum(y1[i], y1[order[1:]])
+        xx2 = torch.minimum(x2[i], x2[order[1:]])
+        yy2 = torch.minimum(y2[i], y2[order[1:]])
+
+        w = torch.maximum(torch.tensor(0.0).to(dets), xx2 - xx1 + 1)
+        h = torch.maximum(torch.tensor(0.0).to(dets), yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = torch.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+class PriorBox:
+    def __init__(self, cfg: dict, image_size: Tuple[int, int]):
+        self.min_sizes = cfg["min_sizes"]
+        self.steps = cfg["steps"]
+        self.clip = cfg["clip"]
+        self.image_size = image_size
+        self.feature_maps = [
+            [ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)]
+            for step in self.steps
+        ]
+
+    def generate_anchors(self, device) -> torch.Tensor:
+        anchors = []
+        for k, f in enumerate(self.feature_maps):
+            min_sizes = self.min_sizes[k]
+            for i, j in product(range(f[0]), range(f[1])):
+                for min_size in min_sizes:
+                    s_kx = min_size / self.image_size[1]
+                    s_ky = min_size / self.image_size[0]
+                    dense_cx = [
+                        x * self.steps[k] / self.image_size[1] for x in [j + 0.5]
+                    ]
+                    dense_cy = [
+                        y * self.steps[k] / self.image_size[0] for y in [i + 0.5]
+                    ]
+                    for cy, cx in product(dense_cy, dense_cx):
+                        anchors += [cx, cy, s_kx, s_ky]
+
+        # back to torch land
+        output = torch.tensor(anchors).view(-1, 4)
+        if self.clip:
+            output.clamp_(max=1, min=0)
+        return output.to(device=device)
+
+
+cfg_mnet = {
+    "name": "mobilenet0.25",
+    "min_sizes": [[16, 32], [64, 128], [256, 512]],
+    "steps": [8, 16, 32],
+    "variance": [0.1, 0.2],
+    "clip": False,
+    "loc_weight": 2.0,
+    "gpu_train": True,
+    "batch_size": 32,
+    "ngpu": 1,
+    "epoch": 250,
+    "decay1": 190,
+    "decay2": 220,
+    "image_size": 640,
+    "pretrain": True,
+    "return_layers": {"stage1": 1, "stage2": 2, "stage3": 3},
+    "in_channel": 32,
+    "out_channel": 64,
+}
+
+cfg_re50 = {
+    "name": "Resnet50",
+    "min_sizes": [[16, 32], [64, 128], [256, 512]],
+    "steps": [8, 16, 32],
+    "variance": [0.1, 0.2],
+    "clip": False,
+    "loc_weight": 2.0,
+    "gpu_train": True,
+    "batch_size": 24,
+    "ngpu": 4,
+    "epoch": 100,
+    "decay1": 70,
+    "decay2": 90,
+    "image_size": 840,
+    "pretrain": False,
+    "return_layers": {"layer2": 1, "layer3": 2, "layer4": 3},
+    "in_channel": 256,
+    "out_channel": 256,
+}
+
+
+def check_keys(model, pretrained_state_dict):
+    ckpt_keys = set(pretrained_state_dict.keys())
+    model_keys = set(model.state_dict().keys())
+    used_pretrained_keys = model_keys & ckpt_keys
+    assert len(used_pretrained_keys) > 0, "load NONE from pretrained checkpoint"
+    return True
+
+
+def remove_prefix(state_dict, prefix):
+    """ Old style model is stored with all names of parameters sharing common prefix 'module.' """
+    def f(x): return x.split(prefix, 1)[-1] if x.startswith(prefix) else x
+    return {f(key): value for key, value in state_dict.items()}
+
+
+def load_model(model, pretrained_path, load_to_cpu, network: str):
+    if pretrained_path is None:
+        url = pretrained_urls[network]
+        if load_to_cpu:
+            pretrained_dict = torch.utils.model_zoo.load_url(
+                url, map_location=lambda storage, loc: storage
+            )
+        else:
+            pretrained_dict = torch.utils.model_zoo.load_url(
+                url, map_location=lambda storage, loc: storage.cuda(device)
+            )
+    else:
+        if load_to_cpu:
+            pretrained_dict = torch.load(
+                pretrained_path, map_location=lambda storage, loc: storage
+            )
+        else:
+            device = torch.cuda.current_device()
+            pretrained_dict = torch.load(
+                pretrained_path, map_location=lambda storage, loc: storage.cuda(
+                    device)
+            )
+    if "state_dict" in pretrained_dict.keys():
+        pretrained_dict = remove_prefix(
+            pretrained_dict["state_dict"], "module.")
+    else:
+        pretrained_dict = remove_prefix(pretrained_dict, "module.")
+    check_keys(model, pretrained_dict)
+    model.load_state_dict(pretrained_dict, strict=False)
+    return model
+
+
+def load_net(model_path, network="mobilenet"):
+    if network == "mobilenet":
+        cfg = cfg_mnet
+    elif network == "resnet50":
+        cfg = cfg_re50
+    else:
+        raise NotImplementedError(network)
+    # net and model
+    net = RetinaFace(cfg=cfg, phase="test")
+    net = load_model(net, model_path, True, network=network)
+    net.eval()
+    cudnn.benchmark = True
+    # net = net.to(device)
+    return net
+
+
+def parse_det(det: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, float]:
+    landmarks = det[5:].reshape(5, 2)
+    box = det[:4]
+    score = det[4]
+    return box, landmarks, score.item()
+
+
+def post_process(
+    loc: torch.Tensor,
+    conf: torch.Tensor,
+    landms: torch.Tensor,
+    prior_data: torch.Tensor,
+    cfg: dict,
+    scale: float,
+    scale1: float,
+    resize,
+    confidence_threshold,
+    top_k,
+    nms_threshold,
+    keep_top_k,
+):
+    boxes = decode(loc, prior_data, cfg["variance"])
+    boxes = boxes * scale / resize
+    # boxes = boxes.cpu().numpy()
+    # scores = conf.cpu().numpy()[:, 1]
+    scores = conf[:, 1]
+    landms_copy = decode_landm(landms, prior_data, cfg["variance"])
+
+    landms_copy = landms_copy * scale1 / resize
+    # landms_copy = landms_copy.cpu().numpy()
+
+    # ignore low scores
+    inds = torch.where(scores > confidence_threshold)[0]
+    boxes = boxes[inds]
+    landms_copy = landms_copy[inds]
+    scores = scores[inds]
+
+    # keep top-K before NMS
+    order = torch.flip(scores.argsort(), [0])[:top_k]
+    boxes = boxes[order]
+    landms_copy = landms_copy[order]
+    scores = scores[order]
+
+    # do NMS
+    dets = torch.hstack((boxes, scores.unsqueeze(-1))).to(
+        dtype=torch.float32, copy=False)
+    keep = nms(dets, nms_threshold)
+    # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
+    dets = dets[keep, :]
+    landms_copy = landms_copy[keep]
+
+    # keep top-K faster NMS
+    dets = dets[:keep_top_k, :]
+    landms_copy = landms_copy[:keep_top_k, :]
+
+    dets = torch.cat((dets, landms_copy), dim=1)
+    # show image
+    dets = sorted(dets, key=lambda x: x[4], reverse=True)
+    dets = [parse_det(x) for x in dets]
+
+    return dets
+
+
+# @torch.no_grad()
+def batch_detect(net: nn.Module, images: torch.Tensor, threshold: float = 0.5):
+    confidence_threshold = threshold
+    cfg = cfg_mnet
+    top_k = 5000
+    nms_threshold = 0.4
+    keep_top_k = 750
+    resize = 1
+
+    img = images.float()
+    mean = torch.as_tensor([104, 117, 123], dtype=img.dtype, device=img.device).view(
+        1, 3, 1, 1
+    )
+    img -= mean
+    (
+        _,
+        _,
+        im_height,
+        im_width,
+    ) = img.shape
+    scale = torch.as_tensor(
+        [im_width, im_height, im_width, im_height],
+        dtype=img.dtype,
+        device=img.device,
+    )
+    scale = scale.to(img.device)
+
+    loc, conf, landms = net(img)  # forward pass
+
+    priorbox = PriorBox(cfg, image_size=(im_height, im_width))
+    prior_data = priorbox.generate_anchors(device=img.device)
+    scale1 = torch.as_tensor(
+        [
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+        ],
+        dtype=img.dtype,
+        device=img.device,
+    )
+    scale1 = scale1.to(img.device)
+
+    all_dets = [
+        post_process(
+            loc_i,
+            conf_i,
+            landms_i,
+            prior_data,
+            cfg,
+            scale,
+            scale1,
+            resize,
+            confidence_threshold,
+            top_k,
+            nms_threshold,
+            keep_top_k,
+        )
+        for loc_i, conf_i, landms_i in zip(loc, conf, landms)
+    ]
+
+    rects = []
+    points = []
+    scores = []
+    image_ids = []
+    for image_id, faces_in_one_image in enumerate(all_dets):
+        for rect, landmarks, score in faces_in_one_image:
+            rects.append(rect)
+            points.append(landmarks)
+            scores.append(score)
+            image_ids.append(image_id)
+
+    if len(rects) == 0:
+        return {
+            'rects': torch.Tensor().to(img.device),
+            'points': torch.Tensor().to(img.device),
+            'scores': torch.Tensor().to(img.device),
+            'image_ids': torch.Tensor().to(img.device),
+        }
+
+    return {
+        'rects': torch.stack(rects, dim=0).to(img.device),
+        'points': torch.stack(points, dim=0).to(img.device),
+        'scores': torch.tensor(scores).to(img.device),
+        'image_ids': torch.tensor(image_ids).to(img.device),
+    }
+
+
+class RetinaFaceDetector(FaceDetector):
+    """RetinaFaceDetector
+
+    Args:
+        images (torch.Tensor): b x c x h x w, uint8, 0~255.
+
+    Returns:
+        faces (Dict[str, torch.Tensor]):
+
+            * image_ids: n, int
+            * rects: n x 4 (x1, y1, x2, y2)
+            * points: n x 5 x 2 (x, y)
+            * scores: n
+    """
+
+    def __init__(self, conf_name: Optional[str] = None,
+                 model_path: Optional[str] = None, threshold=0.8) -> None:
+        super().__init__()
+        if conf_name is None:
+            conf_name = 'mobilenet'
+        self.net = load_net(model_path, conf_name)
+        self.threshold = threshold
+        self.eval()
+
+    def forward(self, images: torch.Tensor) -> Dict[str, torch.Tensor]:
+        return batch_detect(self.net, images.clone(), threshold=self.threshold)

src/pixel3dmm/preprocessing/facer/facer/face_parsing/__init__.py

@@ -0,0 +1,2 @@
+from .base import FaceParser
+from .farl import FaRLFaceParser

src/pixel3dmm/preprocessing/facer/facer/face_parsing/base.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+
+class FaceParser(nn.Module):
+    """ face parser
+
+    Args:
+        images (torch.Tensor): b x c x h x w
+
+        data (Dict[str, Any]):
+
+            * image_ids (torch.Tensor): nfaces
+            * rects (torch.Tensor): nfaces x 4 (x1, y1, x2, y2)
+            * points (torch.Tensor): nfaces x 5 x 2 (x, y)
+
+    Returns:
+        data (Dict[str, Any]):
+
+            * image_ids (torch.Tensor): nfaces
+            * rects (torch.Tensor): nfaces x 4 (x1, y1, x2, y2)
+            * points (torch.Tensor): nfaces x 5 x 2 (x, y)
+            * seg (Dict[str, Any]):
+
+                * logits (torch.Tensor): nfaces x nclasses x h x w
+                * label_names (List[str]): nclasses
+    """
+    pass

src/pixel3dmm/preprocessing/facer/facer/face_parsing/farl.py

@@ -0,0 +1,174 @@
+from typing import Optional, Dict, Any
+import functools
+import torch
+import torch.nn.functional as F
+
+from ..util import download_jit
+from ..transform import (get_crop_and_resize_matrix, get_face_align_matrix, get_face_align_matrix_celebm,
+                         make_inverted_tanh_warp_grid, make_tanh_warp_grid)
+from .base import FaceParser
+import numpy as np
+
+pretrain_settings = {
+    'lapa/448': {
+        'url': [
+            'https://github.com/FacePerceiver/facer/releases/download/models-v1/face_parsing.farl.lapa.main_ema_136500_jit191.pt',
+        ],
+        'matrix_src_tag': 'points',
+        'get_matrix_fn': functools.partial(get_face_align_matrix,
+                                           target_shape=(448, 448), target_face_scale=1.0),
+        'get_grid_fn': functools.partial(make_tanh_warp_grid,
+                                         warp_factor=0.8, warped_shape=(448, 448)),
+        'get_inv_grid_fn': functools.partial(make_inverted_tanh_warp_grid,
+                                             warp_factor=0.8, warped_shape=(448, 448)),
+        'label_names': ['background', 'face', 'rb', 'lb', 're',
+                        'le', 'nose',  'ulip', 'imouth', 'llip', 'hair']
+    },
+    'celebm/448': {
+        'url': [
+            'https://github.com/FacePerceiver/facer/releases/download/models-v1/face_parsing.farl.celebm.main_ema_181500_jit.pt',
+        ],
+        'matrix_src_tag': 'points',
+        'get_matrix_fn': functools.partial(get_face_align_matrix_celebm,
+                                           target_shape=(448, 448)),
+        'get_grid_fn': functools.partial(make_tanh_warp_grid,
+                                         warp_factor=0, warped_shape=(448, 448)),
+        'get_inv_grid_fn': functools.partial(make_inverted_tanh_warp_grid,
+                                             warp_factor=0, warped_shape=(448, 448)),
+        'label_names':  [
+                    'background', 'neck', 'face', 'cloth', 'rr', 'lr', 'rb', 'lb', 're',
+                    'le', 'nose', 'imouth', 'llip', 'ulip', 'hair',
+                    'eyeg', 'hat', 'earr', 'neck_l']
+    }
+}
+
+
+class FaRLFaceParser(FaceParser):
+    """ The face parsing models from [FaRL](https://github.com/FacePerceiver/FaRL).
+
+    Please consider citing 
+    ```bibtex
+        @article{zheng2021farl,
+            title={General Facial Representation Learning in a Visual-Linguistic Manner},
+            author={Zheng, Yinglin and Yang, Hao and Zhang, Ting and Bao, Jianmin and Chen, 
+                Dongdong and Huang, Yangyu and Yuan, Lu and Chen, 
+                Dong and Zeng, Ming and Wen, Fang},
+            journal={arXiv preprint arXiv:2112.03109},
+            year={2021}
+        }
+    ```
+    """
+
+    def __init__(self, conf_name: Optional[str] = None, model_path: Optional[str] = None, device=None) -> None:
+        super().__init__()
+        if conf_name is None:
+            conf_name = 'lapa/448'
+        if model_path is None:
+            model_path = pretrain_settings[conf_name]['url']
+        self.conf_name = conf_name
+        self.net = download_jit(model_path, map_location=device)
+        self.eval()
+        self.device = device
+        self.setting = pretrain_settings[conf_name]
+        self.label_names = self.setting['label_names']
+
+    
+    def get_warp_grid(self, images: torch.Tensor, matrix_src):
+        _, _, h, w = images.shape
+        matrix = self.setting['get_matrix_fn'](matrix_src)
+        grid = self.setting['get_grid_fn'](matrix=matrix, orig_shape=(h, w))
+        inv_grid = self.setting['get_inv_grid_fn'](matrix=matrix, orig_shape=(h, w))
+        return grid, inv_grid
+
+    def warp_images(self, images: torch.Tensor, data: Dict[str, Any]):
+        simages = self.unify_image_dtype(images)
+        simages = simages[data['image_ids']]
+        matrix_src = data[self.setting['matrix_src_tag']]
+        grid, inv_grid = self.get_warp_grid(simages, matrix_src)
+
+        w_images = F.grid_sample(
+            simages, grid, mode='bilinear', align_corners=False)
+        return w_images, grid, inv_grid
+    
+
+    def decode_image_to_cv2(self, images: torch.Tensor):
+        '''
+        output: b x 3 x h x w, torch.uint8, [0, 255]
+        '''
+        assert images.ndim == 4
+        assert images.shape[1] == 3
+        images = images.permute(0, 2, 3, 1).cpu().numpy() * 255
+        images = images.astype(np.uint8)
+        return images
+
+    def unify_image_dtype(self, images: torch.Tensor|np.ndarray|list):
+        '''
+        output: b x 3 x h x w, torch.float32, [0, 1]
+        '''
+        if isinstance(images, np.ndarray):
+            images = torch.from_numpy(images)
+        elif isinstance(images, torch.Tensor):
+            pass
+        elif isinstance(images, list):
+            assert len(images) > 0, "images is empty"
+            first_image = images[0]
+            if isinstance(first_image, np.ndarray):
+                images = [torch.from_numpy(image).permute(2, 0, 1) for image in images]
+                images = torch.stack(images)
+            elif isinstance(first_image, torch.Tensor):
+                images = torch.stack(images)
+            else:
+                raise ValueError(f"Unsupported image type: {type(first_image)}")
+
+        else:
+            raise ValueError(f"Unsupported image type: {type(images)}")
+        
+        assert images.ndim == 4
+        assert images.shape[1] == 3
+
+        max_val = images.max()
+        if max_val <= 1:
+            assert images.dtype == torch.float32 or images.dtype == torch.float16
+        elif max_val <= 255:
+            assert images.dtype == torch.uint8
+            images = images.float() / 255.0
+        else:
+            raise ValueError(f"Unsupported image type: {images.dtype}")
+        if images.device != self.device:
+            images = images.to(device=self.device)
+        return images
+    
+    @torch.no_grad()
+    @torch.inference_mode()
+    def forward(self, images: torch.Tensor, data: Dict[str, Any]):
+        '''
+        images: b x 3 x h x w , torch.uint8, [0, 255]
+        data: {'rects': rects, 'points': points, 'scores': scores, 'image_ids': image_ids}
+        '''
+        w_images, grid, inv_grid = self.warp_images(images, data)
+        w_seg_logits = self.forward_warped(w_images, return_preds=False)
+
+        seg_logits = F.grid_sample(
+            w_seg_logits, inv_grid, mode='bilinear', align_corners=False)
+
+        data['seg'] = {'logits': seg_logits, 'label_names': self.label_names}
+        return data
+    
+
+    def logits2predictions(self, logits: torch.Tensor):
+        return logits.argmax(dim=1)
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    def forward_warped(self, images: torch.Tensor, return_preds: bool = True):
+        '''
+        images: b x 3 x h x w , torch.uint8, [0, 255]
+        '''
+        images = self.unify_image_dtype(images)
+        seg_logits, _ = self.net(images)  # nfaces x c x h x w
+        # seg_probs = seg_logits.softmax(dim=1)  # nfaces x nclasses x h x w
+        if return_preds:
+            seg_preds = self.logits2predictions(seg_logits)
+            return seg_logits, seg_preds, self.label_names
+        else:
+            return seg_logits

src/pixel3dmm/preprocessing/facer/facer/farl/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from .model import load_farl, VisualTransformer
+from .classification import farl_classification

src/pixel3dmm/preprocessing/facer/facer/farl/classification.py

@@ -0,0 +1,149 @@
+from torch import nn
+import torch
+import torch.utils.checkpoint as checkpoint
+from .model import VisualTransformer
+
+
+class VITClassificationHeadV0(nn.Module):
+    def __init__(
+        self,
+        num_features: int,
+        channel: int,
+        num_labels: int,
+        norm=False,
+        dropout=0.0,
+        ret_feat=False,
+    ):
+        super().__init__()
+        self.weights = nn.Parameter(
+            torch.ones(1, num_features * 3, 1, dtype=torch.float32)
+        )
+        self.final_fc = nn.Linear(channel, num_labels)
+        self.norm = norm
+        if self.norm:
+            for i in range(num_features * 3):
+                setattr(self, f"norm_{i}", nn.LayerNorm(channel))
+        self.dropout = nn.Dropout(p=dropout)
+        self.ret_feat = ret_feat
+
+    def forward(self, features, cls_tokens):
+        xs = []
+        for feature, cls_token in zip(features, cls_tokens):
+            # feature: b x c x s x s
+            # cls_token: b x c
+            xs.append(feature.mean([2, 3]))
+            xs.append(feature.max(-1).values.max(-1).values)
+            xs.append(cls_token)
+        if self.norm:
+            xs = [getattr(self, f"norm_{i}")(x) for i, x in enumerate(xs)]
+        xs = torch.stack(xs, dim=1)  # b x 3N x c
+        feat = (xs * self.weights.softmax(dim=1)).sum(1)  # b x c
+        x = self.dropout(feat)
+        x = self.final_fc(x)  # b x num_labels
+        if self.ret_feat:
+            return x, feat
+        else:
+            return x
+
+
+class FACTransformer(nn.Module):
+    """A face attribute classification transformer leveraging multiple cls_tokens.
+    Args:
+        image (torch.Tensor): Float32 tensor with shape [b, 3, h, w], normalized to [0, 1].
+    Returns:
+        logits (torch.Tensor): Float32 tensor with shape [b, n_classes].
+        aux_outputs:
+    """
+
+    def __init__(self, backbone: nn.Module, head: nn.Module):
+        super().__init__()
+        self.backbone = backbone
+        self.head = head
+        self.cuda().float()
+
+    def forward(self, image):
+        logits = self.head(*self.backbone(image))
+        return logits
+
+
+def add_method(obj, name, method):
+    import types
+
+    setattr(obj, name, types.MethodType(method, obj))
+
+
+def get_clip_encode_func(layers):
+    def func(self, x):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        # shape = [*, width, grid ** 2]
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        extra_tokens = getattr(self, "extra_tokens", [])
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        class_token = self.class_embedding.to(x.dtype) + torch.zeros(
+            x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device
+        )
+        special_tokens = [
+            getattr(self, name).to(x.dtype)
+            + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device)
+            for name in extra_tokens
+        ]
+        x = torch.cat(
+            [class_token, *special_tokens, x], dim=1
+        )  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        outs = []
+        max_layer = max(layers)
+        use_checkpoint = self.transformer.use_checkpoint
+        for layer_i, blk in enumerate(self.transformer.resblocks):
+            if layer_i > max_layer:
+                break
+            if self.training and use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+            outs.append(x)
+
+        outs = torch.stack(outs).permute(0, 2, 1, 3)
+        cls_tokens = outs[layers, :, 0, :]
+
+        extra_token_feats = {}
+        for i, name in enumerate(extra_tokens):
+            extra_token_feats[name] = outs[layers, :, i + 1, :]
+        L, B, N, C = outs.shape
+        import math
+
+        W = int(math.sqrt(N - 1 - len(extra_tokens)))
+        features = (
+            outs[layers, :, 1 + len(extra_tokens) :, :]
+            .reshape(len(layers), B, W, W, C)
+            .permute(0, 1, 4, 2, 3)
+        )
+        if getattr(self, "ret_special", False):
+            return features, cls_tokens, extra_token_feats
+        else:
+            return features, cls_tokens
+
+    return func
+
+
+def farl_classification(num_classes=2, layers=list(range(12))):
+    model = VisualTransformer(
+        input_resolution=224,
+        patch_size=16,
+        width=768,
+        layers=12,
+        heads=12,
+        output_dim=512,
+    )
+    channel = 768
+    model = model.cuda()
+    del model.proj
+    del model.ln_post
+    add_method(model, "forward", get_clip_encode_func(layers))
+    head = VITClassificationHeadV0(
+        num_features=len(layers), channel=channel, num_labels=num_classes, norm=True
+    )
+    model = FACTransformer(model, head)
+    return model

src/pixel3dmm/preprocessing/facer/facer/farl/model.py

@@ -0,0 +1,419 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from collections import OrderedDict
+import logging
+import torch
+import torch.nn.functional as F
+from torch import nn
+import torch.utils.checkpoint as checkpoint
+import numpy as np
+from timm.models.layers import trunc_normal_, DropPath
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        self.stride = stride
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes *
+                                self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(
+            torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5
+        )
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2]
+                      * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat(
+                [self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]
+            ),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3,
+                               stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.conv2 = nn.Conv2d(width // 2, width // 2,
+                               kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.conv3 = nn.Conv2d(
+            width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.avgpool = nn.AvgPool2d(2)
+        self.relu = nn.ReLU(inplace=True)
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(
+            input_resolution // 32, embed_dim, heads, output_dim
+        )
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.BatchNorm2d, LayerNorm)):
+            nn.init.constant_(m.weight, 1)
+            nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.Linear, nn.Conv2d)):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            for conv, bn in [
+                (self.conv1, self.bn1),
+                (self.conv2, self.bn2),
+                (self.conv3, self.bn3)
+            ]:
+                x = self.relu(bn(conv(x)))
+            x = self.avgpool(x)
+            return x
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+        return x
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        """Construct a layernorm module in the TF style (epsilon inside the square root).
+        """
+        super(LayerNorm, self).__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        pdtype = x.dtype
+        x = x.float()
+        u = x.mean(-1, keepdim=True)
+        s = (x - u).pow(2).mean(-1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
+        return self.weight * x.to(pdtype) + self.bias
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, drop_path=0.):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+
+    def add_drop_path(self, drop_path):
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) \
+            if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.drop_path(self.attention(self.ln_1(x)))
+        x = x + self.drop_path(self.mlp(self.ln_2(x)))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self,
+                 width: int,
+                 layers: int,
+                 heads: int,
+                 attn_mask: torch.Tensor = None,
+                 use_checkpoint=True,
+                 drop_rate=0.,
+                 attn_drop_rate=0.,
+                 drop_path_rate=0.,
+                 ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.use_checkpoint = use_checkpoint
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, layers)]
+        self.resblocks = nn.ModuleList([
+            ResidualAttentionBlock(width, heads, attn_mask, drop_path=dpr[i])
+            for i in range(layers)
+        ])
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Linear, nn.Conv2d)):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor):
+        for i, blk in enumerate(self.resblocks):
+            x = blk(x)
+        return x
+        
+
+
+class VisualTransformer(nn.Module):
+    positional_embedding: nn.Parameter
+
+    def __init__(self,
+                 input_resolution: int,
+                 patch_size: int,
+                 width: int,
+                 layers: int,
+                 heads: int,
+                 output_dim: int,
+                 pool_type: str = 'default',
+                 skip_cls: bool = False,
+                 drop_path_rate=0.,
+                 **kwargs):
+        super().__init__()
+        self.pool_type = pool_type
+        self.skip_cls = skip_cls
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(
+            in_channels=3,
+            out_channels=width,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=False
+        )
+        self.config = kwargs.get("config", None)
+        self.sequence_length = (input_resolution // patch_size) ** 2 + 1
+        self.conv_pool = nn.Identity()
+        if (self.pool_type == 'linear'):
+            if (not self.skip_cls):
+                self.conv_pool = nn.Conv1d(
+                    width, width, self.sequence_length, stride=self.sequence_length, groups=width)
+            else:
+                self.conv_pool = nn.Conv1d(
+                    width, width, self.sequence_length-1, stride=self.sequence_length, groups=width)
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(
+            scale * torch.randn(
+                self.sequence_length, width
+            )
+        )
+        self.ln_pre = LayerNorm(width)
+        self.transformer = Transformer(
+            width, layers, heads, drop_path_rate=drop_path_rate)
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+        if self.config is not None and self.config.MIM.ENABLE:
+            logging.info("MIM ENABLED")
+            self.mim = True
+            self.lm_transformer = Transformer(
+                width, self.config.MIM.LAYERS, heads)
+            self.ln_lm = LayerNorm(width)
+            self.lm_head = nn.Linear(width, self.config.MIM.VOCAB_SIZE)
+            self.mask_token = nn.Parameter(scale * torch.randn(width))
+        else:
+            self.mim = False
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Linear, nn.Conv2d, nn.Conv1d)):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        # shape = [*, width, grid ** 2]
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1],
+                                                                      dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        if (self.pool_type == 'average'):
+            if self.skip_cls:
+                x = x[:, 1:, :]
+            x = torch.mean(x, dim=1)
+        elif (self.pool_type == 'linear'):
+            if self.skip_cls:
+                x = x[:, 1:, :]
+            x = x.permute(0, 2, 1)
+            x = self.conv_pool(x)
+            x = x.permute(0, 2, 1).squeeze()
+        else:
+            x = x[:, 0, :]
+        x = self.ln_post(x)
+        if self.proj is not None:
+            x = x @ self.proj
+        return x
+
+    def forward_mim(self, x: torch.Tensor, bool_masked_pos, return_all_tokens=False, disable_vlc=False):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        # shape = [*, width, grid ** 2]
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        batch_size, seq_len, _ = x.size()
+        mask_token = self.mask_token.unsqueeze(
+            0).unsqueeze(0).expand(batch_size, seq_len, -1)
+        w = bool_masked_pos.unsqueeze(-1).type_as(mask_token)
+        masked_x = x * (1 - w) + mask_token * w
+        if disable_vlc:
+            x = masked_x
+            masked_start = 0
+        else:
+            x = torch.cat([x, masked_x], 0)
+            masked_start = batch_size
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(
+            x.shape[0], 1, x.shape[-1],
+            dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        masked_x = x[:, masked_start:]
+        masked_x = self.lm_transformer(masked_x)
+        masked_x = masked_x.permute(1, 0, 2)
+        masked_x = masked_x[:, 1:]
+        masked_x = self.ln_lm(masked_x)
+        if not return_all_tokens:
+            masked_x = masked_x[bool_masked_pos]
+        logits = self.lm_head(masked_x)
+        assert self.pool_type == "default"
+        result = {"logits": logits}
+        if not disable_vlc:
+            x = x[0, :batch_size]
+            x = self.ln_post(x)
+            if self.proj is not None:
+                x = x @ self.proj
+            result["feature"] = x
+        return result
+
+
+def load_farl(model_type, model_file=None) -> VisualTransformer:
+    if model_type == "base":
+        model = VisualTransformer(
+            input_resolution=224, patch_size=16, width=768, layers=12, heads=12, output_dim=512)
+    elif model_type == "large":
+        model = VisualTransformer(
+            input_resolution=224, patch_size=16, width=1024, layers=24, heads=16, output_dim=512)
+    elif model_type == "huge":
+        model = VisualTransformer(
+            input_resolution=224, patch_size=14, width=1280, layers=32, heads=16, output_dim=512)
+    else:
+        raise
+    model.transformer.use_checkpoint = False
+    if model_file is not None:
+        checkpoint = torch.load(model_file, map_location='cpu')
+        state_dict = {}
+        for name, weight in checkpoint["state_dict"].items():
+            if name.startswith("visual"):
+                state_dict[name[7:]] = weight
+        inco = model.load_state_dict(state_dict, strict=False)
+        # print(inco.missing_keys)
+        assert len(inco.missing_keys) == 0
+    return model

src/pixel3dmm/preprocessing/facer/facer/io.py

@@ -0,0 +1,28 @@
+import torch
+import numpy as np
+from PIL import Image
+
+
+def read_hwc(path: str) -> torch.Tensor:
+    """Read an image from a given path.
+
+    Args:
+        path (str): The given path.
+    """
+    image = Image.open(path)
+    np_image = np.array(image.convert('RGB'))
+    return torch.from_numpy(np_image)
+
+
+def write_hwc(image: torch.Tensor, path: str):
+    """Write an image to a given path.
+
+    Args:
+        image (torch.Tensor): The image.
+        path (str): The given path.
+    """
+
+    Image.fromarray(image.cpu().numpy()).save(path)
+
+
+

src/pixel3dmm/preprocessing/facer/facer/show.py

@@ -0,0 +1,36 @@
+from typing import Optional, Tuple
+import torch
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from .util import bchw2hwc
+
+
+def set_figsize(*args):
+    if len(args) == 0:
+        plt.rcParams["figure.figsize"] = plt.rcParamsDefault["figure.figsize"]
+    elif len(args) == 1:
+        plt.rcParams["figure.figsize"] = (args[0], args[0])
+    elif len(args) == 2:
+        plt.rcParams["figure.figsize"] = tuple(args)
+    else:
+        raise RuntimeError(
+            f'Supported argument types: set_figsize() or set_figsize(int) or set_figsize(int, int)')
+
+
+def show_hwc(image: torch.Tensor):
+    if image.dtype != torch.uint8:
+        image = image.to(torch.uint8)
+    if image.size(2) == 1:
+        image = image.repeat(1, 1, 3)
+    pimage = Image.fromarray(image.cpu().numpy())
+    plt.imshow(pimage)
+    plt.show()
+
+
+def show_bchw(image: torch.Tensor):
+    show_hwc(bchw2hwc(image))
+
+
+def show_bhw(image: torch.Tensor):
+    show_bchw(image.unsqueeze(1))

src/pixel3dmm/preprocessing/facer/facer/transform.py

@@ -0,0 +1,384 @@
+from typing import List, Dict, Callable, Tuple, Optional
+import torch
+import torch.nn.functional as F
+import functools
+import numpy as np
+
+
+def get_crop_and_resize_matrix(
+        box: torch.Tensor, target_shape: Tuple[int, int],
+        target_face_scale: float = 1.0, make_square_crop: bool = True,
+        offset_xy: Optional[Tuple[float, float]] = None, align_corners: bool = True,
+        offset_box_coords: bool = False) -> torch.Tensor:
+    """
+    Args:
+        box: b x 4(x1, y1, x2, y2)
+        align_corners (bool): Set this to `True` only if the box you give has coordinates
+            ranging from `0` to `h-1` or `w-1`.
+
+        offset_box_coords (bool): Set this to `True` if the box you give has coordinates
+            ranging from `0` to `h` or `w`. 
+
+            Set this to `False` if the box coordinates range from `-0.5` to `h-0.5` or `w-0.5`.
+
+            If the box coordinates range from `0` to `h-1` or `w-1`, set `align_corners=True`.
+
+    Returns:
+        torch.Tensor: b x 3 x 3.
+    """
+    if offset_xy is None:
+        offset_xy = (0.0, 0.0)
+
+    x1, y1, x2, y2 = box.split(1, dim=1)  # b x 1
+    cx = (x1 + x2) / 2 + offset_xy[0]
+    cy = (y1 + y2) / 2 + offset_xy[1]
+    rx = (x2 - x1) / 2 / target_face_scale
+    ry = (y2 - y1) / 2 / target_face_scale
+    if make_square_crop:
+        rx = ry = torch.maximum(rx, ry)
+
+    x1, y1, x2, y2 = cx - rx, cy - ry, cx + rx, cy + ry
+
+    h, w, *_ = target_shape
+
+    zeros_pl = torch.zeros_like(x1)
+    ones_pl = torch.ones_like(x1)
+
+    if align_corners:
+        # x -> (x - x1) / (x2 - x1) * (w - 1)
+        # y -> (y - y1) / (y2 - y1) * (h - 1)
+        ax = 1.0 / (x2 - x1) * (w - 1)
+        ay = 1.0 / (y2 - y1) * (h - 1)
+        matrix = torch.cat([
+            ax, zeros_pl, -x1 * ax,
+            zeros_pl, ay, -y1 * ay,
+            zeros_pl, zeros_pl, ones_pl
+        ], dim=1).reshape(-1, 3, 3)  # b x 3 x 3
+    else:
+        if offset_box_coords:
+            # x1, x2 \in [0, w], y1, y2 \in [0, h]
+            # first we should offset x1, x2, y1, y2 to be ranging in
+            # [-0.5, w-0.5] and [-0.5, h-0.5]
+            # so to convert these pixel coordinates into boundary coordinates.
+            x1, x2, y1, y2 = x1-0.5, x2-0.5, y1-0.5, y2-0.5
+
+        # x -> (x - x1) / (x2 - x1) * w - 0.5
+        # y -> (y - y1) / (y2 - y1) * h - 0.5
+        ax = 1.0 / (x2 - x1) * w
+        ay = 1.0 / (y2 - y1) * h
+        matrix = torch.cat([
+            ax, zeros_pl, -x1 * ax - 0.5*ones_pl,
+            zeros_pl, ay, -y1 * ay - 0.5*ones_pl,
+            zeros_pl, zeros_pl, ones_pl
+        ], dim=1).reshape(-1, 3, 3)  # b x 3 x 3
+    return matrix
+
+
+def get_similarity_transform_matrix(
+        from_pts: torch.Tensor, to_pts: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        from_pts, to_pts: b x n x 2
+
+    Returns:
+        torch.Tensor: b x 3 x 3
+    """
+    mfrom = from_pts.mean(dim=1, keepdim=True)  # b x 1 x 2
+    mto = to_pts.mean(dim=1, keepdim=True)  # b x 1 x 2
+
+    a1 = (from_pts - mfrom).square().sum([1, 2], keepdim=False)  # b
+    c1 = ((to_pts - mto) * (from_pts - mfrom)).sum([1, 2], keepdim=False)  # b
+
+    to_delta = to_pts - mto
+    from_delta = from_pts - mfrom
+    c2 = (to_delta[:, :, 0] * from_delta[:, :, 1] - to_delta[:,
+          :, 1] * from_delta[:, :, 0]).sum([1], keepdim=False)  # b
+
+    a = c1 / a1
+    b = c2 / a1
+    dx = mto[:, 0, 0] - a * mfrom[:, 0, 0] - b * mfrom[:, 0, 1]  # b
+    dy = mto[:, 0, 1] + b * mfrom[:, 0, 0] - a * mfrom[:, 0, 1]  # b
+
+    ones_pl = torch.ones_like(a1)
+    zeros_pl = torch.zeros_like(a1)
+
+    return torch.stack([
+        a, b, dx,
+        -b, a, dy,
+        zeros_pl, zeros_pl, ones_pl,
+    ], dim=-1).reshape(-1, 3, 3)
+
+
+@functools.lru_cache()
+def _standard_face_pts():
+    pts = torch.tensor([
+        196.0, 226.0,
+        316.0, 226.0,
+        256.0, 286.0,
+        220.0, 360.4,
+        292.0, 360.4], dtype=torch.float32) / 256.0 - 1.0
+    return torch.reshape(pts, (5, 2))
+
+
+def get_face_align_matrix(
+        face_pts: torch.Tensor, target_shape: Tuple[int, int],
+        target_face_scale: float = 1.0, offset_xy: Optional[Tuple[float, float]] = None,
+        target_pts: Optional[torch.Tensor] = None):
+
+    if target_pts is None:
+        with torch.no_grad():
+            std_pts = _standard_face_pts().to(face_pts)  # [-1 1]
+            h, w, *_ = target_shape
+            target_pts = (std_pts * target_face_scale + 1) * \
+                torch.tensor([w-1, h-1]).to(face_pts) / 2.0
+            if offset_xy is not None:
+                target_pts[:, 0] += offset_xy[0]
+                target_pts[:, 1] += offset_xy[1]
+    else:
+        target_pts = target_pts.to(face_pts)
+
+    if target_pts.dim() == 2:
+        target_pts = target_pts.unsqueeze(0)
+    if target_pts.size(0) == 1:
+        target_pts = target_pts.broadcast_to(face_pts.shape)
+
+    assert target_pts.shape == face_pts.shape
+
+    return get_similarity_transform_matrix(face_pts, target_pts)
+
+
+def rot90(v):
+    return np.array([-v[1], v[0]])
+
+
+def get_quad(lm: torch.Tensor):
+    # N,2
+    lm = lm.detach().cpu().numpy()
+    # Choose oriented crop rectangle.
+    eye_avg = (lm[0] + lm[1]) * 0.5 + 0.5
+    mouth_avg = (lm[3] + lm[4]) * 0.5 + 0.5
+    eye_to_eye = lm[1] - lm[0]
+    eye_to_mouth = mouth_avg - eye_avg
+    x = eye_to_eye - rot90(eye_to_mouth)
+    x /= np.hypot(*x)
+    x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
+    y = rot90(x)
+    c = eye_avg + eye_to_mouth * 0.1
+    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+    quad_for_coeffs = quad[[0,3, 2,1]] #  顺序改一下
+    return torch.from_numpy(quad_for_coeffs).float()
+
+
+def get_face_align_matrix_celebm(
+        face_pts: torch.Tensor, target_shape: Tuple[int, int]):
+
+    face_pts = torch.stack([get_quad(pts) for pts in face_pts], dim=0).to(face_pts)
+
+    assert target_shape[0] == target_shape[1]
+    target_size  = target_shape[0]
+    target_pts = torch.as_tensor([[0, 0], [target_size,0], [target_size, target_size], [0, target_size]]).to(face_pts)
+
+    if target_pts.dim() == 2:
+        target_pts = target_pts.unsqueeze(0)
+    if target_pts.size(0) == 1:
+        target_pts = target_pts.broadcast_to(face_pts.shape)
+
+    assert target_pts.shape == face_pts.shape
+
+    return get_similarity_transform_matrix(face_pts, target_pts)
+
+@functools.lru_cache(maxsize=128)
+def _meshgrid(h, w) -> Tuple[torch.Tensor, torch.Tensor]:
+    yy, xx = torch.meshgrid(torch.arange(h).float(),
+                            torch.arange(w).float(),
+                            indexing='ij')
+    return yy + 0.5, xx + 0.5
+
+
+def _forge_grid(batch_size: int, device: torch.device,
+                output_shape: Tuple[int, int],
+                fn: Callable[[torch.Tensor], torch.Tensor]
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+    """ Forge transform maps with a given function `fn`.
+
+    Args:
+        output_shape (tuple): (b, h, w, ...).
+        fn (Callable[[torch.Tensor], torch.Tensor]): The function that accepts 
+            a bxnx2 array and outputs the transformed bxnx2 array. Both input 
+            and output store (x, y) coordinates.
+
+    Note: 
+        both input and output arrays of `fn` should store (y, x) coordinates.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Two maps `X` and `Y`, where for each 
+            pixel (y, x) or coordinate (x, y),
+            `(X[y, x], Y[y, x]) = fn([x, y])`
+    """
+    h, w, *_ = output_shape
+    yy, xx = _meshgrid(h, w)  # h x w
+    yy = yy.unsqueeze(0).broadcast_to(batch_size, h, w).to(device)
+    xx = xx.unsqueeze(0).broadcast_to(batch_size, h, w).to(device)
+
+    in_xxyy = torch.stack(
+        [xx, yy], dim=-1).reshape([batch_size, h*w, 2])  # (h x w) x 2
+    out_xxyy: torch.Tensor = fn(in_xxyy)  # (h x w) x 2
+    return out_xxyy.reshape(batch_size, h, w, 2)
+
+
+def _safe_arctanh(x: torch.Tensor, eps: float = 0.001) -> torch.Tensor:
+    return torch.clamp(x, -1+eps, 1-eps).arctanh()
+
+
+def inverted_tanh_warp_transform(coords: torch.Tensor, matrix: torch.Tensor,
+                                 warp_factor: float, warped_shape: Tuple[int, int]):
+    """ Inverted tanh-warp function.
+
+    Args:
+        coords (torch.Tensor): b x n x 2 (x, y). The transformed coordinates.
+        matrix: b x 3 x 3. A matrix that transforms un-normalized coordinates 
+            from the original image to the aligned yet not-warped image.
+        warp_factor (float): The warp factor. 
+            0 means linear transform, 1 means full tanh warp.
+        warped_shape (tuple): [height, width].
+
+    Returns:
+        torch.Tensor: b x n x 2 (x, y). The original coordinates.
+    """
+    h, w, *_ = warped_shape
+    # h -= 1
+    # w -= 1
+
+    w_h = torch.tensor([[w, h]]).to(coords)
+
+    if warp_factor > 0:
+        # normalize coordinates to [-1, +1]
+        coords = coords / w_h * 2 - 1
+
+        nl_part1 = coords > 1.0 - warp_factor
+        nl_part2 = coords < -1.0 + warp_factor
+
+        ret_nl_part1 = _safe_arctanh(
+            (coords - 1.0 + warp_factor) /
+            warp_factor) * warp_factor + \
+            1.0 - warp_factor
+        ret_nl_part2 = _safe_arctanh(
+            (coords + 1.0 - warp_factor) /
+            warp_factor) * warp_factor - \
+            1.0 + warp_factor
+
+        coords = torch.where(nl_part1, ret_nl_part1,
+                             torch.where(nl_part2, ret_nl_part2, coords))
+
+        # denormalize
+        coords = (coords + 1) / 2 * w_h
+
+    coords_homo = torch.cat(
+        [coords, torch.ones_like(coords[:, :, [0]])], dim=-1)  # b x n x 3
+
+    inv_matrix = torch.linalg.inv(matrix)  # b x 3 x 3
+    # inv_matrix = np.linalg.inv(matrix)
+    coords_homo = torch.bmm(
+        coords_homo, inv_matrix.permute(0, 2, 1))  # b x n x 3
+    return coords_homo[:, :, :2] / coords_homo[:, :, [2, 2]]
+
+
+def tanh_warp_transform(
+        coords: torch.Tensor, matrix: torch.Tensor,
+        warp_factor: float, warped_shape: Tuple[int, int]):
+    """ Tanh-warp function.
+
+    Args:
+        coords (torch.Tensor): b x n x 2 (x, y). The original coordinates.
+        matrix: b x 3 x 3. A matrix that transforms un-normalized coordinates 
+            from the original image to the aligned yet not-warped image.
+        warp_factor (float): The warp factor. 
+            0 means linear transform, 1 means full tanh warp.
+        warped_shape (tuple): [height, width].
+
+    Returns:
+        torch.Tensor: b x n x 2 (x, y). The transformed coordinates.
+    """
+    h, w, *_ = warped_shape
+    # h -= 1
+    # w -= 1
+    w_h = torch.tensor([[w, h]]).to(coords)
+
+    coords_homo = torch.cat(
+        [coords, torch.ones_like(coords[:, :, [0]])], dim=-1)  # b x n x 3
+
+    coords_homo = torch.bmm(coords_homo, matrix.transpose(2, 1))  # b x n x 3
+    coords = (coords_homo[:, :, :2] / coords_homo[:, :, [2, 2]])  # b x n x 2
+
+    if warp_factor > 0:
+        # normalize coordinates to [-1, +1]
+        coords = coords / w_h * 2 - 1
+
+        nl_part1 = coords > 1.0 - warp_factor
+        nl_part2 = coords < -1.0 + warp_factor
+
+        ret_nl_part1 = torch.tanh(
+            (coords - 1.0 + warp_factor) /
+            warp_factor) * warp_factor + \
+            1.0 - warp_factor
+        ret_nl_part2 = torch.tanh(
+            (coords + 1.0 - warp_factor) /
+            warp_factor) * warp_factor - \
+            1.0 + warp_factor
+
+        coords = torch.where(nl_part1, ret_nl_part1,
+                             torch.where(nl_part2, ret_nl_part2, coords))
+
+        # denormalize
+        coords = (coords + 1) / 2 * w_h
+
+    return coords
+
+
+def make_tanh_warp_grid(matrix: torch.Tensor, warp_factor: float,
+                        warped_shape: Tuple[int, int],
+                        orig_shape: Tuple[int, int]):
+    """
+    Args:
+        matrix: bx3x3 matrix.
+        warp_factor: The warping factor. `warp_factor=1.0` represents a vannila Tanh-warping, 
+           `warp_factor=0.0` represents a cropping.
+        warped_shape: The target image shape to transform to.
+
+    Returns:
+        torch.Tensor: b x h x w x 2 (x, y).
+    """
+    orig_h, orig_w, *_ = orig_shape
+    w_h = torch.tensor([orig_w, orig_h]).to(matrix).reshape(1, 1, 1, 2)
+    return _forge_grid(
+        matrix.size(0), matrix.device,
+        warped_shape,
+        functools.partial(inverted_tanh_warp_transform,
+                          matrix=matrix,
+                          warp_factor=warp_factor,
+                          warped_shape=warped_shape)) / w_h*2-1
+
+
+def make_inverted_tanh_warp_grid(matrix: torch.Tensor, warp_factor: float,
+                                 warped_shape: Tuple[int, int],
+                                 orig_shape: Tuple[int, int]):
+    """
+    Args:
+        matrix: bx3x3 matrix.
+        warp_factor: The warping factor. `warp_factor=1.0` represents a vannila Tanh-warping, 
+           `warp_factor=0.0` represents a cropping.
+        warped_shape: The target image shape to transform to.
+        orig_shape: The original image shape that is transformed from.
+
+    Returns:
+        torch.Tensor: b x h x w x 2 (x, y).
+    """
+    h, w, *_ = warped_shape
+    w_h = torch.tensor([w, h]).to(matrix).reshape(1, 1, 1, 2)
+    return _forge_grid(
+        matrix.size(0), matrix.device,
+        orig_shape,
+        functools.partial(tanh_warp_transform,
+                          matrix=matrix,
+                          warp_factor=warp_factor,
+                          warped_shape=warped_shape)) / w_h * 2-1

src/pixel3dmm/preprocessing/facer/facer/util.py

@@ -0,0 +1,169 @@
+import torch
+from typing import Any, Optional, Union, List, Dict
+import math
+import os
+from urllib.parse import urlparse
+import errno
+import sys
+import validators
+import requests
+import json
+
+
+def hwc2bchw(images: torch.Tensor) -> torch.Tensor:
+    return images.unsqueeze(0).permute(0, 3, 1, 2)
+
+
+def bchw2hwc(images: torch.Tensor, nrows: Optional[int] = None, border: int = 2,
+             background_value: float = 0) -> torch.Tensor:
+    """ make a grid image from an image batch.
+
+    Args:
+        images (torch.Tensor): input image batch.
+        nrows: rows of grid.
+        border: border size in pixel.
+        background_value: color value of background.
+    """
+    assert images.ndim == 4  # n x c x h x w
+    images = images.permute(0, 2, 3, 1)  # n x h x w x c
+    n, h, w, c = images.shape
+    if nrows is None:
+        nrows = max(int(math.sqrt(n)), 1)
+    ncols = (n + nrows - 1) // nrows
+    result = torch.full([(h + border) * nrows - border,
+                         (w + border) * ncols - border, c], background_value,
+                        device=images.device,
+                        dtype=images.dtype)
+
+    for i, single_image in enumerate(images):
+        row = i // ncols
+        col = i % ncols
+        yy = (h + border) * row
+        xx = (w + border) * col
+        result[yy:(yy + h), xx:(xx + w), :] = single_image
+    return result
+
+
+def bchw2bhwc(images: torch.Tensor) -> torch.Tensor:
+    return images.permute(0, 2, 3, 1)
+
+
+def bhwc2bchw(images: torch.Tensor) -> torch.Tensor:
+    return images.permute(0, 3, 1, 2)
+
+
+def bhwc2hwc(images: torch.Tensor, *kargs, **kwargs) -> torch.Tensor:
+    return bchw2hwc(bhwc2bchw(images), *kargs, **kwargs)
+
+
+def select_data(selection, data):
+    if isinstance(data, dict):
+        return {name: select_data(selection, val) for name, val in data.items()}
+    elif isinstance(data, (list, tuple)):
+        return [select_data(selection, val) for val in data]
+    elif isinstance(data, torch.Tensor):
+        return data[selection]
+    return data
+
+
+def download_from_github(to_path, organisation, repository, file_path, branch='main', username=None, access_token=None):
+    """ download files (including LFS files) from github.
+
+    For example, in order to downlod https://github.com/FacePerceiver/facer/blob/main/README.md, call with
+    ```
+    download_from_github(
+        to_path='README.md', organisation='FacePerceiver', 
+        repository='facer', file_path='README.md', branch='main')
+    ```
+    """
+    if username is not None:
+        assert access_token is not None
+        auth = (username, access_token)
+    else:
+        auth = None
+    r = requests.get(f'https://api.github.com/repos/{organisation}/{repository}/contents/{file_path}?ref={branch}',
+                     auth=auth)
+    data = json.loads(r.content)
+    torch.hub.download_url_to_file(data['download_url'], to_path)
+
+
+def is_github_url(url: str):
+    """
+    A typical github url should be like 
+        https://github.com/FacePerceiver/facer/blob/main/facer/util.py or 
+        https://github.com/FacePerceiver/facer/raw/main/facer/util.py.
+    """
+    return ('blob' in url or 'raw' in url) and url.startswith('https://github.com/')
+
+
+def get_github_components(url: str):
+    assert is_github_url(url)
+    organisation, repository, blob_or_raw, branch, * \
+        path = url[len('https://github.com/'):].split('/')
+    assert blob_or_raw in {'blob', 'raw'}
+    return organisation, repository, branch, '/'.join(path)
+
+
+def download_url_to_file(url, dst, **kwargs):
+    if is_github_url(url):
+        org, rep, branch, path = get_github_components(url)
+        download_from_github(dst, org, rep, path, branch, kwargs.get(
+            'username', None), kwargs.get('access_token', None))
+    else:
+        torch.hub.download_url_to_file(url, dst)
+
+
+def select_data(selection, data):
+    if isinstance(data, dict):
+        return {name: select_data(selection, val) for name, val in data.items()}
+    elif isinstance(data, (list, tuple)):
+        return [select_data(selection, val) for val in data]
+    elif isinstance(data, torch.Tensor):
+        return data[selection]
+    return data
+
+
+def download_jit(url_or_paths: Union[str, List[str]], model_dir=None, map_location=None, jit=True, **kwargs):
+    if isinstance(url_or_paths, str):
+        url_or_paths = [url_or_paths]
+
+    for url_or_path in url_or_paths:
+        try:
+            if validators.url(url_or_path):
+                url = url_or_path
+                if model_dir is None:
+                    if hasattr(torch.hub, 'get_dir'):
+                        hub_dir = torch.hub.get_dir()
+                    else:
+                        hub_dir = os.path.join(os.path.expanduser(
+                            '~'), '.cache', 'torch', 'hub')
+                    model_dir = os.path.join(hub_dir, 'checkpoints')
+
+                try:
+                    os.makedirs(model_dir)
+                except OSError as e:
+                    if e.errno == errno.EEXIST:
+                        # Directory already exists, ignore.
+                        pass
+                    else:
+                        # Unexpected OSError, re-raise.
+                        raise
+
+                parts = urlparse(url)
+                filename = os.path.basename(parts.path)
+                cached_file = os.path.join(model_dir, filename)
+                if not os.path.exists(cached_file):
+                    sys.stderr.write(
+                        'Downloading: "{}" to {}\n'.format(url, cached_file))
+                    download_url_to_file(url, cached_file)
+            else:
+                cached_file = url_or_path
+            if jit:
+                return torch.jit.load(cached_file, map_location=map_location, **kwargs)
+            else:
+                return torch.load(cached_file, map_location=map_location, **kwargs)
+        except:
+            sys.stderr.write(f'failed downloading from {url_or_path}\n')
+            raise
+
+    raise RuntimeError('failed to download jit models from all given urls')

src/pixel3dmm/preprocessing/facer/facer/version.py

@@ -0,0 +1 @@
+__version__="0.0.5"

src/pixel3dmm/preprocessing/facer/requirements.txt

@@ -0,0 +1,11 @@
+torch >= 1.9.1
+torchvision
+timm
+pillow
+numpy
+ipywidgets
+scikit-image
+matplotlib
+validators
+requests
+opencv-python

src/pixel3dmm/preprocessing/facer/samples/download.ipynb

@@ -0,0 +1,66 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.append('..')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "03bf1e12ed8a4b4ebf4ee9c5acda4a4f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0.00/1.25k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import facer\n",
+    "facer.util.download_from_github(\n",
+    "    to_path='.downloaded/download.ipynb', organisation='FacePerceiver',\n",
+    "    repository='facer', file_path='samples/download.ipynb', branch='main')\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "f55a4f7b5691d52d9535fa31bc30af4661bee88c7bc930914ceb21aeaa908798"
+  },
+  "kernelspec": {
+   "display_name": "Python 3.8.12 ('haya38')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/pixel3dmm/preprocessing/facer/scripts/build.sh

@@ -0,0 +1 @@
+python setup.py bdist_wheel