Add V1.0

app.py

@@ -0,0 +1,797 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# pip uninstall nvidia_cublas_cu11
+
+
+import sys
+sys.path.append('..')
+import os
+os.system(f'pip install dlib')
+import torch
+import numpy as np
+from PIL import Image
+import models_mae
+from torch.nn import functional as F
+import dlib
+
+import gradio as gr
+
+
+# loading model
+model = getattr(models_mae, 'mae_vit_base_patch16')()
+
+
+class ITEM:
+    def __init__(self, img, parsing_map):
+        self.image = img
+        self.parsing_map = parsing_map
+
+face_to_show = ITEM(None, None)
+
+
+check_region = {'Eyebrows': [2, 3], 
+                'Eyes': [4, 5], 
+                'Nose': [6], 
+                'Mouth': [7, 8, 9], 
+                'Face Boundaries': [10, 1, 0],
+                'Hair': [10],
+                'Skin': [1],
+                'Background': [0]}
+
+
+def get_boundingbox(face, width, height, minsize=None):
+    """
+    Expects a dlib face to generate a quadratic bounding box.
+    :param face: dlib face class
+    :param width: frame width
+    :param height: frame height
+    :param cfg.face_scale: bounding box size multiplier to get a bigger face region
+    :param minsize: set minimum bounding box size
+    :return: x, y, bounding_box_size in opencv form
+    """
+    x1 = face.left()
+    y1 = face.top()
+    x2 = face.right()
+    y2 = face.bottom()
+    size_bb = int(max(x2 - x1, y2 - y1) * 1.3)
+    if minsize:
+        if size_bb < minsize:
+            size_bb = minsize
+    center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
+
+    # Check for out of bounds, x-y top left corner
+    x1 = max(int(center_x - size_bb // 2), 0)
+    y1 = max(int(center_y - size_bb // 2), 0)
+    # Check for too big bb size for given x, y
+    size_bb = min(width - x1, size_bb)
+    size_bb = min(height - y1, size_bb)
+
+    return x1, y1, size_bb
+
+
+def extract_face(frame):
+    face_detector = dlib.get_frontal_face_detector()
+    image = np.array(frame.convert('RGB'))
+    faces = face_detector(image, 1)
+    if len(faces) > 0:
+        # For now only take the biggest face
+        face = faces[0]
+        # Face crop and rescale(follow FF++)
+        x, y, size = get_boundingbox(face, image.shape[1], image.shape[0])
+        # Get the landmarks/parts for the face in box d only with the five key points
+        cropped_face = image[y:y + size, x:x + size]
+        # cropped_face = cv2.resize(cropped_face, (224, 224), interpolation=cv2.INTER_CUBIC)
+        return Image.fromarray(cropped_face)
+    else:
+        return None
+        
+
+from torchvision.transforms import transforms
+def show_one_img_patchify(img, model): 
+    x = torch.tensor(img)
+
+    # make it a batch-like
+    x = x.unsqueeze(dim=0)
+    x = torch.einsum('nhwc->nchw', x)
+    x_patches = model.patchify(x)
+    
+    # visualize the img_patchify
+    n = int(np.sqrt(x_patches.shape[1]))
+    image_size = int(224/n)  
+    padding = 3
+    new_img = Image.new('RGB', (n * image_size + padding*(n-1), n * image_size + padding*(n-1)), 'white')
+    for i, patch in enumerate(x_patches[0]):
+        ax = i % n
+        ay = int(i / n)
+        patch_img_tensor = torch.reshape(patch, (model.patch_embed.patch_size[0], model.patch_embed.patch_size[1], 3))
+        patch_img_tensor = torch.einsum('hwc->chw', patch_img_tensor)
+        patch_img = transforms.ToPILImage()(patch_img_tensor)
+        new_img.paste(patch_img, (ax * image_size + padding * ax, ay * image_size + padding * ay))
+
+    new_img = new_img.resize((224, 224), Image.BICUBIC)
+    return new_img
+
+
+def show_one_img_parchify_mask(img, parsing_map, mask, model):
+    mask = mask.detach()
+    mask_patches = mask.unsqueeze(-1).repeat(1, 1, model.patch_embed.patch_size[0]**2 *3)  # (N, H*W, p*p*3)
+    mask = model.unpatchify(mask_patches)  # 1 is removing, 0 is keeping
+    mask = torch.einsum('nchw->nhwc', mask).detach().cpu()
+
+    # visualize mask
+    vis_mask = mask[0].clone()
+    vis_mask[vis_mask == 1] = 1  # gray for masked
+    vis_mask[vis_mask == 2] = -1  # black for highlight masked facial region
+    vis_mask[vis_mask == 0] = 2  # white for visible
+    vis_mask = torch.clip(vis_mask * 127, 0, 255).int()
+    fasking_mask = vis_mask.numpy().astype(np.uint8)
+    fasking_mask = Image.fromarray(fasking_mask)
+    
+    # visualize the masked image
+    im_masked = img
+    im_masked[mask[0] == 1] = 127
+    im_masked[mask[0] == 2] = 0
+    im_masked = Image.fromarray(im_masked)
+    
+    # visualize the masked image_patchify
+    parsing_map_masked = parsing_map
+    parsing_map_masked[mask[0] == 1] = 127
+    parsing_map_masked[mask[0] == 2] = 0
+    
+    return [show_one_img_patchify(parsing_map_masked, model), fasking_mask, im_masked]
+
+
+# Random
+class CollateFn_Random:
+    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
+        self.img_size = input_size
+        self.patch_size = patch_size
+        self.num_patches_axis = input_size // patch_size
+        self.num_patches = (input_size // patch_size) ** 2
+        self.mask_ratio = mask_ratio
+
+    def __call__(self, image, parsing_map):
+        random_mask = torch.zeros(parsing_map.size(0), self.num_patches, dtype=torch.float32)  # torch.Size([BS, 14, 14])
+        random_mask = self.masking(parsing_map, random_mask)
+
+        return {'image': image, 'random_mask': random_mask}
+
+    def masking(self, parsing_map, random_mask):
+        """
+        :return:
+        """
+        for i in range(random_mask.size(0)):
+            # normalize the masking to strictly target percentage for batch computation.
+            num_mask_to_change = int(self.mask_ratio * self.num_patches)
+            mask_change_to = 1 if num_mask_to_change >= 0 else 0
+            change_indices = torch.randperm(self.num_patches)
+            for idx in range(num_mask_to_change):
+                random_mask[i, change_indices[idx]] = mask_change_to
+
+        return random_mask
+
+
+def do_random_masking(image, parsing_map_vis, ratio):
+    img = torch.from_numpy(image)
+    img = img.unsqueeze(0).permute(0, 3, 1, 2)
+    parsing_map = face_to_show.parsing_map
+    parsing_map = torch.tensor(parsing_map)
+    
+    mask_method = CollateFn_Random(input_size=224, patch_size=16, mask_ratio=ratio)
+    mask = mask_method(img, parsing_map)['random_mask']
+    
+    random_patch_on_parsing, random_mask, random_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask, model)
+    
+    return random_patch_on_parsing, random_mask, random_mask_on_image
+
+ 
+# Fasking
+class CollateFn_Fasking:
+    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
+        self.img_size = input_size
+        self.patch_size = patch_size
+        self.num_patches_axis = input_size // patch_size
+        self.num_patches = (input_size // patch_size) ** 2
+        self.mask_ratio = mask_ratio
+        # --------------------------------------------------------------------------
+        self.facial_region_group = [
+            [2, 4],  # right eye
+            [3, 5],  # left eye
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10],  # hair
+            [1],  # skin
+            [0]  # background
+        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
+
+    def __call__(self, image, parsing_map):
+        # image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
+        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples])  # torch.Size([bs, 1, 224, 224])
+        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])
+
+        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
+        fasking_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)  # torch.Size([BS, 14, 14])
+        fasking_mask = self.fasking(parsing_map, fasking_mask)
+
+        return {'image': image, 'fasking_mask': fasking_mask}
+
+    def fasking(self, parsing_map, fasking_mask):
+        """
+        :return:
+        """
+        for i in range(parsing_map.size(0)):
+            terminate = False
+            for seg_group in self.facial_region_group[:-2]:
+                if terminate:
+                    break
+                for comp_value in seg_group:
+                    fasking_mask[i] = torch.maximum(
+                        fasking_mask[i], F.max_pool2d((parsing_map[i].unsqueeze(0) == comp_value).float(), kernel_size=self.patch_size))
+                    if fasking_mask[i].mean() >= ((self.mask_ratio * self.num_patches) / self.num_patches):
+                        terminate = True
+                        break
+
+        fasking_mask = fasking_mask.view(parsing_map.size(0), -1)
+        for i in range(fasking_mask.size(0)):
+            # normalize the masking to strictly target percentage for batch computation.
+            num_mask_to_change = (self.mask_ratio * self.num_patches - fasking_mask[i].sum(dim=-1)).int()
+            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+            select_indices = (fasking_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
+            change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+            fasking_mask[i, select_indices[change_indices]] = mask_change_to
+
+        return fasking_mask
+
+
+def do_fasking_masking(image, parsing_map_vis, ratio):
+    img = torch.from_numpy(image)
+    img = img.unsqueeze(0).permute(0, 3, 1, 2)
+    parsing_map = face_to_show.parsing_map
+    parsing_map = torch.tensor(parsing_map)
+    
+    mask_method = CollateFn_Fasking(input_size=224, patch_size=16, mask_ratio=ratio)
+    mask = mask_method(img, parsing_map)['fasking_mask']
+    
+    fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask, model)
+    
+    return fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image
+
+
+# FRP
+class CollateFn_FR_P_Masking:
+    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
+        self.img_size = input_size
+        self.patch_size = patch_size
+        self.num_patches_axis = input_size // patch_size
+        self.num_patches = (input_size // patch_size) ** 2
+        self.mask_ratio = mask_ratio
+        self.facial_region_group = [
+            [2, 3],  # eyebrows
+            [4, 5],  # eyes
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10, 1, 0],  # face boundaries
+            [10],  # hair
+            [1],  # facial skin
+            [0]  # background
+        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
+
+    def __call__(self, image, parsing_map):
+        # image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
+        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples])  # torch.Size([bs, 1, 224, 224])
+        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])
+
+        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
+        P_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)  # torch.Size([BS, 14, 14])
+        P_mask = self.random_variable_facial_semantics_masking(parsing_map, P_mask)
+
+        return {'image': image, 'P_mask': P_mask}
+
+    def random_variable_facial_semantics_masking(self, parsing_map, P_mask):
+        """
+        :return:
+        """
+        P_mask = P_mask.view(P_mask.size(0), -1)
+        for i in range(parsing_map.size(0)):
+
+            for seg_group in self.facial_region_group[:-2]:
+                mask_in_seg_group = torch.zeros(1, self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)
+                if seg_group == [10, 1, 0]:
+                    patch_hair_bg = F.max_pool2d(
+                        ((parsing_map[i].unsqueeze(0) == 10) + (parsing_map[i].unsqueeze(0) == 0)).float(),
+                        kernel_size=self.patch_size)
+                    patch_skin = F.max_pool2d((parsing_map[i].unsqueeze(0) == 1).float(), kernel_size=self.patch_size)
+                    # skin&hair or skin&bg defined as facial boundaries：
+                    mask_in_seg_group = torch.maximum(mask_in_seg_group,
+                                                      (patch_hair_bg.bool() & patch_skin.bool()).float())
+                else:
+                    for comp_value in seg_group:
+                        mask_in_seg_group = torch.maximum(mask_in_seg_group,
+                                                          F.max_pool2d(
+                                                              (parsing_map[i].unsqueeze(0) == comp_value).float(),
+                                                              kernel_size=self.patch_size))
+
+                mask_in_seg_group = mask_in_seg_group.view(-1)
+                # to_mask_patches_in_seg_group = mask_in_seg_group - (mask_in_seg_group & P_mask[i])
+                to_mask_patches_in_seg_group = (mask_in_seg_group - P_mask[i]) > 0
+                mask_num = (mask_in_seg_group.sum(dim=-1) * self.mask_ratio -
+                            (mask_in_seg_group.sum(dim=-1)-to_mask_patches_in_seg_group.sum(dim=-1))).int()
+                if mask_num > 0:
+                    select_indices = (to_mask_patches_in_seg_group == 1).nonzero(as_tuple=False).view(-1)
+                    change_indices = torch.randperm(len(select_indices))[:mask_num]
+                    P_mask[i, select_indices[change_indices]] = 1
+
+            num_mask_to_change = (self.mask_ratio * self.num_patches - P_mask[i].sum(dim=-1)).int()
+            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+            select_indices = (P_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
+            change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+            P_mask[i, select_indices[change_indices]] = mask_change_to
+
+        return P_mask
+    
+
+def do_FRP_masking(image, parsing_map_vis, ratio):
+    img = torch.from_numpy(image)
+    img = img.unsqueeze(0).permute(0, 3, 1, 2)
+    parsing_map = face_to_show.parsing_map
+    parsing_map = torch.tensor(parsing_map)
+    
+    mask_method = CollateFn_FR_P_Masking(input_size=224, patch_size=16, mask_ratio=ratio)
+    masks = mask_method(img, parsing_map)
+    mask = masks['P_mask']
+    
+    FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask, model)
+    
+    return FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image
+
+
+# CRFR_R
+class CollateFn_CRFR_R_Masking:
+    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75, region='Nose'):
+        self.img_size = input_size
+        self.patch_size = patch_size
+        self.num_patches_axis = input_size // patch_size
+        self.num_patches = (input_size // patch_size) ** 2
+        self.mask_ratio = mask_ratio
+        self.facial_region_group = [
+            [2, 3],  # eyebrows
+            [4, 5],  # eyes
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10, 1, 0],  # face boundaries
+            [10],  # hair
+            [1],  # facial skin
+            [0]  # background
+        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
+        self.random_specific_facial_region = check_region[region]
+
+    def __call__(self, image, parsing_map):
+        # mage = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
+        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples]) # torch.Size([bs, 1, 224, 224])
+        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])
+
+        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
+        facial_region_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis, dtype=torch.float32)  # torch.Size([1, H/P, W/P])
+        facial_region_mask, random_specific_facial_region = self.masking_all_patches_in_random_specific_facial_region(parsing_map, facial_region_mask)
+        # torch.Size([num_patches,]), list
+
+        CRFR_R_mask, facial_region_mask = self.random_variable_masking(facial_region_mask)
+        # torch.Size([num_patches,]), torch.Size([num_patches,])
+
+        return {'image': image, 'CRFR_R_mask': CRFR_R_mask, 'fr_mask': facial_region_mask}
+
+    def masking_all_patches_in_random_specific_facial_region(self, parsing_map, facial_region_mask):
+        """
+        :param parsing_map: [1, img_size, img_size])
+        :param facial_region_mask: [1, num_patches ** .5, num_patches ** .5]
+        :return: facial_region_mask, random_specific_facial_region
+        """
+        # random_specific_facial_region = random.choice(self.facial_region_group[:-2])
+        # random_specific_facial_region = [6]  # for test: nose
+        if self.random_specific_facial_region == [10, 1, 0]:  # facial boundaries, 10-hair 1-skin 0-background
+            # True for hair(10) or bg(0) patches:
+            patch_hair_bg = F.max_pool2d(((parsing_map == 10) + (parsing_map == 0)).float(),
+                                         kernel_size=self.patch_size)
+            # True for skin(1) patches:
+            patch_skin = F.max_pool2d((parsing_map == 1).float(), kernel_size=self.patch_size)
+            # skin&hair or skin&bg is defined as facial boundaries：
+            facial_region_mask = (patch_hair_bg.bool() & patch_skin.bool()).float()
+        else:
+            for facial_region_index in self.random_specific_facial_region:
+                facial_region_mask = torch.maximum(facial_region_mask,
+                                                   F.max_pool2d((parsing_map == facial_region_index).float(),
+                                                                kernel_size=self.patch_size))
+
+        return facial_region_mask.view(parsing_map.size(0), -1), self.random_specific_facial_region
+
+    def random_variable_masking(self, facial_region_mask):
+        CRFR_R_mask = facial_region_mask.clone()
+
+        for i in range(facial_region_mask.size(0)):
+            num_mask_to_change = (self.mask_ratio * self.num_patches - facial_region_mask[i].sum(dim=-1)).int()
+            mask_change_to = 1 if num_mask_to_change >= 0 else 0
+
+            select_indices = (facial_region_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
+            change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+            CRFR_R_mask[i, select_indices[change_indices]] = mask_change_to
+
+            facial_region_mask[i] = CRFR_R_mask[i] if num_mask_to_change < 0 else facial_region_mask[i]
+
+        return CRFR_R_mask, facial_region_mask
+    
+
+def do_CRFR_R_masking(image, parsing_map_vis, ratio, region):
+    img = torch.from_numpy(image)
+    img = img.unsqueeze(0).permute(0, 3, 1, 2)
+    parsing_map = face_to_show.parsing_map
+    parsing_map = torch.tensor(parsing_map)
+    
+    mask_method = CollateFn_CRFR_R_Masking(input_size=224, patch_size=16, mask_ratio=ratio, region=region)
+    masks = mask_method(img, parsing_map)
+    mask = masks['CRFR_R_mask']
+    fr_mask = masks['fr_mask']
+    
+    CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask+fr_mask, model)
+    
+    return CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image
+
+
+# CRFR_P
+class CollateFn_CRFR_P_Masking:
+    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75, region='Nose'):
+        self.img_size = input_size
+        self.patch_size = patch_size
+        self.num_patches_axis = input_size // patch_size
+        self.num_patches = (input_size // patch_size) ** 2
+        self.mask_ratio = mask_ratio
+
+        self.facial_region_group = [
+            [2, 3],  # eyebrows
+            [4, 5],  # eyes
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10, 1, 0],  # face boundaries
+            [10],  # hair
+            [1],  # facial skin
+            [0]  # background
+        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
+        self.random_specific_facial_region = check_region[region]
+
+    def __call__(self, image, parsing_map):
+        # image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
+        # parsing_map = torch.stack([sample['parsing_map'] for sample in samples]) # torch.Size([bs, 1, 224, 224])
+        # parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])
+
+        # random select a facial semantic region and get corresponding mask(masking all patches include this region)
+        facial_region_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis,
+                                         dtype=torch.float32)  # torch.Size([1, H/P, W/P])
+        facial_region_mask, random_specific_facial_region = self.masking_all_patches_in_random_specific_facial_region(parsing_map, facial_region_mask)
+        # torch.Size([num_patches,]), list
+
+        CRFR_P_mask, facial_region_mask = self.random_variable_masking(parsing_map, facial_region_mask, random_specific_facial_region)
+        # torch.Size([num_patches,]), torch.Size([num_patches,])
+
+        return {'image': image, 'CRFR_P_mask': CRFR_P_mask, 'fr_mask': facial_region_mask}
+
+    def masking_all_patches_in_random_specific_facial_region(self, parsing_map, facial_region_mask):
+        """
+        :param parsing_map: [1, img_size, img_size])
+        :param facial_region_mask: [1, num_patches ** .5, num_patches ** .5]
+        :return: facial_region_mask, random_specific_facial_region
+        """
+        # random_specific_facial_region = random.choice(self.facial_region_group[:-2])
+        # random_specific_facial_region = [4, 5]   # for test: eyes
+        if self.random_specific_facial_region == [10, 1, 0]:  # facial boundaries, 10-hair 1-skin 0-background
+            # True for hair(10) or bg(0) patches:
+            patch_hair_bg = F.max_pool2d(((parsing_map == 10) + (parsing_map == 0)).float(), kernel_size=self.patch_size)
+            # True for skin(1) patches:
+            patch_skin = F.max_pool2d((parsing_map == 1).float(), kernel_size=self.patch_size)
+            # skin&hair or skin&bg is defined as facial boundaries：
+            facial_region_mask = (patch_hair_bg.bool() & patch_skin.bool()).float()
+
+            # # True for hair(10) or skin(1) patches:
+            # patch_hair_face = F.max_pool2d(((parsing_map == 10) + (parsing_map == 1)).float(),
+            #                                kernel_size=self.patch_size)
+            # # True for bg(0) patches:
+            # patch_bg = F.max_pool2d((parsing_map == 0).float(), kernel_size=self.patch_size)
+            # # skin&bg or hair&bg defined as facial boundaries：
+            # facial_region_mask = (patch_hair_face.bool() & patch_bg.bool()).float()
+
+        else:
+            for facial_region_index in self.random_specific_facial_region:
+                facial_region_mask = torch.maximum(facial_region_mask,
+                                                   F.max_pool2d((parsing_map == facial_region_index).float(),
+                                                                kernel_size=self.patch_size))
+
+        return facial_region_mask.view(parsing_map.size(0), -1), self.random_specific_facial_region
+
+    def random_variable_masking(self, parsing_map, facial_region_mask, random_specific_facial_region):
+        CRFR_P_mask = facial_region_mask.clone()
+        other_facial_region_group = [region for region in self.facial_region_group if
+                                     region != random_specific_facial_region]
+        # print(other_facial_region_group)
+        for i in range(facial_region_mask.size(0)):  # iterate each map in BS
+            num_mask_to_change = (self.mask_ratio * self.num_patches - facial_region_mask[i].sum(dim=-1)).int()
+            # mask_change_to = 1 if num_mask_to_change >= 0 else 0
+            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+
+            # masking patches in other facial regions according to the corresponding ratio
+            if mask_change_to == 1:
+                # mask_ratio_other_fr = remain(unmasked) patches should be masked / remain(unmasked) patches
+                mask_ratio_other_fr = (
+                        num_mask_to_change / (self.num_patches - facial_region_mask[i].sum(dim=-1)))
+
+                masked_patches = facial_region_mask[i].clone()
+                for other_fr in other_facial_region_group:
+                    to_mask_patches = torch.zeros(1, self.num_patches_axis, self.num_patches_axis,
+                                                  dtype=torch.float32)
+                    if other_fr == [10, 1, 0]:
+                        patch_hair_bg = F.max_pool2d(
+                            ((parsing_map[i].unsqueeze(0) == 10) + (parsing_map[i].unsqueeze(0) == 0)).float(),
+                            kernel_size=self.patch_size)
+                        patch_skin = F.max_pool2d((parsing_map[i].unsqueeze(0) == 1).float(), kernel_size=self.patch_size)
+                        # skin&hair or skin&bg defined as facial boundaries：
+                        to_mask_patches = (patch_hair_bg.bool() & patch_skin.bool()).float()
+                    else:
+                        for facial_region_index in other_fr:
+                            to_mask_patches = torch.maximum(to_mask_patches,
+                                                            F.max_pool2d((parsing_map[i].unsqueeze(0) == facial_region_index).float(),
+                                                                         kernel_size=self.patch_size))
+
+                    # ignore already masked patches:
+                    to_mask_patches = (to_mask_patches.view(-1) - masked_patches) > 0
+                    # to_mask_patches = to_mask_patches.view(-1) - (to_mask_patches.view(-1) & masked_patches)
+                    select_indices = to_mask_patches.nonzero(as_tuple=False).view(-1)
+                    change_indices = torch.randperm(len(select_indices))[
+                                     :torch.round(to_mask_patches.sum() * mask_ratio_other_fr).int()]
+                    CRFR_P_mask[i, select_indices[change_indices]] = mask_change_to
+                    # prevent overlap
+                    masked_patches = masked_patches + to_mask_patches.float()
+
+                # mask/unmask patch from other facial regions to get CRFR_P_mask with fixed size
+                num_mask_to_change = (self.mask_ratio * self.num_patches - CRFR_P_mask[i].sum(dim=-1)).int()
+                # mask_change_to = 1 if num_mask_to_change >= 0 else 0
+                mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+                # prevent unmasking facial_region_mask
+                select_indices = ((CRFR_P_mask[i] + facial_region_mask[i]) == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
+                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+                CRFR_P_mask[i, select_indices[change_indices]] = mask_change_to
+
+            else:
+                # if the num of facial_region_mask is over (num_patches*mask_ratio),
+                # unmask it to get CRFR_P_mask with fixed size
+                select_indices = (facial_region_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
+                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+                CRFR_P_mask[i, select_indices[change_indices]] = mask_change_to
+                facial_region_mask[i] = CRFR_P_mask[i]
+
+        return CRFR_P_mask, facial_region_mask
+    
+
+def do_CRFR_P_masking(image, parsing_map_vis, ratio, region):
+    img = torch.from_numpy(image)
+    img = img.unsqueeze(0).permute(0, 3, 1, 2)
+    parsing_map = face_to_show.parsing_map
+    parsing_map = torch.tensor(parsing_map)
+    
+    mask_method = CollateFn_CRFR_P_Masking(input_size=224, patch_size=16, mask_ratio=ratio, region=region)
+    masks = mask_method(img, parsing_map)
+    mask = masks['CRFR_P_mask']
+    fr_mask = masks['fr_mask']
+    
+    CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image = show_one_img_parchify_mask(image, parsing_map_vis, mask+fr_mask, model)
+    
+    return CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image
+
+
+def vis_parsing_maps(parsing_anno):
+    part_colors = [[255, 255, 255],
+                   [0, 0, 255], [255, 128, 0], [255, 255, 0],
+                   [0, 255, 0], [0, 255, 128],
+                   [0, 255, 255], [255, 0, 255], [255, 0, 128],
+                   [128, 0, 255], [255, 0, 0]]
+    vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
+    vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255
+
+    num_of_class = np.max(vis_parsing_anno)
+
+    for pi in range(1, num_of_class + 1):
+        index = np.where(vis_parsing_anno == pi)
+        vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]
+
+    vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
+    return vis_parsing_anno_color
+
+
+#from facer import facer
+import facer
+def do_face_parsing(img):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    face_detector = facer.face_detector('retinaface/mobilenet', device=device, threshold=0.3)  # 0.3 for FF++
+    face_parser = facer.face_parser('farl/lapa/448', device=device)  # celebm parser
+    
+    img = extract_face(img)
+    with torch.inference_mode():
+        img = img.resize((224, 224), Image.BICUBIC)
+        image = torch.from_numpy(np.array(img.convert('RGB')))
+        image = image.unsqueeze(0).permute(0, 3, 1, 2).to(device=device)
+        try:
+            faces = face_detector(image)
+            faces = face_parser(image, faces)
+
+            seg_logits = faces['seg']['logits']
+            seg_probs = seg_logits.softmax(dim=1)  # nfaces x nclasses x h x w
+            seg_probs = seg_probs.data  # torch.Size([1, 11, 224, 224])
+            parsing = seg_probs.argmax(1)  # [1, 224, 224]
+
+            parsing_map = parsing.data.cpu().numpy()  # [1, 224, 224] int64
+            parsing_map = parsing_map.astype(np.int8)  # smaller space
+            parsing_map_vis = vis_parsing_maps(parsing_map.squeeze(0))
+            
+        except KeyError:
+            return gr.update()
+    
+    face_to_show.image = img
+    face_to_show.parsing_map = parsing_map
+    return img, parsing_map_vis, show_one_img_patchify(parsing_map_vis, model)
+
+
+# WebUI
+with gr.Blocks() as demo:
+    gr.Markdown("# Demo of Facial Masking Strategies")
+    gr.Markdown(
+        "This is a demo of different facial masking strategies for MIM that are introduced in [FSFM-3C](https://fsfm-3c.github.io/)"
+    )
+    gr.Markdown(
+        "- <b>Random Masking</b>: Random masking all patches."
+    )
+    gr.Markdown(
+        "- <b>Fasking-I</b>: Use a face parser to divide facial regions and priority masking non-skin and non-background regions."
+    )
+    gr.Markdown(
+        "- <b>FRP</b>: Facial Region Proportional masking, which masks an equal portion of patches in each facial region to the overall masking ratio."
+    )
+    gr.Markdown(
+        "- <b>CRFR-R</b>: (1) Covering a Random Facial Region followed by (2) Random masking other patche."
+    )
+    gr.Markdown(
+        "- <b>CRFR-P (_suggested in FSFM-3C_)</b>: (1) Covering a Random Facial Region followed by (2) Proportional masking masking other regions."
+    )
+
+    with gr.Column():
+        image = gr.Image(label="Upload/Capture/Paste a facial image", type="pil")
+
+        image_submit_btn = gr.Button("🖱️ Face Parsing")
+        with gr.Row():  
+            ori_image = gr.Image(interactive=False, label="Detected Face")
+            parsing_map_vis = gr.Image(interactive=False, label="Face Parsing")
+            patch_parsing_map = gr.Image(interactive=False, label="Patchify")
+        gr.HTML('<div class="spacer-20"></div>')
+    
+    with gr.Column():  # Random
+        random_submit_btn = gr.Button("🖱️ Random Masking")
+        ratio_random = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for Random Masking")
+        with gr.Row():
+            random_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
+            random_mask = gr.Image(interactive=False, label="Mask")
+            random_mask_on_image = gr.Image(interactive=False, label="Masked Face")
+        gr.HTML('<div class="spacer-20"></div>')
+
+    with gr.Column():  # Fasking-I
+        fasking_submit_btn = gr.Button("🖱️ Fasking-I")
+        ratio_fasking = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for Fasking")
+        with gr.Row():
+            fasking_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
+            fasking_mask = gr.Image(interactive=False, label="Mask")
+            fasking_mask_on_image = gr.Image(interactive=False, label="Masked Face")
+        gr.HTML('<div class="spacer-20"></div>')
+        
+    with gr.Column():  # FRP
+        FRP_submit_btn = gr.Button("🖱️ FRP")
+        ratio_FRP = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for FRP")
+        with gr.Row():
+            FRP_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
+            FRP_mask = gr.Image(interactive=False, label="Mask")
+            FRP_mask_on_image = gr.Image(interactive=False, label="Masked Face")
+        gr.HTML('<div class="spacer-20"></div>')
+               
+    with gr.Column():  # CRFR-R
+        CRFR_R_submit_btn = gr.Button("🖱️ CRFR-R")
+        ratio_CRFR_R = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for CRFR-R")
+        mask_region_CRFR_R = gr.Radio(choices=['Eyebrows', 'Eyes', 'Nose', 'Mouth', 'Face Boundaries', 'Hair','Skin','Background'],
+                                      value='Eyes',
+                                      label="Facial Region (for CRFR, highlighted by black)")
+        with gr.Row():
+            CRFR_R_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
+            CRFR_R_mask = gr.Image(interactive=False, label="Mask")
+            CRFR_R_mask_on_image = gr.Image(interactive=False, label="Masked Face")
+        gr.HTML('<div class="spacer-20"></div>')
+        
+    with gr.Column():  # CRFR-P
+        CRFR_P_submit_btn = gr.Button("🖱️ CRFR-P (suggested in FSFM-3C)")
+        ratio_CRFR_P = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Masking Ratio for CRFR-P")
+        mask_region_CRFR_P = gr.Radio(choices=['Eyebrows', 'Eyes', 'Nose', 'Mouth', 'Face Boundaries', 'Hair', 'Skin', 'Background'],
+                                      value='Eyes',
+                                      label="Facial Region (for CRFR, highlighted by black)")
+        with gr.Row():
+            CRFR_P_patch_on_parsing = gr.Image(interactive=False, label="Mask/Parsing")
+            CRFR_P_mask = gr.Image(interactive=False, label="Mask")
+            CRFR_P_mask_on_image = gr.Image(interactive=False, label="Masked Face")
+
+    
+    parseing_map = []
+    image_submit_btn.click(
+        fn = do_face_parsing,
+        inputs=image,
+        outputs=[ori_image, parsing_map_vis, patch_parsing_map]
+    )
+    random_submit_btn.click(
+        fn = do_random_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_random],
+        outputs=[random_patch_on_parsing, random_mask, random_mask_on_image],
+    )
+    ratio_random.change(
+        fn = do_random_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_random],
+        outputs=[random_patch_on_parsing, random_mask, random_mask_on_image],
+    )
+    fasking_submit_btn.click(
+        fn = do_fasking_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_fasking],
+        outputs=[fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image],
+    )
+    ratio_fasking.change(
+        fn = do_fasking_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_fasking],
+        outputs=[fasking_patch_on_parsing, fasking_mask, fasking_mask_on_image],
+    )
+    FRP_submit_btn.click(
+        fn = do_FRP_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_FRP],
+        outputs=[FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image],
+    )
+    ratio_FRP.change(
+        fn = do_FRP_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_FRP],
+        outputs=[FRP_patch_on_parsing, FRP_mask, FRP_mask_on_image],
+    )
+    CRFR_R_submit_btn.click(
+        fn = do_CRFR_R_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_CRFR_R, mask_region_CRFR_R],
+        outputs=[CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image],
+    )
+    ratio_CRFR_R.change(
+        fn = do_CRFR_R_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_CRFR_R, mask_region_CRFR_R],
+        outputs=[CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image],
+    )
+    mask_region_CRFR_R.change(
+        fn = do_CRFR_R_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_CRFR_R, mask_region_CRFR_R],
+        outputs=[CRFR_R_patch_on_parsing, CRFR_R_mask, CRFR_R_mask_on_image],
+    )
+    CRFR_P_submit_btn.click(
+        fn = do_CRFR_P_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_CRFR_P, mask_region_CRFR_P],
+        outputs=[CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image],
+    )
+    ratio_CRFR_P.change(
+        fn=do_CRFR_P_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_CRFR_P, mask_region_CRFR_P],
+        outputs=[CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image],
+    )
+    mask_region_CRFR_P.change(
+        fn=do_CRFR_P_masking,
+        inputs=[ori_image, parsing_map_vis, ratio_CRFR_P, mask_region_CRFR_P],
+        outputs=[CRFR_P_patch_on_parsing, CRFR_P_mask, CRFR_P_mask_on_image],
+    )
+
+if __name__ == "__main__":
+    gr.close_all()
+    demo.queue()
+    demo.launch()

facer/__init__.py

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

facer/draw.py

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

facer/face_alignment/__init__.py

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

facer/face_alignment/base.py

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

facer/face_alignment/farl.py

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

facer/face_alignment/network/__init__.py

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

facer/face_alignment/network/common.py

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

facer/face_alignment/network/geometry.py

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

facer/face_alignment/network/mmseg.py

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

facer/face_alignment/network/transformers.py

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

facer/face_attribute/__init__.py

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

facer/face_attribute/base.py

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

facer/face_attribute/farl.py

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

facer/face_detection/__init__.py

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

facer/face_detection/base.py

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

facer/face_detection/retinaface.py

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

facer/face_parsing/__init__.py

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

facer/face_parsing/base.py

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

facer/face_parsing/farl.py

@@ -0,0 +1,92 @@
+from typing import Optional, Dict, Any
+import functools
+import torch
+import torch.nn.functional as F
+
+from ..util import download_jit
+from ..transform import (get_crop_and_resize_matrix, get_face_align_matrix, get_face_align_matrix_celebm,
+                         make_inverted_tanh_warp_grid, make_tanh_warp_grid)
+from .base import FaceParser
+
+pretrain_settings = {
+    'lapa/448': {
+        'url': [
+            'https://github.com/FacePerceiver/facer/releases/download/models-v1/face_parsing.farl.lapa.main_ema_136500_jit191.pt',
+        ],
+        'matrix_src_tag': 'points',
+        'get_matrix_fn': functools.partial(get_face_align_matrix,
+                                           target_shape=(448, 448), target_face_scale=1.0),
+        'get_grid_fn': functools.partial(make_tanh_warp_grid,
+                                         warp_factor=0.8, warped_shape=(448, 448)),
+        'get_inv_grid_fn': functools.partial(make_inverted_tanh_warp_grid,
+                                             warp_factor=0.8, warped_shape=(448, 448)),
+        'label_names': ['background', 'face', 'rb', 'lb', 're',
+                        'le', 'nose',  'ulip', 'imouth', 'llip', 'hair']
+    },
+    'celebm/448': {
+        'url': [
+            'https://github.com/FacePerceiver/facer/releases/download/models-v1/face_parsing.farl.celebm.main_ema_181500_jit.pt',
+        ],
+        'matrix_src_tag': 'points',
+        'get_matrix_fn': functools.partial(get_face_align_matrix_celebm,
+                                           target_shape=(448, 448)),
+        'get_grid_fn': functools.partial(make_tanh_warp_grid,
+                                         warp_factor=0, warped_shape=(448, 448)),
+        'get_inv_grid_fn': functools.partial(make_inverted_tanh_warp_grid,
+                                             warp_factor=0, warped_shape=(448, 448)),
+        'label_names':  [
+                    'background', 'neck', 'face', 'cloth', 'rr', 'lr', 'rb', 'lb', 're',
+                    'le', 'nose', 'imouth', 'llip', 'ulip', 'hair',
+                    'eyeg', 'hat', 'earr', 'neck_l']
+    }
+}
+
+
+class FaRLFaceParser(FaceParser):
+    """ The face parsing models from [FaRL](https://github.com/FacePerceiver/FaRL).
+
+    Please consider citing 
+    ```bibtex
+        @article{zheng2021farl,
+            title={General Facial Representation Learning in a Visual-Linguistic Manner},
+            author={Zheng, Yinglin and Yang, Hao and Zhang, Ting and Bao, Jianmin and Chen, 
+                Dongdong and Huang, Yangyu and Yuan, Lu and Chen, 
+                Dong and Zeng, Ming and Wen, Fang},
+            journal={arXiv preprint arXiv:2112.03109},
+            year={2021}
+        }
+    ```
+    """
+
+    def __init__(self, conf_name: Optional[str] = None,
+                 model_path: Optional[str] = None, device=None) -> None:
+        super().__init__()
+        if conf_name is None:
+            conf_name = 'lapa/448'
+        if model_path is None:
+            model_path = pretrain_settings[conf_name]['url']
+        self.conf_name = conf_name
+        self.net = download_jit(model_path, map_location=device)
+        self.eval()
+
+    def forward(self, images: torch.Tensor, data: Dict[str, Any]):
+        setting = pretrain_settings[self.conf_name]
+        images = images.float() / 255.0
+        _, _, h, w = images.shape
+
+        simages = images[data['image_ids']]
+        matrix = setting['get_matrix_fn'](data[setting['matrix_src_tag']])
+        grid = setting['get_grid_fn'](matrix=matrix, orig_shape=(h, w))
+        inv_grid = setting['get_inv_grid_fn'](matrix=matrix, orig_shape=(h, w))
+
+        w_images = F.grid_sample(
+            simages, grid, mode='bilinear', align_corners=False)
+
+        w_seg_logits, _ = self.net(w_images)  # (b*n) x c x h x w
+
+        seg_logits = F.grid_sample(
+            w_seg_logits, inv_grid, mode='bilinear', align_corners=False)
+
+        data['seg'] = {'logits': seg_logits,
+                       'label_names': setting['label_names']}
+        return data

facer/farl/__init__.py

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

facer/farl/classification.py

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

facer/farl/model.py

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

facer/io.py

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

facer/show.py

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

facer/transform.py

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

facer/util.py

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

facer/version.py

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

models_mae.py

@@ -0,0 +1,251 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from timm.models.vision_transformer import PatchEmbed, Block
+
+from util.pos_embed import get_2d_sincos_pos_embed
+
+
+class MaskedAutoencoderViT(nn.Module):
+    """ Masked Autoencoder with VisionTransformer backbone
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False):
+        super().__init__()
+
+        # --------------------------------------------------------------------------
+        # MAE encoder specifics
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim),
+                                      requires_grad=False)  # fixed sin-cos embedding
+
+        self.blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer) # qk_scale=None
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+        # --------------------------------------------------------------------------
+
+        # --------------------------------------------------------------------------
+        # MAE decoder specifics
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+
+        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim),
+                                              requires_grad=False)  # fixed sin-cos embedding
+
+        self.decoder_blocks = nn.ModuleList([
+            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer) # qk_scale=None
+            for i in range(decoder_depth)])
+
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+        self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size ** 2 * in_chans, bias=True)  # decoder to patch
+        # --------------------------------------------------------------------------
+
+        self.norm_pix_loss = norm_pix_loss
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # initialization
+        # initialize (and freeze) pos_embed by sin-cos embedding
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches ** .5),
+                                            cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1],
+                                                    int(self.patch_embed.num_patches ** .5), cls_token=True)
+        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
+
+        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.patch_embed.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
+        torch.nn.init.normal_(self.cls_token, std=.02)
+        torch.nn.init.normal_(self.mask_token, std=.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def patchify(self, imgs):
+        """
+        imgs: (N, 3, H, W)
+        x: (N, L, patch_size**2 *3)
+        """
+        p = self.patch_embed.patch_size[0]
+        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
+
+        h = w = imgs.shape[2] // p
+        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
+        x = torch.einsum('nchpwq->nhwpqc', x)
+        x = x.reshape(shape=(imgs.shape[0], h * w, p ** 2 * 3))
+        return x
+
+    def unpatchify(self, x):
+        """
+        x: (N, L, patch_size**2 *3)
+        imgs: (N, 3, H, W)
+        """
+        p = self.patch_embed.patch_size[0]
+        h = w = int(x.shape[1] ** .5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
+        return imgs
+
+    def random_masking(self, x, mask_ratio):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, mask, ids_restore
+
+    def forward_encoder(self, x, mask_ratio):
+        # embed patches
+        x = self.patch_embed(x)
+
+        # add pos embed w/o cls token
+        x = x + self.pos_embed[:, 1:, :]
+
+        # masking: length -> length * mask_ratio
+        x, mask, ids_restore = self.random_masking(x, mask_ratio)
+
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x, mask, ids_restore
+
+    def forward_decoder(self, x, ids_restore):
+        # embed tokens
+        x = self.decoder_embed(x)
+
+        # append mask tokens to sequence
+        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
+        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token
+        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
+        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token
+
+        # add pos embed
+        x = x + self.decoder_pos_embed
+
+        # apply Transformer blocks
+        for blk in self.decoder_blocks:
+            x = blk(x)
+        x = self.decoder_norm(x)
+
+        # predictor projection
+        x = self.decoder_pred(x)
+
+        # remove cls token
+        x = x[:, 1:, :]
+
+        return x
+
+    def forward_loss(self, imgs, pred, mask):
+        """
+        imgs: [N, 3, H, W]
+        pred: [N, L, p*p*3]
+        mask: [N, L], 0 is keep, 1 is remove, 
+        """
+        target = self.patchify(imgs)
+        if self.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.e-6) ** .5
+
+        loss = (pred - target) ** 2
+        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
+
+        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
+        return loss
+
+    def forward(self, imgs, mask_ratio=0.75):
+        latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        pred = self.forward_decoder(latent, ids_restore)  # [N, L, p*p*3]
+        loss = self.forward_loss(imgs, pred, mask)
+        return loss, pred, mask
+
+
+def mae_vit_base_patch16_dec512d8b(**kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def mae_vit_large_patch16_dec512d8b(**kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def mae_vit_huge_patch14_dec512d8b(**kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+# set recommended archs
+mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
+mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
+mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b  # decoder: 512 dim, 8 blocks

requirements.txt

@@ -0,0 +1,80 @@
+aiofiles==23.2.1
+annotated-types==0.7.0
+anyio==4.6.2.post1
+certifi==2024.8.30
+charset-normalizer==3.4.0
+click==8.1.7
+cmake==3.31.0.1
+contourpy==1.3.0
+cycler==0.12.1
+# dlib==19.24.6
+exceptiongroup==1.2.2
+fastapi==0.115.4
+ffmpy==0.4.0
+filelock==3.16.1
+fonttools==4.54.1
+fsspec==2024.10.0
+gradio==4.44.1
+gradio_client==1.3.0
+h11==0.14.0
+httpcore==1.0.6
+httpx==0.27.2
+huggingface-hub==0.26.2
+idna==3.10
+imageio==2.36.0
+importlib_resources==6.4.5
+Jinja2==3.1.4
+kiwisolver==1.4.7
+lazy_loader==0.4
+lit==18.1.8
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.2
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.24.4
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+opencv-python==4.10.0.84
+orjson==3.10.11
+packaging==24.2
+pandas==2.2.3
+pillow==10.4.0
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.2.0
+python-dateutil==2.9.0.post0
+python-multipart==0.0.17
+pytz==2024.2
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.9.4
+ruff==0.7.3
+safetensors==0.4.5
+scikit-image==0.24.0
+scipy==1.13.1
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+starlette==0.41.2
+sympy==1.13.1
+tifffile==2024.8.30
+timm==1.0.11
+tomlkit==0.12.0
+torch==2.0.0
+torchvision==0.15.1
+tqdm==4.67.0
+triton==2.0.0
+typer==0.13.0
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+uvicorn==0.32.0
+validators==0.34.0
+websockets==12.0
+zipp==3.21.0

util/crop.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import math
+
+import torch
+
+from torchvision import transforms
+from torchvision.transforms import functional as F
+
+
+class RandomResizedCrop(transforms.RandomResizedCrop):
+    """
+    RandomResizedCrop for matching TF/TPU implementation: no for-loop is used.
+    This may lead to results different with torchvision's version.
+    Following BYOL's TF code:
+    https://github.com/deepmind/deepmind-research/blob/master/byol/utils/dataset.py#L206
+    """
+    @staticmethod
+    def get_params(img, scale, ratio):
+        width, height = F._get_image_size(img)
+        area = height * width
+
+        target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
+        log_ratio = torch.log(torch.tensor(ratio))
+        aspect_ratio = torch.exp(
+            torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
+        ).item()
+
+        w = int(round(math.sqrt(target_area * aspect_ratio)))
+        h = int(round(math.sqrt(target_area / aspect_ratio)))
+
+        w = min(w, width)
+        h = min(h, height)
+
+        i = torch.randint(0, height - h + 1, size=(1,)).item()
+        j = torch.randint(0, width - w + 1, size=(1,)).item()
+
+        return i, j, h, w

util/datasets.py

@@ -0,0 +1,349 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import os
+import json
+import shutil
+
+from torchvision import datasets, transforms
+
+from timm.data import create_transform
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+import numpy as np
+from PIL import Image
+import random
+import torch
+from torch.utils.data import DataLoader, Dataset, ConcatDataset
+from torchvision import transforms
+from torch.nn import functional as F
+
+
+class collate_fn_crfrp:
+    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
+        self.img_size = input_size
+        self.patch_size = patch_size
+        self.num_patches_axis = input_size // patch_size
+        self.num_patches = (input_size // patch_size) ** 2
+        self.mask_ratio = mask_ratio
+
+        # --------------------------------------------------------------------------
+        # self.facial_region_group = [
+        #     [2],  # right eyebrow
+        #     [3],  # left eyebrow
+        #     [4],  # right eye
+        #     [5],  # left eye
+        #     [6],  # nose
+        #     [7, 8],  # upper mouth
+        #     [8, 9],  # lower mouth
+        #     [10, 1, 0],  # facial boundaries
+        #     [10],  # hair
+        #     [1],  # facial skin
+        #     [0]  # background
+        # ]
+        self.facial_region_group = [
+            [2, 3],  # eyebrows
+            [4, 5],  # eyes
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10, 1, 0],  # face boundaries
+            [10],  # hair
+            [1],  # facial skin
+            [0]  # background
+        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
+
+    def __call__(self, samples):
+        image, img_mask, facial_region_mask, random_specific_facial_region \
+            = self.CRFR_P_masking(samples, specified_facial_region=None)
+
+        return {'image': image, 'img_mask': img_mask, 'specific_facial_region_mask': facial_region_mask}
+
+        # # using following code if using different data augmentation for target view
+        # image, img_mask, facial_region_mask, random_specific_facial_region \
+        #     = self.CRFR_P_masking(samples, specified_facial_region=None)
+        # image_cl, img_mask_cl, facial_region_mask_cl, random_specific_facial_region_cl \
+        #     = self.CRFR_P_masking(samples, specified_facial_region=random_specific_facial_region)
+        #
+        # return {'image': image, 'img_mask': img_mask, 'specific_facial_region_mask': facial_region_mask,
+        #         'image_cl': image_cl, 'img_mask_cl': img_mask_cl, 'specific_facial_region_mask_cl': facial_region_mask_cl}
+
+    def CRFR_P_masking(self, samples, specified_facial_region=None):
+        image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
+        parsing_map = torch.stack([sample['parsing_map'] for sample in samples])  # torch.Size([bs, 1, 224, 224])
+        parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])
+
+        # covering a randomly select facial_region_group and get fr_mask(masking all patches include this region)
+        facial_region_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis,
+                                         dtype=torch.float32)  # torch.Size([BS, H/P, W/P])
+        facial_region_mask, random_specific_facial_region \
+            = self.masking_all_patches_in_random_specific_facial_region(parsing_map, facial_region_mask)
+        # torch.Size([num_patches,]), list
+
+        img_mask, facial_region_mask \
+            = self.variable_proportional_masking(parsing_map, facial_region_mask, random_specific_facial_region)
+        # torch.Size([num_patches,]), torch.Size([num_patches,])
+
+        del parsing_map
+        return image, img_mask, facial_region_mask, random_specific_facial_region
+
+    def masking_all_patches_in_random_specific_facial_region(self, parsing_map, facial_region_mask,
+                                                             # specified_facial_region=None
+                                                             ):
+        # while True:
+        #     random_specific_facial_region = random.choice(self.facial_region_group[:-2])
+        #     if random_specific_facial_region != specified_facial_region:
+        #         break
+        random_specific_facial_region = random.choice(self.facial_region_group[:-2])
+        if random_specific_facial_region == [10, 1, 0]:  # facial boundaries, 10-hair 1-skin 0-background
+            # True for hair(10) or bg(0) patches:
+            patch_hair_bg = F.max_pool2d(((parsing_map == 10) + (parsing_map == 0)).float(),
+                                         kernel_size=self.patch_size)
+            # True for skin(1) patches:
+            patch_skin = F.max_pool2d((parsing_map == 1).float(), kernel_size=self.patch_size)
+            # skin&hair or skin&bg is defined as facial boundaries：
+            facial_region_mask = (patch_hair_bg.bool() & patch_skin.bool()).float()
+        else:
+            for facial_region_index in random_specific_facial_region:
+                facial_region_mask = torch.maximum(facial_region_mask,
+                                                   F.max_pool2d((parsing_map == facial_region_index).float(),
+                                                                kernel_size=self.patch_size))
+
+        return facial_region_mask.view(parsing_map.size(0), -1), random_specific_facial_region
+
+    def variable_proportional_masking(self, parsing_map, facial_region_mask, random_specific_facial_region):
+        img_mask = facial_region_mask.clone()
+
+        # proportional masking patches in other regions
+        other_facial_region_group = [region for region in self.facial_region_group if
+                                     region != random_specific_facial_region]
+        # print(other_facial_region_group)
+        for i in range(facial_region_mask.size(0)):  # iterate each map in BS
+            num_mask_to_change = (self.mask_ratio * self.num_patches - facial_region_mask[i].sum(dim=-1)).int()
+            # mask_change_to = 1 if num_mask_to_change >= 0 else 0
+            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+
+            if mask_change_to == 1:
+                # proportional masking patches in other facial regions according to the corresponding ratio
+                mask_ratio_other_fr = (
+                        num_mask_to_change / (self.num_patches - facial_region_mask[i].sum(dim=-1)))
+
+                masked_patches = facial_region_mask[i].clone()
+                for other_fr in other_facial_region_group:
+                    to_mask_patches = torch.zeros(1, self.num_patches_axis, self.num_patches_axis,
+                                                  dtype=torch.float32)
+                    if other_fr == [10, 1, 0]:
+                        patch_hair_bg = F.max_pool2d(
+                            ((parsing_map[i].unsqueeze(0) == 10) + (parsing_map[i].unsqueeze(0) == 0)).float(),
+                            kernel_size=self.patch_size)
+                        patch_skin = F.max_pool2d((parsing_map[i].unsqueeze(0) == 1).float(),
+                                                  kernel_size=self.patch_size)
+                        # skin&hair or skin&bg defined as facial boundaries：
+                        to_mask_patches = (patch_hair_bg.bool() & patch_skin.bool()).float()
+                    else:
+                        for facial_region_index in other_fr:
+                            to_mask_patches = torch.maximum(to_mask_patches,
+                                                            F.max_pool2d((parsing_map[i].unsqueeze(
+                                                                0) == facial_region_index).float(),
+                                                                         kernel_size=self.patch_size))
+
+                    # ignore already masked patches:
+                    to_mask_patches = (to_mask_patches.view(-1) - masked_patches) > 0
+                    select_indices = to_mask_patches.nonzero(as_tuple=False).view(-1)
+                    change_indices = torch.randperm(len(select_indices))[
+                                     :torch.round(to_mask_patches.sum() * mask_ratio_other_fr).int()]
+                    img_mask[i, select_indices[change_indices]] = mask_change_to
+                    # prevent overlap
+                    masked_patches = masked_patches + to_mask_patches.float()
+
+                # mask/unmask patch from other facial regions to get img_mask with fixed size
+                num_mask_to_change = (self.mask_ratio * self.num_patches - img_mask[i].sum(dim=-1)).int()
+                # mask_change_to = 1 if num_mask_to_change >= 0 else 0
+                mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+                # prevent unmasking facial_region_mask
+                select_indices = ((img_mask[i] + facial_region_mask[i]) == (1 - mask_change_to)).nonzero(
+                    as_tuple=False).view(-1)
+                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+                img_mask[i, select_indices[change_indices]] = mask_change_to
+
+            else:
+                # Extreme situations:
+                # if fr_mask is already over(>=) num_patches*mask_ratio, unmask it to get img_mask with fixed ratio
+                select_indices = (facial_region_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
+                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+                img_mask[i, select_indices[change_indices]] = mask_change_to
+                facial_region_mask[i] = img_mask[i]
+
+        return img_mask, facial_region_mask
+
+
+class FaceParsingDataset(Dataset):
+    def __init__(self, root, transform=None):
+        self.root_dir = root
+        self.transform = transform
+        self.image_folder = os.path.join(root, 'images')
+        self.parsing_map_folder = os.path.join(root, 'parsing_maps')
+        self.image_names = os.listdir(self.image_folder)
+
+    def __len__(self):
+        return len(self.image_names)
+
+    def __getitem__(self, idx):
+        img_name = os.path.join(self.image_folder, self.image_names[idx])
+        parsing_map_name = os.path.join(self.parsing_map_folder, self.image_names[idx].replace('.png', '.npy'))
+
+        image = Image.open(img_name).convert("RGB")
+        parsing_map_np = np.load(parsing_map_name)
+
+        if self.transform:
+            image = self.transform(image)
+
+        # Convert mask to tensor
+        parsing_map = torch.from_numpy(parsing_map_np)
+        del parsing_map_np  # may save mem
+
+        return {'image': image, 'parsing_map': parsing_map}
+
+
+class TestImageFolder(datasets.ImageFolder):
+    def __init__(self, root, transform=None, target_transform=None):
+        super(TestImageFolder, self).__init__(root, transform, target_transform)
+
+    def __getitem__(self, index):
+        # Call the parent class method to load image and label
+        original_tuple = super(TestImageFolder, self).__getitem__(index)
+
+        # Get the video name
+        video_name = self.imgs[index][0].split('/')[-1].split('_frame_')[0]  # the separator of video name
+
+        # Extend the tuple to include video name
+        extended_tuple = (original_tuple + (video_name,))
+
+        return extended_tuple
+
+
+def get_mean_std(args):
+    print('dataset_paths:', args.data_path)
+    transform = transforms.Compose([transforms.ToTensor(),
+                                    transforms.Resize((args.input_size, args.input_size),
+                                                      interpolation=transforms.InterpolationMode.BICUBIC)])
+
+    if len(args.data_path) > 1:
+        pretrain_datasets = [FaceParsingDataset(root=path, transform=transform) for path in args.data_path]
+        dataset_pretrain = ConcatDataset(pretrain_datasets)
+    else:
+        pretrain_datasets = args.data_path[0]
+        dataset_pretrain = FaceParsingDataset(root=pretrain_datasets, transform=transform)
+
+    print('Compute mean and variance for pretraining data.')
+    print('len(dataset_train): ', len(dataset_pretrain))
+
+    loader = DataLoader(
+        dataset_pretrain,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+    )
+
+    channels_sum, channels_squared_sum, num_batches = 0, 0, 0
+    for sample in loader:
+        data = sample['image']
+        channels_sum += torch.mean(data, dim=[0, 2, 3])
+        channels_squared_sum += torch.mean(data ** 2, dim=[0, 2, 3])
+        num_batches += 1
+
+    mean = channels_sum / num_batches
+    std = (channels_squared_sum / num_batches - mean ** 2) ** 0.5
+
+    print(f'train dataset mean%: {mean.numpy()} std: %{std.numpy()} ')
+    del pretrain_datasets, dataset_pretrain, loader
+    return mean.numpy(), std.numpy()
+
+
+def build_dataset(is_train, args):
+    transform = build_transform(is_train, args)
+    if args.eval:
+        # no loading training set
+        root = os.path.join(args.data_path, 'test' if is_train else 'test')
+        dataset = TestImageFolder(root, transform=transform)
+    else:
+        root = os.path.join(args.data_path, 'train' if is_train else 'val')
+        dataset = datasets.ImageFolder(root, transform=transform)
+        print(dataset)
+
+    return dataset
+
+
+def build_transform(is_train, args):
+    if args.normalize_from_IMN:
+        mean = IMAGENET_DEFAULT_MEAN
+        std = IMAGENET_DEFAULT_STD
+        # print(f'mean:{mean}, std:{std}')
+    else:
+        if not os.path.exists(os.path.join(args.output_dir, "/pretrain_ds_mean_std.txt")) and not args.eval:
+            shutil.copyfile(os.path.dirname(args.finetune) + '/pretrain_ds_mean_std.txt',
+                            os.path.join(args.output_dir) + '/pretrain_ds_mean_std.txt')
+        with open(os.path.join(os.path.dirname(args.resume)) + '/pretrain_ds_mean_std.txt' if args.eval
+                  else os.path.join(args.output_dir) + '/pretrain_ds_mean_std.txt', 'r') as file:
+            ds_stat = json.loads(file.readline())
+            mean = ds_stat['mean']
+            std = ds_stat['std']
+            # print(f'mean:{mean}, std:{std}')
+
+    if args.apply_simple_augment:
+        if is_train:
+            # this should always dispatch to transforms_imagenet_train
+            transform = create_transform(
+                input_size=args.input_size,
+                is_training=True,
+                color_jitter=args.color_jitter,
+                auto_augment=args.aa,
+                interpolation=transforms.InterpolationMode.BICUBIC,
+                re_prob=args.reprob,
+                re_mode=args.remode,
+                re_count=args.recount,
+                mean=mean,
+                std=std,
+            )
+            return transform
+
+        # no augment / eval transform
+        t = []
+        if args.input_size <= 224:
+            crop_pct = 224 / 256
+        else:
+            crop_pct = 1.0
+        size = int(args.input_size / crop_pct)  # 256
+        t.append(
+            transforms.Resize(size, interpolation=transforms.InterpolationMode.BICUBIC),
+            # to maintain same ratio w.r.t. 224 images
+        )
+        t.append(transforms.CenterCrop(args.input_size))  # 224
+
+        t.append(transforms.ToTensor())
+        t.append(transforms.Normalize(mean, std))
+        return transforms.Compose(t)
+
+    else:
+        t = []
+        if args.input_size < 224:
+            crop_pct = input_size / 224
+        else:
+            crop_pct = 1.0
+        size = int(args.input_size / crop_pct)  # size = 224
+        t.append(
+            transforms.Resize(size, interpolation=transforms.InterpolationMode.BICUBIC),
+            # to maintain same ratio w.r.t. 224 images
+        )
+        # t.append(
+        #     transforms.Resize((224, 224), interpolation=transforms.InterpolationMode.BICUBIC),
+        #     # to maintain same ratio w.r.t. 224 images
+        # )
+
+        t.append(transforms.ToTensor())
+        t.append(transforms.Normalize(mean, std))
+        return transforms.Compose(t)

util/lars.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import torch
+
+
+class LARS(torch.optim.Optimizer):
+    """
+    LARS optimizer, no rate scaling or weight decay for parameters <= 1D.
+    """
+    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, trust_coefficient=0.001):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, trust_coefficient=trust_coefficient)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for g in self.param_groups:
+            for p in g['params']:
+                dp = p.grad
+
+                if dp is None:
+                    continue
+
+                if p.ndim > 1: # if not normalization gamma/beta or bias
+                    dp = dp.add(p, alpha=g['weight_decay'])
+                    param_norm = torch.norm(p)
+                    update_norm = torch.norm(dp)
+                    one = torch.ones_like(param_norm)
+                    q = torch.where(param_norm > 0.,
+                                    torch.where(update_norm > 0,
+                                    (g['trust_coefficient'] * param_norm / update_norm), one),
+                                    one)
+                    dp = dp.mul(q)
+
+                param_state = self.state[p]
+                if 'mu' not in param_state:
+                    param_state['mu'] = torch.zeros_like(p)
+                mu = param_state['mu']
+                mu.mul_(g['momentum']).add_(dp)
+                p.add_(mu, alpha=-g['lr'])

util/loss_contrastive.py

@@ -0,0 +1,360 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+from __future__ import print_function
+
+import torch
+import torch.nn as nn
+import math
+
+
+class SimSiamLoss(nn.Module):
+    def __init__(self):
+        super(SimSiamLoss, self).__init__()
+        self.criterion = nn.CosineSimilarity(dim=1)
+
+    def forward(self, cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+        loss = -(self.criterion(cl_features_1, cl_features_2).mean()) * 0.5
+
+        # if not math.isfinite(loss):
+        #     print(cl_features_1, '\n', cl_features_2)
+        #     print(self.criterion(cl_features_1, cl_features_2))
+
+        return loss
+
+
+class BYOLLoss(nn.Module):
+    def __init__(self):
+        super(BYOLLoss, self).__init__()
+
+    @staticmethod
+    def forward(cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+        loss = 2 - 2 * (cl_features_1 * cl_features_2).sum(dim=-1)
+        # loss = 1 - (cl_features_1 * cl_features_2).sum(dim=-1)
+        loss = loss.mean()
+
+        if not math.isfinite(loss):
+            print(cl_features_1, '\n', cl_features_2)
+            print(2 - 2 * (cl_features_1 * cl_features_2).sum(dim=-1))
+
+        return loss
+
+
+# different implementation of InfoNCELoss, including MOCOV3Loss; SupConLoss
+class InfoNCELoss(nn.Module):
+    def __init__(self, temperature=0.1, contrast_sample='all'):
+        """
+        from CMAE: https://github.com/ZhichengHuang/CMAE/issues/5
+        :param temperature: 0.1 0.5 1.0, 1.5 2.0
+        """
+        super(InfoNCELoss, self).__init__()
+        self.temperature = temperature
+        self.criterion = nn.CrossEntropyLoss()
+        self.contrast_sample = contrast_sample
+
+    def forward(self, cl_features):
+        """
+        Args:
+            :param cl_features: : hidden vector of shape [bsz, n_views, ...]
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if cl_features.is_cuda
+                  else torch.device('cpu'))
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+        score_all = torch.matmul(cl_features_1, cl_features_2.transpose(1, 0))  # [BS, BS]
+        score_all = score_all / self.temperature
+        bs = score_all.size(0)
+
+        if self.contrast_sample == 'all':
+            score = score_all
+        elif self.contrast_sample == 'positive':
+            mask = torch.eye(bs, dtype=torch.float).to(device)  # torch.Size([BS, BS])
+            score = score_all * mask
+        else:
+            raise ValueError('Contrastive sample: all{pos&neg} or positive(positive)')
+
+        # label = (torch.arange(bs, dtype=torch.long) +
+        #          bs * torch.distributed.get_rank()).to(device)
+        label = torch.arange(bs, dtype=torch.long).to(device)
+
+        loss = 2 * self.temperature * self.criterion(score, label)
+
+        if not math.isfinite(loss):
+            print(cl_features_1, '\n', cl_features_2)
+            print(score_all, '\n', score, '\n', mask)
+
+        return loss
+
+
+class MOCOV3Loss(nn.Module):
+    def __init__(self, temperature=0.1):
+        super(MOCOV3Loss, self).__init__()
+        self.temperature = temperature
+
+    def forward(self, cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+
+        # normalize
+        cl_features_1 = nn.functional.normalize(cl_features_1, dim=1)
+        cl_features_2 = nn.functional.normalize(cl_features_2, dim=1)
+        # Einstein sum is more intuitive
+        logits = torch.einsum('nc,mc->nm', [cl_features_1, cl_features_2]) / self.temperature
+        N = logits.shape[0]
+        labels = (torch.arange(N, dtype=torch.long)).cuda()
+        return nn.CrossEntropyLoss()(logits, labels) * (2 * self.temperature)
+
+
+class SupConLoss(nn.Module):
+    """
+    from: https://github.com/HobbitLong/SupContrast
+    Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
+    It also supports the unsupervised contrastive loss in SimCLR"""
+    def __init__(self, temperature=0.1, contrast_mode='all', contrast_sample='all',
+                 base_temperature=0.1):
+        super(SupConLoss, self).__init__()
+        self.temperature = temperature
+        self.contrast_mode = contrast_mode
+        self.contrast_sample = contrast_sample
+        self.base_temperature = base_temperature
+
+    def forward(self, features, labels=None, mask=None):
+        """Compute loss for model. If both `labels` and `mask` are None,
+        it degenerates to SimCLR unsupervised loss:
+        https://arxiv.org/pdf/2002.05709.pdf
+
+        Args:
+            features: hidden vector of shape [bsz, n_views, ...].
+            labels: ground truth of shape [bsz].
+            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
+                has the same class as sample i. Can be asymmetric.
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if features.is_cuda
+                  else torch.device('cpu'))
+
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        batch_size = features.shape[0]
+        if labels is not None and mask is not None:
+            raise ValueError('Cannot define both `labels` and `mask`')
+        elif labels is None and mask is None:
+            mask = torch.eye(batch_size, dtype=torch.float32).to(device)  # torch.Size([BS, BS])
+        elif labels is not None:
+            labels = labels.contiguous().view(-1, 1)
+            if labels.shape[0] != batch_size:
+                raise ValueError('Num of labels does not match num of features')
+            mask = torch.eq(labels, labels.T).float().to(device)
+        else:
+            mask = mask.float().to(device)
+
+        contrast_count = features.shape[1]  # contrast_count(2)
+        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)  # [BS*contrast_count, D]
+        if self.contrast_mode == 'one':
+            anchor_feature = features[:, 0]  # [BS, D]
+            anchor_count = 1
+        elif self.contrast_mode == 'all':
+            anchor_feature = contrast_feature  # [BS*contrast_count, D]
+            anchor_count = contrast_count
+        else:
+            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))
+
+        # compute logits
+        anchor_dot_contrast = torch.div(
+            torch.matmul(anchor_feature, contrast_feature.T),
+            self.temperature)  # [BS*contrast_count, BS*contrast_count]
+        # for numerical stability
+        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)  # [BS*contrast_count, 1]
+        logits = anchor_dot_contrast - logits_max.detach()  # [BS*contrast_count, BS*contrast_count]
+
+        # tile mask
+        mask = mask.repeat(anchor_count, contrast_count)  # [BS*anchor_count, BS*contrast_count]
+        # mask-out self-contrast cases
+        logits_mask = torch.scatter(
+            torch.ones_like(mask),
+            1,
+            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
+            0
+        )  # [BS*anchor_count, BS*contrast_count]
+        mask = mask * logits_mask  # [BS*anchor_count, BS*contrast_count]
+
+        """
+        logits_mask is used to get the denominator(positives and negatives).
+        mask is used to get the numerator(positives). mask is applied to log_prob.
+        """
+
+        # compute log_prob，logits_mask is contrast anchor with both positives and negatives
+        exp_logits = torch.exp(logits) * logits_mask  # [BS*anchor_count, BS*contrast_count]
+        # compute log_prob，logits_mask is contrast anchor with negatives, i.e., denominator only negatives contrast:
+        # exp_logits = torch.exp(logits) * (logits_mask-mask)
+
+        if self.contrast_sample == 'all':
+            log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))  # [BS*anchor_count, BS*anchor_count]
+            # compute mean of log-likelihood over positive
+            mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)  # [BS*anchor_count]
+        elif self.contrast_sample == 'positive':
+            mean_log_prob_pos = (mask * logits).sum(1) / mask.sum(1)
+        else:
+            raise ValueError('Contrastive sample: all{pos&neg} or positive(positive)')
+
+        # loss
+        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+        loss = loss.view(anchor_count, batch_size).mean()
+
+        return loss
+
+
+class InfoNCELossPatchLevel(nn.Module):
+    """
+    test: ref ConMIM: https://github.com/TencentARC/ConMIM.
+    """
+    def __init__(self, temperature=0.1, contrast_sample='all'):
+        """
+        :param temperature: 0.1 0.5 1.0, 1.5 2.0
+        """
+        super(InfoNCELossPatchLevel, self).__init__()
+        self.temperature = temperature
+        self.criterion = nn.CrossEntropyLoss()
+        self.contrast_sample = contrast_sample
+
+        self.facial_region_group = [
+            [2, 3],  # eyebrows
+            [4, 5],  # eyes
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10, 1, 0],  # face boundaries
+            [10],  # hair
+            [1],  # facial skin
+            [0]  # background
+        ]
+
+    def forward(self, cl_features, parsing_map=None):
+        """
+        Args:
+            :param parsing_map:
+            :param cl_features: : hidden vector of shape [bsz, n_views, ...]
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if cl_features.is_cuda
+                  else torch.device('cpu'))
+
+        if len(cl_features.shape) < 4:
+            raise ValueError('`features` needs to be [bsz, n_views, n_cl_patches, ...],'
+                             'at least 4 dimensions are required')
+        if len(cl_features.shape) > 4:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], cl_features.shape[2], -1)
+            # [BS, 2, num_cl_patches, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]
+        cl_features_2 = cl_features[:, 1]
+        score = torch.matmul(cl_features_1, cl_features_2.permute(0, 2, 1))  # [BS, num_cl_patches, num_cl_patches]
+        score = score / self.temperature
+        bs = score.size(0)
+        num_cl_patches = score.size(1)
+
+        if self.contrast_sample == 'all':
+            score = score
+        elif self.contrast_sample == 'positive':
+            mask = torch.eye(num_cl_patches, dtype=torch.float32)  # torch.Size([num_cl_patches, num_cl_patches])
+            mask_batch = mask.unsqueeze(0).expand(bs, -1).to(device)  # [bs, num_cl_patches, num_cl_patches]
+            score = score*mask_batch
+        elif self.contrast_sample == 'region':
+            cl_features_1_fr = []
+            cl_features_2_fr = []
+            for facial_region_index in self.facial_region_group:
+                fr_mask = (parsing_map == facial_region_index).unsqueeze(2).expand(-1, -1, cl_features_1.size(-1))
+                cl_features_1_fr.append((cl_features_1 * fr_mask).mean(dim=1, keepdim=False))
+                cl_features_2_fr.append((cl_features_1 * fr_mask).mean(dim=1, keepdim=False))
+            cl_features_1_fr = torch.stack(cl_features_1_fr, dim=1)
+            cl_features_2_fr = torch.stack(cl_features_2_fr, dim=1)
+            score = torch.matmul(cl_features_1_fr, cl_features_2_fr.permute(0, 2, 1))  # [BS, 8, 8]
+            score = score / self.temperature
+            # mask = torch.eye(cl_features_1_fr.size(1), dtype=torch.bool)
+            # torch.Size([cl_features_1_fr.size(1), cl_features_1_fr.size(1)])
+            # mask_batch = mask.unsqueeze(0).expand(bs, -1).to(device)
+            # [bs, cl_features_1_fr.size(1), cl_features_1_fr.size(1)]
+            # score = score*mask_batch
+            label = torch.arange(cl_features_1_fr.size(1), dtype=torch.long).to(device)
+            labels_batch = label.unsqueeze(0).expand(bs, -1)
+            loss = 2 * self.temperature * self.criterion(score, labels_batch)
+            return loss
+        else:
+            raise ValueError('Contrastive sample: all{pos&neg} or positive(positive)')
+
+        # label = (torch.arange(bs, dtype=torch.long) +
+        #          bs * torch.distributed.get_rank()).to(device)
+        label = torch.arange(num_cl_patches, dtype=torch.long).to(device)
+        labels_batch = label.unsqueeze(0).expand(bs, -1)
+
+        loss = 2 * self.temperature * self.criterion(score, labels_batch)
+
+        return loss
+
+
+class MSELoss(nn.Module):
+    """
+    test: unused
+    """
+    def __init__(self):
+        super(MSELoss, self).__init__()
+
+    @staticmethod
+    def forward(cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, n_patches, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0].float()  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1].float()  # [BS, feat_cl_dim]
+
+        return torch.nn.functional.mse_loss(cl_features_1, cl_features_2, reduction='mean')

util/lr_decay.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import json
+
+
+def param_groups_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
+    """
+    Parameter groups for layer-wise lr decay
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
+    """
+    param_group_names = {}
+    param_groups = {}
+
+    num_layers = len(model.blocks) + 1
+
+    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))
+
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+
+        # no decay: all 1D parameters and model specific ones
+        if p.ndim == 1 or n in no_weight_decay_list:
+            g_decay = "no_decay"
+            this_decay = 0.
+        else:
+            g_decay = "decay"
+            this_decay = weight_decay
+            
+        layer_id = get_layer_id_for_vit(n, num_layers)
+        group_name = "layer_%d_%s" % (layer_id, g_decay)
+
+        if group_name not in param_group_names:
+            this_scale = layer_scales[layer_id]
+
+            param_group_names[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+            param_groups[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+
+        param_group_names[group_name]["params"].append(n)
+        param_groups[group_name]["params"].append(p)
+
+    # print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
+
+    return list(param_groups.values())
+
+
+def get_layer_id_for_vit(name, num_layers):
+    """
+    Assign a parameter with its layer id
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
+    """
+    if name in ['cls_token', 'pos_embed']:
+        return 0
+    elif name.startswith('patch_embed'):
+        return 0
+    elif name.startswith('blocks'):
+        return int(name.split('.')[1]) + 1
+    else:
+        return num_layers

util/lr_sched.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import math
+import numpy as np
+
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < args.warmup_epochs:
+        lr = args.lr * epoch / args.warmup_epochs 
+    else:
+        lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
+            (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr
+    return lr
+
+
+def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0,
+                     start_warmup_value=0, warmup_steps=-1):
+    warmup_schedule = np.array([])
+    warmup_iters = warmup_epochs * niter_per_ep
+    if warmup_steps > 0:
+        warmup_iters = warmup_steps
+    print("Set warmup steps = %d" % warmup_iters)
+    if warmup_epochs > 0:
+        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+    iters = np.arange(epochs * niter_per_ep - warmup_iters)
+    schedule = np.array(
+        [final_value + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters])
+
+    schedule = np.concatenate((warmup_schedule, schedule))
+
+    assert len(schedule) == epochs * niter_per_ep
+    return schedule

util/metrics.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+from sklearn.metrics import roc_auc_score
+from sklearn.metrics import roc_curve
+from sklearn.metrics import auc, accuracy_score, balanced_accuracy_score
+from scipy.optimize import brentq
+from scipy.interpolate import interp1d
+
+
+def frame_level_acc(labels, y_preds):
+    return accuracy_score(labels, y_preds) * 100.
+
+
+def frame_level_balanced_acc(labels, y_preds):
+    return balanced_accuracy_score(labels, y_preds) * 100.
+
+
+def frame_level_auc(labels, preds):
+    return roc_auc_score(labels, preds) * 100.
+
+
+def frame_level_eer(labels, preds):
+    # 推荐；更正确的，MaskRelation（TIFS23也是）
+    fpr, tpr, thresholds = roc_curve(labels, preds, pos_label=1)  # 如果标签不是二进制的，则应显式地给出pos_标签
+    eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
+    # eer_thresh = interp1d(fpr, thresholds)(eer)
+    return eer
+
+
+# def frame_level_eer(labels, preds):
+#     fpr, tpr, thresholds = roc_curve(labels, preds, pos_label=1)
+#     eer_threshold = thresholds[(fpr + (1 - tpr)).argmin()]
+#     fpr_eer = fpr[thresholds == eer_threshold][0]
+#     fnr_eer = 1 - tpr[thresholds == eer_threshold][0]
+#     eer = (fpr_eer + fnr_eer) / 2
+#     metric_logger.meters['eer'].update(eer)
+#     return eer, eer_thresh
+
+
+def get_video_level_label_pred(f_label_list, v_name_list, f_pred_list):
+    """
+    References:
+    CADDM: https://github.com/megvii-research/CADDM
+    """
+    video_res_dict = dict()
+    video_pred_list = list()
+    video_y_pred_list = list()
+    video_label_list = list()
+    # summarize all the results for each video
+    for label, video, score in zip(f_label_list, v_name_list, f_pred_list):
+        if video not in video_res_dict.keys():
+            video_res_dict[video] = {"scores": [score], "label": label}
+        else:
+            video_res_dict[video]["scores"].append(score)
+    # get the score and label for each video
+    for video, res in video_res_dict.items():
+        score = sum(res['scores']) / len(res['scores'])
+        label = res['label']
+        video_pred_list.append(score)
+        video_label_list.append(label)
+        video_y_pred_list.append(score >= 0.5)
+
+    return video_label_list, video_pred_list, video_y_pred_list
+
+
+def video_level_acc(video_label_list, video_y_pred_list):
+    return accuracy_score(video_label_list, video_y_pred_list) * 100.
+
+
+def video_level_balanced_acc(video_label_list, video_y_pred_list):
+    return balanced_accuracy_score(video_label_list, video_y_pred_list) * 100.
+
+
+def video_level_auc(video_label_list, video_pred_list):
+    return roc_auc_score(video_label_list, video_pred_list) * 100.
+
+
+def video_level_eer(video_label_list, video_pred_list):
+    # 推荐；更正确的，MaskRelation（TIFS23也是）
+    fpr, tpr, thresholds = roc_curve(video_label_list, video_pred_list, pos_label=1)  # 如果标签不是二进制的，则应显式地给出pos_标签
+    v_eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
+    # eer_thresh = interp1d(fpr, thresholds)(eer)
+    return v_eer

util/misc.py

@@ -0,0 +1,390 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import os
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch._six import inf
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('max mem: {memory:.0f}')
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print('[{}] '.format(now), end='')  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if args.dist_on_itp:
+        args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+        args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
+        args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+        args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+        os.environ['LOCAL_RANK'] = str(args.gpu)
+        os.environ['RANK'] = str(args.rank)
+        os.environ['WORLD_SIZE'] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}, gpu {}'.format(
+        args.rank, args.dist_url, args.gpu), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]),
+                                norm_type)
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, tag=None):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        if tag is None:
+            checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+        else:
+            checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % tag)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        if tag is None:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name,
+                                  client_state=client_state)
+        else:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % tag,
+                                  client_state=client_state)
+
+
+def save_model_target_encoder(args, epoch, model, model_target_encoder_without_ddp, optimizer, loss_scaler, tag=None):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        if tag is None:
+            checkpoint_paths = [output_dir / ('checkpoint-te-%s.pth' % epoch_name)]
+        else:
+            checkpoint_paths = [output_dir / ('checkpoint-te-%s.pth' % tag)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_target_encoder_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        if tag is None:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-te-%s" % epoch_name,
+                                  client_state=client_state)
+        else:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-te-%s" % tag,
+                                  client_state=client_state)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        print("Resume checkpoint %s" % args.resume)
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
+
+def load_model_target_encoder(args, model_target_encoder_without_ddp, optimizer, loss_scaler):
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume_target_encoder, map_location='cpu')
+        model_target_encoder_without_ddp.load_state_dict(checkpoint['model'])
+        print("Resume checkpoint %s" % args.resume_target_encoder)
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
+
+def all_reduce_mean(x):
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x

util/pos_embed.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# -------------------------------------------------------------
+
+import numpy as np
+import torch
+
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed
+
+
+def interpolate_pos_embed_ema(model, checkpoint_model):
+    if checkpoint_model.pos_embed is not None:
+        pos_embed_checkpoint = checkpoint_model.pos_embed
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model.pos_embed = new_pos_embed