Upload 29 files

.gitattributes

@@ -2,34 +2,27 @@
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,160 @@
+# 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/
+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/
+cover/
+
+# 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
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .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
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__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/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

Dockerfile

@@ -0,0 +1,29 @@
+# Stage 1: Build
+FROM python:3.10.12-slim as builder
+
+RUN apt update && \
+    apt install --no-install-recommends -y build-essential gcc git
+
+COPY requirements-min.txt /requirements.txt
+
+RUN pip install --no-cache-dir --no-warn-script-location --user -r requirements.txt
+
+# Stage 2: Runtime
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
+ENV GRADIO_SERVER_NAME=0.0.0.0
+
+RUN apt update && \
+    apt install --no-install-recommends -y python3 python3-pip libgl1-mesa-glx libglib2.0-0 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY models/ /models
+COPY util/ /util
+COPY backbone/ /backbone
+COPY align/ /align
+COPY app.py app.py
+
+COPY --from=builder /root/.local/lib/python3.10/site-packages /root/.local/lib/python3.10/site-packages
+
+CMD [ "python3" , "-u", "app.py" ]
+EXPOSE 7860

README.md

@@ -1,12 +1,12 @@
 ---
-title: Lowkey V 0.1
-emoji: 📊
-colorFrom: green
-colorTo: red
+title: LowKey
+emoji: 😒
+colorFrom: pink
+colorTo: blue
 sdk: gradio
-sdk_version: 4.38.1
+sdk_version: 3.0.20
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Needs git lfs for the model weights

align/__init__.py

@@ -0,0 +1 @@
+

align/align_trans.py

@@ -0,0 +1,296 @@
+import numpy as np
+import cv2
+from align.matlab_cp2tform import get_similarity_transform_for_cv2
+
+
+# reference facial points, a list of coordinates (x,y)
+REFERENCE_FACIAL_POINTS = [        # default reference facial points for crop_size = (112, 112); should adjust REFERENCE_FACIAL_POINTS accordingly for other crop_size
+    [30.29459953,  51.69630051], 
+    [65.53179932,  51.50139999],
+    [48.02519989,  71.73660278],
+    [33.54930115,  92.3655014],
+    [62.72990036,  92.20410156]
+]
+
+DEFAULT_CROP_SIZE = (96, 112)
+
+
+class FaceWarpException(Exception):
+    def __str__(self):
+        return 'In File {}:{}'.format(
+            __file__, super.__str__(self))
+
+
+def get_reference_facial_points(output_size = None,
+                                inner_padding_factor = 0.0,
+                                outer_padding=(0, 0),
+                                default_square = False):
+    """
+    Function:
+    ----------
+        get reference 5 key points according to crop settings:
+        0. Set default crop_size:
+            if default_square: 
+                crop_size = (112, 112)
+            else: 
+                crop_size = (96, 112)
+        1. Pad the crop_size by inner_padding_factor in each side;
+        2. Resize crop_size into (output_size - outer_padding*2),
+            pad into output_size with outer_padding;
+        3. Output reference_5point;
+    Parameters:
+    ----------
+        @output_size: (w, h) or None
+            size of aligned face image
+        @inner_padding_factor: (w_factor, h_factor)
+            padding factor for inner (w, h)
+        @outer_padding: (w_pad, h_pad)
+            each row is a pair of coordinates (x, y)
+        @default_square: True or False
+            if True:
+                default crop_size = (112, 112)
+            else:
+                default crop_size = (96, 112);
+        !!! make sure, if output_size is not None:
+                (output_size - outer_padding) 
+                = some_scale * (default crop_size * (1.0 + inner_padding_factor))
+    Returns:
+    ----------
+        @reference_5point: 5x2 np.array
+            each row is a pair of transformed coordinates (x, y)
+    """
+    #print('\n===> get_reference_facial_points():')
+
+    #print('---> Params:')
+    #print('            output_size: ', output_size)
+    #print('            inner_padding_factor: ', inner_padding_factor)
+    #print('            outer_padding:', outer_padding)
+    #print('            default_square: ', default_square)
+
+    tmp_5pts = np.array(REFERENCE_FACIAL_POINTS)
+    tmp_crop_size = np.array(DEFAULT_CROP_SIZE)
+
+    # 0) make the inner region a square
+    if default_square:
+        size_diff = max(tmp_crop_size) - tmp_crop_size
+        tmp_5pts += size_diff / 2
+        tmp_crop_size += size_diff
+
+    #print('---> default:')
+    #print('              crop_size = ', tmp_crop_size)
+    #print('              reference_5pts = ', tmp_5pts)
+
+    if (output_size and
+            output_size[0] == tmp_crop_size[0] and
+            output_size[1] == tmp_crop_size[1]):
+        #print('output_size == DEFAULT_CROP_SIZE {}: return default reference points'.format(tmp_crop_size))
+        return tmp_5pts
+
+    if (inner_padding_factor == 0 and
+            outer_padding == (0, 0)):
+        if output_size is None:
+            #print('No paddings to do: return default reference points')
+            return tmp_5pts
+        else:
+            raise FaceWarpException(
+                'No paddings to do, output_size must be None or {}'.format(tmp_crop_size))
+
+    # check output size
+    if not (0 <= inner_padding_factor <= 1.0):
+        raise FaceWarpException('Not (0 <= inner_padding_factor <= 1.0)')
+
+    if ((inner_padding_factor > 0 or outer_padding[0] > 0 or outer_padding[1] > 0)
+            and output_size is None):
+        output_size = tmp_crop_size * \
+            (1 + inner_padding_factor * 2).astype(np.int32)
+        output_size += np.array(outer_padding)
+        #print('              deduced from paddings, output_size = ', output_size)
+
+    if not (outer_padding[0] < output_size[0]
+            and outer_padding[1] < output_size[1]):
+        raise FaceWarpException('Not (outer_padding[0] < output_size[0]'
+                                'and outer_padding[1] < output_size[1])')
+
+    # 1) pad the inner region according inner_padding_factor
+    #print('---> STEP1: pad the inner region according inner_padding_factor')
+    if inner_padding_factor > 0:
+        size_diff = tmp_crop_size * inner_padding_factor * 2
+        tmp_5pts += size_diff / 2
+        tmp_crop_size += np.round(size_diff).astype(np.int32)
+
+    #print('              crop_size = ', tmp_crop_size)
+    #print('              reference_5pts = ', tmp_5pts)
+
+    # 2) resize the padded inner region
+    #print('---> STEP2: resize the padded inner region')
+    size_bf_outer_pad = np.array(output_size) - np.array(outer_padding) * 2
+    #print('              crop_size = ', tmp_crop_size)
+    #print('              size_bf_outer_pad = ', size_bf_outer_pad)
+
+    if size_bf_outer_pad[0] * tmp_crop_size[1] != size_bf_outer_pad[1] * tmp_crop_size[0]:
+        raise FaceWarpException('Must have (output_size - outer_padding)'
+                                '= some_scale * (crop_size * (1.0 + inner_padding_factor)')
+
+    scale_factor = size_bf_outer_pad[0].astype(np.float32) / tmp_crop_size[0]
+    #print('              resize scale_factor = ', scale_factor)
+    tmp_5pts = tmp_5pts * scale_factor
+#    size_diff = tmp_crop_size * (scale_factor - min(scale_factor))
+#    tmp_5pts = tmp_5pts + size_diff / 2
+    tmp_crop_size = size_bf_outer_pad
+    #print('              crop_size = ', tmp_crop_size)
+    #print('              reference_5pts = ', tmp_5pts)
+
+    # 3) add outer_padding to make output_size
+    reference_5point = tmp_5pts + np.array(outer_padding)
+    tmp_crop_size = output_size
+    #print('---> STEP3: add outer_padding to make output_size')
+    #print('              crop_size = ', tmp_crop_size)
+    #print('              reference_5pts = ', tmp_5pts)
+
+    #print('===> end get_reference_facial_points\n')
+
+    return reference_5point
+
+
+def get_affine_transform_matrix(src_pts, dst_pts):
+    """
+    Function:
+    ----------
+        get affine transform matrix 'tfm' from src_pts to dst_pts
+    Parameters:
+    ----------
+        @src_pts: Kx2 np.array
+            source points matrix, each row is a pair of coordinates (x, y)
+        @dst_pts: Kx2 np.array
+            destination points matrix, each row is a pair of coordinates (x, y)
+    Returns:
+    ----------
+        @tfm: 2x3 np.array
+            transform matrix from src_pts to dst_pts
+    """
+
+    tfm = np.float32([[1, 0, 0], [0, 1, 0]])
+    n_pts = src_pts.shape[0]
+    ones = np.ones((n_pts, 1), src_pts.dtype)
+    src_pts_ = np.hstack([src_pts, ones])
+    dst_pts_ = np.hstack([dst_pts, ones])
+
+#    #print(('src_pts_:\n' + str(src_pts_))
+#    #print(('dst_pts_:\n' + str(dst_pts_))
+
+    A, res, rank, s = np.linalg.lstsq(src_pts_, dst_pts_)
+
+#    #print(('np.linalg.lstsq return A: \n' + str(A))
+#    #print(('np.linalg.lstsq return res: \n' + str(res))
+#    #print(('np.linalg.lstsq return rank: \n' + str(rank))
+#    #print(('np.linalg.lstsq return s: \n' + str(s))
+
+    if rank == 3:
+        tfm = np.float32([
+            [A[0, 0], A[1, 0], A[2, 0]],
+            [A[0, 1], A[1, 1], A[2, 1]]
+        ])
+    elif rank == 2:
+        tfm = np.float32([
+            [A[0, 0], A[1, 0], 0],
+            [A[0, 1], A[1, 1], 0]
+        ])
+
+    return tfm
+
+
+def warp_and_crop_face(src_img,
+                       facial_pts,
+                       reference_pts = None,
+                       crop_size=(96, 112),
+                       align_type = 'smilarity'):
+    """
+    Function:
+    ----------
+        apply affine transform 'trans' to uv
+    Parameters:
+    ----------
+        @src_img: 3x3 np.array
+            input image
+        @facial_pts: could be
+            1)a list of K coordinates (x,y)
+        or
+            2) Kx2 or 2xK np.array
+            each row or col is a pair of coordinates (x, y)
+        @reference_pts: could be
+            1) a list of K coordinates (x,y)
+        or
+            2) Kx2 or 2xK np.array
+            each row or col is a pair of coordinates (x, y)
+        or
+            3) None
+            if None, use default reference facial points
+        @crop_size: (w, h)
+            output face image size
+        @align_type: transform type, could be one of
+            1) 'similarity': use similarity transform
+            2) 'cv2_affine': use the first 3 points to do affine transform,
+                    by calling cv2.getAffineTransform()
+            3) 'affine': use all points to do affine transform
+    Returns:
+    ----------
+        @face_img: output face image with size (w, h) = @crop_size
+    """
+
+    if reference_pts is None:
+        if crop_size[0] == 96 and crop_size[1] == 112:
+            reference_pts = REFERENCE_FACIAL_POINTS
+        else:
+            default_square = False
+            inner_padding_factor = 0
+            outer_padding = (0, 0)
+            output_size = crop_size
+
+            reference_pts = get_reference_facial_points(output_size,
+                                                        inner_padding_factor,
+                                                        outer_padding,
+                                                        default_square)
+
+    ref_pts = np.float32(reference_pts)
+    ref_pts_shp = ref_pts.shape
+    if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2:
+        raise FaceWarpException(
+            'reference_pts.shape must be (K,2) or (2,K) and K>2')
+
+    if ref_pts_shp[0] == 2:
+        ref_pts = ref_pts.T
+
+    src_pts = np.float32(facial_pts)
+    src_pts_shp = src_pts.shape
+    if max(src_pts_shp) < 3 or min(src_pts_shp) != 2:
+        raise FaceWarpException(
+            'facial_pts.shape must be (K,2) or (2,K) and K>2')
+
+    if src_pts_shp[0] == 2:
+        src_pts = src_pts.T
+
+#    #print('--->src_pts:\n', src_pts
+#    #print('--->ref_pts\n', ref_pts
+
+    if src_pts.shape != ref_pts.shape:
+        raise FaceWarpException(
+            'facial_pts and reference_pts must have the same shape')
+
+    if align_type is 'cv2_affine':
+        tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3])
+#        #print(('cv2.getAffineTransform() returns tfm=\n' + str(tfm))
+    elif align_type is 'affine':
+        tfm = get_affine_transform_matrix(src_pts, ref_pts)
+#        #print(('get_affine_transform_matrix() returns tfm=\n' + str(tfm))
+    else:
+        tfm,tfm_inv = get_similarity_transform_for_cv2(src_pts, ref_pts)
+#        #print(('get_similarity_transform_for_cv2() returns tfm=\n' + str(tfm))
+
+#    #print('--->Transform matrix: '
+#    #print(('type(tfm):' + str(type(tfm)))
+#    #print(('tfm.dtype:' + str(tfm.dtype))
+#    #print( tfm
+
+    face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1]))
+
+    return face_img, tfm

align/box_utils.py

@@ -0,0 +1,238 @@
+import numpy as np
+from PIL import Image
+
+
+def nms(boxes, overlap_threshold = 0.5, mode = 'union'):
+    """Non-maximum suppression.
+
+    Arguments:
+        boxes: a float numpy array of shape [n, 5],
+            where each row is (xmin, ymin, xmax, ymax, score).
+        overlap_threshold: a float number.
+        mode: 'union' or 'min'.
+
+    Returns:
+        list with indices of the selected boxes
+    """
+
+    # if there are no boxes, return the empty list
+    if len(boxes) == 0:
+        return []
+
+    # list of picked indices
+    pick = []
+
+    # grab the coordinates of the bounding boxes
+    x1, y1, x2, y2, score = [boxes[:, i] for i in range(5)]
+
+    area = (x2 - x1 + 1.0)*(y2 - y1 + 1.0)
+    ids = np.argsort(score)  # in increasing order
+
+    while len(ids) > 0:
+
+        # grab index of the largest value
+        last = len(ids) - 1
+        i = ids[last]
+        pick.append(i)
+
+        # compute intersections
+        # of the box with the largest score
+        # with the rest of boxes
+
+        # left top corner of intersection boxes
+        ix1 = np.maximum(x1[i], x1[ids[:last]])
+        iy1 = np.maximum(y1[i], y1[ids[:last]])
+
+        # right bottom corner of intersection boxes
+        ix2 = np.minimum(x2[i], x2[ids[:last]])
+        iy2 = np.minimum(y2[i], y2[ids[:last]])
+
+        # width and height of intersection boxes
+        w = np.maximum(0.0, ix2 - ix1 + 1.0)
+        h = np.maximum(0.0, iy2 - iy1 + 1.0)
+
+        # intersections' areas
+        inter = w * h
+        if mode == 'min':
+            overlap = inter/np.minimum(area[i], area[ids[:last]])
+        elif mode == 'union':
+            # intersection over union (IoU)
+            overlap = inter/(area[i] + area[ids[:last]] - inter)
+
+        # delete all boxes where overlap is too big
+        ids = np.delete(
+            ids,
+            np.concatenate([[last], np.where(overlap > overlap_threshold)[0]])
+        )
+
+    return pick
+
+
+def convert_to_square(bboxes):
+    """Convert bounding boxes to a square form.
+
+    Arguments:
+        bboxes: a float numpy array of shape [n, 5].
+
+    Returns:
+        a float numpy array of shape [n, 5],
+            squared bounding boxes.
+    """
+
+    square_bboxes = np.zeros_like(bboxes)
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    h = y2 - y1 + 1.0
+    w = x2 - x1 + 1.0
+    max_side = np.maximum(h, w)
+    square_bboxes[:, 0] = x1 + w*0.5 - max_side*0.5
+    square_bboxes[:, 1] = y1 + h*0.5 - max_side*0.5
+    square_bboxes[:, 2] = square_bboxes[:, 0] + max_side - 1.0
+    square_bboxes[:, 3] = square_bboxes[:, 1] + max_side - 1.0
+    return square_bboxes
+
+
+def calibrate_box(bboxes, offsets):
+    """Transform bounding boxes to be more like true bounding boxes.
+    'offsets' is one of the outputs of the nets.
+
+    Arguments:
+        bboxes: a float numpy array of shape [n, 5].
+        offsets: a float numpy array of shape [n, 4].
+
+    Returns:
+        a float numpy array of shape [n, 5].
+    """
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    w = x2 - x1 + 1.0
+    h = y2 - y1 + 1.0
+    w = np.expand_dims(w, 1)
+    h = np.expand_dims(h, 1)
+
+    # this is what happening here:
+    # tx1, ty1, tx2, ty2 = [offsets[:, i] for i in range(4)]
+    # x1_true = x1 + tx1*w
+    # y1_true = y1 + ty1*h
+    # x2_true = x2 + tx2*w
+    # y2_true = y2 + ty2*h
+    # below is just more compact form of this
+
+    # are offsets always such that
+    # x1 < x2 and y1 < y2 ?
+
+    translation = np.hstack([w, h, w, h])*offsets
+    bboxes[:, 0:4] = bboxes[:, 0:4] + translation
+    return bboxes
+
+
+def get_image_boxes(bounding_boxes, img, size = 24):
+    """Cut out boxes from the image.
+
+    Arguments:
+        bounding_boxes: a float numpy array of shape [n, 5].
+        img: an instance of PIL.Image.
+        size: an integer, size of cutouts.
+
+    Returns:
+        a float numpy array of shape [n, 3, size, size].
+    """
+
+    num_boxes = len(bounding_boxes)
+    width, height = img.size
+
+    [dy, edy, dx, edx, y, ey, x, ex, w, h] = correct_bboxes(bounding_boxes, width, height)
+    img_boxes = np.zeros((num_boxes, 3, size, size), 'float32')
+
+    for i in range(num_boxes):
+        img_box = np.zeros((h[i], w[i], 3), 'uint8')
+
+        img_array = np.asarray(img, 'uint8')
+        img_box[dy[i]:(edy[i] + 1), dx[i]:(edx[i] + 1), :] =\
+            img_array[y[i]:(ey[i] + 1), x[i]:(ex[i] + 1), :]
+
+        # resize
+        img_box = Image.fromarray(img_box)
+        img_box = img_box.resize((size, size), Image.BILINEAR)
+        img_box = np.asarray(img_box, 'float32')
+
+        img_boxes[i, :, :, :] = _preprocess(img_box)
+
+    return img_boxes
+
+
+def correct_bboxes(bboxes, width, height):
+    """Crop boxes that are too big and get coordinates
+    with respect to cutouts.
+
+    Arguments:
+        bboxes: a float numpy array of shape [n, 5],
+            where each row is (xmin, ymin, xmax, ymax, score).
+        width: a float number.
+        height: a float number.
+
+    Returns:
+        dy, dx, edy, edx: a int numpy arrays of shape [n],
+            coordinates of the boxes with respect to the cutouts.
+        y, x, ey, ex: a int numpy arrays of shape [n],
+            corrected ymin, xmin, ymax, xmax.
+        h, w: a int numpy arrays of shape [n],
+            just heights and widths of boxes.
+
+        in the following order:
+            [dy, edy, dx, edx, y, ey, x, ex, w, h].
+    """
+
+    x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
+    w, h = x2 - x1 + 1.0,  y2 - y1 + 1.0
+    num_boxes = bboxes.shape[0]
+
+    # 'e' stands for end
+    # (x, y) -> (ex, ey)
+    x, y, ex, ey = x1, y1, x2, y2
+
+    # we need to cut out a box from the image.
+    # (x, y, ex, ey) are corrected coordinates of the box
+    # in the image.
+    # (dx, dy, edx, edy) are coordinates of the box in the cutout
+    # from the image.
+    dx, dy = np.zeros((num_boxes,)), np.zeros((num_boxes,))
+    edx, edy = w.copy() - 1.0, h.copy() - 1.0
+
+    # if box's bottom right corner is too far right
+    ind = np.where(ex > width - 1.0)[0]
+    edx[ind] = w[ind] + width - 2.0 - ex[ind]
+    ex[ind] = width - 1.0
+
+    # if box's bottom right corner is too low
+    ind = np.where(ey > height - 1.0)[0]
+    edy[ind] = h[ind] + height - 2.0 - ey[ind]
+    ey[ind] = height - 1.0
+
+    # if box's top left corner is too far left
+    ind = np.where(x < 0.0)[0]
+    dx[ind] = 0.0 - x[ind]
+    x[ind] = 0.0
+
+    # if box's top left corner is too high
+    ind = np.where(y < 0.0)[0]
+    dy[ind] = 0.0 - y[ind]
+    y[ind] = 0.0
+
+    return_list = [dy, edy, dx, edx, y, ey, x, ex, w, h]
+    return_list = [i.astype('int32') for i in return_list]
+
+    return return_list
+
+
+def _preprocess(img):
+    """Preprocessing step before feeding the network.
+
+    Arguments:
+        img: a float numpy array of shape [h, w, c].
+
+    Returns:
+        a float numpy array of shape [1, c, h, w].
+    """
+    img = img.transpose((2, 0, 1))
+    img = np.expand_dims(img, 0)
+    img = (img - 127.5) * 0.0078125
+    return img

align/detector.py

@@ -0,0 +1,126 @@
+import numpy as np
+import torch
+from torch.autograd import Variable
+import sys
+sys.path.append('./')
+from align.get_nets import PNet, RNet, ONet
+from align.box_utils import nms, calibrate_box, get_image_boxes, convert_to_square
+from align.first_stage import run_first_stage
+
+
+
+def detect_faces(image, min_face_size = 20.0,
+                 thresholds=[0.6, 0.7, 0.8],
+                 nms_thresholds=[0.7, 0.7, 0.7]):
+    """
+    Arguments:
+        image: an instance of PIL.Image.
+        min_face_size: a float number.
+        thresholds: a list of length 3.
+        nms_thresholds: a list of length 3.
+
+    Returns:
+        two float numpy arrays of shapes [n_boxes, 4] and [n_boxes, 10],
+        bounding boxes and facial landmarks.
+    """
+    # LOAD MODELS
+    pnet = PNet()
+    rnet = RNet()
+    onet = ONet()
+    onet.eval()
+
+    # BUILD AN IMAGE PYRAMID
+    width, height = image.size
+    min_length = min(height, width)
+
+    min_detection_size = 12
+    factor = 0.707  # sqrt(0.5)
+
+    # scales for scaling the image
+    scales = []
+
+    # scales the image so that
+    # minimum size that we can detect equals to
+    # minimum face size that we want to detect
+    m = min_detection_size/min_face_size
+    min_length *= m
+
+    factor_count = 0
+    while min_length > min_detection_size:
+        scales.append(m*factor**factor_count)
+        min_length *= factor
+        factor_count += 1
+
+    # STAGE 1
+
+    # it will be returned
+    bounding_boxes = []
+
+    # run P-Net on different scales
+    for s in scales:
+        boxes = run_first_stage(image, pnet, scale = s, threshold = thresholds[0])
+        bounding_boxes.append(boxes)
+
+    # collect boxes (and offsets, and scores) from different scales
+    bounding_boxes = [i for i in bounding_boxes if i is not None]
+    bounding_boxes = np.vstack(bounding_boxes)
+
+    keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
+    bounding_boxes = bounding_boxes[keep]
+
+    # use offsets predicted by pnet to transform bounding boxes
+    bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
+    # shape [n_boxes, 5]
+
+    bounding_boxes = convert_to_square(bounding_boxes)
+    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+    # STAGE 2
+
+    img_boxes = get_image_boxes(bounding_boxes, image, size = 24)
+    img_boxes = Variable(torch.FloatTensor(img_boxes), volatile = True)
+    output = rnet(img_boxes)
+    offsets = output[0].data.numpy()  # shape [n_boxes, 4]
+    probs = output[1].data.numpy()  # shape [n_boxes, 2]
+
+    keep = np.where(probs[:, 1] > thresholds[1])[0]
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1, ))
+    offsets = offsets[keep]
+
+    keep = nms(bounding_boxes, nms_thresholds[1])
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
+    bounding_boxes = convert_to_square(bounding_boxes)
+    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
+
+    # STAGE 3
+
+    img_boxes = get_image_boxes(bounding_boxes, image, size = 48)
+    if len(img_boxes) == 0: 
+        return [], []
+    img_boxes = Variable(torch.FloatTensor(img_boxes), volatile = True)
+    output = onet(img_boxes)
+    landmarks = output[0].data.numpy()  # shape [n_boxes, 10]
+    offsets = output[1].data.numpy()  # shape [n_boxes, 4]
+    probs = output[2].data.numpy()  # shape [n_boxes, 2]
+
+    keep = np.where(probs[:, 1] > thresholds[2])[0]
+    bounding_boxes = bounding_boxes[keep]
+    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1, ))
+    offsets = offsets[keep]
+    landmarks = landmarks[keep]
+
+    # compute landmark points
+    width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
+    height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
+    xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
+    landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1)*landmarks[:, 0:5]
+    landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1)*landmarks[:, 5:10]
+
+    bounding_boxes = calibrate_box(bounding_boxes, offsets)
+    keep = nms(bounding_boxes, nms_thresholds[2], mode = 'min')
+    bounding_boxes = bounding_boxes[keep]
+    landmarks = landmarks[keep]
+
+    return bounding_boxes, landmarks

align/first_stage.py

@@ -0,0 +1,97 @@
+import torch
+from torch.autograd import Variable
+import math
+from PIL import Image
+import numpy as np
+from align.box_utils import nms, _preprocess
+
+
+def run_first_stage(image, net, scale, threshold):
+    """Run P-Net, generate bounding boxes, and do NMS.
+
+    Arguments:
+        image: an instance of PIL.Image.
+        net: an instance of pytorch's nn.Module, P-Net.
+        scale: a float number,
+            scale width and height of the image by this number.
+        threshold: a float number,
+            threshold on the probability of a face when generating
+            bounding boxes from predictions of the net.
+
+    Returns:
+        a float numpy array of shape [n_boxes, 9],
+            bounding boxes with scores and offsets (4 + 1 + 4).
+    """
+
+    # scale the image and convert it to a float array
+    width, height = image.size
+    sw, sh = math.ceil(width*scale), math.ceil(height*scale)
+    img = image.resize((sw, sh), Image.BILINEAR)
+    img = np.asarray(img, 'float32')
+
+    img = Variable(torch.FloatTensor(_preprocess(img)), volatile = True)
+    output = net(img)
+    probs = output[1].data.numpy()[0, 1, :, :]
+    offsets = output[0].data.numpy()
+    # probs: probability of a face at each sliding window
+    # offsets: transformations to true bounding boxes
+
+    boxes = _generate_bboxes(probs, offsets, scale, threshold)
+    if len(boxes) == 0:
+        return None
+
+    keep = nms(boxes[:, 0:5], overlap_threshold = 0.5)
+    return boxes[keep]
+
+
+def _generate_bboxes(probs, offsets, scale, threshold):
+    """Generate bounding boxes at places
+    where there is probably a face.
+
+    Arguments:
+        probs: a float numpy array of shape [n, m].
+        offsets: a float numpy array of shape [1, 4, n, m].
+        scale: a float number,
+            width and height of the image were scaled by this number.
+        threshold: a float number.
+
+    Returns:
+        a float numpy array of shape [n_boxes, 9]
+    """
+
+    # applying P-Net is equivalent, in some sense, to
+    # moving 12x12 window with stride 2
+    stride = 2
+    cell_size = 12
+
+    # indices of boxes where there is probably a face
+    inds = np.where(probs > threshold)
+
+    if inds[0].size == 0:
+        return np.array([])
+
+    # transformations of bounding boxes
+    tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
+    # they are defined as:
+    # w = x2 - x1 + 1
+    # h = y2 - y1 + 1
+    # x1_true = x1 + tx1*w
+    # x2_true = x2 + tx2*w
+    # y1_true = y1 + ty1*h
+    # y2_true = y2 + ty2*h
+
+    offsets = np.array([tx1, ty1, tx2, ty2])
+    score = probs[inds[0], inds[1]]
+
+    # P-Net is applied to scaled images
+    # so we need to rescale bounding boxes back
+    bounding_boxes = np.vstack([
+        np.round((stride*inds[1] + 1.0)/scale),
+        np.round((stride*inds[0] + 1.0)/scale),
+        np.round((stride*inds[1] + 1.0 + cell_size)/scale),
+        np.round((stride*inds[0] + 1.0 + cell_size)/scale),
+        score, offsets
+    ])
+    # why one is added?
+
+    return bounding_boxes.T

align/get_nets.py

@@ -0,0 +1,169 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+import numpy as np
+
+
+class Flatten(nn.Module):
+
+    def __init__(self):
+        super(Flatten, self).__init__()
+
+    def forward(self, x):
+        """
+        Arguments:
+            x: a float tensor with shape [batch_size, c, h, w].
+        Returns:
+            a float tensor with shape [batch_size, c*h*w].
+        """
+
+        # without this pretrained model isn't working
+        x = x.transpose(3, 2).contiguous()
+
+        return x.view(x.size(0), -1)
+
+
+class PNet(nn.Module):
+
+    def __init__(self):
+
+        super(PNet, self).__init__()
+
+        # suppose we have input with size HxW, then
+        # after first layer: H - 2,
+        # after pool: ceil((H - 2)/2),
+        # after second conv: ceil((H - 2)/2) - 2,
+        # after last conv: ceil((H - 2)/2) - 4,
+        # and the same for W
+
+        self.features = nn.Sequential(OrderedDict([
+            ('conv1', nn.Conv2d(3, 10, 3, 1)),
+            ('prelu1', nn.PReLU(10)),
+            ('pool1', nn.MaxPool2d(2, 2, ceil_mode = True)),
+
+            ('conv2', nn.Conv2d(10, 16, 3, 1)),
+            ('prelu2', nn.PReLU(16)),
+
+            ('conv3', nn.Conv2d(16, 32, 3, 1)),
+            ('prelu3', nn.PReLU(32))
+        ]))
+
+        self.conv4_1 = nn.Conv2d(32, 2, 1, 1)
+        self.conv4_2 = nn.Conv2d(32, 4, 1, 1)
+
+        weights = np.load("align/pnet.npy", allow_pickle=True)[()]
+        for n, p in self.named_parameters():
+            p.data = torch.FloatTensor(weights[n])
+
+    def forward(self, x):
+        """
+        Arguments:
+            x: a float tensor with shape [batch_size, 3, h, w].
+        Returns:
+            b: a float tensor with shape [batch_size, 4, h', w'].
+            a: a float tensor with shape [batch_size, 2, h', w'].
+        """
+        x = self.features(x)
+        a = self.conv4_1(x)
+        b = self.conv4_2(x)
+        a = F.softmax(a)
+        return b, a
+
+
+class RNet(nn.Module):
+
+    def __init__(self):
+
+        super(RNet, self).__init__()
+
+        self.features = nn.Sequential(OrderedDict([
+            ('conv1', nn.Conv2d(3, 28, 3, 1)),
+            ('prelu1', nn.PReLU(28)),
+            ('pool1', nn.MaxPool2d(3, 2, ceil_mode = True)),
+
+            ('conv2', nn.Conv2d(28, 48, 3, 1)),
+            ('prelu2', nn.PReLU(48)),
+            ('pool2', nn.MaxPool2d(3, 2, ceil_mode = True)),
+
+            ('conv3', nn.Conv2d(48, 64, 2, 1)),
+            ('prelu3', nn.PReLU(64)),
+
+            ('flatten', Flatten()),
+            ('conv4', nn.Linear(576, 128)),
+            ('prelu4', nn.PReLU(128))
+        ]))
+
+        self.conv5_1 = nn.Linear(128, 2)
+        self.conv5_2 = nn.Linear(128, 4)
+
+        weights = np.load("align/rnet.npy", allow_pickle=True)[()]
+        for n, p in self.named_parameters():
+            p.data = torch.FloatTensor(weights[n])
+
+    def forward(self, x):
+        """
+        Arguments:
+            x: a float tensor with shape [batch_size, 3, h, w].
+        Returns:
+            b: a float tensor with shape [batch_size, 4].
+            a: a float tensor with shape [batch_size, 2].
+        """
+        x = self.features(x)
+        a = self.conv5_1(x)
+        b = self.conv5_2(x)
+        a = F.softmax(a)
+        return b, a
+
+
+class ONet(nn.Module):
+
+    def __init__(self):
+
+        super(ONet, self).__init__()
+
+        self.features = nn.Sequential(OrderedDict([
+            ('conv1', nn.Conv2d(3, 32, 3, 1)),
+            ('prelu1', nn.PReLU(32)),
+            ('pool1', nn.MaxPool2d(3, 2, ceil_mode = True)),
+
+            ('conv2', nn.Conv2d(32, 64, 3, 1)),
+            ('prelu2', nn.PReLU(64)),
+            ('pool2', nn.MaxPool2d(3, 2, ceil_mode = True)),
+
+            ('conv3', nn.Conv2d(64, 64, 3, 1)),
+            ('prelu3', nn.PReLU(64)),
+            ('pool3', nn.MaxPool2d(2, 2, ceil_mode = True)),
+
+            ('conv4', nn.Conv2d(64, 128, 2, 1)),
+            ('prelu4', nn.PReLU(128)),
+
+            ('flatten', Flatten()),
+            ('conv5', nn.Linear(1152, 256)),
+            ('drop5', nn.Dropout(0.25)),
+            ('prelu5', nn.PReLU(256)),
+        ]))
+
+        self.conv6_1 = nn.Linear(256, 2)
+        self.conv6_2 = nn.Linear(256, 4)
+        self.conv6_3 = nn.Linear(256, 10)
+
+        weights = np.load("align/onet.npy", allow_pickle=True)[()]
+        for n, p in self.named_parameters():
+            p.data = torch.FloatTensor(weights[n])
+
+    def forward(self, x):
+        """
+        Arguments:
+            x: a float tensor with shape [batch_size, 3, h, w].
+        Returns:
+            c: a float tensor with shape [batch_size, 10].
+            b: a float tensor with shape [batch_size, 4].
+            a: a float tensor with shape [batch_size, 2].
+        """
+        x = self.features(x)
+        a = self.conv6_1(x)
+        b = self.conv6_2(x)
+        c = self.conv6_3(x)
+        a = F.softmax(a)
+        return c, b, a

align/matlab_cp2tform.py

@@ -0,0 +1,345 @@
+import numpy as np
+from numpy.linalg import inv, norm, lstsq
+from numpy.linalg import matrix_rank as rank
+
+
+class MatlabCp2tormException(Exception):
+    def __str__(self):
+        return "In File {}:{}".format(
+                __file__, super.__str__(self))
+
+def tformfwd(trans, uv):
+    """
+    Function:
+    ----------
+        apply affine transform 'trans' to uv
+
+    Parameters:
+    ----------
+        @trans: 3x3 np.array
+            transform matrix
+        @uv: Kx2 np.array
+            each row is a pair of coordinates (x, y)
+
+    Returns:
+    ----------
+        @xy: Kx2 np.array
+            each row is a pair of transformed coordinates (x, y)
+    """
+    uv = np.hstack((
+        uv, np.ones((uv.shape[0], 1))
+    ))
+    xy = np.dot(uv, trans)
+    xy = xy[:, 0:-1]
+    return xy
+
+
+def tforminv(trans, uv):
+    """
+    Function:
+    ----------
+        apply the inverse of affine transform 'trans' to uv
+
+    Parameters:
+    ----------
+        @trans: 3x3 np.array
+            transform matrix
+        @uv: Kx2 np.array
+            each row is a pair of coordinates (x, y)
+
+    Returns:
+    ----------
+        @xy: Kx2 np.array
+            each row is a pair of inverse-transformed coordinates (x, y)
+    """
+    Tinv = inv(trans)
+    xy = tformfwd(Tinv, uv)
+    return xy
+
+
+def findNonreflectiveSimilarity(uv, xy, options=None):
+
+    options = {'K': 2}
+
+    K = options['K']
+    M = xy.shape[0]
+    x = xy[:, 0].reshape((-1, 1))  # use reshape to keep a column vector
+    y = xy[:, 1].reshape((-1, 1))  # use reshape to keep a column vector
+    # print('--->x, y:\n', x, y
+
+    tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1))))
+    tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1))))
+    X = np.vstack((tmp1, tmp2))
+    # print('--->X.shape: ', X.shape
+    # print('X:\n', X
+
+    u = uv[:, 0].reshape((-1, 1))  # use reshape to keep a column vector
+    v = uv[:, 1].reshape((-1, 1))  # use reshape to keep a column vector
+    U = np.vstack((u, v))
+    # print('--->U.shape: ', U.shape
+    # print('U:\n', U
+
+    # We know that X * r = U
+    if rank(X) >= 2 * K:
+        r, _, _, _ = lstsq(X, U)
+        r = np.squeeze(r)
+    else:
+        raise Exception("cp2tform: two Unique Points Req")
+
+    # print('--->r:\n', r
+
+    sc = r[0]
+    ss = r[1]
+    tx = r[2]
+    ty = r[3]
+
+    Tinv = np.array([
+        [sc, -ss, 0],
+        [ss,  sc, 0],
+        [tx,  ty, 1]
+    ])
+
+    # print('--->Tinv:\n', Tinv
+
+    T = inv(Tinv)
+    # print('--->T:\n', T
+
+    T[:, 2] = np.array([0, 0, 1])
+
+    return T, Tinv
+
+
+def findSimilarity(uv, xy, options=None):
+
+    options = {'K': 2}
+
+#    uv = np.array(uv)
+#    xy = np.array(xy)
+
+    # Solve for trans1
+    trans1, trans1_inv = findNonreflectiveSimilarity(uv, xy, options)
+
+    # Solve for trans2
+
+    # manually reflect the xy data across the Y-axis
+    xyR = xy
+    xyR[:, 0] = -1 * xyR[:, 0]
+
+    trans2r, trans2r_inv = findNonreflectiveSimilarity(uv, xyR, options)
+
+    # manually reflect the tform to undo the reflection done on xyR
+    TreflectY = np.array([
+        [-1, 0, 0],
+        [0, 1, 0],
+        [0, 0, 1]
+    ])
+
+    trans2 = np.dot(trans2r, TreflectY)
+
+    # Figure out if trans1 or trans2 is better
+    xy1 = tformfwd(trans1, uv)
+    norm1 = norm(xy1 - xy)
+
+    xy2 = tformfwd(trans2, uv)
+    norm2 = norm(xy2 - xy)
+
+    if norm1 <= norm2:
+        return trans1, trans1_inv
+    else:
+        trans2_inv = inv(trans2)
+        return trans2, trans2_inv
+
+
+def get_similarity_transform(src_pts, dst_pts, reflective = True):
+    """
+    Function:
+    ----------
+        Find Similarity Transform Matrix 'trans':
+            u = src_pts[:, 0]
+            v = src_pts[:, 1]
+            x = dst_pts[:, 0]
+            y = dst_pts[:, 1]
+            [x, y, 1] = [u, v, 1] * trans
+
+    Parameters:
+    ----------
+        @src_pts: Kx2 np.array
+            source points, each row is a pair of coordinates (x, y)
+        @dst_pts: Kx2 np.array
+            destination points, each row is a pair of transformed
+            coordinates (x, y)
+        @reflective: True or False
+            if True:
+                use reflective similarity transform
+            else:
+                use non-reflective similarity transform
+
+    Returns:
+    ----------
+       @trans: 3x3 np.array
+            transform matrix from uv to xy
+        trans_inv: 3x3 np.array
+            inverse of trans, transform matrix from xy to uv
+    """
+
+    if reflective:
+        trans, trans_inv = findSimilarity(src_pts, dst_pts)
+    else:
+        trans, trans_inv = findNonreflectiveSimilarity(src_pts, dst_pts)
+
+    return trans, trans_inv
+
+
+def cvt_tform_mat_for_cv2(trans):
+    """
+    Function:
+    ----------
+        Convert Transform Matrix 'trans' into 'cv2_trans' which could be
+        directly used by cv2.warpAffine():
+            u = src_pts[:, 0]
+            v = src_pts[:, 1]
+            x = dst_pts[:, 0]
+            y = dst_pts[:, 1]
+            [x, y].T = cv_trans * [u, v, 1].T
+
+    Parameters:
+    ----------
+        @trans: 3x3 np.array
+            transform matrix from uv to xy
+
+    Returns:
+    ----------
+        @cv2_trans: 2x3 np.array
+            transform matrix from src_pts to dst_pts, could be directly used
+            for cv2.warpAffine()
+    """
+    cv2_trans = trans[:, 0:2].T
+
+    return cv2_trans
+
+
+def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective = True):
+    """
+    Function:
+    ----------
+        Find Similarity Transform Matrix 'cv2_trans' which could be
+        directly used by cv2.warpAffine():
+            u = src_pts[:, 0]
+            v = src_pts[:, 1]
+            x = dst_pts[:, 0]
+            y = dst_pts[:, 1]
+            [x, y].T = cv_trans * [u, v, 1].T
+
+    Parameters:
+    ----------
+        @src_pts: Kx2 np.array
+            source points, each row is a pair of coordinates (x, y)
+        @dst_pts: Kx2 np.array
+            destination points, each row is a pair of transformed
+            coordinates (x, y)
+        reflective: True or False
+            if True:
+                use reflective similarity transform
+            else:
+                use non-reflective similarity transform
+
+    Returns:
+    ----------
+        @cv2_trans: 2x3 np.array
+            transform matrix from src_pts to dst_pts, could be directly used
+            for cv2.warpAffine()
+    """
+    trans, trans_inv = get_similarity_transform(src_pts, dst_pts, reflective)
+    cv2_trans = cvt_tform_mat_for_cv2(trans)
+    cv2_trans_inv = cvt_tform_mat_for_cv2(trans_inv)
+
+    return cv2_trans, cv2_trans_inv
+
+
+if __name__ == '__main__':
+    """
+    u = [0, 6, -2]
+    v = [0, 3, 5]
+    x = [-1, 0, 4]
+    y = [-1, -10, 4]
+
+    # In Matlab, run:
+    #
+    #   uv = [u'; v'];
+    #   xy = [x'; y'];
+    #   tform_sim=cp2tform(uv,xy,'similarity');
+    #
+    #   trans = tform_sim.tdata.T
+    #   ans =
+    #       -0.0764   -1.6190         0
+    #        1.6190   -0.0764         0
+    #       -3.2156    0.0290    1.0000
+    #   trans_inv = tform_sim.tdata.Tinv
+    #    ans =
+    #
+    #       -0.0291    0.6163         0
+    #       -0.6163   -0.0291         0
+    #       -0.0756    1.9826    1.0000
+    #    xy_m=tformfwd(tform_sim, u,v)
+    #
+    #    xy_m =
+    #
+    #       -3.2156    0.0290
+    #        1.1833   -9.9143
+    #        5.0323    2.8853
+    #    uv_m=tforminv(tform_sim, x,y)
+    #
+    #    uv_m =
+    #
+    #        0.5698    1.3953
+    #        6.0872    2.2733
+    #       -2.6570    4.3314
+    """
+    u = [0, 6, -2]
+    v = [0, 3, 5]
+    x = [-1, 0, 4]
+    y = [-1, -10, 4]
+
+    uv = np.array((u, v)).T
+    xy = np.array((x, y)).T
+
+    print("\n--->uv:")
+    print(uv)
+    print("\n--->xy:")
+    print(xy)
+
+    trans, trans_inv = get_similarity_transform(uv, xy)
+
+    print("\n--->trans matrix:")
+    print(trans)
+
+    print("\n--->trans_inv matrix:")
+    print(trans_inv)
+
+    print("\n---> apply transform to uv")
+    print("\nxy_m = uv_augmented * trans")
+    uv_aug = np.hstack((
+        uv, np.ones((uv.shape[0], 1))
+    ))
+    xy_m = np.dot(uv_aug, trans)
+    print(xy_m)
+
+    print("\nxy_m = tformfwd(trans, uv)")
+    xy_m = tformfwd(trans, uv)
+    print(xy_m)
+
+    print("\n---> apply inverse transform to xy")
+    print("\nuv_m = xy_augmented * trans_inv")
+    xy_aug = np.hstack((
+        xy, np.ones((xy.shape[0], 1))
+    ))
+    uv_m = np.dot(xy_aug, trans_inv)
+    print(uv_m)
+
+    print("\nuv_m = tformfwd(trans_inv, xy)")
+    uv_m = tformfwd(trans_inv, xy)
+    print(uv_m)
+
+    uv_m = tforminv(trans, xy)
+    print("\nuv_m = tforminv(trans, xy)")
+    print(uv_m)

align/onet.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09b2ec49f8e6210dc4203e898a6fcf9f9f1330d6750f3c8c7a1d009aa75b3c6e
+size 132

align/pnet.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8364064b9199f3df5a01c72b88e6f01024d1c76657e55d2524ca8a87b97bea5d
+size 130

align/rnet.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80917ce9d27de824859196469e37680b81504e99310b581e96d524c533476803
+size 131

app.py

@@ -0,0 +1,80 @@
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+from util.feature_extraction_utils import normalize_transforms
+from util.attack_utils import Attack
+from util.prepare_utils import prepare_models, prepare_dir_vec, get_ensemble
+from align.detector import detect_faces
+from align.align_trans import get_reference_facial_points, warp_and_crop_face
+import torchvision.transforms as transforms
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(device)
+to_tensor = transforms.ToTensor()
+
+eps = 0.05
+n_iters = 50
+input_size = [112, 112]
+attack_type = 'lpips'
+c_tv = None
+c_sim = 0.05
+lr = 0.0025
+net_type = 'alex'
+noise_size = 0.005
+n_starts = 1
+kernel_size_gf = 7
+sigma_gf = 3
+combination = True
+using_subspace = False
+V_reduction_root = './'
+model_backbones = ['IR_152', 'IR_152', 'ResNet_152', 'ResNet_152']
+model_roots = ['models/Backbone_IR_152_Arcface_Epoch_112.pth', 'models/Backbone_IR_152_Cosface_Epoch_70.pth', \
+    'models/Backbone_ResNet_152_Arcface_Epoch_65.pth', 'models/Backbone_ResNet_152_Cosface_Epoch_68.pth'] 
+direction = 1
+crop_size = 112
+scale = crop_size / 112.
+
+models_attack, V_reduction, dim = prepare_models(model_backbones,
+              input_size,
+              model_roots,
+              kernel_size_gf,
+              sigma_gf,
+              combination,
+              using_subspace,
+              V_reduction_root)
+
+def protect(img):
+  img = Image.fromarray(img)
+  reference = get_reference_facial_points(default_square=True) * scale
+  h,w,c = np.array(img).shape
+
+  _, landmarks = detect_faces(img)
+  facial5points = [[landmarks[0][j], landmarks[0][j + 5]] for j in range(5)]
+
+  _, tfm = warp_and_crop_face(np.array(img), facial5points, reference, crop_size=(crop_size, crop_size))
+
+  # pytorch transform
+  theta = normalize_transforms(tfm, w, h)
+  tensor_img = to_tensor(img).unsqueeze(0).to(device)
+
+  V_reduction = None
+  dim = 512
+
+  # Find gradient direction vector
+  dir_vec_extractor = get_ensemble(models = models_attack, sigma_gf=None, kernel_size_gf=None, combination=False, V_reduction=V_reduction, warp=True, theta_warp=theta)
+  dir_vec = prepare_dir_vec(dir_vec_extractor, tensor_img, dim, combination)
+
+  img_attacked = tensor_img.clone()
+  attack = Attack(models_attack, dim, attack_type, 
+                  eps, c_sim, net_type, lr, n_iters, 
+                  noise_size, n_starts, c_tv, sigma_gf, 
+                  kernel_size_gf, combination, warp=True, 
+                  theta_warp=theta, V_reduction=V_reduction)
+  img_attacked = attack.execute(tensor_img, dir_vec, direction).detach().cpu()
+
+  img_attacked_pil = transforms.ToPILImage()(img_attacked[0])
+  return img_attacked_pil
+
+gr.Interface(fn=protect, inputs=gr.components.Image(shape=(512,512)),
+                         outputs=gr.components.Image(type="pil"), allow_flagging="never").launch(show_error=True, quiet=False, share=True)

backbone/__init__.py

@@ -0,0 +1 @@
+

backbone/model_irse.py

@@ -0,0 +1,237 @@
+import torch
+import torch.nn as nn
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, ReLU, Sigmoid, Dropout, MaxPool2d, \
+    AdaptiveAvgPool2d, Sequential, Module
+from collections import namedtuple
+
+
+# Support: ['IR_50', 'IR_101', 'IR_152', 'IR_SE_50', 'IR_SE_101', 'IR_SE_152']
+
+
+class Flatten(Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+
+def l2_norm(input, axis=1):
+    norm = torch.norm(input, 2, axis, True)
+    output = torch.div(input, norm)
+
+    return output
+
+
+class SEModule(Module):
+    def __init__(self, channels, reduction):
+        super(SEModule, self).__init__()
+        self.avg_pool = AdaptiveAvgPool2d(1)
+        self.fc1 = Conv2d(
+            channels, channels // reduction, kernel_size=1, padding=0, bias=False)
+
+        nn.init.xavier_uniform_(self.fc1.weight.data)
+
+        self.relu = ReLU(inplace=True)
+        self.fc2 = Conv2d(
+            channels // reduction, channels, kernel_size=1, padding=0, bias=False)
+
+        self.sigmoid = Sigmoid()
+
+    def forward(self, x):
+        module_input = x
+        x = self.avg_pool(x)
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.sigmoid(x)
+
+        return module_input * x
+
+
+class bottleneck_IR(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False), BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth))
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class bottleneck_IR_SE(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR_SE, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+            PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+            BatchNorm2d(depth),
+            SEModule(depth, 16)
+        )
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
+    '''A named tuple describing a ResNet block.'''
+
+
+def get_block(in_channel, depth, num_units, stride=2):
+
+    return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
+
+
+def get_blocks(num_layers):
+    if num_layers == 50:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=4),
+            get_block(in_channel=128, depth=256, num_units=14),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 100:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=13),
+            get_block(in_channel=128, depth=256, num_units=30),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 152:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=8),
+            get_block(in_channel=128, depth=256, num_units=36),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+
+    return blocks
+
+
+class Backbone(Module):
+    def __init__(self, input_size, num_layers, mode='ir'):
+        super(Backbone, self).__init__()
+        assert input_size[0] in [112, 224], "input_size should be [112, 112] or [224, 224]"
+        assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152"
+        assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se"
+        blocks = get_blocks(num_layers)
+        if mode == 'ir':
+            unit_module = bottleneck_IR
+        elif mode == 'ir_se':
+            unit_module = bottleneck_IR_SE
+        self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+                                      BatchNorm2d(64),
+                                      PReLU(64))
+        if input_size[0] == 112:
+            self.output_layer = Sequential(BatchNorm2d(512),
+                                           Dropout(),
+                                           Flatten(),
+                                           Linear(512 * 7 * 7, 512),
+                                           BatchNorm1d(512))
+        else:
+            self.output_layer = Sequential(BatchNorm2d(512),
+                                           Dropout(),
+                                           Flatten(),
+                                           Linear(512 * 14 * 14, 512),
+                                           BatchNorm1d(512))
+
+        modules = []
+        for block in blocks:
+            for bottleneck in block:
+                modules.append(
+                    unit_module(bottleneck.in_channel,
+                                bottleneck.depth,
+                                bottleneck.stride))
+        self.body = Sequential(*modules)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.input_layer(x)
+        x = self.body(x)
+        x = self.output_layer(x)
+
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight.data)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm1d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight.data)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+
+
+def IR_50(input_size):
+    """Constructs a ir-50 model.
+    """
+    model = Backbone(input_size, 50, 'ir')
+
+    return model
+
+
+def IR_101(input_size):
+    """Constructs a ir-101 model.
+    """
+    model = Backbone(input_size, 100, 'ir')
+
+    return model
+
+
+def IR_152(input_size):
+    """Constructs a ir-152 model.
+    """
+    model = Backbone(input_size, 152, 'ir')
+
+    return model
+
+
+def IR_SE_50(input_size):
+    """Constructs a ir_se-50 model.
+    """
+    model = Backbone(input_size, 50, 'ir_se')
+
+    return model
+
+
+def IR_SE_101(input_size):
+    """Constructs a ir_se-101 model.
+    """
+    model = Backbone(input_size, 100, 'ir_se')
+
+    return model
+
+
+def IR_SE_152(input_size):
+    """Constructs a ir_se-152 model.
+    """
+    model = Backbone(input_size, 152, 'ir_se')
+
+    return model

backbone/model_resnet.py

@@ -0,0 +1,195 @@
+import torch.nn as nn
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, ReLU, Dropout, MaxPool2d, Sequential, Module
+
+
+# Support: ['ResNet_50', 'ResNet_101', 'ResNet_152']
+
+
+def conv3x3(in_planes, out_planes, stride = 1):
+    """3x3 convolution with padding"""
+
+    return Conv2d(in_planes, out_planes, kernel_size = 3, stride = stride,
+                     padding = 1, bias = False)
+
+
+def conv1x1(in_planes, out_planes, stride = 1):
+    """1x1 convolution"""
+
+    return Conv2d(in_planes, out_planes, kernel_size = 1, stride = stride, bias = False)
+
+
+class BasicBlock(Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride = 1, downsample = None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = BatchNorm2d(planes)
+        self.relu = ReLU(inplace = True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride = 1, downsample = None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = BatchNorm2d(planes * self.expansion)
+        self.relu = ReLU(inplace = True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(Module):
+
+    def __init__(self, input_size, block, layers, zero_init_residual = True):
+        super(ResNet, self).__init__()
+        assert input_size[0] in [112, 224], "input_size should be [112, 112] or [224, 224]"
+        self.inplanes = 64
+        self.conv1 = Conv2d(3, 64, kernel_size = 7, stride = 2, padding = 3, bias = False)
+        self.bn1 = BatchNorm2d(64)
+        self.relu = ReLU(inplace = True)
+        self.maxpool = MaxPool2d(kernel_size = 3, stride = 2, padding = 1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride = 2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride = 2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride = 2)
+
+        self.bn_o1 = BatchNorm2d(2048)
+        self.dropout = Dropout()
+        if input_size[0] == 112:
+            self.fc = Linear(2048 * 4 * 4, 512)
+        else:
+            self.fc = Linear(2048 * 8 * 8, 512)
+        self.bn_o2 = BatchNorm1d(512)
+
+        for m in self.modules():
+            if isinstance(m, Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode = 'fan_out', nonlinearity = 'relu')
+            elif isinstance(m, BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride = 1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.bn_o1(x)
+        x = self.dropout(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        x = self.bn_o2(x)
+
+        return x
+
+def ResNet_18(input_size, **kwargs):
+    """Constructs a ResNet-50 model.
+    """
+    model = ResNet(input_size, Bottleneck, [2, 2, 2, 2], **kwargs)
+
+    return model
+
+
+def ResNet_50(input_size, **kwargs):
+    """Constructs a ResNet-50 model.
+    """
+    model = ResNet(input_size, Bottleneck, [3, 4, 6, 3], **kwargs)
+
+    return model
+
+
+def ResNet_101(input_size, **kwargs):
+    """Constructs a ResNet-101 model.
+    """
+    model = ResNet(input_size, Bottleneck, [3, 4, 23, 3], **kwargs)
+
+    return model
+
+
+def ResNet_152(input_size, **kwargs):
+    """Constructs a ResNet-152 model.
+    """
+    model = ResNet(input_size, Bottleneck, [3, 8, 36, 3], **kwargs)
+
+    return model

backbone/models2.py

@@ -0,0 +1,305 @@
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, ReLU, Sigmoid, Dropout2d, Dropout, AvgPool2d, MaxPool2d, AdaptiveAvgPool2d, Sequential, Module, Parameter
+import torch.nn.functional as F
+import torch
+from collections import namedtuple
+import math
+import pdb
+
+##################################  Original Arcface Model #############################################################
+
+class Flatten(Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+def l2_norm(input,axis=1):
+    norm = torch.norm(input,2,axis,True)
+    output = torch.div(input, norm)
+    return output
+
+class SEModule(Module):
+    def __init__(self, channels, reduction):
+        super(SEModule, self).__init__()
+        self.avg_pool = AdaptiveAvgPool2d(1)
+        self.fc1 = Conv2d(
+            channels, channels // reduction, kernel_size=1, padding=0 ,bias=False)
+        self.relu = ReLU(inplace=True)
+        self.fc2 = Conv2d(
+            channels // reduction, channels, kernel_size=1, padding=0 ,bias=False)
+        self.sigmoid = Sigmoid()
+
+    def forward(self, x):
+        module_input = x
+        x = self.avg_pool(x)
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.sigmoid(x)
+        return module_input * x
+
+class bottleneck_IR(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride ,bias=False), BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 1 ,bias=False), PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 1 ,bias=False), BatchNorm2d(depth))
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+        return res + shortcut
+
+class bottleneck_IR_SE(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR_SE, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride ,bias=False), 
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3,3), (1,1),1 ,bias=False),
+            PReLU(depth),
+            Conv2d(depth, depth, (3,3), stride, 1 ,bias=False),
+            BatchNorm2d(depth),
+            SEModule(depth,16)
+            )
+    def forward(self,x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+        return res + shortcut
+
+class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
+    '''A named tuple describing a ResNet block.'''
+    
+def get_block(in_channel, depth, num_units, stride = 2):
+  return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units-1)]
+
+def get_blocks(num_layers):
+    if num_layers == 50:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units = 3),
+            get_block(in_channel=64, depth=128, num_units=4),
+            get_block(in_channel=128, depth=256, num_units=14),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 100:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=13),
+            get_block(in_channel=128, depth=256, num_units=30),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 152:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=8),
+            get_block(in_channel=128, depth=256, num_units=36),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    return blocks
+
+class Backbone(Module):
+    def __init__(self, num_layers, drop_ratio, mode='ir'):
+        super(Backbone, self).__init__()
+        assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
+        assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
+        blocks = get_blocks(num_layers)
+        if mode == 'ir':
+            unit_module = bottleneck_IR
+        elif mode == 'ir_se':
+            unit_module = bottleneck_IR_SE
+        self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1 ,bias=False), 
+                                      BatchNorm2d(64), 
+                                      PReLU(64))
+        self.output_layer = Sequential(BatchNorm2d(512), 
+                                       Dropout(drop_ratio),
+                                       Flatten(),
+                                       Linear(512 * 7 * 7, 512),
+                                       BatchNorm1d(512))
+        modules = []
+        for block in blocks:
+            for bottleneck in block:
+                modules.append(
+                    unit_module(bottleneck.in_channel,
+                                bottleneck.depth,
+                                bottleneck.stride))
+        self.body = Sequential(*modules)
+    
+    def forward(self,x):
+        x = self.input_layer(x)
+        x = self.body(x)
+        x = self.output_layer(x)
+        return l2_norm(x)
+
+##################################  MobileFaceNet #############################################################
+    
+class Conv_block(Module):
+    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
+        super(Conv_block, self).__init__()
+        self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding, bias=False)
+        self.bn = BatchNorm2d(out_c)
+        self.prelu = PReLU(out_c)
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.prelu(x)
+        return x
+
+class Linear_block(Module):
+    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
+        super(Linear_block, self).__init__()
+        self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding, bias=False)
+        self.bn = BatchNorm2d(out_c)
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+class Depth_Wise(Module):
+     def __init__(self, in_c, out_c, residual = False, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=1):
+        super(Depth_Wise, self).__init__()
+        self.conv = Conv_block(in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1))
+        self.conv_dw = Conv_block(groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride)
+        self.project = Linear_block(groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1))
+        self.residual = residual
+     def forward(self, x):
+        if self.residual:
+            short_cut = x
+        x = self.conv(x)
+        x = self.conv_dw(x)
+        x = self.project(x)
+        if self.residual:
+            output = short_cut + x
+        else:
+            output = x
+        return output
+
+class Residual(Module):
+    def __init__(self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)):
+        super(Residual, self).__init__()
+        modules = []
+        for _ in range(num_block):
+            modules.append(Depth_Wise(c, c, residual=True, kernel=kernel, padding=padding, stride=stride, groups=groups))
+        self.model = Sequential(*modules)
+    def forward(self, x):
+        return self.model(x)
+
+class MobileFaceNet(Module):
+    def __init__(self, embedding_size):
+        super(MobileFaceNet, self).__init__()
+        self.conv1 = Conv_block(3, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1))
+        self.conv2_dw = Conv_block(64, 64, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64)
+        self.conv_23 = Depth_Wise(64, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128)
+        self.conv_3 = Residual(64, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv_34 = Depth_Wise(64, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256)
+        self.conv_4 = Residual(128, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv_45 = Depth_Wise(128, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512)
+        self.conv_5 = Residual(128, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv_6_sep = Conv_block(128, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0))
+        self.conv_6_dw = Linear_block(512, 512, groups=512, kernel=(7,7), stride=(1, 1), padding=(0, 0))
+        self.conv_6_flatten = Flatten()
+        self.linear = Linear(512, embedding_size, bias=False)
+        self.bn = BatchNorm1d(embedding_size)
+    
+    def forward(self, x):
+        out = self.conv1(x)
+
+        out = self.conv2_dw(out)
+
+        out = self.conv_23(out)
+
+        out = self.conv_3(out)
+        
+        out = self.conv_34(out)
+
+        out = self.conv_4(out)
+
+        out = self.conv_45(out)
+
+        out = self.conv_5(out)
+
+        out = self.conv_6_sep(out)
+
+        out = self.conv_6_dw(out)
+
+        out = self.conv_6_flatten(out)
+
+        out = self.linear(out)
+
+        out = self.bn(out)
+        return l2_norm(out)
+
+##################################  Arcface head #############################################################
+
+class Arcface(Module):
+    # implementation of additive margin softmax loss in https://arxiv.org/abs/1801.05599    
+    def __init__(self, embedding_size=512, classnum=51332,  s=64., m=0.5):
+        super(Arcface, self).__init__()
+        self.classnum = classnum
+        self.kernel = Parameter(torch.Tensor(embedding_size,classnum))
+        # initial kernel
+        self.kernel.data.uniform_(-1, 1).renorm_(2,1,1e-5).mul_(1e5)
+        self.m = m # the margin value, default is 0.5
+        self.s = s # scalar value default is 64, see normface https://arxiv.org/abs/1704.06369
+        self.cos_m = math.cos(m)
+        self.sin_m = math.sin(m)
+        self.mm = self.sin_m * m  # issue 1
+        self.threshold = math.cos(math.pi - m)
+    def forward(self, embbedings, label):
+        # weights norm
+        nB = len(embbedings)
+        kernel_norm = l2_norm(self.kernel,axis=0)
+        # cos(theta+m)
+        cos_theta = torch.mm(embbedings,kernel_norm)
+#         output = torch.mm(embbedings,kernel_norm)
+        cos_theta = cos_theta.clamp(-1,1) # for numerical stability
+        cos_theta_2 = torch.pow(cos_theta, 2)
+        sin_theta_2 = 1 - cos_theta_2
+        sin_theta = torch.sqrt(sin_theta_2)
+        cos_theta_m = (cos_theta * self.cos_m - sin_theta * self.sin_m)
+        # this condition controls the theta+m should in range [0, pi]
+        #      0<=theta+m<=pi
+        #     -m<=theta<=pi-m
+        cond_v = cos_theta - self.threshold
+        cond_mask = cond_v <= 0
+        keep_val = (cos_theta - self.mm) # when theta not in [0,pi], use cosface instead
+        cos_theta_m[cond_mask] = keep_val[cond_mask]
+        output = cos_theta * 1.0 # a little bit hacky way to prevent in_place operation on cos_theta
+        idx_ = torch.arange(0, nB, dtype=torch.long)
+        output[idx_, label] = cos_theta_m[idx_, label]
+        output *= self.s # scale up in order to make softmax work, first introduced in normface
+        return output
+
+##################################  Cosface head #############################################################    
+    
+class Am_softmax(Module):
+    # implementation of additive margin softmax loss in https://arxiv.org/abs/1801.05599    
+    def __init__(self,embedding_size=512,classnum=51332):
+        super(Am_softmax, self).__init__()
+        self.classnum = classnum
+        self.kernel = Parameter(torch.Tensor(embedding_size,classnum))
+        # initial kernel
+        self.kernel.data.uniform_(-1, 1).renorm_(2,1,1e-5).mul_(1e5)
+        self.m = 0.35 # additive margin recommended by the paper
+        self.s = 30. # see normface https://arxiv.org/abs/1704.06369
+    def forward(self,embbedings,label):
+        kernel_norm = l2_norm(self.kernel,axis=0)
+        cos_theta = torch.mm(embbedings,kernel_norm)
+        cos_theta = cos_theta.clamp(-1,1) # for numerical stability
+        phi = cos_theta - self.m
+        label = label.view(-1,1) #size=(B,1)
+        index = cos_theta.data * 0.0 #size=(B,Classnum)
+        index.scatter_(1,label.data.view(-1,1),1)
+        index = index.byte()
+        output = cos_theta * 1.0
+        output[index] = phi[index] #only change the correct predicted output
+        output *= self.s # scale up in order to make softmax work, first introduced in normface
+        return output

models/Backbone_IR_152_Arcface_Epoch_112.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e906b5405eba33e2b90887832ebfc72369df4e349fef65d354d3e6ca9e709fb
+size 134

models/Backbone_IR_152_Cosface_Epoch_70.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c89ffc23eec6c3adb66067a74ffb42b3d90a54ee5e8da2d058190bc65b4333b
+size 134

models/Backbone_ResNet_152_Arcface_Epoch_65.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64fd31b1bb52e89b10662d4d2a33c246b7058311b5eb4065509584f15322460b
+size 134

models/Backbone_ResNet_152_Cosface_Epoch_68.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e13a1c8d6df7f6999a54e56f3e6398ed4da17ca8d63b665d0d749ee9dc3aedad
+size 134

requirements-min.txt

@@ -0,0 +1,5 @@
+scikit-learn
+git+https://github.com/logasja/lpips-pytorch.git
+gradio
+pillow
+opencv-python

requirements.txt

@@ -0,0 +1,103 @@
+aiofiles==23.1.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+altair==4.2.2
+analytics-python==1.4.0
+anyio==3.6.1
+asgiref==3.5.2
+async-timeout==4.0.2
+attrs==21.4.0
+backoff==1.10.0
+bcrypt==3.2.2
+blis==0.7.7
+catalogue==2.0.7
+certifi==2022.6.15
+cffi==1.15.0
+charset-normalizer==2.0.12
+click==8.1.3
+confection==0.0.4
+contourpy==1.0.7
+cryptography==37.0.2
+cycler==0.11.0
+cymem==2.0.6
+entrypoints==0.4
+fastapi==0.78.0
+fastcore==1.4.5
+fastdownload==0.0.6
+fastprogress==1.0.2
+ffmpy==0.3.0
+filelock==3.9.1
+fonttools==4.33.3
+frozenlist==1.3.0
+fsspec==2022.5.0
+gradio==3.21.0
+h11==0.13.0
+httpcore==0.16.3
+httpx==0.23.3
+huggingface-hub==0.13.2
+idna==3.3
+Jinja2==3.1.2
+joblib==1.1.0
+jsonschema==4.17.3
+kiwisolver==1.4.3
+langcodes==3.3.0
+linkify-it-py==1.0.3
+git+https://github.com/logasja/lpips-pytorch.git
+markdown-it-py==2.1.0
+MarkupSafe==2.1.1
+matplotlib==3.5.2
+mdit-py-plugins==0.3.0
+mdurl==0.1.1
+monotonic==1.6
+multidict==6.0.2
+murmurhash==1.0.7
+numpy==1.22.4
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+opencv-python==4.6.0.66
+orjson==3.7.2
+packaging==21.3
+pandas==1.4.2
+paramiko==2.11.0
+pathy==0.6.1
+Pillow==9.1.1
+preshed==3.0.6
+pycparser==2.21
+pycryptodome==3.14.1
+pydantic==1.8.2
+pydub==0.25.1
+PyNaCl==1.5.0
+pyparsing==3.0.9
+pyrsistent==0.19.3
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytz==2022.1
+PyYAML==6.0
+requests==2.28.0
+rfc3986==1.5.0
+scikit-learn==1.1.1
+scipy==1.8.1
+six==1.16.0
+smart-open==5.2.1
+sniffio==1.2.0
+spacy==3.3.1
+spacy-legacy==3.0.9
+spacy-loggers==1.0.2
+srsly==2.4.3
+starlette==0.19.1
+thinc==8.0.17
+threadpoolctl==3.1.0
+toolz==0.12.0
+torch==1.11.0
+torchvision==0.12.0
+tqdm==4.64.0
+typer==0.4.1
+typing_extensions==4.2.0
+uc-micro-py==1.0.1
+urllib3==1.26.9
+uvicorn==0.17.6
+wasabi==0.9.1
+websockets==10.4
+yarl==1.7.2

util/__init__.py

@@ -0,0 +1 @@
+

util/attack_utils.py

@@ -0,0 +1,127 @@
+# Helper function for extracting features from pre-trained models
+import os
+import sys
+import math
+import numbers
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+import torch.nn as nn
+from PIL import Image
+import numpy as np
+from util.feature_extraction_utils import  warp_image, normalize_batch
+from util.prepare_utils import get_ensemble, extract_features
+from lpips_pytorch import LPIPS, lpips
+from tqdm import tqdm 
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+tensor_transform = transforms.ToTensor()
+pil_transform = transforms.ToPILImage()
+
+
+
+
+
+
+class Attack(nn.Module):
+
+    def __init__(self, models, dim, attack_type, eps, c_sim=0.5, net_type='alex', lr=0.05,
+                 n_iters=100, noise_size=0.001, n_starts=10, c_tv=None, sigma_gf=None, kernel_size_gf=None,
+                 combination=False, warp=False, theta_warp=None, V_reduction=None):
+        super(Attack, self).__init__()
+        self.extractor_ens = get_ensemble(models, sigma_gf, kernel_size_gf, combination, V_reduction, warp, theta_warp)
+        #print("There are '{}'' models in the attack ensemble".format(len(self.extractor_ens)))
+        self.dim = dim
+        self.eps = eps
+        self.c_sim = c_sim
+        self.net_type = net_type
+        self.lr = lr
+        self.n_iters = n_iters
+        self.noise_size = noise_size
+        self.n_starts = n_starts
+        self.c_tv=None
+        self.attack_type = attack_type
+        self.warp = warp
+        self.theta_warp = theta_warp
+        if self.attack_type == 'lpips':
+            self.lpips_loss = LPIPS(self.net_type).to(device)
+
+
+    def execute(self, images, dir_vec, direction):
+
+        images = Variable(images).to(device)
+        dir_vec = dir_vec.to(device)
+        # take norm wrt dim
+        dir_vec_norm = dir_vec.norm(dim = 2).unsqueeze(2).to(device)
+        dist = torch.zeros(images.shape[0]).to(device)
+        adv_images = images.detach().clone()
+
+        if self.warp:
+            self.face_img = warp_image(images, self.theta_warp)
+
+        for start in range(self.n_starts):
+            # update adversarial images old and distance old
+            adv_images_old = adv_images.detach().clone()
+            dist_old = dist.clone()
+            # add noise to initialize ( - noise_size, noise_size)
+            noise_uniform = Variable(2 * self.noise_size * torch.rand(images.size()) - self.noise_size).to(device)
+            adv_images = Variable(images.detach().clone() + noise_uniform, requires_grad=True).to(device)
+
+            for i in range(self.n_iters):
+                adv_features = extract_features(adv_images, self.extractor_ens, self.dim).to(device)
+                # normalize feature vectors in ensembles 
+                loss = direction*torch.mean((adv_features - dir_vec)**2/dir_vec_norm)
+
+                if self.c_tv != None:
+                    tv_out = self.total_var_reg(images, adv_images)
+                    loss -= self.c_tv * tv_out
+
+                if self.attack_type == 'lpips':
+                    lpips_out = self.lpips_reg(images, adv_images)
+                    loss -= self.c_sim * lpips_out
+
+                grad = torch.autograd.grad(loss, [adv_images])
+                adv_images = adv_images + self.lr * grad[0].sign()
+                perturbation = adv_images - images
+
+                if self.attack_type == 'sgd':
+                    perturbation = torch.clamp(perturbation, min=-self.eps, max=self.eps)
+                    adv_images = images + perturbation
+
+            adv_images = torch.clamp(adv_images, min=0, max=1)
+            adv_features = extract_features(adv_images, self.extractor_ens, self.dim).to(device)
+            dist = torch.mean((adv_features - dir_vec)**2/dir_vec_norm, dim = [1,2])
+
+            if direction ==1:
+                adv_images[dist < dist_old] = adv_images_old[dist < dist_old]
+                dist[dist < dist_old] = dist_old[dist < dist_old]
+            else: 
+                adv_images[dist > dist_old] = adv_images_old[dist > dist_old]
+                dist[dist > dist_old] = dist_old[dist > dist_old]
+
+        return adv_images.detach().cpu()
+
+
+    def lpips_reg(self, images, adv_images):
+        if self.warp:
+            face_adv = warp_image(adv_images, self.theta_warp)
+            lpips_out = self.lpips_loss(normalize_batch(self.face_img).to(device), normalize_batch(face_adv).to(device))[0][0][0][0] /(2*adv_images.shape[0])
+            lpips_out += self.lpips_loss(normalize_batch(images).to(device), normalize_batch(adv_images).to(device))[0][0][0][0] / (2*adv_images.shape[0])
+
+        else:
+            lpips_out = self.lpips_loss(normalize_batch(images).to(device), normalize_batch(adv_images).to(device))[0][0][0][0] / adv_images.shape[0]
+
+        return lpips_out
+
+    def total_var_reg(images, adv_images):
+        perturbation = adv_images - images
+        tv = torch.mean(torch.abs(perturbation[:, :, :, :-1] - perturbation[:, :, :, 1:])) + \
+         torch.mean(torch.abs(perturbation[:, :, :-1, :] - perturbation[:, :, 1:, :]))
+
+        return tv
+
+
+
+

util/feature_extraction_utils.py

@@ -0,0 +1,112 @@
+# Helper function for extracting features from pre-trained models
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import torch.nn as nn
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def warp_image(tensor_img, theta_warp, crop_size=112):
+    # applies affine transform theta to image and crops it 
+
+    theta_warp = torch.Tensor(theta_warp).unsqueeze(0).to(device)
+    grid = F.affine_grid(theta_warp, tensor_img.size())
+    img_warped = F.grid_sample(tensor_img, grid)
+    img_cropped = img_warped[:,:,0:crop_size, 0:crop_size]
+    return(img_cropped)
+
+def normalize_transforms(tfm, W,H):
+    # normalizes affine transform from cv2 for pytorch
+    tfm_t = np.concatenate((tfm, np.array([[0,0,1]])), axis = 0)
+    transforms = np.linalg.inv(tfm_t)[0:2,:]
+    transforms[0,0] = transforms[0,0]
+    transforms[0,1] = transforms[0,1]*H/W
+    transforms[0,2] = transforms[0,2]*2/W + transforms[0,0] + transforms[0,1] - 1
+
+    transforms[1,0] = transforms[1,0]*W/H
+    transforms[1,1] = transforms[1,1]
+    transforms[1,2] = transforms[1,2]*2/H + transforms[1,0] + transforms[1,1] - 1
+
+    return transforms
+
+def l2_norm(input, axis = 1):
+    # normalizes input with respect to second norm
+    norm = torch.norm(input, 2, axis, True)
+    output = torch.div(input, norm)
+    return output
+
+
+def de_preprocess(tensor):
+    # normalize images from [-1,1] to [0,1]
+    return tensor * 0.5 + 0.5
+
+# normalize image to [-1,1]
+normalize = transforms.Compose([
+            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+        ])
+
+def normalize_batch(imgs_tensor):
+    normalized_imgs = torch.empty_like(imgs_tensor)
+    for i, img_ten in enumerate(imgs_tensor):
+        normalized_imgs[i] = normalize(img_ten)
+
+    return normalized_imgs
+
+def resize2d(img, size):
+    # resizes image
+    return (F.adaptive_avg_pool2d(img, size))
+
+
+
+class face_extractor(nn.Module):
+    def __init__(self, crop_size = 112, warp = False, theta_warp = None):
+        super(face_extractor, self).__init__()
+        self.crop_size = crop_size
+        self.warp = warp
+        self.theta_warp = theta_warp
+
+    def forward(self, input):
+
+        if self.warp:
+            assert(input.shape[0] == 1)
+            input = warp_image(input, self.theta_warp, self.crop_size)
+
+
+        return input
+
+
+
+class feature_extractor(nn.Module):
+    def __init__(self, model, crop_size = 112, tta = True, warp = False, theta_warp = None):
+        super(feature_extractor, self).__init__()
+        self.model = model
+        self.crop_size = crop_size
+        self.tta = tta
+        self.warp = warp
+        self.theta_warp = theta_warp
+
+        self.model = model
+
+
+    def forward(self, input):
+
+        if self.warp:
+            assert(input.shape[0] == 1)
+            input = warp_image(input, self.theta_warp, self.crop_size)
+
+
+        batch_normalized = normalize_batch(input)
+        batch_flipped = torch.flip(batch_normalized, [3])
+        # extract features
+        self.model.eval() # set to evaluation mode
+        if self.tta:
+            embed = self.model(batch_normalized) + self.model(batch_flipped)
+            features = l2_norm(embed)
+        else:
+            features = l2_norm(self.model(batch_normalized))
+        return features
+

util/prepare_utils.py

@@ -0,0 +1,254 @@
+# Helper function for extracting features from pre-trained models
+import os
+import sys
+import math
+import numbers
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+import torch.nn as nn
+from PIL import Image
+import numpy as np
+import torchvision.datasets as datasets
+import copy
+import time
+from util.feature_extraction_utils import feature_extractor, face_extractor, warp_image, de_preprocess, normalize_batch
+from lpips_pytorch import LPIPS, lpips
+from backbone.model_irse import IR_50, IR_101, IR_152, IR_SE_50, IR_SE_101, IR_SE_152
+from backbone.model_resnet import ResNet_50, ResNet_101, ResNet_152
+from tqdm import tqdm 
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+tensor_transform = transforms.ToTensor()
+pil_transform = transforms.ToPILImage()
+
+
+
+class ImageFolderWithPaths(datasets.ImageFolder):
+    """Custom dataset that includes image file paths. Extends
+    torchvision.datasets.ImageFolder
+    """
+
+    # override the __getitem__ method. this is the method that dataloader calls
+    def __getitem__(self, index):
+        # this is what ImageFolder normally returns
+        original_tuple = super(ImageFolderWithPaths, self).__getitem__(index)
+        # the image file path
+        path = self.imgs[index][0]
+        # make a new tuple that includes original and the path
+        tuple_with_path = (original_tuple + (path,))
+        return tuple_with_path
+
+
+
+class GaussianSmoothing(nn.Module):
+    """
+    Apply gaussian smoothing on a
+    1d, 2d or 3d tensor. Filtering is performed seperately for each channel
+    in the input using a depthwise convolution.
+    Arguments:
+        channels (int, sequence): Number of channels of the input tensors. Output will
+            have this number of channels as well.
+        kernel_size (int, sequence): Size of the gaussian kernel.
+        sigma (float, sequence): Standard deviation of the gaussian kernel.
+        dim (int, optional): The number of dimensions of the data.
+            Default value is 2 (spatial).
+    """
+
+    def __init__(self, channels, kernel_size, sigma, dim=2):
+        super(GaussianSmoothing, self).__init__()
+        if isinstance(kernel_size, numbers.Number):
+            kernel_size = [kernel_size] * dim
+        if isinstance(sigma, numbers.Number):
+            sigma = [sigma] * dim
+
+        # The gaussian kernel is the product of the
+        # gaussian function of each dimension.
+        kernel = 1
+        meshgrids = torch.meshgrid(
+            [
+                torch.arange(size, dtype=torch.float32)
+                for size in kernel_size
+            ]
+        )
+        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
+            mean = (size - 1) / 2
+            kernel *= 1 / (std * math.sqrt(2 * math.pi)) * \
+                      torch.exp(-((mgrid - mean) / std) ** 2 / 2)
+
+        # Make sure sum of values in gaussian kernel equals 1.
+        kernel = kernel / torch.sum(kernel)
+
+        # Reshape to depthwise convolutional weight
+        kernel = kernel.view(1, 1, *kernel.size())
+        kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
+
+        self.register_buffer('weight', kernel)
+        self.groups = channels
+
+        if dim == 1:
+            self.conv = F.conv1d
+        elif dim == 2:
+            self.conv = F.conv2d
+        elif dim == 3:
+            self.conv = F.conv3d
+        else:
+            raise RuntimeError(
+                'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim)
+            )
+        self.pad_size = int(kernel_size[0] / 2)
+
+    def forward(self, input):
+        """
+        Apply gaussian filter to input.
+        Arguments:
+            input (torch.Tensor): Input to apply gaussian filter on.
+        Returns:
+            filtered (torch.Tensor): Filtered output.
+        """
+        input = F.pad(input, (self.pad_size, self.pad_size, self.pad_size, self.pad_size), mode='reflect')
+        return self.conv(input, weight=self.weight, groups=self.groups)
+
+
+
+
+class dim_reduction(nn.Module):
+    def __init__(self, V):
+        super(dim_reduction, self).__init__()
+        self.V = V
+
+    def forward(self, input):
+        return torch.matmul(input, self.V.to(input.device))
+
+
+
+def get_ensemble(models, sigma_gf, kernel_size_gf, combination, V_reduction, warp=False,
+                                       theta_warp=None):
+    # function prepares ensemble of feature extractors
+    # outputs list of pytorch nn models 
+    feature_extractor_ensemble = []
+    if sigma_gf != None:
+        # if apply gaussian filterng during attack
+        gaussian_filtering = GaussianSmoothing(3, kernel_size_gf, sigma_gf)
+        if V_reduction == None:
+            for model in models:
+                feature_extractor_model = nn.DataParallel(nn.Sequential(gaussian_filtering,
+                                                                        feature_extractor(model=model, warp=warp,
+                                                                                          theta_warp=theta_warp))).to(device)
+                feature_extractor_ensemble.append(feature_extractor_model)
+                if combination:
+                    feature_extractor_model = nn.DataParallel(
+                        feature_extractor(model=model, warp=warp, theta_warp=theta_warp)).to(device)
+                    feature_extractor_ensemble.append(feature_extractor_model)
+
+        else:
+            for i, model in enumerate(models):
+                feature_extractor_model = nn.DataParallel(
+                    nn.Sequential(gaussian_filtering, feature_extractor(model=model, warp=warp, theta_warp=theta_warp),
+                                  dim_reduction(V_reduction[i]))).to(device)
+                feature_extractor_ensemble.append(feature_extractor_model)
+                if combination:
+                    feature_extractor_model = nn.DataParallel(
+                        nn.Sequential(feature_extractor(model=model, warp=warp, theta_warp=theta_warp),
+                                      dim_reduction(V_reduction[i]))).to(device)
+                    feature_extractor_ensemble.append(feature_extractor_model)
+
+    else:
+        if V_reduction == None:
+            for model in models:
+                feature_extractor_model = nn.DataParallel(
+                    feature_extractor(model=model, warp=warp, theta_warp=theta_warp)).to(device)
+                feature_extractor_ensemble.append(feature_extractor_model)
+        else:
+            for i, model in enumerate(models):
+                feature_extractor_model = nn.DataParallel(
+                    nn.Sequential(feature_extractor(model=model, warp=warp, theta_warp=theta_warp),
+                                  dim_reduction(V_reduction[i]))).to(device)
+                feature_extractor_ensemble.append(feature_extractor_model)
+
+    return feature_extractor_ensemble
+
+
+def extract_features(imgs, feature_extractor_ensemble, dim):
+    # function computes mean feature vector of images with ensemble of feature extractors 
+
+    features = torch.zeros(imgs.shape[0], len(feature_extractor_ensemble), dim)
+    for i, feature_extractor_model in enumerate(feature_extractor_ensemble):
+        # batch size, model in ensemble, dim
+        features_model = feature_extractor_model(imgs)
+        features[:, i, :] = features_model
+
+
+    return features
+
+
+
+
+
+def prepare_models(model_backbones,
+                 input_size,
+                 model_roots,
+                 kernel_size_attack,
+                 sigma_attack,
+                 combination,
+                 using_subspace,
+                 V_reduction_root):
+
+    backbone_dict = {'IR_50': IR_50(input_size), 'IR_152': IR_152(input_size), 'ResNet_50': ResNet_50(input_size),
+                     'ResNet_152': ResNet_152(input_size)}
+
+    print("Loading Attack Backbone Checkpoint '{}'".format(model_roots))
+    print('=' * 20)
+
+    models_attack = []
+    for i in range(len(model_backbones)):
+        model = backbone_dict[model_backbones[i]]
+        model.load_state_dict(torch.load(model_roots[i], map_location=device))
+        models_attack.append(model)
+
+    if using_subspace:
+
+        V_reduction = []
+        for i in range(len(model_backbones)):
+            V_reduction.append(torch.tensor(np.load(V_reduction_root[i])))
+
+        dim = V_reduction[0].shape[1]
+    else:
+        V_reduction = None
+        dim = 512
+
+    return models_attack, V_reduction, dim
+
+def prepare_data(query_data_root, target_data_root, freq, batch_size, warp = False, theta_warp = None):
+
+    data = datasets.ImageFolder(query_data_root, tensor_transform)
+
+    subset_query = list(range(0, len(data), freq))
+    subset_gallery = [x for x in list(range(0, len(data))) if x not in subset_query]
+    query_set = torch.utils.data.Subset(data, subset_query)
+    gallery_set = torch.utils.data.Subset(data, subset_gallery)
+
+    if target_data_root != None:
+        target_data =  datasets.ImageFolder(target_data_root, tensor_transform)
+        target_loader = torch.utils.data.DataLoader(
+            target_data, batch_size = batch_size)
+    else:
+        target_loader = None
+
+    query_loader = torch.utils.data.DataLoader(
+        query_set, batch_size = batch_size)
+    gallery_loader = torch.utils.data.DataLoader(
+        gallery_set, batch_size = batch_size)
+
+    return query_loader, gallery_loader, target_loader
+
+
+def prepare_dir_vec(dir_vec_extractor, imgs, dim, combination):
+    dir_vec = extract_features(imgs, dir_vec_extractor, dim).detach().cpu()
+    if combination:
+        dir_vec = torch.repeat_interleave(dir_vec,2,1)
+    return dir_vec 
+
+