add: PIPNet, arcface

update: gpu or cpu mode

app.py

@@ -15,156 +15,156 @@ from PIL import Image
 import tqdm
 
 from modules.networks.faceshifter import FSGenerator
-# from inference.alignment import norm_crop, norm_crop_with_M, paste_back
-# from inference.utils import save, get_5_from_98, get_detector, get_lmk
-# from inference.PIPNet.lib.tools import get_lmk_model, demo_image
-# from inference.landmark_smooth import kalman_filter_landmark, savgol_filter_landmark
+from inference.alignment import norm_crop, norm_crop_with_M, paste_back
+from inference.utils import save, get_5_from_98, get_detector, get_lmk
+from third_party.PIPNet.lib.tools import get_lmk_model, demo_image
+from inference.landmark_smooth import kalman_filter_landmark, savgol_filter_landmark
 from inference.tricks import Trick
 
-# make_abs_path = lambda fn: os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), fn))
-#
-#
-# fs_model_name = 'faceshifter'
-# in_size = 512
-#
-# mouth_net_param = {
-#     "use": True,
-#     "feature_dim": 128,
-#     "crop_param": (28, 56, 84, 112),
-#     "weight_path": "../../modules/third_party/arcface/weights/mouth_net_28_56_84_112.pth",
-# }
-# trick = Trick()
-#
-# T = transforms.Compose(
-#         [
-#             transforms.ToTensor(),
-#             transforms.Normalize(0.5, 0.5),
-#         ]
-#     )
-# tensor2pil_transform = transforms.ToPILImage()
-#
-#
-# def extract_generator(ckpt: str, pt: str):
-#     print(f'[extract_generator] loading ckpt...')
-#     from trainer.faceshifter.faceshifter_pl import FaceshifterPL512, FaceshifterPL
-#     import yaml
-#     with open(make_abs_path('../../trainer/faceshifter/config.yaml'), 'r') as f:
-#         config = yaml.load(f, Loader=yaml.FullLoader)
-#     config['mouth_net'] = mouth_net_param
-#
-#     if in_size == 256:
-#         net = FaceshifterPL(n_layers=3, num_D=3, config=config)
-#     elif in_size == 512:
-#         net = FaceshifterPL512(n_layers=3, num_D=3, config=config, verbose=False)
-#     else:
-#         raise ValueError('Not supported in_size.')
-#     checkpoint = torch.load(ckpt, map_location="cpu", )
-#     net.load_state_dict(checkpoint["state_dict"], strict=False)
-#     net.eval()
-#
-#     G = net.generator
-#     torch.save(G.state_dict(), pt)
-#     print(f'[extract_generator] extracted from {ckpt}, pth saved to {pt}')
-#
-#
-# ''' load model '''
-# if fs_model_name == 'faceshifter':
-#     # pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t38.pth")
-#     # pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t512_6.pth")
-#     # ckpt_path = "/apdcephfs/share_1290939/gavinyuan/out/triplet512_6/epoch=3-step=128999.ckpt"
-#     pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t512_4.pth")
-#     ckpt_path = "/apdcephfs/share_1290939/gavinyuan/out/triplet512_4/epoch=2-step=185999.ckpt"
-#     if not os.path.exists(pt_path) or 't512' in pt_path:
-#         extract_generator(ckpt_path, pt_path)
-#     fs_model = FSGenerator(
-#         make_abs_path("../../modules/third_party/arcface/weights/ms1mv3_arcface_r100_fp16/backbone.pth"),
-#         mouth_net_param=mouth_net_param,
-#         in_size=in_size,
-#         downup=in_size == 512,
-#     )
-#     fs_model.load_state_dict(torch.load(pt_path, "cpu"), strict=True)
-#     fs_model.eval()
-#
-#     @torch.no_grad()
-#     def infer_batch_to_img(i_s, i_t, post: bool = False):
-#         i_r = fs_model(i_s, i_t)[0]  # x, id_vector, att
-#
-#         if post:
-#             target_hair_mask = trick.get_any_mask(i_t, par=[0, 17])
-#             target_hair_mask = trick.smooth_mask(target_hair_mask)
-#             i_r = target_hair_mask * i_t + (target_hair_mask * (-1) + 1) * i_r
-#             i_r = trick.finetune_mouth(i_s, i_t, i_r) if in_size == 256 else i_r
-#
-#         img_r = trick.tensor_to_arr(i_r)[0]
-#         return img_r
-#
-# elif fs_model_name == 'simswap_triplet' or fs_model_name == 'simswap_vanilla':
-#     from modules.networks.simswap import Generator_Adain_Upsample
-#     sw_model = Generator_Adain_Upsample(
-#         input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=False,
-#         mouth_net_param=mouth_net_param
-#     )
-#     if fs_model_name == 'simswap_triplet':
-#         pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_st5.pth")
-#         ckpt_path = make_abs_path("/apdcephfs/share_1290939/gavinyuan/out/"
-#                                   "simswap_triplet_5/epoch=12-step=782999.ckpt")
-#     elif fs_model_name == 'simswap_vanilla':
-#         pt_path = make_abs_path("../ffplus/extracted_ckpt/G_tmp_sv4_off.pth")
-#         ckpt_path = make_abs_path("/apdcephfs/share_1290939/gavinyuan/out/"
-#                                   "simswap_vanilla_4/epoch=694-step=1487999.ckpt")
-#     else:
-#         pt_path = None
-#         ckpt_path = None
-#     sw_model.load_state_dict(torch.load(pt_path, "cpu"), strict=False)
-#     sw_model.eval()
-#     fs_model = sw_model
-#
-#     from trainer.simswap.simswap_pl import SimSwapPL
-#     import yaml
-#     with open(make_abs_path('../../trainer/simswap/config.yaml'), 'r') as f:
-#         config = yaml.load(f, Loader=yaml.FullLoader)
-#     config['mouth_net'] = mouth_net_param
-#     net = SimSwapPL(config=config, use_official_arc='off' in pt_path)
-#
-#     checkpoint = torch.load(ckpt_path, map_location="cpu")
-#     net.load_state_dict(checkpoint["state_dict"], strict=False)
-#     net.eval()
-#     sw_mouth_net = net.mouth_net  # maybe None
-#     sw_netArc = net.netArc
-#     fs_model = fs_model.cuda()
-#     sw_mouth_net = sw_mouth_net.cuda() if sw_mouth_net is not None else sw_mouth_net
-#     sw_netArc = sw_netArc.cuda()
-#
-#     @torch.no_grad()
-#     def infer_batch_to_img(i_s, i_t, post: bool = False):
-#         i_r = fs_model(source=i_s, target=i_t, net_arc=sw_netArc, mouth_net=sw_mouth_net,)
-#         if post:
-#             target_hair_mask = trick.get_any_mask(i_t, par=[0, 17])
-#             target_hair_mask = trick.smooth_mask(target_hair_mask)
-#             i_r = target_hair_mask * i_t + (target_hair_mask * (-1) + 1) * i_r
-#         i_r = i_r.clamp(-1, 1)
-#         i_r = trick.tensor_to_arr(i_r)[0]
-#         return i_r
-#
-# elif fs_model_name == 'simswap_official':
-#     from simswap.image_infer import SimSwapOfficialImageInfer
-#     fs_model = SimSwapOfficialImageInfer()
-#     pt_path = 'Simswap Official'
-#     mouth_net_param = {
-#         "use": False
-#     }
-#
-#     @torch.no_grad()
-#     def infer_batch_to_img(i_s, i_t):
-#         i_r = fs_model.image_infer(source_tensor=i_s, target_tensor=i_t)
-#         i_r = i_r.clamp(-1, 1)
-#         return i_r
-#
-# else:
-#     raise ValueError('Not supported fs_model_name.')
-#
-#
-# print(f'[demo] model loaded from {pt_path}')
+make_abs_path = lambda fn: os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), fn))
+
+
+fs_model_name = 'faceshifter'
+in_size = 256
+
+mouth_net_param = {
+    "use": True,
+    "feature_dim": 128,
+    "crop_param": (28, 56, 84, 112),
+    "weight_path": "../../modules/third_party/arcface/weights/mouth_net_28_56_84_112.pth",
+}
+trick = Trick()
+
+T = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize(0.5, 0.5),
+        ]
+    )
+tensor2pil_transform = transforms.ToPILImage()
+
+
+def extract_generator(ckpt: str, pt: str):
+    print(f'[extract_generator] loading ckpt...')
+    from trainer.faceshifter.faceshifter_pl import FaceshifterPL512, FaceshifterPL
+    import yaml
+    with open(make_abs_path('../../trainer/faceshifter/config.yaml'), 'r') as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    config['mouth_net'] = mouth_net_param
+
+    if in_size == 256:
+        net = FaceshifterPL(n_layers=3, num_D=3, config=config)
+    elif in_size == 512:
+        net = FaceshifterPL512(n_layers=3, num_D=3, config=config, verbose=False)
+    else:
+        raise ValueError('Not supported in_size.')
+    checkpoint = torch.load(ckpt, map_location="cpu", )
+    net.load_state_dict(checkpoint["state_dict"], strict=False)
+    net.eval()
+
+    G = net.generator
+    torch.save(G.state_dict(), pt)
+    print(f'[extract_generator] extracted from {ckpt}, pth saved to {pt}')
+
+
+''' load model '''
+if fs_model_name == 'faceshifter':
+    pt_path = make_abs_path("./weights/extracted/G_mouth1_t38.pth")
+    # pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t512_6.pth")
+    # ckpt_path = "/apdcephfs/share_1290939/gavinyuan/out/triplet512_6/epoch=3-step=128999.ckpt"
+    # pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_t512_4.pth")
+    # ckpt_path = "/apdcephfs/share_1290939/gavinyuan/out/triplet512_4/epoch=2-step=185999.ckpt"
+    if not os.path.exists(pt_path) or 't512' in pt_path:
+        extract_generator(ckpt_path, pt_path)
+    fs_model = FSGenerator(
+        make_abs_path("./weights/arcface/ms1mv3_arcface_r100_fp16/backbone.pth"),
+        mouth_net_param=mouth_net_param,
+        in_size=in_size,
+        downup=in_size == 512,
+    )
+    fs_model.load_state_dict(torch.load(pt_path, "cpu"), strict=True)
+    fs_model.eval()
+
+    @torch.no_grad()
+    def infer_batch_to_img(i_s, i_t, post: bool = False):
+        i_r = fs_model(i_s, i_t)[0]  # x, id_vector, att
+
+        if post:
+            target_hair_mask = trick.get_any_mask(i_t, par=[0, 17])
+            target_hair_mask = trick.smooth_mask(target_hair_mask)
+            i_r = target_hair_mask * i_t + (target_hair_mask * (-1) + 1) * i_r
+            i_r = trick.finetune_mouth(i_s, i_t, i_r) if in_size == 256 else i_r
+
+        img_r = trick.tensor_to_arr(i_r)[0]
+        return img_r
+
+elif fs_model_name == 'simswap_triplet' or fs_model_name == 'simswap_vanilla':
+    from modules.networks.simswap import Generator_Adain_Upsample
+    sw_model = Generator_Adain_Upsample(
+        input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=False,
+        mouth_net_param=mouth_net_param
+    )
+    if fs_model_name == 'simswap_triplet':
+        pt_path = make_abs_path("../ffplus/extracted_ckpt/G_mouth1_st5.pth")
+        ckpt_path = make_abs_path("/apdcephfs/share_1290939/gavinyuan/out/"
+                                  "simswap_triplet_5/epoch=12-step=782999.ckpt")
+    elif fs_model_name == 'simswap_vanilla':
+        pt_path = make_abs_path("../ffplus/extracted_ckpt/G_tmp_sv4_off.pth")
+        ckpt_path = make_abs_path("/apdcephfs/share_1290939/gavinyuan/out/"
+                                  "simswap_vanilla_4/epoch=694-step=1487999.ckpt")
+    else:
+        pt_path = None
+        ckpt_path = None
+    sw_model.load_state_dict(torch.load(pt_path, "cpu"), strict=False)
+    sw_model.eval()
+    fs_model = sw_model
+
+    from trainer.simswap.simswap_pl import SimSwapPL
+    import yaml
+    with open(make_abs_path('../../trainer/simswap/config.yaml'), 'r') as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    config['mouth_net'] = mouth_net_param
+    net = SimSwapPL(config=config, use_official_arc='off' in pt_path)
+
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+    net.load_state_dict(checkpoint["state_dict"], strict=False)
+    net.eval()
+    sw_mouth_net = net.mouth_net  # maybe None
+    sw_netArc = net.netArc
+    fs_model = fs_model.cuda()
+    sw_mouth_net = sw_mouth_net.cuda() if sw_mouth_net is not None else sw_mouth_net
+    sw_netArc = sw_netArc.cuda()
+
+    @torch.no_grad()
+    def infer_batch_to_img(i_s, i_t, post: bool = False):
+        i_r = fs_model(source=i_s, target=i_t, net_arc=sw_netArc, mouth_net=sw_mouth_net,)
+        if post:
+            target_hair_mask = trick.get_any_mask(i_t, par=[0, 17])
+            target_hair_mask = trick.smooth_mask(target_hair_mask)
+            i_r = target_hair_mask * i_t + (target_hair_mask * (-1) + 1) * i_r
+        i_r = i_r.clamp(-1, 1)
+        i_r = trick.tensor_to_arr(i_r)[0]
+        return i_r
+
+elif fs_model_name == 'simswap_official':
+    from simswap.image_infer import SimSwapOfficialImageInfer
+    fs_model = SimSwapOfficialImageInfer()
+    pt_path = 'Simswap Official'
+    mouth_net_param = {
+        "use": False
+    }
+
+    @torch.no_grad()
+    def infer_batch_to_img(i_s, i_t):
+        i_r = fs_model.image_infer(source_tensor=i_s, target_tensor=i_t)
+        i_r = i_r.clamp(-1, 1)
+        return i_r
+
+else:
+    raise ValueError('Not supported fs_model_name.')
+
+
+print(f'[demo] model loaded from {pt_path}')
 
 
 def swap_image(
@@ -435,6 +435,7 @@ def swap_video_gr(img1, target_path, use_gpu=True, frames=9999999):
 
 
 if __name__ == "__main__":
+    use_gpu = torch.cuda.is_available()
     with gr.Blocks() as demo:
         gr.Markdown("SuperSwap")
 
@@ -458,12 +459,12 @@ if __name__ == "__main__":
                     video_button = gr.Button("换脸")
         image_button.click(
             swap_image_gr,
-            inputs=[image1_input, image2_input, use_post, use_gpen],
+            inputs=[image1_input, image2_input, use_post, use_gpen, use_gpu],
             outputs=image_output,
         )
         video_button.click(
             swap_video_gr,
-            inputs=[image3_input, video_input],
+            inputs=[image3_input, video_input, use_gpu],
             outputs=video_output,
         )
 

inference/alignment.py

@@ -0,0 +1,245 @@
+import cv2
+import numpy as np
+from skimage import transform as trans
+
+
+def get_center(points):
+    x = [p[0] for p in points]
+    y = [p[1] for p in points]
+    centroid = (sum(x) / len(points), sum(y) / len(points))
+    return np.array([centroid])
+
+
+def extract_five_lmk(lmk):
+    x = lmk[..., :2]
+    left_eye = get_center(x[36:42])
+    right_eye = get_center(x[42:48])
+    nose = x[30:31]
+    left_mouth = x[48:49]
+    right_mouth = x[54:55]
+    x = np.concatenate([left_eye, right_eye, nose, left_mouth, right_mouth], axis=0)
+    return x
+
+
+set1 = np.array(
+    [
+        [41.125, 50.75],
+        [71.75, 49.4375],
+        [49.875, 73.0625],
+        [45.9375, 87.9375],
+        [70.4375, 87.9375],
+    ],
+    dtype=np.float32,
+)
+
+arcface_src = np.array(
+    [
+        [38.2946, 51.6963],
+        [73.5318, 51.5014],
+        [56.0252, 71.7366],
+        [41.5493, 92.3655],
+        [70.7299, 92.2041],
+    ],
+    dtype=np.float32,
+)
+
+
+ffhq = np.array(
+    [
+        [192.98138, 239.94708],
+        [318.90277, 240.1936],
+        [256.63416, 314.01935],
+        [201.26117, 371.41043],
+        [313.08905, 371.15118],
+    ],
+    dtype=np.float32,
+)
+
+mtcnn = np.array(
+    [
+        [40.95041, 52.341854],
+        [70.90203, 52.17619],
+        [56.02142, 69.376114],
+        [43.716904, 86.910675],
+        [68.52042, 86.77348],
+    ],
+    dtype=np.float32,
+)
+
+arcface_src = np.expand_dims(arcface_src, axis=0)
+set1 = np.expand_dims(set1, axis=0)
+ffhq = np.expand_dims(ffhq, axis=0)
+mtcnn = np.expand_dims(mtcnn, axis=0)
+
+
+# lmk is prediction; src is template
+def estimate_norm(lmk, image_size=112, mode="set1"):
+    assert lmk.shape == (5, 2)
+    tform = trans.SimilarityTransform()
+    lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
+    min_M = []
+    min_index = []
+    min_error = float("inf")
+    if mode == "arcface":
+        if image_size == 112:
+            src = arcface_src
+        else:
+            src = float(image_size) / 112 * arcface_src
+    elif mode == "set1":
+        if image_size == 112:
+            src = set1
+        else:
+            src = float(image_size) / 112 * set1
+    elif mode == "ffhq":
+        if image_size == 512:
+            src = ffhq
+        else:
+            src = float(image_size) / 512 * ffhq
+    elif mode == "mtcnn":
+        if image_size == 112:
+            src = mtcnn
+        else:
+            src = float(image_size) / 112 * mtcnn
+    else:
+        print("no mode like {}".format(mode))
+        exit()
+    for i in np.arange(src.shape[0]):
+        tform.estimate(lmk, src[i])
+        M = tform.params[0:2, :]
+        results = np.dot(M, lmk_tran.T)
+        results = results.T
+        error = np.sum(np.sqrt(np.sum((results - src[i]) ** 2, axis=1)))
+        #         print(error)
+        if error < min_error:
+            min_error = error
+            min_M = M
+            min_index = i
+    return min_M, min_index
+
+
+def estimate_norm_any(lmk_from, lmk_to, image_size=112):
+    tform = trans.SimilarityTransform()
+    lmk_tran = np.insert(lmk_from, 2, values=np.ones(5), axis=1)
+    min_M = []
+    min_index = []
+    min_error = float("inf")
+    src = lmk_to[np.newaxis, ...]
+    for i in np.arange(src.shape[0]):
+        tform.estimate(lmk_from, src[i])
+        M = tform.params[0:2, :]
+        results = np.dot(M, lmk_tran.T)
+        results = results.T
+        error = np.sum(np.sqrt(np.sum((results - src[i]) ** 2, axis=1)))
+        #         print(error)
+        if error < min_error:
+            min_error = error
+            min_M = M
+            min_index = i
+    return min_M, min_index
+
+
+def norm_crop(img, landmark, image_size=112, mode="arcface", borderValue=0.0):
+    M, pose_index = estimate_norm(landmark, image_size, mode)
+    warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=borderValue)
+    return warped
+
+
+def norm_crop_with_M(img, landmark, image_size=112, mode="arcface", borderValue=0.0):
+    M, pose_index = estimate_norm(landmark, image_size, mode)
+    warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=borderValue)
+    return warped, M
+
+
+def square_crop(im, S):
+    if im.shape[0] > im.shape[1]:
+        height = S
+        width = int(float(im.shape[1]) / im.shape[0] * S)
+        scale = float(S) / im.shape[0]
+    else:
+        width = S
+        height = int(float(im.shape[0]) / im.shape[1] * S)
+        scale = float(S) / im.shape[1]
+    resized_im = cv2.resize(im, (width, height))
+    det_im = np.zeros((S, S, 3), dtype=np.uint8)
+    det_im[: resized_im.shape[0], : resized_im.shape[1], :] = resized_im
+    return det_im, scale
+
+
+def transform(data, center, output_size, scale, rotation):
+    scale_ratio = scale
+    rot = float(rotation) * np.pi / 180.0
+    # translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)
+    t1 = trans.SimilarityTransform(scale=scale_ratio)
+    cx = center[0] * scale_ratio
+    cy = center[1] * scale_ratio
+    t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy))
+    t3 = trans.SimilarityTransform(rotation=rot)
+    t4 = trans.SimilarityTransform(translation=(output_size / 2, output_size / 2))
+    t = t1 + t2 + t3 + t4
+    M = t.params[0:2]
+    cropped = cv2.warpAffine(data, M, (output_size, output_size), borderValue=0.0)
+    return cropped, M
+
+
+def trans_points2d(pts, M):
+    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
+    for i in range(pts.shape[0]):
+        pt = pts[i]
+        new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
+        new_pt = np.dot(M, new_pt)
+        # print('new_pt', new_pt.shape, new_pt)
+        new_pts[i] = new_pt[0:2]
+
+    return new_pts
+
+
+def trans_points3d(pts, M):
+    scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1])
+    # print(scale)
+    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
+    for i in range(pts.shape[0]):
+        pt = pts[i]
+        new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
+        new_pt = np.dot(M, new_pt)
+        # print('new_pt', new_pt.shape, new_pt)
+        new_pts[i][0:2] = new_pt[0:2]
+        new_pts[i][2] = pts[i][2] * scale
+
+    return new_pts
+
+
+def trans_points(pts, M):
+    if pts.shape[1] == 2:
+        return trans_points2d(pts, M)
+    else:
+        return trans_points3d(pts, M)
+
+
+def paste_back(img, mat, ori_img):
+    mat_rev = np.zeros([2, 3])
+    div1 = mat[0][0] * mat[1][1] - mat[0][1] * mat[1][0]
+    mat_rev[0][0] = mat[1][1] / div1
+    mat_rev[0][1] = -mat[0][1] / div1
+    mat_rev[0][2] = -(mat[0][2] * mat[1][1] - mat[0][1] * mat[1][2]) / div1
+    div2 = mat[0][1] * mat[1][0] - mat[0][0] * mat[1][1]
+    mat_rev[1][0] = mat[1][0] / div2
+    mat_rev[1][1] = -mat[0][0] / div2
+    mat_rev[1][2] = -(mat[0][2] * mat[1][0] - mat[0][0] * mat[1][2]) / div2
+
+    img_shape = (ori_img.shape[1], ori_img.shape[0])
+
+    img = cv2.warpAffine(img, mat_rev, img_shape)
+    img_white = np.full((256, 256), 255, dtype=float)
+    img_white = cv2.warpAffine(img_white, mat_rev, img_shape)
+    img_white[img_white > 20] = 255
+    img_mask = img_white
+    kernel = np.ones((40, 40), np.uint8)
+    img_mask = cv2.erode(img_mask, kernel, iterations=2)
+    kernel_size = (20, 20)
+    blur_size = tuple(2 * j + 1 for j in kernel_size)
+    img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
+    img_mask /= 255
+    img_mask = np.reshape(img_mask, [img_mask.shape[0], img_mask.shape[1], 1])
+    ori_img = img_mask * img + (1 - img_mask) * ori_img
+    ori_img = ori_img.astype(np.uint8)
+    return ori_img

inference/landmark_smooth.py

@@ -0,0 +1,117 @@
+import cv2
+import numpy as np
+from scipy.signal import savgol_filter
+
+def kalman_filter(inputs: np.array,
+                  process_noise: float = 0.03,
+                  measure_noise: float = 0.01,
+                  ):
+    """ OpenCV - Kalman Filter
+    https://blog.csdn.net/angelfish91/article/details/61768575
+    https://blog.csdn.net/qq_23981335/article/details/82968422
+    """
+    assert inputs.ndim == 2, "inputs should be 2-dim np.array"
+
+    '''
+    它有3个输入参数，
+    dynam_params：状态空间的维数，这里为2；
+    measure_param：测量值的维数，这里也为2; 
+    control_params：控制向量的维数，默认为0。由于这里该模型中并没有控制变量，因此也为0。
+    '''
+    kalman = cv2.KalmanFilter(2,2)
+
+    kalman.measurementMatrix = np.array([[1,0],[0,1]],np.float32)
+    kalman.transitionMatrix = np.array([[1,0],[0,1]], np.float32)
+    kalman.processNoiseCov = np.array([[1,0],[0,1]], np.float32) * process_noise
+    kalman.measurementNoiseCov = np.array([[1,0],[0,1]], np.float32) * measure_noise
+    '''
+    kalman.measurementNoiseCov为测量系统的协方差矩阵，方差越小，预测结果越接近测量值，
+    kalman.processNoiseCov为模型系统的噪声，噪声越大，预测结果越不稳定，越容易接近模型系统预测值，且单步变化越大，
+    相反，若噪声小，则预测结果与上个计算结果相差不大。
+    '''
+
+    kalman.statePre = np.array([[inputs[0][0]],
+                                [inputs[0][1]]])
+
+    '''
+    Kalman Filtering
+    '''
+    outputs = np.zeros_like(inputs)
+    for i in range(len(inputs)):
+        mes = np.reshape(inputs[i,:],(2,1))
+
+        x = kalman.correct(mes)
+
+        y = kalman.predict()
+        outputs[i] = np.squeeze(y)
+        # print (kalman.statePost[0],kalman.statePost[1])
+        # print (kalman.statePre[0],kalman.statePre[1])
+        # print ('measurement:\t',mes[0],mes[1])
+        # print ('correct:\t',x[0],x[1])
+        # print ('predict:\t',y[0],y[1])
+        # print ('='*30)
+
+    return outputs
+
+
+def kalman_filter_landmark(landmarks: np.array,
+                           process_noise: float = 0.03,
+                           measure_noise: float = 0.01,
+                           ):
+    """ Kalman Filter for Landmarks
+    :param process_noise: large means unstable and close to model predictions
+    :param measure_noise: small means close to measurement
+    """
+    print('[Using Kalman Filter for Landmark Smoothing, process_noise=%f, measure_noise=%f]' %
+          (process_noise, measure_noise))
+
+    ''' 
+    landmarks: (#frames, key, xy)
+    '''
+    assert landmarks.ndim == 3, 'landmarks should be 3-dim np.array'
+    assert landmarks.dtype == 'float32', 'landmarks dtype should be float32'
+
+    for s1 in range(landmarks.shape[1]):
+        landmarks[:, s1] = kalman_filter(landmarks[:, s1],
+                                         process_noise,
+                                         measure_noise)
+    return landmarks
+
+
+def savgol_filter_landmark(landmarks: np.array,
+                           window_length: int = 25,
+                           poly_order: int = 2,
+                           ):
+    """ Savgol Filter for Landmarks
+    https://blog.csdn.net/kaever/article/details/105520941
+    """
+    print('[Using Savgol Filter for Landmark Smoothing, window_length=%d, poly_order=%d]' %
+          (window_length, poly_order))
+
+    ''' 
+    landmarks: (#frames, key, xy)
+    '''
+    assert landmarks.ndim == 3, 'landmarks should be 3-dim np.array'
+    assert landmarks.dtype == 'float32', 'landmarks dtype should be float32'
+    assert window_length % 2 == 1, 'window_length should be odd'
+
+    for s1 in range(landmarks.shape[1]):
+        for s2 in range(landmarks.shape[2]):
+            landmarks[:, s1, s2] = savgol_filter(landmarks[:, s1, s2],
+                                                 window_length,
+                                                 poly_order)
+    return landmarks
+
+if __name__ == '__main__':
+
+    pos = np.array([
+            [10,    50],
+            [12,    49],
+            [11,    52],
+            [13,    52.2],
+            [12.9,  50]], np.float32)
+
+    print(pos)
+    pos_filtered = kalman_filter(pos)
+    print(pos)
+    print(pos_filtered)

inference/tricks.py

@@ -74,7 +74,7 @@ class Trick(object):
         if not use_gpen:
             return img_np
         if self.gpen_model is None:
-            self.gpen_model = GPENImageInfer()
+            self.gpen_model = GPENImageInfer(device=global_device)
         img_np = self.gpen_model.image_infer(img_np)
         return img_np
 
@@ -139,22 +139,22 @@ class SoftErosion(nn.Module):
 
 
 if torch.cuda.is_available():
-    device = torch.device(0)
+    global_device = torch.device(0)
 else:
-    device = torch.device('cpu')
+    global_device = torch.device('cpu')
 vgg_mean = torch.tensor([[[0.485]], [[0.456]], [[0.406]]],
-                             requires_grad=False, device=device)
+                             requires_grad=False, device=global_device)
 vgg_std = torch.tensor([[[0.229]], [[0.224]], [[0.225]]],
-                            requires_grad=False, device=device)
+                            requires_grad=False, device=global_device)
 def load_bisenet():
     bisenet_model = BiSeNet(n_classes=19)
     bisenet_model.load_state_dict(
-        torch.load(make_abs_path("../weights/79999_iter.pth",), map_location="cpu")
+        torch.load(make_abs_path("../weights/bisenet/79999_iter.pth",), map_location="cpu")
     )
     bisenet_model.eval()
-    bisenet_model = bisenet_model.to(device)
+    bisenet_model = bisenet_model.to(global_device)
 
-    smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=7).to(device)
+    smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=7).to(global_device)
     print('[Global] bisenet loaded.')
     return bisenet_model, smooth_mask
 

inference/utils.py

@@ -0,0 +1,133 @@
+import numpy as np
+import cv2
+from PIL import Image
+# from TDDFA_V2.FaceBoxes import FaceBoxes
+# from TDDFA_V2.TDDFA import TDDFA
+
+def get_5_from_98(lmk):
+    lefteye = (lmk[60] + lmk[64] + lmk[96]) / 3  # lmk[96]
+    righteye = (lmk[68] + lmk[72] + lmk[97]) / 3  # lmk[97]
+    nose = lmk[54]
+    leftmouth = lmk[76]
+    rightmouth = lmk[82]
+    return np.array([lefteye, righteye, nose, leftmouth, rightmouth])
+
+
+def get_center(points):
+    x = [p[0] for p in points]
+    y = [p[1] for p in points]
+    centroid = (sum(x) / len(points), sum(y) / len(points))
+    return np.array([centroid])
+
+
+def get_lmk(img, tddfa, face_boxes):
+    # 仅接受一个人的图像
+    boxes = face_boxes(img)
+    n = len(boxes)
+    if n < 1:
+        return None
+    param_lst, roi_box_lst = tddfa(img, boxes)
+    ver_lst = tddfa.recon_vers(param_lst, roi_box_lst, dense_flag=False)
+    x = ver_lst[0].transpose(1, 0)[..., :2]
+    left_eye = get_center(x[36:42])
+    right_eye = get_center(x[42:48])
+    nose = x[30:31]
+    left_mouth = x[48:49]
+    right_mouth = x[54:55]
+    x = np.concatenate([left_eye, right_eye, nose, left_mouth, right_mouth], axis=0)
+    return x
+
+
+def get_landmark_once(img, gpu_mode=False):
+    tddfa = TDDFA(
+        gpu_mode=gpu_mode,
+        arch="resnet",
+        checkpoint_fp="./TDDFA_V2/weights/resnet22.pth",
+        bfm_fp="TDDFA_V2/configs/bfm_noneck_v3.pkl",
+        size=120,
+        num_params=62,
+    )
+    face_boxes = FaceBoxes()
+    boxes = face_boxes(img)
+    n = len(boxes)
+    if n < 1:
+        return None
+    param_lst, roi_box_lst = tddfa(img, boxes)
+    ver_lst = tddfa.recon_vers(param_lst, roi_box_lst, dense_flag=False)
+    x = ver_lst[0].transpose(1, 0)[..., :2]
+    left_eye = get_center(x[36:42])
+    right_eye = get_center(x[42:48])
+    nose = x[30:31]
+    left_mouth = x[48:49]
+    right_mouth = x[54:55]
+    x = np.concatenate([left_eye, right_eye, nose, left_mouth, right_mouth], axis=0)
+    return x
+
+
+def get_detector(gpu_mode=False):
+    tddfa = TDDFA(
+        gpu_mode=gpu_mode,
+        arch="resnet",
+        checkpoint_fp="./TDDFA_V2/weights/resnet22.pth",
+        bfm_fp="TDDFA_V2/configs/bfm_noneck_v3.pkl",
+        size=120,
+        num_params=62,
+    )
+    face_boxes = FaceBoxes()
+    return tddfa, face_boxes
+
+
+def save(x, trick=None, use_post=False):
+    """ Paste img to ori_img """
+    img, mat, ori_img, save_path, img_mask = x
+    if mat is None:
+        print('[Warning] mat is None.')
+        ori_img = ori_img.astype(np.uint8)
+        Image.fromarray(ori_img).save(save_path)
+        return
+
+    H, W = img.shape[0], img.shape[1]  # (256,256) or (512,512)
+    mat_rev = np.zeros([2, 3])
+    div1 = mat[0][0] * mat[1][1] - mat[0][1] * mat[1][0]
+    mat_rev[0][0] = mat[1][1] / div1
+    mat_rev[0][1] = -mat[0][1] / div1
+    mat_rev[0][2] = -(mat[0][2] * mat[1][1] - mat[0][1] * mat[1][2]) / div1
+    div2 = mat[0][1] * mat[1][0] - mat[0][0] * mat[1][1]
+    mat_rev[1][0] = mat[1][0] / div2
+    mat_rev[1][1] = -mat[0][0] / div2
+    mat_rev[1][2] = -(mat[0][2] * mat[1][0] - mat[0][0] * mat[1][2]) / div2
+
+    img_shape = (ori_img.shape[1], ori_img.shape[0])  # (h,w)
+
+    img = cv2.warpAffine(img, mat_rev, img_shape)
+
+    if img_mask is None:
+        ''' hanbang version of paste masks '''
+        img_white = np.full((H, W), 255, dtype=float)
+        img_white = cv2.warpAffine(img_white, mat_rev, img_shape)
+        img_white[img_white > 20] = 255
+        img_mask = img_white
+
+        kernel = np.ones((40, 40), np.uint8)
+        img_mask = cv2.erode(img_mask, kernel, iterations=2)
+
+        kernel_size = (20, 20)
+        blur_size = tuple(2 * j + 1 for j in kernel_size)
+        img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
+        img_mask /= 255
+        img_mask = np.reshape(img_mask, [img_mask.shape[0], img_mask.shape[1], 1])
+    else:
+        ''' yuange version of paste masks '''
+        img_mask = cv2.warpAffine(img_mask, mat_rev, img_shape)
+        img_mask = np.expand_dims(img_mask, axis=-1)
+
+    ori_img = img_mask * img + (1 - img_mask) * ori_img
+    ori_img = ori_img.astype(np.uint8)
+
+    if trick is not None:
+        ori_img = trick.gpen(ori_img, use_post)
+
+    Image.fromarray(ori_img).save(save_path)
+
+    # img_mask = np.array((img_mask * 255), dtype=np.uint8).squeeze()
+    # Image.fromarray(img_mask).save('img_mask.jpg')

third_party/GPEN/infer_image.py

@@ -14,7 +14,7 @@ make_abs_path = lambda fn: os.path.abspath(os.path.join(os.path.dirname(os.path.
 
 
 class GPENImageInfer(object):
-    def __init__(self):
+    def __init__(self, device):
         super(GPENImageInfer, self).__init__()
 
         model = {
@@ -32,6 +32,7 @@ class GPENImageInfer(object):
             model=model["name"],
             channel_multiplier=model["channel_multiplier"],
             narrow=model["narrow"],
+            device=device,
         )
         self.faceenhancer = faceenhancer
 
@@ -77,6 +78,7 @@ class GPENImageInfer(object):
         :return: out_batch: (N,RGB,H,W), in [-1,1]
         """
         B, C, H, W = in_batch.shape
+        device = in_batch.device
 
         in_batch = ((in_batch + 1.) * 127.5).permute(0, 2, 3, 1)
         in_batch = in_batch.cpu().numpy().astype(np.uint8)  # (N,H,W,RGB), in [0,255]
@@ -89,7 +91,7 @@ class GPENImageInfer(object):
             out_batch[b_idx] = out_img[:, :, ::-1]
             if save_batch_idx is not None and b_idx == save_batch_idx:
                 cv2.imwrite(os.path.join(save_folder, save_name), out_img)
-        out_batch = torch.FloatTensor(out_batch).cuda()
+        out_batch = torch.FloatTensor(out_batch).to(device)
         out_batch = out_batch / 127.5 - 1.  # (N,H,W,RGB)
         out_batch = out_batch.permute(0, 3, 1, 2)  # (N,RGB,H,W)
         out_batch = out_batch.clamp(-1, 1)

third_party/PIPNet/FaceBoxesV2/detector.py

@@ -0,0 +1,39 @@
+import cv2
+
+class Detector(object):
+    def __init__(self, model_arch, model_weights):
+        self.model_arch = model_arch
+        self.model_weights = model_weights
+
+    def detect(self, image, thresh):
+        raise NotImplementedError
+
+    def crop(self, image, detections):
+        crops = []
+        for det in detections:
+            xmin = max(det[2], 0)
+            ymin = max(det[3], 0)
+            width = det[4]
+            height = det[5]
+            xmax = min(xmin+width, image.shape[1])
+            ymax = min(ymin+height, image.shape[0])
+            cut = image[ymin:ymax, xmin:xmax,:]
+            crops.append(cut)
+
+        return crops
+
+    def draw(self, image, detections, im_scale=None):
+        if im_scale is not None:
+            image = cv2.resize(image, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR)
+            detections = [[det[0],det[1],int(det[2]*im_scale),int(det[3]*im_scale),int(det[4]*im_scale),int(det[5]*im_scale)] for det in detections]
+
+        for det in detections:
+            xmin = det[2]
+            ymin = det[3]
+            width = det[4]
+            height = det[5]
+            xmax = xmin + width
+            ymax = ymin + height
+            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
+
+        return image

third_party/PIPNet/FaceBoxesV2/faceboxes_detector.py

@@ -0,0 +1,124 @@
+from third_party.PIPNet.FaceBoxesV2.detector import Detector
+import cv2, os
+import numpy as np
+import torch
+import torch.nn as nn
+from third_party.PIPNet.FaceBoxesV2.utils.config import cfg
+from third_party.PIPNet.FaceBoxesV2.utils.prior_box import PriorBox
+from third_party.PIPNet.FaceBoxesV2.utils.nms_wrapper import nms
+from third_party.PIPNet.FaceBoxesV2.utils.faceboxes import FaceBoxesV2
+from third_party.PIPNet.FaceBoxesV2.utils.box_utils import decode
+import time
+
+
+class FaceBoxesDetector(Detector):
+    def __init__(self, model_arch, model_weights, use_gpu, device):
+        super().__init__(model_arch, model_weights)
+        self.name = "FaceBoxesDetector"
+        self.net = FaceBoxesV2(
+            phase="test", size=None, num_classes=2
+        )  # initialize detector
+        self.use_gpu = use_gpu
+        self.device = device
+
+        state_dict = torch.load(self.model_weights, map_location=self.device)
+        # create new OrderedDict that does not contain `module.`
+        from collections import OrderedDict
+
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            name = k[7:]  # remove `module.`
+            new_state_dict[name] = v
+        # load params
+        self.net.load_state_dict(new_state_dict)
+        self.net = self.net.to(self.device)
+        self.net.eval()
+
+    def detect(self, image, thresh=0.6, im_scale=None):
+        # auto resize for large images
+        if im_scale is None:
+            height, width, _ = image.shape
+            if min(height, width) > 600:
+                im_scale = 600.0 / min(height, width)
+            else:
+                im_scale = 1
+        image_scale = cv2.resize(
+            image, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR
+        )
+
+        scale = torch.Tensor(
+            [
+                image_scale.shape[1],
+                image_scale.shape[0],
+                image_scale.shape[1],
+                image_scale.shape[0],
+            ]
+        )
+        image_scale = (
+            torch.from_numpy(image_scale.transpose(2, 0, 1)).to(self.device).int()
+        )
+        mean_tmp = torch.IntTensor([104, 117, 123]).to(self.device)
+        mean_tmp = mean_tmp.unsqueeze(1).unsqueeze(2)
+        image_scale -= mean_tmp
+        image_scale = image_scale.float().unsqueeze(0)
+        scale = scale.to(self.device)
+
+        with torch.no_grad():
+            out = self.net(image_scale)
+            # priorbox = PriorBox(cfg, out[2], (image_scale.size()[2], image_scale.size()[3]), phase='test')
+            priorbox = PriorBox(
+                cfg, image_size=(image_scale.size()[2], image_scale.size()[3])
+            )
+            priors = priorbox.forward()
+            priors = priors.to(self.device)
+            loc, conf = out
+            prior_data = priors.data
+            boxes = decode(loc.data.squeeze(0), prior_data, cfg["variance"])
+            boxes = boxes * scale
+            boxes = boxes.cpu().numpy()
+            scores = conf.data.cpu().numpy()[:, 1]
+
+            # ignore low scores
+            inds = np.where(scores > thresh)[0]
+            boxes = boxes[inds]
+            scores = scores[inds]
+
+            # keep top-K before NMS
+            order = scores.argsort()[::-1][:5000]
+            boxes = boxes[order]
+            scores = scores[order]
+
+            # do NMS
+            dets = np.hstack((boxes, scores[:, np.newaxis])).astype(
+                np.float32, copy=False
+            )
+            keep = nms(dets, 0.3)
+            dets = dets[keep, :]
+
+            dets = dets[:750, :]
+            detections_scale = []
+            for i in range(dets.shape[0]):
+                xmin = int(dets[i][0])
+                ymin = int(dets[i][1])
+                xmax = int(dets[i][2])
+                ymax = int(dets[i][3])
+                score = dets[i][4]
+                width = xmax - xmin
+                height = ymax - ymin
+                detections_scale.append(["face", score, xmin, ymin, width, height])
+
+        # adapt bboxes to the original image size
+        if len(detections_scale) > 0:
+            detections_scale = [
+                [
+                    det[0],
+                    det[1],
+                    int(det[2] / im_scale),
+                    int(det[3] / im_scale),
+                    int(det[4] / im_scale),
+                    int(det[5] / im_scale),
+                ]
+                for det in detections_scale
+            ]
+
+        return detections_scale, im_scale

third_party/PIPNet/FaceBoxesV2/utils/box_utils.py

@@ -0,0 +1,276 @@
+import torch
+import numpy as np
+
+
+def point_form(boxes):
+    """ Convert prior_boxes to (xmin, ymin, xmax, ymax)
+    representation for comparison to point form ground truth data.
+    Args:
+        boxes: (tensor) center-size default boxes from priorbox layers.
+    Return:
+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
+    """
+    return torch.cat((boxes[:, :2] - boxes[:, 2:]/2,     # xmin, ymin
+                     boxes[:, :2] + boxes[:, 2:]/2), 1)  # xmax, ymax
+
+
+def center_size(boxes):
+    """ Convert prior_boxes to (cx, cy, w, h)
+    representation for comparison to center-size form ground truth data.
+    Args:
+        boxes: (tensor) point_form boxes
+    Return:
+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
+    """
+    return torch.cat((boxes[:, 2:] + boxes[:, :2])/2,  # cx, cy
+                     boxes[:, 2:] - boxes[:, :2], 1)  # w, h
+
+
+def intersect(box_a, box_b):
+    """ We resize both tensors to [A,B,2] without new malloc:
+    [A,2] -> [A,1,2] -> [A,B,2]
+    [B,2] -> [1,B,2] -> [A,B,2]
+    Then we compute the area of intersect between box_a and box_b.
+    Args:
+      box_a: (tensor) bounding boxes, Shape: [A,4].
+      box_b: (tensor) bounding boxes, Shape: [B,4].
+    Return:
+      (tensor) intersection area, Shape: [A,B].
+    """
+    A = box_a.size(0)
+    B = box_b.size(0)
+    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
+    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
+    inter = torch.clamp((max_xy - min_xy), min=0)
+    return inter[:, :, 0] * inter[:, :, 1]
+
+
+def jaccard(box_a, box_b):
+    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
+    is simply the intersection over union of two boxes.  Here we operate on
+    ground truth boxes and default boxes.
+    E.g.:
+        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
+    Args:
+        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
+        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
+    Return:
+        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
+    """
+    inter = intersect(box_a, box_b)
+    area_a = ((box_a[:, 2]-box_a[:, 0]) *
+              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
+    area_b = ((box_b[:, 2]-box_b[:, 0]) *
+              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
+    union = area_a + area_b - inter
+    return inter / union  # [A,B]
+
+
+def matrix_iou(a, b):
+    """
+    return iou of a and b, numpy version for data augenmentation
+    """
+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
+
+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
+    area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
+    return area_i / (area_a[:, np.newaxis] + area_b - area_i)
+
+
+def matrix_iof(a, b):
+    """
+    return iof of a and b, numpy version for data augenmentation
+    """
+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
+
+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
+    return area_i / np.maximum(area_a[:, np.newaxis], 1)
+
+
+def match(threshold, truths, priors, variances, labels, loc_t, conf_t, idx):
+    """Match each prior box with the ground truth box of the highest jaccard
+    overlap, encode the bounding boxes, then return the matched indices
+    corresponding to both confidence and location preds.
+    Args:
+        threshold: (float) The overlap threshold used when mathing boxes.
+        truths: (tensor) Ground truth boxes, Shape: [num_obj, num_priors].
+        priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
+        variances: (tensor) Variances corresponding to each prior coord,
+            Shape: [num_priors, 4].
+        labels: (tensor) All the class labels for the image, Shape: [num_obj].
+        loc_t: (tensor) Tensor to be filled w/ endcoded location targets.
+        conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
+        idx: (int) current batch index
+    Return:
+        The matched indices corresponding to 1)location and 2)confidence preds.
+    """
+    # jaccard index
+    overlaps = jaccard(
+        truths,
+        point_form(priors)
+    )
+    # (Bipartite Matching)
+    # [1,num_objects] best prior for each ground truth
+    best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
+
+    # ignore hard gt
+    valid_gt_idx = best_prior_overlap[:, 0] >= 0.2
+    best_prior_idx_filter = best_prior_idx[valid_gt_idx, :]
+    if best_prior_idx_filter.shape[0] <= 0:
+        loc_t[idx] = 0
+        conf_t[idx] = 0
+        return
+
+    # [1,num_priors] best ground truth for each prior
+    best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
+    best_truth_idx.squeeze_(0)
+    best_truth_overlap.squeeze_(0)
+    best_prior_idx.squeeze_(1)
+    best_prior_idx_filter.squeeze_(1)
+    best_prior_overlap.squeeze_(1)
+    best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2)  # ensure best prior
+    # NoTODO refactor: index  best_prior_idx with long tensor
+    # ensure every gt matches with its prior of max overlap
+    for j in range(best_prior_idx.size(0)):
+        best_truth_idx[best_prior_idx[j]] = j
+    matches = truths[best_truth_idx]          # Shape: [num_priors,4]
+    conf = labels[best_truth_idx]          # Shape: [num_priors]
+    conf[best_truth_overlap < threshold] = 0  # label as background
+    loc = encode(matches, priors, variances)
+    loc_t[idx] = loc    # [num_priors,4] encoded offsets to learn
+    conf_t[idx] = conf  # [num_priors] top class label for each prior
+
+
+def encode(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth boxes
+    we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 4].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded boxes (tensor), Shape: [num_priors, 4]
+    """
+
+    # dist b/t match center and prior's center
+    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, 2:])
+    # match wh / prior wh
+    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
+    g_wh = torch.log(g_wh) / variances[1]
+    # return target for smooth_l1_loss
+    return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]
+
+
+# Adapted from https://github.com/Hakuyume/chainer-ssd
+def decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    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 log_sum_exp(x):
+    """Utility function for computing log_sum_exp while determining
+    This will be used to determine unaveraged confidence loss across
+    all examples in a batch.
+    Args:
+        x (Variable(tensor)): conf_preds from conf layers
+    """
+    x_max = x.data.max()
+    return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max
+
+
+# Original author: Francisco Massa:
+# https://github.com/fmassa/object-detection.torch
+# Ported to PyTorch by Max deGroot (02/01/2017)
+def nms(boxes, scores, overlap=0.5, top_k=200):
+    """Apply non-maximum suppression at test time to avoid detecting too many
+    overlapping bounding boxes for a given object.
+    Args:
+        boxes: (tensor) The location preds for the img, Shape: [num_priors,4].
+        scores: (tensor) The class predscores for the img, Shape:[num_priors].
+        overlap: (float) The overlap thresh for suppressing unnecessary boxes.
+        top_k: (int) The Maximum number of box preds to consider.
+    Return:
+        The indices of the kept boxes with respect to num_priors.
+    """
+
+    keep = torch.Tensor(scores.size(0)).fill_(0).long()
+    if boxes.numel() == 0:
+        return keep
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    area = torch.mul(x2 - x1, y2 - y1)
+    v, idx = scores.sort(0)  # sort in ascending order
+    # I = I[v >= 0.01]
+    idx = idx[-top_k:]  # indices of the top-k largest vals
+    xx1 = boxes.new()
+    yy1 = boxes.new()
+    xx2 = boxes.new()
+    yy2 = boxes.new()
+    w = boxes.new()
+    h = boxes.new()
+
+    # keep = torch.Tensor()
+    count = 0
+    while idx.numel() > 0:
+        i = idx[-1]  # index of current largest val
+        # keep.append(i)
+        keep[count] = i
+        count += 1
+        if idx.size(0) == 1:
+            break
+        idx = idx[:-1]  # remove kept element from view
+        # load bboxes of next highest vals
+        torch.index_select(x1, 0, idx, out=xx1)
+        torch.index_select(y1, 0, idx, out=yy1)
+        torch.index_select(x2, 0, idx, out=xx2)
+        torch.index_select(y2, 0, idx, out=yy2)
+        # store element-wise max with next highest score
+        xx1 = torch.clamp(xx1, min=x1[i])
+        yy1 = torch.clamp(yy1, min=y1[i])
+        xx2 = torch.clamp(xx2, max=x2[i])
+        yy2 = torch.clamp(yy2, max=y2[i])
+        w.resize_as_(xx2)
+        h.resize_as_(yy2)
+        w = xx2 - xx1
+        h = yy2 - yy1
+        # check sizes of xx1 and xx2.. after each iteration
+        w = torch.clamp(w, min=0.0)
+        h = torch.clamp(h, min=0.0)
+        inter = w*h
+        # IoU = i / (area(a) + area(b) - i)
+        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
+        union = (rem_areas - inter) + area[i]
+        IoU = inter/union  # store result in iou
+        # keep only elements with an IoU <= overlap
+        idx = idx[IoU.le(overlap)]
+    return keep, count
+
+

third_party/PIPNet/FaceBoxesV2/utils/build.py

@@ -0,0 +1,57 @@
+# coding: utf-8
+
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import os
+from os.path import join as pjoin
+import numpy as np
+from distutils.core import setup
+from distutils.extension import Extension
+from Cython.Distutils import build_ext
+
+
+def find_in_path(name, path):
+    "Find a file in a search path"
+    # adapted fom http://code.activestate.com/recipes/52224-find-a-file-given-a-search-path/
+    for dir in path.split(os.pathsep):
+        binpath = pjoin(dir, name)
+        if os.path.exists(binpath):
+            return os.path.abspath(binpath)
+    return None
+
+
+# Obtain the numpy include directory.  This logic works across numpy versions.
+try:
+    numpy_include = np.get_include()
+except AttributeError:
+    numpy_include = np.get_numpy_include()
+
+
+# run the customize_compiler
+class custom_build_ext(build_ext):
+    def build_extensions(self):
+        # customize_compiler_for_nvcc(self.compiler)
+        build_ext.build_extensions(self)
+
+
+ext_modules = [
+    Extension(
+        "nms.cpu_nms",
+        ["nms/cpu_nms.pyx"],
+        # extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]},
+        extra_compile_args=["-Wno-cpp", "-Wno-unused-function"],
+        include_dirs=[numpy_include]
+    )
+]
+
+setup(
+    name='mot_utils',
+    ext_modules=ext_modules,
+    # inject our custom trigger
+    cmdclass={'build_ext': custom_build_ext},
+)

third_party/PIPNet/FaceBoxesV2/utils/config.py

@@ -0,0 +1,14 @@
+# config.py
+
+cfg = {
+    'name': 'FaceBoxes',
+    #'min_dim': 1024,
+    #'feature_maps': [[32, 32], [16, 16], [8, 8]],
+    # 'aspect_ratios': [[1], [1], [1]],
+    'min_sizes': [[32, 64, 128], [256], [512]],
+    'steps': [32, 64, 128],
+    'variance': [0.1, 0.2],
+    'clip': False,
+    'loc_weight': 2.0,
+    'gpu_train': True
+}

third_party/PIPNet/FaceBoxesV2/utils/faceboxes.py

@@ -0,0 +1,239 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicConv2d(nn.Module):
+
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
+        self.bn = nn.BatchNorm2d(out_channels, eps=1e-5)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return F.relu(x, inplace=True)
+
+
+class Inception(nn.Module):
+
+  def __init__(self):
+    super(Inception, self).__init__()
+    self.branch1x1 = BasicConv2d(128, 32, kernel_size=1, padding=0)
+    self.branch1x1_2 = BasicConv2d(128, 32, kernel_size=1, padding=0)
+    self.branch3x3_reduce = BasicConv2d(128, 24, kernel_size=1, padding=0)
+    self.branch3x3 = BasicConv2d(24, 32, kernel_size=3, padding=1)
+    self.branch3x3_reduce_2 = BasicConv2d(128, 24, kernel_size=1, padding=0)
+    self.branch3x3_2 = BasicConv2d(24, 32, kernel_size=3, padding=1)
+    self.branch3x3_3 = BasicConv2d(32, 32, kernel_size=3, padding=1)
+
+  def forward(self, x):
+    branch1x1 = self.branch1x1(x)
+
+    branch1x1_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
+    branch1x1_2 = self.branch1x1_2(branch1x1_pool)
+
+    branch3x3_reduce = self.branch3x3_reduce(x)
+    branch3x3 = self.branch3x3(branch3x3_reduce)
+
+    branch3x3_reduce_2 = self.branch3x3_reduce_2(x)
+    branch3x3_2 = self.branch3x3_2(branch3x3_reduce_2)
+    branch3x3_3 = self.branch3x3_3(branch3x3_2)
+
+    outputs = [branch1x1, branch1x1_2, branch3x3, branch3x3_3]
+    return torch.cat(outputs, 1)
+
+
+class CRelu(nn.Module):
+
+  def __init__(self, in_channels, out_channels, **kwargs):
+    super(CRelu, self).__init__()
+    self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
+    self.bn = nn.BatchNorm2d(out_channels, eps=1e-5)
+
+  def forward(self, x):
+    x = self.conv(x)
+    x = self.bn(x)
+    x = torch.cat([x, -x], 1)
+    x = F.relu(x, inplace=True)
+    return x
+
+
+class FaceBoxes(nn.Module):
+
+  def __init__(self, phase, size, num_classes):
+    super(FaceBoxes, self).__init__()
+    self.phase = phase
+    self.num_classes = num_classes
+    self.size = size
+
+    self.conv1 = CRelu(3, 24, kernel_size=7, stride=4, padding=3)
+    self.conv2 = CRelu(48, 64, kernel_size=5, stride=2, padding=2)
+
+    self.inception1 = Inception()
+    self.inception2 = Inception()
+    self.inception3 = Inception()
+
+    self.conv3_1 = BasicConv2d(128, 128, kernel_size=1, stride=1, padding=0)
+    self.conv3_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
+
+    self.conv4_1 = BasicConv2d(256, 128, kernel_size=1, stride=1, padding=0)
+    self.conv4_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
+
+    self.loc, self.conf = self.multibox(self.num_classes)
+
+    if self.phase == 'test':
+        self.softmax = nn.Softmax(dim=-1)
+
+    if self.phase == 'train':
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                if m.bias is not None:
+                    nn.init.xavier_normal_(m.weight.data)
+                    m.bias.data.fill_(0.02)
+                else:
+                    m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+  def multibox(self, num_classes):
+    loc_layers = []
+    conf_layers = []
+    loc_layers += [nn.Conv2d(128, 21 * 4, kernel_size=3, padding=1)]
+    conf_layers += [nn.Conv2d(128, 21 * num_classes, kernel_size=3, padding=1)]
+    loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
+    conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
+    loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
+    conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
+    return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers)
+
+  def forward(self, x):
+
+    detection_sources = list()
+    loc = list()
+    conf = list()
+
+    x = self.conv1(x)
+    x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
+    x = self.conv2(x)
+    x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
+    x = self.inception1(x)
+    x = self.inception2(x)
+    x = self.inception3(x)
+    detection_sources.append(x)
+
+    x = self.conv3_1(x)
+    x = self.conv3_2(x)
+    detection_sources.append(x)
+
+    x = self.conv4_1(x)
+    x = self.conv4_2(x)
+    detection_sources.append(x)
+
+    for (x, l, c) in zip(detection_sources, self.loc, self.conf):
+        loc.append(l(x).permute(0, 2, 3, 1).contiguous())
+        conf.append(c(x).permute(0, 2, 3, 1).contiguous())
+
+    loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
+    conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
+
+    if self.phase == "test":
+      output = (loc.view(loc.size(0), -1, 4),
+                self.softmax(conf.view(conf.size(0), -1, self.num_classes)))
+    else:
+      output = (loc.view(loc.size(0), -1, 4),
+                conf.view(conf.size(0), -1, self.num_classes))
+
+    return output
+
+class FaceBoxesV2(nn.Module):
+
+  def __init__(self, phase, size, num_classes):
+    super(FaceBoxesV2, self).__init__()
+    self.phase = phase
+    self.num_classes = num_classes
+    self.size = size
+    
+    self.conv1 = BasicConv2d(3, 8, kernel_size=3, stride=2, padding=1)
+    self.conv2 = BasicConv2d(8, 16, kernel_size=3, stride=2, padding=1)
+    self.conv3 = BasicConv2d(16, 32, kernel_size=3, stride=2, padding=1)
+    self.conv4 = BasicConv2d(32, 64, kernel_size=3, stride=2, padding=1)
+    self.conv5 = BasicConv2d(64, 128, kernel_size=3, stride=2, padding=1)
+    
+    self.inception1 = Inception()
+    self.inception2 = Inception()
+    self.inception3 = Inception()
+    
+    self.conv6_1 = BasicConv2d(128, 128, kernel_size=1, stride=1, padding=0)
+    self.conv6_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
+    
+    self.conv7_1 = BasicConv2d(256, 128, kernel_size=1, stride=1, padding=0)
+    self.conv7_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
+    
+    self.loc, self.conf = self.multibox(self.num_classes)
+    
+    if self.phase == 'test':
+        self.softmax = nn.Softmax(dim=-1)
+
+    if self.phase == 'train':
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                if m.bias is not None:
+                    nn.init.xavier_normal_(m.weight.data)
+                    m.bias.data.fill_(0.02)
+                else:
+                    m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+  def multibox(self, num_classes):
+    loc_layers = []
+    conf_layers = []
+    loc_layers += [nn.Conv2d(128, 21 * 4, kernel_size=3, padding=1)]
+    conf_layers += [nn.Conv2d(128, 21 * num_classes, kernel_size=3, padding=1)]
+    loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
+    conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
+    loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
+    conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
+    return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers)
+    
+  def forward(self, x):
+  
+    sources = list()
+    loc = list()
+    conf = list()
+
+    x = self.conv1(x)
+    x = self.conv2(x)
+    x = self.conv3(x)
+    x = self.conv4(x)
+    x = self.conv5(x)
+    x = self.inception1(x)
+    x = self.inception2(x)
+    x = self.inception3(x)
+    sources.append(x)
+    x = self.conv6_1(x)
+    x = self.conv6_2(x)
+    sources.append(x)
+    x = self.conv7_1(x)
+    x = self.conv7_2(x)
+    sources.append(x)
+
+    for (x, l, c) in zip(sources, self.loc, self.conf):
+        loc.append(l(x).permute(0, 2, 3, 1).contiguous())
+        conf.append(c(x).permute(0, 2, 3, 1).contiguous())
+        
+    loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
+    conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
+
+    if self.phase == "test":
+      output = (loc.view(loc.size(0), -1, 4),
+                self.softmax(conf.view(-1, self.num_classes)))
+    else:
+      output = (loc.view(loc.size(0), -1, 4),
+                conf.view(conf.size(0), -1, self.num_classes))
+  
+    return output

third_party/PIPNet/FaceBoxesV2/utils/make.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+python3 build.py build_ext --inplace
+

third_party/PIPNet/FaceBoxesV2/utils/nms/cpu_nms.pyx

@@ -0,0 +1,163 @@
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import numpy as np
+cimport numpy as np
+
+cdef inline np.float32_t max(np.float32_t a, np.float32_t b):
+    return a if a >= b else b
+
+cdef inline np.float32_t min(np.float32_t a, np.float32_t b):
+    return a if a <= b else b
+
+def cpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh):
+    cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0]
+    cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1]
+    cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2]
+    cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:, 3]
+    cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4]
+
+    cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1]
+
+    cdef int ndets = dets.shape[0]
+    cdef np.ndarray[np.int_t, ndim=1] suppressed = \
+            np.zeros((ndets), dtype=np.int)
+
+    # nominal indices
+    cdef int _i, _j
+    # sorted indices
+    cdef int i, j
+    # temp variables for box i's (the box currently under consideration)
+    cdef np.float32_t ix1, iy1, ix2, iy2, iarea
+    # variables for computing overlap with box j (lower scoring box)
+    cdef np.float32_t xx1, yy1, xx2, yy2
+    cdef np.float32_t w, h
+    cdef np.float32_t inter, ovr
+
+    keep = []
+    for _i in range(ndets):
+        i = order[_i]
+        if suppressed[i] == 1:
+            continue
+        keep.append(i)
+        ix1 = x1[i]
+        iy1 = y1[i]
+        ix2 = x2[i]
+        iy2 = y2[i]
+        iarea = areas[i]
+        for _j in range(_i + 1, ndets):
+            j = order[_j]
+            if suppressed[j] == 1:
+                continue
+            xx1 = max(ix1, x1[j])
+            yy1 = max(iy1, y1[j])
+            xx2 = min(ix2, x2[j])
+            yy2 = min(iy2, y2[j])
+            w = max(0.0, xx2 - xx1 + 1)
+            h = max(0.0, yy2 - yy1 + 1)
+            inter = w * h
+            ovr = inter / (iarea + areas[j] - inter)
+            if ovr >= thresh:
+                suppressed[j] = 1
+
+    return keep
+
+def cpu_soft_nms(np.ndarray[float, ndim=2] boxes, float sigma=0.5, float Nt=0.3, float threshold=0.001, unsigned int method=0):
+    cdef unsigned int N = boxes.shape[0]
+    cdef float iw, ih, box_area
+    cdef float ua
+    cdef int pos = 0
+    cdef float maxscore = 0
+    cdef int maxpos = 0
+    cdef float x1,x2,y1,y2,tx1,tx2,ty1,ty2,ts,area,weight,ov
+
+    for i in range(N):
+        maxscore = boxes[i, 4]
+        maxpos = i
+
+        tx1 = boxes[i,0]
+        ty1 = boxes[i,1]
+        tx2 = boxes[i,2]
+        ty2 = boxes[i,3]
+        ts = boxes[i,4]
+
+        pos = i + 1
+	# get max box
+        while pos < N:
+            if maxscore < boxes[pos, 4]:
+                maxscore = boxes[pos, 4]
+                maxpos = pos
+            pos = pos + 1
+
+	# add max box as a detection 
+        boxes[i,0] = boxes[maxpos,0]
+        boxes[i,1] = boxes[maxpos,1]
+        boxes[i,2] = boxes[maxpos,2]
+        boxes[i,3] = boxes[maxpos,3]
+        boxes[i,4] = boxes[maxpos,4]
+
+	# swap ith box with position of max box
+        boxes[maxpos,0] = tx1
+        boxes[maxpos,1] = ty1
+        boxes[maxpos,2] = tx2
+        boxes[maxpos,3] = ty2
+        boxes[maxpos,4] = ts
+
+        tx1 = boxes[i,0]
+        ty1 = boxes[i,1]
+        tx2 = boxes[i,2]
+        ty2 = boxes[i,3]
+        ts = boxes[i,4]
+
+        pos = i + 1
+	# NMS iterations, note that N changes if detection boxes fall below threshold
+        while pos < N:
+            x1 = boxes[pos, 0]
+            y1 = boxes[pos, 1]
+            x2 = boxes[pos, 2]
+            y2 = boxes[pos, 3]
+            s = boxes[pos, 4]
+
+            area = (x2 - x1 + 1) * (y2 - y1 + 1)
+            iw = (min(tx2, x2) - max(tx1, x1) + 1)
+            if iw > 0:
+                ih = (min(ty2, y2) - max(ty1, y1) + 1)
+                if ih > 0:
+                    ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih)
+                    ov = iw * ih / ua #iou between max box and detection box
+
+                    if method == 1: # linear
+                        if ov > Nt: 
+                            weight = 1 - ov
+                        else:
+                            weight = 1
+                    elif method == 2: # gaussian
+                        weight = np.exp(-(ov * ov)/sigma)
+                    else: # original NMS
+                        if ov > Nt: 
+                            weight = 0
+                        else:
+                            weight = 1
+
+                    boxes[pos, 4] = weight*boxes[pos, 4]
+		    
+		    # if box score falls below threshold, discard the box by swapping with last box
+		    # update N
+                    if boxes[pos, 4] < threshold:
+                        boxes[pos,0] = boxes[N-1, 0]
+                        boxes[pos,1] = boxes[N-1, 1]
+                        boxes[pos,2] = boxes[N-1, 2]
+                        boxes[pos,3] = boxes[N-1, 3]
+                        boxes[pos,4] = boxes[N-1, 4]
+                        N = N - 1
+                        pos = pos - 1
+
+            pos = pos + 1
+
+    keep = [i for i in range(N)]
+    return keep

third_party/PIPNet/FaceBoxesV2/utils/nms/gpu_nms.hpp

@@ -0,0 +1,2 @@
+void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num,
+          int boxes_dim, float nms_overlap_thresh, int device_id);

third_party/PIPNet/FaceBoxesV2/utils/nms/gpu_nms.pyx

@@ -0,0 +1,31 @@
+# --------------------------------------------------------
+# Faster R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import numpy as np
+cimport numpy as np
+
+assert sizeof(int) == sizeof(np.int32_t)
+
+cdef extern from "gpu_nms.hpp":
+    void _nms(np.int32_t*, int*, np.float32_t*, int, int, float, int)
+
+def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh,
+            np.int32_t device_id=0):
+    cdef int boxes_num = dets.shape[0]
+    cdef int boxes_dim = dets.shape[1]
+    cdef int num_out
+    cdef np.ndarray[np.int32_t, ndim=1] \
+        keep = np.zeros(boxes_num, dtype=np.int32)
+    cdef np.ndarray[np.float32_t, ndim=1] \
+        scores = dets[:, 4]
+    cdef np.ndarray[np.int_t, ndim=1] \
+        order = scores.argsort()[::-1]
+    cdef np.ndarray[np.float32_t, ndim=2] \
+        sorted_dets = dets[order, :]
+    _nms(&keep[0], &num_out, &sorted_dets[0, 0], boxes_num, boxes_dim, thresh, device_id)
+    keep = keep[:num_out]
+    return list(order[keep])

third_party/PIPNet/FaceBoxesV2/utils/nms/nms_kernel.cu

@@ -0,0 +1,144 @@
+// ------------------------------------------------------------------
+// Faster R-CNN
+// Copyright (c) 2015 Microsoft
+// Licensed under The MIT License [see fast-rcnn/LICENSE for details]
+// Written by Shaoqing Ren
+// ------------------------------------------------------------------
+
+#include "gpu_nms.hpp"
+#include <vector>
+#include <iostream>
+
+#define CUDA_CHECK(condition) \
+  /* Code block avoids redefinition of cudaError_t error */ \
+  do { \
+    cudaError_t error = condition; \
+    if (error != cudaSuccess) { \
+      std::cout << cudaGetErrorString(error) << std::endl; \
+    } \
+  } while (0)
+
+#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
+int const threadsPerBlock = sizeof(unsigned long long) * 8;
+
+__device__ inline float devIoU(float const * const a, float const * const b) {
+  float left = max(a[0], b[0]), right = min(a[2], b[2]);
+  float top = max(a[1], b[1]), bottom = min(a[3], b[3]);
+  float width = max(right - left + 1, 0.f), height = max(bottom - top + 1, 0.f);
+  float interS = width * height;
+  float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1);
+  float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1);
+  return interS / (Sa + Sb - interS);
+}
+
+__global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh,
+                           const float *dev_boxes, unsigned long long *dev_mask) {
+  const int row_start = blockIdx.y;
+  const int col_start = blockIdx.x;
+
+  // if (row_start > col_start) return;
+
+  const int row_size =
+        min(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
+  const int col_size =
+        min(n_boxes - col_start * threadsPerBlock, threadsPerBlock);
+
+  __shared__ float block_boxes[threadsPerBlock * 5];
+  if (threadIdx.x < col_size) {
+    block_boxes[threadIdx.x * 5 + 0] =
+        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0];
+    block_boxes[threadIdx.x * 5 + 1] =
+        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1];
+    block_boxes[threadIdx.x * 5 + 2] =
+        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2];
+    block_boxes[threadIdx.x * 5 + 3] =
+        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3];
+    block_boxes[threadIdx.x * 5 + 4] =
+        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4];
+  }
+  __syncthreads();
+
+  if (threadIdx.x < row_size) {
+    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
+    const float *cur_box = dev_boxes + cur_box_idx * 5;
+    int i = 0;
+    unsigned long long t = 0;
+    int start = 0;
+    if (row_start == col_start) {
+      start = threadIdx.x + 1;
+    }
+    for (i = start; i < col_size; i++) {
+      if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) {
+        t |= 1ULL << i;
+      }
+    }
+    const int col_blocks = DIVUP(n_boxes, threadsPerBlock);
+    dev_mask[cur_box_idx * col_blocks + col_start] = t;
+  }
+}
+
+void _set_device(int device_id) {
+  int current_device;
+  CUDA_CHECK(cudaGetDevice(&current_device));
+  if (current_device == device_id) {
+    return;
+  }
+  // The call to cudaSetDevice must come before any calls to Get, which
+  // may perform initialization using the GPU.
+  CUDA_CHECK(cudaSetDevice(device_id));
+}
+
+void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num,
+          int boxes_dim, float nms_overlap_thresh, int device_id) {
+  _set_device(device_id);
+
+  float* boxes_dev = NULL;
+  unsigned long long* mask_dev = NULL;
+
+  const int col_blocks = DIVUP(boxes_num, threadsPerBlock);
+
+  CUDA_CHECK(cudaMalloc(&boxes_dev,
+                        boxes_num * boxes_dim * sizeof(float)));
+  CUDA_CHECK(cudaMemcpy(boxes_dev,
+                        boxes_host,
+                        boxes_num * boxes_dim * sizeof(float),
+                        cudaMemcpyHostToDevice));
+
+  CUDA_CHECK(cudaMalloc(&mask_dev,
+                        boxes_num * col_blocks * sizeof(unsigned long long)));
+
+  dim3 blocks(DIVUP(boxes_num, threadsPerBlock),
+              DIVUP(boxes_num, threadsPerBlock));
+  dim3 threads(threadsPerBlock);
+  nms_kernel<<<blocks, threads>>>(boxes_num,
+                                  nms_overlap_thresh,
+                                  boxes_dev,
+                                  mask_dev);
+
+  std::vector<unsigned long long> mask_host(boxes_num * col_blocks);
+  CUDA_CHECK(cudaMemcpy(&mask_host[0],
+                        mask_dev,
+                        sizeof(unsigned long long) * boxes_num * col_blocks,
+                        cudaMemcpyDeviceToHost));
+
+  std::vector<unsigned long long> remv(col_blocks);
+  memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);
+
+  int num_to_keep = 0;
+  for (int i = 0; i < boxes_num; i++) {
+    int nblock = i / threadsPerBlock;
+    int inblock = i % threadsPerBlock;
+
+    if (!(remv[nblock] & (1ULL << inblock))) {
+      keep_out[num_to_keep++] = i;
+      unsigned long long *p = &mask_host[0] + i * col_blocks;
+      for (int j = nblock; j < col_blocks; j++) {
+        remv[j] |= p[j];
+      }
+    }
+  }
+  *num_out = num_to_keep;
+
+  CUDA_CHECK(cudaFree(boxes_dev));
+  CUDA_CHECK(cudaFree(mask_dev));
+}

third_party/PIPNet/FaceBoxesV2/utils/nms/py_cpu_nms.py

@@ -0,0 +1,38 @@
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import numpy as np
+
+def py_cpu_nms(dets, thresh):
+    """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 = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    return keep

third_party/PIPNet/FaceBoxesV2/utils/nms_wrapper.py

@@ -0,0 +1,15 @@
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+from .nms.cpu_nms import cpu_nms, cpu_soft_nms
+
+def nms(dets, thresh):
+    """Dispatch to either CPU or GPU NMS implementations."""
+
+    if dets.shape[0] == 0:
+        return []
+    return cpu_nms(dets, thresh)

third_party/PIPNet/FaceBoxesV2/utils/prior_box.py

@@ -0,0 +1,43 @@
+import torch
+from itertools import product as product
+import numpy as np
+from math import ceil
+
+
+class PriorBox(object):
+    def __init__(self, cfg, image_size=None, phase='train'):
+        super(PriorBox, self).__init__()
+        #self.aspect_ratios = cfg['aspect_ratios']
+        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 forward(self):
+        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]
+                    if min_size == 32:
+                        dense_cx = [x*self.steps[k]/self.image_size[1] for x in [j+0, j+0.25, j+0.5, j+0.75]]
+                        dense_cy = [y*self.steps[k]/self.image_size[0] for y in [i+0, i+0.25, i+0.5, i+0.75]]
+                        for cy, cx in product(dense_cy, dense_cx):
+                            anchors += [cx, cy, s_kx, s_ky]
+                    elif min_size == 64:
+                        dense_cx = [x*self.steps[k]/self.image_size[1] for x in [j+0, j+0.5]]
+                        dense_cy = [y*self.steps[k]/self.image_size[0] for y in [i+0, i+0.5]]
+                        for cy, cx in product(dense_cy, dense_cx):
+                            anchors += [cx, cy, s_kx, s_ky]
+                    else:
+                        cx = (j + 0.5) * self.steps[k] / self.image_size[1]
+                        cy = (i + 0.5) * self.steps[k] / self.image_size[0]
+                        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

third_party/PIPNet/FaceBoxesV2/utils/timer.py

@@ -0,0 +1,40 @@
+# --------------------------------------------------------
+# Fast R-CNN
+# Copyright (c) 2015 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ross Girshick
+# --------------------------------------------------------
+
+import time
+
+
+class Timer(object):
+    """A simple timer."""
+    def __init__(self):
+        self.total_time = 0.
+        self.calls = 0
+        self.start_time = 0.
+        self.diff = 0.
+        self.average_time = 0.
+
+    def tic(self):
+        # using time.time instead of time.clock because time time.clock
+        # does not normalize for multithreading
+        self.start_time = time.time()
+
+    def toc(self, average=True):
+        self.diff = time.time() - self.start_time
+        self.total_time += self.diff
+        self.calls += 1
+        self.average_time = self.total_time / self.calls
+        if average:
+            return self.average_time
+        else:
+            return self.diff
+
+    def clear(self):
+        self.total_time = 0.
+        self.calls = 0
+        self.start_time = 0.
+        self.diff = 0.
+        self.average_time = 0.

third_party/PIPNet/LICENSE

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

third_party/PIPNet/README.md

@@ -0,0 +1,153 @@
+# Pixel-in-Pixel Net: Towards Efficient Facial Landmark Detection in the Wild
+## Introduction
+This is the code of paper [Pixel-in-Pixel Net: Towards Efficient Facial Landmark Detection in the Wild](https://arxiv.org/abs/2003.03771). We propose a novel facial landmark detector, PIPNet, that is **fast**, **accurate**, and **robust**. PIPNet can be trained under two settings: (1) supervised learning; (2) generalizable semi-supervised learning (GSSL). With GSSL, PIPNet has better cross-domain generalization performance by utilizing massive amounts of unlabeled data across domains. 
+
+<img src="images/speed.png" alt="speed" width="640px">
+Figure 1. Comparison to existing methods on speed-accuracy tradeoff, tested on WFLW full test set (closer to bottom-right corner is better).<br><br>
+
+<img src="images/detection_heads.png" alt="det_heads" width="512px">
+Figure 2. Comparison of different detection heads.<br>
+
+## Installation
+1. Install Python3 and PyTorch >= v1.1
+2. Clone this repository.
+```Shell
+git clone https://github.com/jhb86253817/PIPNet.git
+```
+3. Install the dependencies in requirements.txt.
+```Shell
+pip install -r requirements.txt
+```
+
+## Demo
+1. We use a [modified version](https://github.com/jhb86253817/FaceBoxesV2) of [FaceBoxes](https://github.com/zisianw/FaceBoxes.PyTorch) as the face detector, so go to folder `FaceBoxesV2/utils`, run `sh make.sh` to build for NMS.
+2. Back to folder `PIPNet`, create two empty folders `logs` and `snapshots`. For PIPNets, you can download our trained models from [here](https://drive.google.com/drive/folders/17OwDgJUfuc5_ymQ3QruD8pUnh5zHreP2?usp=sharing), and put them under folder `snapshots/DATA_NAME/EXPERIMENT_NAME/`. 
+3. Edit `run_demo.sh` to choose the config file and input source you want, then run `sh run_demo.sh`. We support image, video, and camera as the input. Some sample predictions can be seen as follows.
+* PIPNet-ResNet18 trained on WFLW, with image `images/1.jpg` as the input:
+<img src="images/1_out_WFLW_model.jpg" alt="1_out_WFLW_model" width="400px">
+
+* PIPNet-ResNet18 trained on WFLW, with a snippet from *Shaolin Soccer* as the input:
+<img src="videos/shaolin_soccer.gif" alt="shaolin_soccer" width="400px">
+
+* PIPNet-ResNet18 trained on WFLW, with video `videos/002.avi` as the input:
+<img src="videos/002_out_WFLW_model.gif" alt="002_out_WFLW_model" width="512px">
+
+* PIPNet-ResNet18 trained on 300W+CelebA (GSSL), with video `videos/007.avi` as the input:
+<img src="videos/007_out_300W_CELEBA_model.gif" alt="007_out_300W_CELEBA_model" width="512px">
+
+## Training
+
+### Supervised Learning
+Datasets: [300W](https://ibug.doc.ic.ac.uk/resources/facial-point-annotations/), [COFW](http://www.vision.caltech.edu/xpburgos/ICCV13/), [WFLW](https://wywu.github.io/projects/LAB/WFLW.html), [AFLW](https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/aflw/)
+
+1. Download the datasets from official sources, then put them under folder `data`. The folder structure should look like this:
+````
+PIPNet
+-- FaceBoxesV2
+-- lib
+-- experiments
+-- logs
+-- snapshots
+-- data
+   |-- data_300W
+       |-- afw
+       |-- helen
+       |-- ibug
+       |-- lfpw
+   |-- COFW
+       |-- COFW_train_color.mat
+       |-- COFW_test_color.mat
+   |-- WFLW
+       |-- WFLW_images
+       |-- WFLW_annotations
+   |-- AFLW
+       |-- flickr
+       |-- AFLWinfo_release.mat
+````
+2. Go to folder `lib`, preprocess a dataset by running ```python preprocess.py DATA_NAME```. For example, to process 300W:
+```
+python preprocess.py data_300W
+```
+3. Back to folder `PIPNet`, edit `run_train.sh` to choose the config file you want. Then, train the model by running:
+```
+sh run_train.sh
+```
+
+### Generalizable Semi-supervised Learning
+Datasets: 
+* data_300W_COFW_WFLW: 300W + COFW-68 (unlabeled) + WFLW-68 (unlabeled) 
+* data_300W_CELEBA: 300W + CelebA (unlabeled)
+
+1. Download 300W, COFW, and WFLW as in the supervised learning setting. Download annotations of COFW-68 test from [here](https://github.com/golnazghiasi/cofw68-benchmark). For 300W+CelebA, you also need to download the in-the-wild CelebA images from [here](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html), and the [face bounding boxes](https://drive.google.com/drive/folders/17OwDgJUfuc5_ymQ3QruD8pUnh5zHreP2?usp=sharing) detected by us. The folder structure should look like this:
+````
+PIPNet
+-- FaceBoxesV2
+-- lib
+-- experiments
+-- logs
+-- snapshots
+-- data
+   |-- data_300W
+       |-- afw
+       |-- helen
+       |-- ibug
+       |-- lfpw
+   |-- COFW
+       |-- COFW_train_color.mat
+       |-- COFW_test_color.mat
+   |-- WFLW
+       |-- WFLW_images
+       |-- WFLW_annotations
+   |-- data_300W_COFW_WFLW
+       |-- cofw68_test_annotations
+       |-- cofw68_test_bboxes.mat
+   |-- CELEBA
+       |-- img_celeba
+       |-- celeba_bboxes.txt
+   |-- data_300W_CELEBA
+       |-- cofw68_test_annotations
+       |-- cofw68_test_bboxes.mat
+````
+2. Go to folder `lib`, preprocess a dataset by running ```python preprocess_gssl.py DATA_NAME```.
+   To process data_300W_COFW_WFLW, run
+   ```
+   python preprocess_gssl.py data_300W_COFW_WFLW
+   ```
+   To process data_300W_CELEBA, run
+   ```
+   python preprocess_gssl.py CELEBA
+   ```
+   and
+   ```
+   python preprocess_gssl.py data_300W_CELEBA
+   ```
+3. Back to folder `PIPNet`, edit `run_train.sh` to choose the config file you want. Then, train the model by running:
+```
+sh run_train.sh
+```
+
+## Evaluation
+1. Edit `run_test.sh` to choose the config file you want. Then, test the model by running:
+```
+sh run_test.sh
+```
+
+## Citation
+````
+@article{JLS21,
+  title={Pixel-in-Pixel Net: Towards Efficient Facial Landmark Detection in the Wild},
+  author={Haibo Jin and Shengcai Liao and Ling Shao},
+  journal={International Journal of Computer Vision},
+  publisher={Springer Science and Business Media LLC},
+  ISSN={1573-1405},
+  url={http://dx.doi.org/10.1007/s11263-021-01521-4},
+  DOI={10.1007/s11263-021-01521-4},
+  year={2021},
+  month={Sep}
+}
+````
+
+## Acknowledgement
+We thank the following great works:
+* [human-pose-estimation.pytorch](https://github.com/microsoft/human-pose-estimation.pytorch)
+* [HRNet-Facial-Landmark-Detection](https://github.com/HRNet/HRNet-Facial-Landmark-Detection)

third_party/PIPNet/lib/data_utils.py

@@ -0,0 +1,166 @@
+import torch.utils.data as data
+import torch
+from PIL import Image, ImageFilter 
+import os, cv2
+import numpy as np
+import random
+from scipy.stats import norm
+from math import floor
+
+def random_translate(image, target):
+    if random.random() > 0.5:
+        image_height, image_width = image.size
+        a = 1
+        b = 0
+        #c = 30 #left/right (i.e. 5/-5)
+        c = int((random.random()-0.5) * 60)
+        d = 0
+        e = 1
+        #f = 30 #up/down (i.e. 5/-5)
+        f = int((random.random()-0.5) * 60)
+        image = image.transform(image.size, Image.AFFINE, (a, b, c, d, e, f))
+        target_translate = target.copy()
+        target_translate = target_translate.reshape(-1, 2)
+        target_translate[:, 0] -= 1.*c/image_width
+        target_translate[:, 1] -= 1.*f/image_height
+        target_translate = target_translate.flatten()
+        target_translate[target_translate < 0] = 0
+        target_translate[target_translate > 1] = 1
+        return image, target_translate
+    else:
+        return image, target
+
+def random_blur(image):
+    if random.random() > 0.7:
+        image = image.filter(ImageFilter.GaussianBlur(random.random()*5))
+    return image
+
+def random_occlusion(image):
+    if random.random() > 0.5:
+        image_np = np.array(image).astype(np.uint8)
+        image_np = image_np[:,:,::-1]
+        image_height, image_width, _ = image_np.shape
+        occ_height = int(image_height*0.4*random.random())
+        occ_width = int(image_width*0.4*random.random())
+        occ_xmin = int((image_width - occ_width - 10) * random.random())
+        occ_ymin = int((image_height - occ_height - 10) * random.random())
+        image_np[occ_ymin:occ_ymin+occ_height, occ_xmin:occ_xmin+occ_width, 0] = int(random.random() * 255)
+        image_np[occ_ymin:occ_ymin+occ_height, occ_xmin:occ_xmin+occ_width, 1] = int(random.random() * 255)
+        image_np[occ_ymin:occ_ymin+occ_height, occ_xmin:occ_xmin+occ_width, 2] = int(random.random() * 255)
+        image_pil = Image.fromarray(image_np[:,:,::-1].astype('uint8'), 'RGB')
+        return image_pil
+    else:
+        return image
+
+def random_flip(image, target, points_flip):
+    if random.random() > 0.5:
+        image = image.transpose(Image.FLIP_LEFT_RIGHT)
+        target = np.array(target).reshape(-1, 2)
+        target = target[points_flip, :]
+        target[:,0] = 1-target[:,0]
+        target = target.flatten()
+        return image, target
+    else:
+        return image, target
+
+def random_rotate(image, target, angle_max):
+    if random.random() > 0.5:
+        center_x = 0.5
+        center_y = 0.5
+        landmark_num= int(len(target) / 2)
+        target_center = np.array(target) - np.array([center_x, center_y]*landmark_num)
+        target_center = target_center.reshape(landmark_num, 2)
+        theta_max = np.radians(angle_max)
+        theta = random.uniform(-theta_max, theta_max)
+        angle = np.degrees(theta)
+        image = image.rotate(angle)
+
+        c, s = np.cos(theta), np.sin(theta)
+        rot = np.array(((c,-s), (s, c)))
+        target_center_rot = np.matmul(target_center, rot)
+        target_rot = target_center_rot.reshape(landmark_num*2) + np.array([center_x, center_y]*landmark_num)
+        return image, target_rot
+    else:
+        return image, target
+
+def gen_target_pip(target, meanface_indices, target_map, target_local_x, target_local_y, target_nb_x, target_nb_y):
+    num_nb = len(meanface_indices[0])
+    map_channel, map_height, map_width = target_map.shape
+    target = target.reshape(-1, 2)
+    assert map_channel == target.shape[0]
+
+    for i in range(map_channel):
+        mu_x = int(floor(target[i][0] * map_width))
+        mu_y = int(floor(target[i][1] * map_height))
+        mu_x = max(0, mu_x)
+        mu_y = max(0, mu_y)
+        mu_x = min(mu_x, map_width-1)
+        mu_y = min(mu_y, map_height-1)
+        target_map[i, mu_y, mu_x] = 1
+        shift_x = target[i][0] * map_width - mu_x
+        shift_y = target[i][1] * map_height - mu_y
+        target_local_x[i, mu_y, mu_x] = shift_x
+        target_local_y[i, mu_y, mu_x] = shift_y
+
+        for j in range(num_nb):
+            nb_x = target[meanface_indices[i][j]][0] * map_width - mu_x
+            nb_y = target[meanface_indices[i][j]][1] * map_height - mu_y
+            target_nb_x[num_nb*i+j, mu_y, mu_x] = nb_x
+            target_nb_y[num_nb*i+j, mu_y, mu_x] = nb_y
+
+    return target_map, target_local_x, target_local_y, target_nb_x, target_nb_y
+
+class ImageFolder_pip(data.Dataset):
+    def __init__(self, root, imgs, input_size, num_lms, net_stride, points_flip, meanface_indices, transform=None, target_transform=None):
+        self.root = root
+        self.imgs = imgs
+        self.num_lms = num_lms
+        self.net_stride = net_stride
+        self.points_flip = points_flip
+        self.meanface_indices = meanface_indices
+        self.num_nb = len(meanface_indices[0])
+        self.transform = transform
+        self.target_transform = target_transform
+        self.input_size = input_size
+
+    def __getitem__(self, index):
+
+        img_name, target = self.imgs[index]
+
+        img = Image.open(os.path.join(self.root, img_name)).convert('RGB')
+        img, target = random_translate(img, target)
+        img = random_occlusion(img)
+        img, target = random_flip(img, target, self.points_flip)
+        img, target = random_rotate(img, target, 30)
+        img = random_blur(img)
+
+        target_map = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_local_x = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_local_y = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_nb_x = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_nb_y = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_map, target_local_x, target_local_y, target_nb_x, target_nb_y = gen_target_pip(target, self.meanface_indices, target_map, target_local_x, target_local_y, target_nb_x, target_nb_y)
+        
+        target_map = torch.from_numpy(target_map).float()
+        target_local_x = torch.from_numpy(target_local_x).float()
+        target_local_y = torch.from_numpy(target_local_y).float()
+        target_nb_x = torch.from_numpy(target_nb_x).float()
+        target_nb_y = torch.from_numpy(target_nb_y).float()
+
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.target_transform is not None:
+            target_map = self.target_transform(target_map)
+            target_local_x = self.target_transform(target_local_x)
+            target_local_y = self.target_transform(target_local_y)
+            target_nb_x = self.target_transform(target_nb_x)
+            target_nb_y = self.target_transform(target_nb_y)
+
+        return img, target_map, target_local_x, target_local_y, target_nb_x, target_nb_y
+
+    def __len__(self):
+        return len(self.imgs)
+
+if __name__ == '__main__':
+    pass
+    

third_party/PIPNet/lib/data_utils_gssl.py

@@ -0,0 +1,290 @@
+import torch.utils.data as data
+import torch
+from PIL import Image, ImageFilter 
+import os, cv2
+import numpy as np
+import random
+from scipy.stats import norm
+from math import floor
+
+def random_translate(image, target):
+    if random.random() > 0.5:
+        image_height, image_width = image.size
+        a = 1
+        b = 0
+        #c = 30 #left/right (i.e. 5/-5)
+        c = int((random.random()-0.5) * 60)
+        d = 0
+        e = 1
+        #f = 30 #up/down (i.e. 5/-5)
+        f = int((random.random()-0.5) * 60)
+        image = image.transform(image.size, Image.AFFINE, (a, b, c, d, e, f))
+        target_translate = target.copy()
+        target_translate = target_translate.reshape(-1, 2)
+        target_translate[:, 0] -= 1.*c/image_width
+        target_translate[:, 1] -= 1.*f/image_height
+        target_translate = target_translate.flatten()
+        target_translate[target_translate < 0] = 0
+        target_translate[target_translate > 1] = 1
+        return image, target_translate
+    else:
+        return image, target
+
+def random_blur(image):
+    if random.random() > 0.7:
+        image = image.filter(ImageFilter.GaussianBlur(random.random()*5))
+    return image
+
+def random_occlusion(image):
+    if random.random() > 0.5:
+        image_np = np.array(image).astype(np.uint8)
+        image_np = image_np[:,:,::-1]
+        image_height, image_width, _ = image_np.shape
+        occ_height = int(image_height*0.4*random.random())
+        occ_width = int(image_width*0.4*random.random())
+        occ_xmin = int((image_width - occ_width - 10) * random.random())
+        occ_ymin = int((image_height - occ_height - 10) * random.random())
+        image_np[occ_ymin:occ_ymin+occ_height, occ_xmin:occ_xmin+occ_width, 0] = int(random.random() * 255)
+        image_np[occ_ymin:occ_ymin+occ_height, occ_xmin:occ_xmin+occ_width, 1] = int(random.random() * 255)
+        image_np[occ_ymin:occ_ymin+occ_height, occ_xmin:occ_xmin+occ_width, 2] = int(random.random() * 255)
+        image_pil = Image.fromarray(image_np[:,:,::-1].astype('uint8'), 'RGB')
+        return image_pil
+    else:
+        return image
+
+def random_flip(image, target, points_flip):
+    if random.random() > 0.5:
+        image = image.transpose(Image.FLIP_LEFT_RIGHT)
+        target = np.array(target).reshape(-1, 2)
+        target = target[points_flip, :]
+        target[:,0] = 1-target[:,0]
+        target = target.flatten()
+        return image, target
+    else:
+        return image, target
+
+def random_rotate(image, target, angle_max):
+    if random.random() > 0.5:
+        center_x = 0.5
+        center_y = 0.5
+        landmark_num= int(len(target) / 2)
+        target_center = np.array(target) - np.array([center_x, center_y]*landmark_num)
+        target_center = target_center.reshape(landmark_num, 2)
+        theta_max = np.radians(angle_max)
+        theta = random.uniform(-theta_max, theta_max)
+        angle = np.degrees(theta)
+        image = image.rotate(angle)
+
+        c, s = np.cos(theta), np.sin(theta)
+        rot = np.array(((c,-s), (s, c)))
+        target_center_rot = np.matmul(target_center, rot)
+        target_rot = target_center_rot.reshape(landmark_num*2) + np.array([center_x, center_y]*landmark_num)
+        return image, target_rot
+    else:
+        return image, target
+
+def gen_target_pip(target, meanface_indices, target_map1, target_map2, target_map3, target_local_x, target_local_y, target_nb_x, target_nb_y):
+    num_nb = len(meanface_indices[0])
+    map_channel1, map_height1, map_width1 = target_map1.shape
+    map_channel2, map_height2, map_width2 = target_map2.shape
+    map_channel3, map_height3, map_width3 = target_map3.shape
+    target = target.reshape(-1, 2)
+    assert map_channel1 == target.shape[0]
+
+    for i in range(map_channel1):
+        mu_x1 = int(floor(target[i][0] * map_width1))
+        mu_y1 = int(floor(target[i][1] * map_height1))
+        mu_x1 = max(0, mu_x1)
+        mu_y1 = max(0, mu_y1)
+        mu_x1 = min(mu_x1, map_width1-1)
+        mu_y1 = min(mu_y1, map_height1-1)
+        target_map1[i, mu_y1, mu_x1] = 1
+
+        shift_x = target[i][0] * map_width1 - mu_x1
+        shift_y = target[i][1] * map_height1 - mu_y1
+        target_local_x[i, mu_y1, mu_x1] = shift_x
+        target_local_y[i, mu_y1, mu_x1] = shift_y
+
+        for j in range(num_nb):
+            nb_x = target[meanface_indices[i][j]][0] * map_width1 - mu_x1
+            nb_y = target[meanface_indices[i][j]][1] * map_height1 - mu_y1
+            target_nb_x[num_nb*i+j, mu_y1, mu_x1] = nb_x
+            target_nb_y[num_nb*i+j, mu_y1, mu_x1] = nb_y
+
+        mu_x2 = int(floor(target[i][0] * map_width2))
+        mu_y2 = int(floor(target[i][1] * map_height2))
+        mu_x2 = max(0, mu_x2)
+        mu_y2 = max(0, mu_y2)
+        mu_x2 = min(mu_x2, map_width2-1)
+        mu_y2 = min(mu_y2, map_height2-1)
+        target_map2[i, mu_y2, mu_x2] = 1
+
+        mu_x3 = int(floor(target[i][0] * map_width3))
+        mu_y3 = int(floor(target[i][1] * map_height3))
+        mu_x3 = max(0, mu_x3)
+        mu_y3 = max(0, mu_y3)
+        mu_x3 = min(mu_x3, map_width3-1)
+        mu_y3 = min(mu_y3, map_height3-1)
+        target_map3[i, mu_y3, mu_x3] = 1
+
+    return target_map1, target_map2, target_map3, target_local_x, target_local_y, target_nb_x, target_nb_y
+
+def gen_target_pip_cls1(target, target_map1):
+    map_channel1, map_height1, map_width1 = target_map1.shape
+    target = target.reshape(-1, 2)
+    assert map_channel1 == target.shape[0]
+
+    for i in range(map_channel1):
+        mu_x1 = int(floor(target[i][0] * map_width1))
+        mu_y1 = int(floor(target[i][1] * map_height1))
+        mu_x1 = max(0, mu_x1)
+        mu_y1 = max(0, mu_y1)
+        mu_x1 = min(mu_x1, map_width1-1)
+        mu_y1 = min(mu_y1, map_height1-1)
+        target_map1[i, mu_y1, mu_x1] = 1
+
+    return target_map1 
+
+def gen_target_pip_cls2(target, target_map2):
+    map_channel2, map_height2, map_width2 = target_map2.shape
+    target = target.reshape(-1, 2)
+    assert map_channel2 == target.shape[0]
+
+    for i in range(map_channel2):
+        mu_x2 = int(floor(target[i][0] * map_width2))
+        mu_y2 = int(floor(target[i][1] * map_height2))
+        mu_x2 = max(0, mu_x2)
+        mu_y2 = max(0, mu_y2)
+        mu_x2 = min(mu_x2, map_width2-1)
+        mu_y2 = min(mu_y2, map_height2-1)
+        target_map2[i, mu_y2, mu_x2] = 1
+
+    return target_map2 
+
+def gen_target_pip_cls3(target, target_map3):
+    map_channel3, map_height3, map_width3 = target_map3.shape
+    target = target.reshape(-1, 2)
+    assert map_channel3 == target.shape[0]
+
+    for i in range(map_channel3):
+        mu_x3 = int(floor(target[i][0] * map_width3))
+        mu_y3 = int(floor(target[i][1] * map_height3))
+        mu_x3 = max(0, mu_x3)
+        mu_y3 = max(0, mu_y3)
+        mu_x3 = min(mu_x3, map_width3-1)
+        mu_y3 = min(mu_y3, map_height3-1)
+        target_map3[i, mu_y3, mu_x3] = 1
+
+    return target_map3 
+
+class ImageFolder_pip(data.Dataset):
+    def __init__(self, root, imgs, input_size, num_lms, net_stride, points_flip, meanface_indices, transform=None, target_transform=None):
+        self.root = root
+        self.imgs = imgs
+        self.num_lms = num_lms
+        self.net_stride = net_stride
+        self.points_flip = points_flip
+        self.meanface_indices = meanface_indices
+        self.num_nb = len(meanface_indices[0])
+        self.transform = transform
+        self.target_transform = target_transform
+        self.input_size = input_size
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (image, target) where target is class_index of the target class.
+        """
+        img_name, target_type, target = self.imgs[index]
+        img = Image.open(os.path.join(self.root, img_name)).convert('RGB')
+
+        img, target = random_translate(img, target)
+        img = random_occlusion(img)
+        img, target = random_flip(img, target, self.points_flip)
+        img, target = random_rotate(img, target, 30)
+        img = random_blur(img)
+
+        target_map1 = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_map2 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/2), int(self.input_size/self.net_stride/2)))
+        target_map3 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/4), int(self.input_size/self.net_stride/4)))
+        target_local_x = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_local_y = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_nb_x = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        target_nb_y = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+
+        mask_map1 = np.ones((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        mask_map2 = np.ones((self.num_lms, int(self.input_size/self.net_stride/2), int(self.input_size/self.net_stride/2)))
+        mask_map3 = np.ones((self.num_lms, int(self.input_size/self.net_stride/4), int(self.input_size/self.net_stride/4)))
+        mask_local_x = np.ones((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        mask_local_y = np.ones((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        mask_nb_x = np.ones((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        mask_nb_y = np.ones((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+
+        if target_type == 'std':
+            target_map1, target_map2, target_map3, target_local_x, target_local_y, target_nb_x, target_nb_y = gen_target_pip(target, self.meanface_indices, target_map1, target_map2, target_map3, target_local_x, target_local_y, target_nb_x, target_nb_y)
+            mask_map2 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/2), int(self.input_size/self.net_stride/2)))
+            mask_map3 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/4), int(self.input_size/self.net_stride/4)))
+        elif target_type == 'cls1':
+            target_map1 = gen_target_pip_cls1(target, target_map1)
+            mask_map2 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/2), int(self.input_size/self.net_stride/2)))
+            mask_map3 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/4), int(self.input_size/self.net_stride/4)))
+            mask_local_x = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_local_y = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_nb_x = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_nb_y = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        elif target_type == 'cls2':
+            target_map2 = gen_target_pip_cls2(target, target_map2)
+            mask_map1 = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_map3 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/4), int(self.input_size/self.net_stride/4)))
+            mask_local_x = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_local_y = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_nb_x = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_nb_y = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        elif target_type == 'cls3':
+            target_map3 = gen_target_pip_cls3(target, target_map3)
+            mask_map1 = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_map2 = np.zeros((self.num_lms, int(self.input_size/self.net_stride/2), int(self.input_size/self.net_stride/2)))
+            mask_local_x = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_local_y = np.zeros((self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_nb_x = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+            mask_nb_y = np.zeros((self.num_nb*self.num_lms, int(self.input_size/self.net_stride), int(self.input_size/self.net_stride)))
+        else:
+            print('No such target type!')
+            exit(0)
+
+        target_map1 = torch.from_numpy(target_map1).float()
+        target_map2 = torch.from_numpy(target_map2).float()
+        target_map3 = torch.from_numpy(target_map3).float()
+        target_local_x = torch.from_numpy(target_local_x).float()
+        target_local_y = torch.from_numpy(target_local_y).float()
+        target_nb_x = torch.from_numpy(target_nb_x).float()
+        target_nb_y = torch.from_numpy(target_nb_y).float()
+        mask_map1 = torch.from_numpy(mask_map1).float()
+        mask_map2 = torch.from_numpy(mask_map2).float()
+        mask_map3 = torch.from_numpy(mask_map3).float()
+        mask_local_x = torch.from_numpy(mask_local_x).float()
+        mask_local_y = torch.from_numpy(mask_local_y).float()
+        mask_nb_x = torch.from_numpy(mask_nb_x).float()
+        mask_nb_y = torch.from_numpy(mask_nb_y).float()
+
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.target_transform is not None:
+            target_map1 = self.target_transform(target_map1)
+            target_map2 = self.target_transform(target_map2)
+            target_map3 = self.target_transform(target_map3)
+            target_local_x = self.target_transform(target_local_x)
+            target_local_y = self.target_transform(target_local_y)
+            target_nb_x = self.target_transform(target_nb_x)
+            target_nb_y = self.target_transform(target_nb_y)
+
+        return img, target_map1, target_map2, target_map3, target_local_x, target_local_y, target_nb_x, target_nb_y, mask_map1, mask_map2, mask_map3, mask_local_x, mask_local_y, mask_nb_x, mask_nb_y
+
+    def __len__(self):
+        return len(self.imgs)
+
+if __name__ == '__main__':
+    pass
+    

third_party/PIPNet/lib/demo.py

@@ -0,0 +1,159 @@
+import cv2
+import sys
+
+sys.path.insert(0, "FaceBoxesV2")
+sys.path.insert(0, "..")
+from math import floor
+from faceboxes_detector import *
+
+import torch
+import torch.nn.parallel
+import torch.utils.data
+import torchvision.transforms as transforms
+import torchvision.models as models
+
+from networks import *
+from functions import *
+from PIPNet.reverse_index import ri1, ri2
+
+
+class Config:
+    def __init__(self):
+        self.det_head = "pip"
+        self.net_stride = 32
+        self.batch_size = 16
+        self.init_lr = 0.0001
+        self.num_epochs = 60
+        self.decay_steps = [30, 50]
+        self.input_size = 256
+        self.backbone = "resnet101"
+        self.pretrained = True
+        self.criterion_cls = "l2"
+        self.criterion_reg = "l1"
+        self.cls_loss_weight = 10
+        self.reg_loss_weight = 1
+        self.num_lms = 98
+        self.save_interval = self.num_epochs
+        self.num_nb = 10
+        self.use_gpu = True
+        self.gpu_id = 3
+
+
+def get_lmk_model():
+
+    cfg = Config()
+
+    resnet101 = models.resnet101(pretrained=cfg.pretrained)
+    net = Pip_resnet101(
+        resnet101,
+        cfg.num_nb,
+        num_lms=cfg.num_lms,
+        input_size=cfg.input_size,
+        net_stride=cfg.net_stride,
+    )
+
+    if cfg.use_gpu:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device("cpu")
+    net = net.to(device)
+
+    weight_file = "/apdcephfs/share_1290939/ahbanliang/codes/PIPNet/snapshots/WFLW/pip_32_16_60_r101_l2_l1_10_1_nb10/epoch59.pth"
+    state_dict = torch.load(weight_file, map_location=device)
+    net.load_state_dict(state_dict)
+
+    detector = FaceBoxesDetector(
+        "FaceBoxes",
+        "FaceBoxesV2/weights/FaceBoxesV2.pth",
+        use_gpu=True,
+        device="cuda:0",
+    )
+    return net, detector
+
+
+def demo_image(
+    image_file,
+    net,
+    detector,
+    input_size=256,
+    net_stride=32,
+    num_nb=10,
+    use_gpu=True,
+    device="cuda:0",
+):
+
+    my_thresh = 0.6
+    det_box_scale = 1.2
+    net.eval()
+    preprocess = transforms.Compose(
+        [
+            transforms.Resize((256, 256)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+    reverse_index1, reverse_index2, max_len = ri1, ri2, 17
+    image = cv2.imread(image_file)
+    image_height, image_width, _ = image.shape
+    detections, _ = detector.detect(image, my_thresh, 1)
+    for i in range(len(detections)):
+        det_xmin = detections[i][2]
+        det_ymin = detections[i][3]
+        det_width = detections[i][4]
+        det_height = detections[i][5]
+        det_xmax = det_xmin + det_width - 1
+        det_ymax = det_ymin + det_height - 1
+
+        det_xmin -= int(det_width * (det_box_scale - 1) / 2)
+        # remove a part of top area for alignment, see paper for details
+        det_ymin += int(det_height * (det_box_scale - 1) / 2)
+        det_xmax += int(det_width * (det_box_scale - 1) / 2)
+        det_ymax += int(det_height * (det_box_scale - 1) / 2)
+        det_xmin = max(det_xmin, 0)
+        det_ymin = max(det_ymin, 0)
+        det_xmax = min(det_xmax, image_width - 1)
+        det_ymax = min(det_ymax, image_height - 1)
+        det_width = det_xmax - det_xmin + 1
+        det_height = det_ymax - det_ymin + 1
+        cv2.rectangle(image, (det_xmin, det_ymin), (det_xmax, det_ymax), (0, 0, 255), 2)
+        det_crop = image[det_ymin:det_ymax, det_xmin:det_xmax, :]
+        det_crop = cv2.resize(det_crop, (input_size, input_size))
+        inputs = Image.fromarray(det_crop[:, :, ::-1].astype("uint8"), "RGB")
+        inputs = preprocess(inputs).unsqueeze(0)
+        inputs = inputs.to(device)
+        (
+            lms_pred_x,
+            lms_pred_y,
+            lms_pred_nb_x,
+            lms_pred_nb_y,
+            outputs_cls,
+            max_cls,
+        ) = forward_pip(net, inputs, preprocess, input_size, net_stride, num_nb)
+        lms_pred = torch.cat((lms_pred_x, lms_pred_y), dim=1).flatten()
+        tmp_nb_x = lms_pred_nb_x[reverse_index1, reverse_index2].view(98, max_len)
+        tmp_nb_y = lms_pred_nb_y[reverse_index1, reverse_index2].view(98, max_len)
+        tmp_x = torch.mean(torch.cat((lms_pred_x, tmp_nb_x), dim=1), dim=1).view(-1, 1)
+        tmp_y = torch.mean(torch.cat((lms_pred_y, tmp_nb_y), dim=1), dim=1).view(-1, 1)
+        lms_pred_merge = torch.cat((tmp_x, tmp_y), dim=1).flatten()
+        lms_pred = lms_pred.cpu().numpy()
+        lms_pred_merge = lms_pred_merge.cpu().numpy()
+        for i in range(98):
+            x_pred = lms_pred_merge[i * 2] * det_width
+            y_pred = lms_pred_merge[i * 2 + 1] * det_height
+            cv2.circle(
+                image,
+                (int(x_pred) + det_xmin, int(y_pred) + det_ymin),
+                1,
+                (0, 0, 255),
+                2,
+            )
+    cv2.imwrite("images/1_out.jpg", image)
+
+
+if __name__ == "__main__":
+    net, detector = get_lmk_model()
+    demo_image(
+        "/apdcephfs/private_ahbanliang/codes/Real-ESRGAN-master/tmp_frames/yanikefu/frame00000046.png",
+        net,
+        detector,
+    )

third_party/PIPNet/lib/demo_video.py

@@ -0,0 +1,141 @@
+import cv2, os
+import sys
+sys.path.insert(0, 'FaceBoxesV2')
+sys.path.insert(0, '..')
+import numpy as np
+import pickle
+import importlib
+from math import floor
+from faceboxes_detector import *
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+import torchvision.models as models
+
+from networks import *
+import data_utils
+from functions import *
+
+if not len(sys.argv) == 3:
+    print('Format:')
+    print('python lib/demo_video.py config_file video_file')
+    exit(0)
+experiment_name = sys.argv[1].split('/')[-1][:-3]
+data_name = sys.argv[1].split('/')[-2]
+config_path = '.experiments.{}.{}'.format(data_name, experiment_name)
+video_file = sys.argv[2]
+
+my_config = importlib.import_module(config_path, package='PIPNet')
+Config = getattr(my_config, 'Config')
+cfg = Config()
+cfg.experiment_name = experiment_name
+cfg.data_name = data_name
+
+save_dir = os.path.join('./snapshots', cfg.data_name, cfg.experiment_name)
+
+meanface_indices, reverse_index1, reverse_index2, max_len = get_meanface(os.path.join('data', cfg.data_name, 'meanface.txt'), cfg.num_nb)
+
+if cfg.backbone == 'resnet18':
+    resnet18 = models.resnet18(pretrained=cfg.pretrained)
+    net = Pip_resnet18(resnet18, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+elif cfg.backbone == 'resnet50':
+    resnet50 = models.resnet50(pretrained=cfg.pretrained)
+    net = Pip_resnet50(resnet50, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+elif cfg.backbone == 'resnet101':
+    resnet101 = models.resnet101(pretrained=cfg.pretrained)
+    net = Pip_resnet101(resnet101, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+else:
+    print('No such backbone!')
+    exit(0)
+
+if cfg.use_gpu:
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+else:
+    device = torch.device("cpu")
+net = net.to(device)
+
+weight_file = os.path.join(save_dir, 'epoch%d.pth' % (cfg.num_epochs-1))
+state_dict = torch.load(weight_file, map_location=device)
+net.load_state_dict(state_dict)
+
+normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+preprocess = transforms.Compose([transforms.Resize((cfg.input_size, cfg.input_size)), transforms.ToTensor(), normalize])
+
+def demo_video(video_file, net, preprocess, input_size, net_stride, num_nb, use_gpu, device):
+    detector = FaceBoxesDetector('FaceBoxes', 'FaceBoxesV2/weights/FaceBoxesV2.pth', use_gpu, device)
+    my_thresh = 0.9
+    det_box_scale = 1.2
+
+    net.eval()
+    if video_file == 'camera':
+        cap = cv2.VideoCapture(0)
+    else:
+        cap = cv2.VideoCapture(video_file)
+    if (cap.isOpened()== False): 
+        print("Error opening video stream or file")
+    frame_width = int(cap.get(3))
+    frame_height = int(cap.get(4))
+    count = 0
+    while(cap.isOpened()):
+        ret, frame = cap.read()
+        if ret == True:
+            detections, _ = detector.detect(frame, my_thresh, 1)
+            for i in range(len(detections)):
+                det_xmin = detections[i][2]
+                det_ymin = detections[i][3]
+                det_width = detections[i][4]
+                det_height = detections[i][5]
+                det_xmax = det_xmin + det_width - 1
+                det_ymax = det_ymin + det_height - 1
+
+                det_xmin -= int(det_width * (det_box_scale-1)/2)
+                # remove a part of top area for alignment, see paper for details
+                det_ymin += int(det_height * (det_box_scale-1)/2)
+                det_xmax += int(det_width * (det_box_scale-1)/2)
+                det_ymax += int(det_height * (det_box_scale-1)/2)
+                det_xmin = max(det_xmin, 0)
+                det_ymin = max(det_ymin, 0)
+                det_xmax = min(det_xmax, frame_width-1)
+                det_ymax = min(det_ymax, frame_height-1)
+                det_width = det_xmax - det_xmin + 1
+                det_height = det_ymax - det_ymin + 1
+                cv2.rectangle(frame, (det_xmin, det_ymin), (det_xmax, det_ymax), (0, 0, 255), 2)
+                det_crop = frame[det_ymin:det_ymax, det_xmin:det_xmax, :]
+                det_crop = cv2.resize(det_crop, (input_size, input_size))
+                inputs = Image.fromarray(det_crop[:,:,::-1].astype('uint8'), 'RGB')
+                inputs = preprocess(inputs).unsqueeze(0)
+                inputs = inputs.to(device)
+                lms_pred_x, lms_pred_y, lms_pred_nb_x, lms_pred_nb_y, outputs_cls, max_cls = forward_pip(net, inputs, preprocess, input_size, net_stride, num_nb)
+                lms_pred = torch.cat((lms_pred_x, lms_pred_y), dim=1).flatten()
+                tmp_nb_x = lms_pred_nb_x[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+                tmp_nb_y = lms_pred_nb_y[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+                tmp_x = torch.mean(torch.cat((lms_pred_x, tmp_nb_x), dim=1), dim=1).view(-1,1)
+                tmp_y = torch.mean(torch.cat((lms_pred_y, tmp_nb_y), dim=1), dim=1).view(-1,1)
+                lms_pred_merge = torch.cat((tmp_x, tmp_y), dim=1).flatten()
+                lms_pred = lms_pred.cpu().numpy()
+                lms_pred_merge = lms_pred_merge.cpu().numpy()
+                for i in range(cfg.num_lms):
+                    x_pred = lms_pred_merge[i*2] * det_width
+                    y_pred = lms_pred_merge[i*2+1] * det_height
+                    cv2.circle(frame, (int(x_pred)+det_xmin, int(y_pred)+det_ymin), 1, (0, 0, 255), 2)
+                
+            count += 1
+            #cv2.imwrite('video_out2/'+str(count)+'.jpg', frame)
+            cv2.imshow('1', frame)
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
+        else:
+            break
+
+    cap.release()
+    cv2.destroyAllWindows()
+
+demo_video(video_file, net, preprocess, cfg.input_size, cfg.net_stride, cfg.num_nb, cfg.use_gpu, device)

third_party/PIPNet/lib/functions.py

@@ -0,0 +1,210 @@
+import os, cv2
+import numpy as np
+from PIL import Image, ImageFilter
+import logging
+import torch
+import torch.nn as nn
+import random
+import time
+from scipy.integrate import simps
+
+
+def get_label(data_name, label_file, task_type=None):
+    label_path = os.path.join('data', data_name, label_file)
+    with open(label_path, 'r') as f:
+        labels = f.readlines()
+    labels = [x.strip().split() for x in labels]
+    if len(labels[0])==1:
+        return labels
+
+    labels_new = []
+    for label in labels:
+        image_name = label[0]
+        target = label[1:]
+        target = np.array([float(x) for x in target])
+        if task_type is None:
+            labels_new.append([image_name, target])
+        else:
+            labels_new.append([image_name, task_type, target])
+    return labels_new
+
+def get_meanface(meanface_file, num_nb):
+    with open(meanface_file) as f:
+        meanface = f.readlines()[0]
+        
+    meanface = meanface.strip().split()
+    meanface = [float(x) for x in meanface]
+    meanface = np.array(meanface).reshape(-1, 2)
+    # each landmark predicts num_nb neighbors
+    meanface_indices = []
+    for i in range(meanface.shape[0]):
+        pt = meanface[i,:]
+        dists = np.sum(np.power(pt-meanface, 2), axis=1)
+        indices = np.argsort(dists)
+        meanface_indices.append(indices[1:1+num_nb])
+    
+    # each landmark predicted by X neighbors, X varies
+    meanface_indices_reversed = {}
+    for i in range(meanface.shape[0]):
+        meanface_indices_reversed[i] = [[],[]]
+    for i in range(meanface.shape[0]):
+        for j in range(num_nb):
+            meanface_indices_reversed[meanface_indices[i][j]][0].append(i)
+            meanface_indices_reversed[meanface_indices[i][j]][1].append(j)
+    
+    max_len = 0
+    for i in range(meanface.shape[0]):
+        tmp_len = len(meanface_indices_reversed[i][0])
+        if tmp_len > max_len:
+            max_len = tmp_len
+    
+    # tricks, make them have equal length for efficient computation
+    for i in range(meanface.shape[0]):
+        tmp_len = len(meanface_indices_reversed[i][0])
+        meanface_indices_reversed[i][0] += meanface_indices_reversed[i][0]*10
+        meanface_indices_reversed[i][1] += meanface_indices_reversed[i][1]*10
+        meanface_indices_reversed[i][0] = meanface_indices_reversed[i][0][:max_len]
+        meanface_indices_reversed[i][1] = meanface_indices_reversed[i][1][:max_len]
+
+    # make the indices 1-dim
+    reverse_index1 = []
+    reverse_index2 = []
+    for i in range(meanface.shape[0]):
+        reverse_index1 += meanface_indices_reversed[i][0]
+        reverse_index2 += meanface_indices_reversed[i][1]
+    return meanface_indices, reverse_index1, reverse_index2, max_len
+
+def compute_loss_pip(outputs_map, outputs_local_x, outputs_local_y, outputs_nb_x, outputs_nb_y, labels_map, labels_local_x, labels_local_y, labels_nb_x, labels_nb_y,  criterion_cls, criterion_reg, num_nb):
+
+    tmp_batch, tmp_channel, tmp_height, tmp_width = outputs_map.size()
+    labels_map = labels_map.view(tmp_batch*tmp_channel, -1)
+    labels_max_ids = torch.argmax(labels_map, 1)
+    labels_max_ids = labels_max_ids.view(-1, 1)
+    labels_max_ids_nb = labels_max_ids.repeat(1, num_nb).view(-1, 1)
+
+    outputs_local_x = outputs_local_x.view(tmp_batch*tmp_channel, -1)
+    outputs_local_x_select = torch.gather(outputs_local_x, 1, labels_max_ids)
+    outputs_local_y = outputs_local_y.view(tmp_batch*tmp_channel, -1)
+    outputs_local_y_select = torch.gather(outputs_local_y, 1, labels_max_ids)
+    outputs_nb_x = outputs_nb_x.view(tmp_batch*num_nb*tmp_channel, -1)
+    outputs_nb_x_select = torch.gather(outputs_nb_x, 1, labels_max_ids_nb)
+    outputs_nb_y = outputs_nb_y.view(tmp_batch*num_nb*tmp_channel, -1)
+    outputs_nb_y_select = torch.gather(outputs_nb_y, 1, labels_max_ids_nb)
+
+    labels_local_x = labels_local_x.view(tmp_batch*tmp_channel, -1)
+    labels_local_x_select = torch.gather(labels_local_x, 1, labels_max_ids)
+    labels_local_y = labels_local_y.view(tmp_batch*tmp_channel, -1)
+    labels_local_y_select = torch.gather(labels_local_y, 1, labels_max_ids)
+    labels_nb_x = labels_nb_x.view(tmp_batch*num_nb*tmp_channel, -1)
+    labels_nb_x_select = torch.gather(labels_nb_x, 1, labels_max_ids_nb)
+    labels_nb_y = labels_nb_y.view(tmp_batch*num_nb*tmp_channel, -1)
+    labels_nb_y_select = torch.gather(labels_nb_y, 1, labels_max_ids_nb)
+
+    labels_map = labels_map.view(tmp_batch, tmp_channel, tmp_height, tmp_width)
+    loss_map = criterion_cls(outputs_map, labels_map)
+    loss_x = criterion_reg(outputs_local_x_select, labels_local_x_select)
+    loss_y = criterion_reg(outputs_local_y_select, labels_local_y_select)
+    loss_nb_x = criterion_reg(outputs_nb_x_select, labels_nb_x_select)
+    loss_nb_y = criterion_reg(outputs_nb_y_select, labels_nb_y_select)
+    return loss_map, loss_x, loss_y, loss_nb_x, loss_nb_y
+
+def train_model(det_head, net, train_loader, criterion_cls, criterion_reg, cls_loss_weight, reg_loss_weight, num_nb, optimizer, num_epochs, scheduler, save_dir, save_interval, device):
+    for epoch in range(num_epochs):
+        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
+        logging.info('Epoch {}/{}'.format(epoch, num_epochs - 1))
+        print('-' * 10)
+        logging.info('-' * 10)
+        net.train()
+        epoch_loss = 0.0
+
+        for i, data in enumerate(train_loader):
+            if det_head == 'pip':
+                inputs, labels_map, labels_x, labels_y, labels_nb_x, labels_nb_y = data
+                inputs = inputs.to(device)
+                labels_map = labels_map.to(device)
+                labels_x = labels_x.to(device)
+                labels_y = labels_y.to(device)
+                labels_nb_x = labels_nb_x.to(device)
+                labels_nb_y = labels_nb_y.to(device)
+                outputs_map, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y = net(inputs)
+                loss_map, loss_x, loss_y, loss_nb_x, loss_nb_y = compute_loss_pip(outputs_map, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y, labels_map, labels_x, labels_y, labels_nb_x, labels_nb_y, criterion_cls, criterion_reg, num_nb)
+                loss = cls_loss_weight*loss_map + reg_loss_weight*loss_x + reg_loss_weight*loss_y + reg_loss_weight*loss_nb_x + reg_loss_weight*loss_nb_y
+            else:
+                print('No such head:', det_head)
+                exit(0)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if i%10 == 0:
+                if det_head == 'pip':
+                    print('[Epoch {:d}/{:d}, Batch {:d}/{:d}] <Total loss: {:.6f}> <map loss: {:.6f}> <x loss: {:.6f}> <y loss: {:.6f}> <nbx loss: {:.6f}> <nby loss: {:.6f}>'.format(
+                        epoch, num_epochs-1, i, len(train_loader)-1, loss.item(), cls_loss_weight*loss_map.item(), reg_loss_weight*loss_x.item(), reg_loss_weight*loss_y.item(), reg_loss_weight*loss_nb_x.item(), reg_loss_weight*loss_nb_y.item()))
+                    logging.info('[Epoch {:d}/{:d}, Batch {:d}/{:d}] <Total loss: {:.6f}> <map loss: {:.6f}> <x loss: {:.6f}> <y loss: {:.6f}> <nbx loss: {:.6f}> <nby loss: {:.6f}>'.format(
+                        epoch, num_epochs-1, i, len(train_loader)-1, loss.item(), cls_loss_weight*loss_map.item(), reg_loss_weight*loss_x.item(), reg_loss_weight*loss_y.item(), reg_loss_weight*loss_nb_x.item(), reg_loss_weight*loss_nb_y.item()))
+                else:
+                    print('No such head:', det_head)
+                    exit(0)
+            epoch_loss += loss.item()
+        epoch_loss /= len(train_loader)
+        if epoch%(save_interval-1) == 0 and epoch > 0:
+            filename = os.path.join(save_dir, 'epoch%d.pth' % epoch)
+            torch.save(net.state_dict(), filename)
+            print(filename, 'saved')
+        scheduler.step()
+    return net
+
+def forward_pip(net, inputs, preprocess, input_size, net_stride, num_nb):
+    net.eval()
+    with torch.no_grad():
+        outputs_cls, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y = net(inputs)
+        tmp_batch, tmp_channel, tmp_height, tmp_width = outputs_cls.size()
+        assert tmp_batch == 1
+
+        outputs_cls = outputs_cls.view(tmp_batch*tmp_channel, -1)
+        max_ids = torch.argmax(outputs_cls, 1)
+        max_cls = torch.max(outputs_cls, 1)[0]
+        max_ids = max_ids.view(-1, 1)
+        max_ids_nb = max_ids.repeat(1, num_nb).view(-1, 1)
+
+        outputs_x = outputs_x.view(tmp_batch*tmp_channel, -1)
+        outputs_x_select = torch.gather(outputs_x, 1, max_ids)
+        outputs_x_select = outputs_x_select.squeeze(1)
+        outputs_y = outputs_y.view(tmp_batch*tmp_channel, -1)
+        outputs_y_select = torch.gather(outputs_y, 1, max_ids)
+        outputs_y_select = outputs_y_select.squeeze(1)
+
+        outputs_nb_x = outputs_nb_x.view(tmp_batch*num_nb*tmp_channel, -1)
+        outputs_nb_x_select = torch.gather(outputs_nb_x, 1, max_ids_nb)
+        outputs_nb_x_select = outputs_nb_x_select.squeeze(1).view(-1, num_nb)
+        outputs_nb_y = outputs_nb_y.view(tmp_batch*num_nb*tmp_channel, -1)
+        outputs_nb_y_select = torch.gather(outputs_nb_y, 1, max_ids_nb)
+        outputs_nb_y_select = outputs_nb_y_select.squeeze(1).view(-1, num_nb)
+
+        tmp_x = (max_ids%tmp_width).view(-1,1).float()+outputs_x_select.view(-1,1)
+        tmp_y = (max_ids//tmp_width).view(-1,1).float()+outputs_y_select.view(-1,1)
+        tmp_x /= 1.0 * input_size / net_stride
+        tmp_y /= 1.0 * input_size / net_stride
+
+        tmp_nb_x = (max_ids%tmp_width).view(-1,1).float()+outputs_nb_x_select
+        tmp_nb_y = (max_ids//tmp_width).view(-1,1).float()+outputs_nb_y_select
+        tmp_nb_x = tmp_nb_x.view(-1, num_nb)
+        tmp_nb_y = tmp_nb_y.view(-1, num_nb)
+        tmp_nb_x /= 1.0 * input_size / net_stride
+        tmp_nb_y /= 1.0 * input_size / net_stride
+
+    return tmp_x, tmp_y, tmp_nb_x, tmp_nb_y, outputs_cls, max_cls
+
+def compute_nme(lms_pred, lms_gt, norm):
+    lms_pred = lms_pred.reshape((-1, 2))
+    lms_gt = lms_gt.reshape((-1, 2))
+    nme = np.mean(np.linalg.norm(lms_pred - lms_gt, axis=1)) / norm 
+    return nme
+
+def compute_fr_and_auc(nmes, thres=0.1, step=0.0001):
+    num_data = len(nmes)
+    xs = np.arange(0, thres + step, step)
+    ys = np.array([np.count_nonzero(nmes <= x) for x in xs]) / float(num_data)
+    fr = 1.0 - ys[-1]
+    auc = simps(ys, x=xs) / thres
+    return fr, auc

third_party/PIPNet/lib/functions_gssl.py

@@ -0,0 +1,241 @@
+import os, cv2
+import numpy as np
+from PIL import Image, ImageFilter
+import logging
+import torch
+import torch.nn as nn
+import random
+
+def get_label(data_name, label_file, task_type=None):
+    label_path = os.path.join('data', data_name, label_file)
+    with open(label_path, 'r') as f:
+        labels = f.readlines()
+    labels = [x.strip().split() for x in labels]
+    if len(labels[0])==1:
+        return labels
+
+    labels_new = []
+    for label in labels:
+        image_name = label[0]
+        target = label[1:]
+        target = np.array([float(x) for x in target])
+        if task_type is None:
+            labels_new.append([image_name, target])
+        else:
+            labels_new.append([image_name, task_type, target])
+    return labels_new
+
+def get_meanface(meanface_file, num_nb):
+    with open(meanface_file) as f:
+        meanface = f.readlines()[0]
+        
+    meanface = meanface.strip().split()
+    meanface = [float(x) for x in meanface]
+    meanface = np.array(meanface).reshape(-1, 2)
+    # each landmark predicts num_nb neighbors
+    meanface_indices = []
+    for i in range(meanface.shape[0]):
+        pt = meanface[i,:]
+        dists = np.sum(np.power(pt-meanface, 2), axis=1)
+        indices = np.argsort(dists)
+        meanface_indices.append(indices[1:1+num_nb])
+    
+    # each landmark predicted by X neighbors, X varies
+    meanface_indices_reversed = {}
+    for i in range(meanface.shape[0]):
+        meanface_indices_reversed[i] = [[],[]]
+    for i in range(meanface.shape[0]):
+        for j in range(num_nb):
+            meanface_indices_reversed[meanface_indices[i][j]][0].append(i)
+            meanface_indices_reversed[meanface_indices[i][j]][1].append(j)
+    
+    max_len = 0
+    for i in range(meanface.shape[0]):
+        tmp_len = len(meanface_indices_reversed[i][0])
+        if tmp_len > max_len:
+            max_len = tmp_len
+    
+    # tricks, make them have equal length for efficient computation
+    for i in range(meanface.shape[0]):
+        tmp_len = len(meanface_indices_reversed[i][0])
+        meanface_indices_reversed[i][0] += meanface_indices_reversed[i][0]*10
+        meanface_indices_reversed[i][1] += meanface_indices_reversed[i][1]*10
+        meanface_indices_reversed[i][0] = meanface_indices_reversed[i][0][:max_len]
+        meanface_indices_reversed[i][1] = meanface_indices_reversed[i][1][:max_len]
+
+    # make the indices 1-dim
+    reverse_index1 = []
+    reverse_index2 = []
+    for i in range(meanface.shape[0]):
+        reverse_index1 += meanface_indices_reversed[i][0]
+        reverse_index2 += meanface_indices_reversed[i][1]
+    return meanface_indices, reverse_index1, reverse_index2, max_len
+
+def compute_loss_pip(outputs_map1, outputs_map2, outputs_map3, outputs_local_x, outputs_local_y, outputs_nb_x, outputs_nb_y, labels_map1, labels_map2, labels_map3, labels_local_x, labels_local_y, labels_nb_x, labels_nb_y, masks_map1, masks_map2, masks_map3, masks_local_x, masks_local_y, masks_nb_x, masks_nb_y, criterion_cls, criterion_reg, num_nb):
+
+    tmp_batch, tmp_channel, tmp_height, tmp_width = outputs_map1.size()
+    labels_map1 = labels_map1.view(tmp_batch*tmp_channel, -1)
+    labels_max_ids = torch.argmax(labels_map1, 1)
+    labels_max_ids = labels_max_ids.view(-1, 1)
+    labels_max_ids_nb = labels_max_ids.repeat(1, num_nb).view(-1, 1)
+
+    outputs_local_x = outputs_local_x.view(tmp_batch*tmp_channel, -1)
+    outputs_local_x_select = torch.gather(outputs_local_x, 1, labels_max_ids)
+    outputs_local_y = outputs_local_y.view(tmp_batch*tmp_channel, -1)
+    outputs_local_y_select = torch.gather(outputs_local_y, 1, labels_max_ids)
+    outputs_nb_x = outputs_nb_x.view(tmp_batch*num_nb*tmp_channel, -1)
+    outputs_nb_x_select = torch.gather(outputs_nb_x, 1, labels_max_ids_nb)
+    outputs_nb_y = outputs_nb_y.view(tmp_batch*num_nb*tmp_channel, -1)
+    outputs_nb_y_select = torch.gather(outputs_nb_y, 1, labels_max_ids_nb)
+
+    labels_local_x = labels_local_x.view(tmp_batch*tmp_channel, -1)
+    labels_local_x_select = torch.gather(labels_local_x, 1, labels_max_ids)
+    labels_local_y = labels_local_y.view(tmp_batch*tmp_channel, -1)
+    labels_local_y_select = torch.gather(labels_local_y, 1, labels_max_ids)
+    labels_nb_x = labels_nb_x.view(tmp_batch*num_nb*tmp_channel, -1)
+    labels_nb_x_select = torch.gather(labels_nb_x, 1, labels_max_ids_nb)
+    labels_nb_y = labels_nb_y.view(tmp_batch*num_nb*tmp_channel, -1)
+    labels_nb_y_select = torch.gather(labels_nb_y, 1, labels_max_ids_nb)
+
+    masks_local_x = masks_local_x.view(tmp_batch*tmp_channel, -1)
+    masks_local_x_select = torch.gather(masks_local_x, 1, labels_max_ids)
+    masks_local_y = masks_local_y.view(tmp_batch*tmp_channel, -1)
+    masks_local_y_select = torch.gather(masks_local_y, 1, labels_max_ids)
+    masks_nb_x = masks_nb_x.view(tmp_batch*num_nb*tmp_channel, -1)
+    masks_nb_x_select = torch.gather(masks_nb_x, 1, labels_max_ids_nb)
+    masks_nb_y = masks_nb_y.view(tmp_batch*num_nb*tmp_channel, -1)
+    masks_nb_y_select = torch.gather(masks_nb_y, 1, labels_max_ids_nb)
+
+    ##########################################
+    outputs_map1 = outputs_map1.view(tmp_batch*tmp_channel, -1)
+    outputs_map2 = outputs_map2.view(tmp_batch*tmp_channel, -1)
+    outputs_map3 = outputs_map3.view(tmp_batch*tmp_channel, -1)
+    labels_map2 = labels_map2.view(tmp_batch*tmp_channel, -1)
+    labels_map3 = labels_map3.view(tmp_batch*tmp_channel, -1)
+    masks_map1 = masks_map1.view(tmp_batch*tmp_channel, -1)
+    masks_map2 = masks_map2.view(tmp_batch*tmp_channel, -1)
+    masks_map3 = masks_map3.view(tmp_batch*tmp_channel, -1)
+    outputs_map = torch.cat([outputs_map1, outputs_map2, outputs_map3], 1)
+    labels_map = torch.cat([labels_map1, labels_map2, labels_map3], 1)
+    masks_map = torch.cat([masks_map1, masks_map2, masks_map3], 1)
+    loss_map = criterion_cls(outputs_map*masks_map, labels_map*masks_map)
+    if not masks_map.sum() == 0:
+        loss_map /= masks_map.sum()
+    ##########################################
+
+    loss_x = criterion_reg(outputs_local_x_select*masks_local_x_select, labels_local_x_select*masks_local_x_select)
+    if not masks_local_x_select.sum() == 0:
+        loss_x /= masks_local_x_select.sum()
+    loss_y = criterion_reg(outputs_local_y_select*masks_local_y_select, labels_local_y_select*masks_local_y_select)
+    if not masks_local_y_select.sum() == 0:
+        loss_y /= masks_local_y_select.sum()
+    loss_nb_x = criterion_reg(outputs_nb_x_select*masks_nb_x_select, labels_nb_x_select*masks_nb_x_select)
+    if not masks_nb_x_select.sum() == 0:
+        loss_nb_x /= masks_nb_x_select.sum()
+    loss_nb_y = criterion_reg(outputs_nb_y_select*masks_nb_y_select, labels_nb_y_select*masks_nb_y_select)
+    if not masks_nb_y_select.sum() == 0:
+        loss_nb_y /= masks_nb_y_select.sum()
+    return loss_map, loss_x, loss_y, loss_nb_x, loss_nb_y
+
+def train_model(det_head, net, train_loader, criterion_cls, criterion_reg, cls_loss_weight, reg_loss_weight, num_nb, optimizer, num_epochs, scheduler, save_dir, save_interval, device):
+    for epoch in range(num_epochs):
+        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
+        logging.info('Epoch {}/{}'.format(epoch, num_epochs - 1))
+        print('-' * 10)
+        logging.info('-' * 10)
+        net.train()
+        epoch_loss = 0.0
+
+        for i, data in enumerate(train_loader):
+            if det_head == 'pip':
+                inputs, labels_map1, labels_map2, labels_map3, labels_x, labels_y, labels_nb_x, labels_nb_y, masks_map1, masks_map2, masks_map3, masks_x, masks_y, masks_nb_x, masks_nb_y = data
+                inputs = inputs.to(device)
+                labels_map1 = labels_map1.to(device)
+                labels_map2 = labels_map2.to(device)
+                labels_map3 = labels_map3.to(device)
+                labels_x = labels_x.to(device)
+                labels_y = labels_y.to(device)
+                labels_nb_x = labels_nb_x.to(device)
+                labels_nb_y = labels_nb_y.to(device)
+                masks_map1 = masks_map1.to(device)
+                masks_map2 = masks_map2.to(device)
+                masks_map3 = masks_map3.to(device)
+                masks_x = masks_x.to(device)
+                masks_y = masks_y.to(device)
+                masks_nb_x = masks_nb_x.to(device)
+                masks_nb_y = masks_nb_y.to(device)
+                outputs_map1, outputs_map2, outputs_map3, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y = net(inputs)
+                loss_map, loss_x, loss_y, loss_nb_x, loss_nb_y = compute_loss_pip(outputs_map1, outputs_map2, outputs_map3, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y, labels_map1, labels_map2, labels_map3, labels_x, labels_y, labels_nb_x, labels_nb_y, masks_map1, masks_map2, masks_map3, masks_x, masks_y, masks_nb_x, masks_nb_y, criterion_cls, criterion_reg, num_nb)
+                loss = cls_loss_weight*loss_map + reg_loss_weight*loss_x + reg_loss_weight*loss_y + reg_loss_weight*loss_nb_x + reg_loss_weight*loss_nb_y
+            else:
+                print('No such head:', det_head)
+                exit(0)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if i%10 == 0:
+                if det_head == 'pip':
+                    print('[Epoch {:d}/{:d}, Batch {:d}/{:d}] <Total loss: {:.6f}> <map loss: {:.6f}> <x loss: {:.6f}> <y loss: {:.6f}> <nbx loss: {:.6f}> <nby loss: {:.6f}>'.format(
+                        epoch, num_epochs-1, i, len(train_loader)-1, loss.item(), cls_loss_weight*loss_map.item(), reg_loss_weight*loss_x.item(), reg_loss_weight*loss_y.item(), reg_loss_weight*loss_nb_x.item(), reg_loss_weight*loss_nb_y.item()))
+                    logging.info('[Epoch {:d}/{:d}, Batch {:d}/{:d}] <Total loss: {:.6f}> <map loss: {:.6f}> <x loss: {:.6f}> <y loss: {:.6f}> <nbx loss: {:.6f}> <nby loss: {:.6f}>'.format(
+                        epoch, num_epochs-1, i, len(train_loader)-1, loss.item(), cls_loss_weight*loss_map.item(), reg_loss_weight*loss_x.item(), reg_loss_weight*loss_y.item(), reg_loss_weight*loss_nb_x.item(), reg_loss_weight*loss_nb_y.item()))
+                else:
+                    print('No such head:', det_head)
+                    exit(0)
+            epoch_loss += loss.item()
+        epoch_loss /= len(train_loader)
+        if epoch%(save_interval-1) == 0 and epoch > 0:
+            filename = os.path.join(save_dir, 'epoch%d.pth' % epoch)
+            torch.save(net.state_dict(), filename)
+            print(filename, 'saved')
+        scheduler.step()
+    return net
+
+def forward_pip(net, inputs, preprocess, input_size, net_stride, num_nb):
+    net.eval()
+    with torch.no_grad():
+        outputs_cls1, outputs_cls2, outputs_cls3, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y = net(inputs)
+        tmp_batch, tmp_channel, tmp_height, tmp_width = outputs_cls1.size()
+        assert tmp_batch == 1
+        
+        outputs_cls1 = outputs_cls1.view(tmp_batch*tmp_channel, -1)
+        max_ids = torch.argmax(outputs_cls1, 1)
+        max_cls = torch.max(outputs_cls1, 1)[0]
+        max_ids = max_ids.view(-1, 1)
+        max_ids_nb = max_ids.repeat(1, num_nb).view(-1, 1)
+
+        outputs_x = outputs_x.view(tmp_batch*tmp_channel, -1)
+        outputs_x_select = torch.gather(outputs_x, 1, max_ids)
+        outputs_x_select = outputs_x_select.squeeze(1)
+        outputs_y = outputs_y.view(tmp_batch*tmp_channel, -1)
+        outputs_y_select = torch.gather(outputs_y, 1, max_ids)
+        outputs_y_select = outputs_y_select.squeeze(1)
+
+        outputs_nb_x = outputs_nb_x.view(tmp_batch*num_nb*tmp_channel, -1)
+        outputs_nb_x_select = torch.gather(outputs_nb_x, 1, max_ids_nb)
+        outputs_nb_x_select = outputs_nb_x_select.squeeze(1).view(-1, num_nb)
+        outputs_nb_y = outputs_nb_y.view(tmp_batch*num_nb*tmp_channel, -1)
+        outputs_nb_y_select = torch.gather(outputs_nb_y, 1, max_ids_nb)
+        outputs_nb_y_select = outputs_nb_y_select.squeeze(1).view(-1, num_nb)
+
+        tmp_x = (max_ids%tmp_width).view(-1,1).float()+outputs_x_select.view(-1,1)
+        tmp_y = (max_ids//tmp_width).view(-1,1).float()+outputs_y_select.view(-1,1)
+        tmp_x /= 1.0 * input_size / net_stride
+        tmp_y /= 1.0 * input_size / net_stride
+
+        tmp_nb_x = (max_ids%tmp_width).view(-1,1).float()+outputs_nb_x_select
+        tmp_nb_y = (max_ids//tmp_width).view(-1,1).float()+outputs_nb_y_select
+        tmp_nb_x = tmp_nb_x.view(-1, num_nb)
+        tmp_nb_y = tmp_nb_y.view(-1, num_nb)
+        tmp_nb_x /= 1.0 * input_size / net_stride
+        tmp_nb_y /= 1.0 * input_size / net_stride
+
+    return tmp_x, tmp_y, tmp_nb_x, tmp_nb_y, [outputs_cls1, outputs_cls2, outputs_cls3], max_cls
+
+def compute_nme(lms_pred, lms_gt, norm):
+    lms_pred = lms_pred.reshape((-1, 2))
+    lms_gt = lms_gt.reshape((-1, 2))
+    nme = np.mean(np.linalg.norm(lms_pred - lms_gt, axis=1)) / norm 
+    return nme
+

third_party/PIPNet/lib/mobilenetv3.py

@@ -0,0 +1,233 @@
+"""
+Creates a MobileNetV3 Model as defined in:
+Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, Quoc V. Le, Hartwig Adam. (2019).
+Searching for MobileNetV3
+arXiv preprint arXiv:1905.02244.
+"""
+
+import torch
+import torch.nn as nn
+import math
+
+
+__all__ = ['mobilenetv3_large', 'mobilenetv3_small']
+
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class h_sigmoid(nn.Module):
+    def __init__(self, inplace=True):
+        super(h_sigmoid, self).__init__()
+        self.relu = nn.ReLU6(inplace=inplace)
+
+    def forward(self, x):
+        return self.relu(x + 3) / 6
+
+
+class h_swish(nn.Module):
+    def __init__(self, inplace=True):
+        super(h_swish, self).__init__()
+        self.sigmoid = h_sigmoid(inplace=inplace)
+
+    def forward(self, x):
+        return x * self.sigmoid(x)
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=4):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+                nn.Linear(channel, _make_divisible(channel // reduction, 8)),
+                nn.ReLU(inplace=True),
+                nn.Linear(_make_divisible(channel // reduction, 8), channel),
+                h_sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+
+
+def conv_3x3_bn(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        h_swish()
+    )
+
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        h_swish()
+    )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, hidden_dim, oup, kernel_size, stride, use_se, use_hs):
+        super(InvertedResidual, self).__init__()
+        assert stride in [1, 2]
+
+        self.identity = stride == 1 and inp == oup
+
+        if inp == hidden_dim:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, (kernel_size - 1) // 2, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                h_swish() if use_hs else nn.ReLU(inplace=True),
+                # Squeeze-and-Excite
+                SELayer(hidden_dim) if use_se else nn.Identity(),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                h_swish() if use_hs else nn.ReLU(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, (kernel_size - 1) // 2, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                # Squeeze-and-Excite
+                SELayer(hidden_dim) if use_se else nn.Identity(),
+                h_swish() if use_hs else nn.ReLU(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.identity:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV3(nn.Module):
+    def __init__(self, cfgs, mode, num_classes=1000, width_mult=1.):
+        super(MobileNetV3, self).__init__()
+        # setting of inverted residual blocks
+        self.cfgs = cfgs
+        assert mode in ['large', 'small']
+
+        # building first layer
+        input_channel = _make_divisible(16 * width_mult, 8)
+        layers = [conv_3x3_bn(3, input_channel, 2)]
+        # building inverted residual blocks
+        block = InvertedResidual
+        for k, t, c, use_se, use_hs, s in self.cfgs:
+            output_channel = _make_divisible(c * width_mult, 8)
+            exp_size = _make_divisible(input_channel * t, 8)
+            layers.append(block(input_channel, exp_size, output_channel, k, s, use_se, use_hs))
+            input_channel = output_channel
+        self.features = nn.Sequential(*layers)
+        # building last several layers
+        self.conv = conv_1x1_bn(input_channel, exp_size)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        output_channel = {'large': 1280, 'small': 1024}
+        output_channel = _make_divisible(output_channel[mode] * width_mult, 8) if width_mult > 1.0 else output_channel[mode]
+        self.classifier = nn.Sequential(
+            nn.Linear(exp_size, output_channel),
+            h_swish(),
+            nn.Dropout(0.2),
+            nn.Linear(output_channel, num_classes),
+        )
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.conv(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+
+
+def mobilenetv3_large(**kwargs):
+    """
+    Constructs a MobileNetV3-Large model
+    """
+    cfgs = [
+        # k, t, c, SE, HS, s 
+        [3,   1,  16, 0, 0, 1],
+        [3,   4,  24, 0, 0, 2],
+        [3,   3,  24, 0, 0, 1],
+        [5,   3,  40, 1, 0, 2],
+        [5,   3,  40, 1, 0, 1],
+        [5,   3,  40, 1, 0, 1],
+        [3,   6,  80, 0, 1, 2],
+        [3, 2.5,  80, 0, 1, 1],
+        [3, 2.3,  80, 0, 1, 1],
+        [3, 2.3,  80, 0, 1, 1],
+        [3,   6, 112, 1, 1, 1],
+        [3,   6, 112, 1, 1, 1],
+        [5,   6, 160, 1, 1, 2],
+        [5,   6, 160, 1, 1, 1],
+        [5,   6, 160, 1, 1, 1]
+    ]
+    return MobileNetV3(cfgs, mode='large', **kwargs)
+
+
+def mobilenetv3_small(**kwargs):
+    """
+    Constructs a MobileNetV3-Small model
+    """
+    cfgs = [
+        # k, t, c, SE, HS, s 
+        [3,    1,  16, 1, 0, 2],
+        [3,  4.5,  24, 0, 0, 2],
+        [3, 3.67,  24, 0, 0, 1],
+        [5,    4,  40, 1, 1, 2],
+        [5,    6,  40, 1, 1, 1],
+        [5,    6,  40, 1, 1, 1],
+        [5,    3,  48, 1, 1, 1],
+        [5,    3,  48, 1, 1, 1],
+        [5,    6,  96, 1, 1, 2],
+        [5,    6,  96, 1, 1, 1],
+        [5,    6,  96, 1, 1, 1],
+    ]
+
+    return MobileNetV3(cfgs, mode='small', **kwargs)
+
+
+

third_party/PIPNet/lib/networks.py

@@ -0,0 +1,415 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import numpy as np
+
+# net_stride output_size
+# 128        2x2
+# 64         4x4
+# 32         8x8
+# pip regression, resnet101
+class Pip_resnet101(nn.Module):
+    def __init__(self, resnet, num_nb, num_lms=68, input_size=256, net_stride=32):
+        super(Pip_resnet101, self).__init__()
+        self.num_nb = num_nb
+        self.num_lms = num_lms
+        self.input_size = input_size
+        self.net_stride = net_stride
+        self.conv1 = resnet.conv1
+        self.bn1 = resnet.bn1
+        self.maxpool = resnet.maxpool
+        self.sigmoid = nn.Sigmoid()
+        self.layer1 = resnet.layer1
+        self.layer2 = resnet.layer2
+        self.layer3 = resnet.layer3
+        self.layer4 = resnet.layer4
+        if self.net_stride == 128:
+            self.layer5 = nn.Conv2d(2048, 512, kernel_size=3, stride=2, padding=1)
+            self.bn5 = nn.BatchNorm2d(512)
+            self.layer6 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)
+            self.bn6 = nn.BatchNorm2d(512)
+            # init
+            nn.init.normal_(self.layer5.weight, std=0.001)
+            if self.layer5.bias is not None:
+                nn.init.constant_(self.layer5.bias, 0)
+            nn.init.constant_(self.bn5.weight, 1)
+            nn.init.constant_(self.bn5.bias, 0)
+
+            nn.init.normal_(self.layer6.weight, std=0.001)
+            if self.layer6.bias is not None:
+                nn.init.constant_(self.layer6.bias, 0)
+            nn.init.constant_(self.bn6.weight, 1)
+            nn.init.constant_(self.bn6.bias, 0)
+        elif self.net_stride == 64:
+            self.layer5 = nn.Conv2d(2048, 512, kernel_size=3, stride=2, padding=1)
+            self.bn5 = nn.BatchNorm2d(512)
+            # init
+            nn.init.normal_(self.layer5.weight, std=0.001)
+            if self.layer5.bias is not None:
+                nn.init.constant_(self.layer5.bias, 0)
+            nn.init.constant_(self.bn5.weight, 1)
+            nn.init.constant_(self.bn5.bias, 0)
+        elif self.net_stride == 32:
+            pass
+        else:
+            print('No such net_stride!')
+            exit(0)
+
+        self.cls_layer = nn.Conv2d(2048, num_lms, kernel_size=1, stride=1, padding=0)
+        self.x_layer = nn.Conv2d(2048, num_lms, kernel_size=1, stride=1, padding=0)
+        self.y_layer = nn.Conv2d(2048, num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_x_layer = nn.Conv2d(2048, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_y_layer = nn.Conv2d(2048, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+
+        nn.init.normal_(self.cls_layer.weight, std=0.001)
+        if self.cls_layer.bias is not None:
+            nn.init.constant_(self.cls_layer.bias, 0)
+
+        nn.init.normal_(self.x_layer.weight, std=0.001)
+        if self.x_layer.bias is not None:
+            nn.init.constant_(self.x_layer.bias, 0)
+
+        nn.init.normal_(self.y_layer.weight, std=0.001)
+        if self.y_layer.bias is not None:
+            nn.init.constant_(self.y_layer.bias, 0)
+
+        nn.init.normal_(self.nb_x_layer.weight, std=0.001)
+        if self.nb_x_layer.bias is not None:
+            nn.init.constant_(self.nb_x_layer.bias, 0)
+
+        nn.init.normal_(self.nb_y_layer.weight, std=0.001)
+        if self.nb_y_layer.bias is not None:
+            nn.init.constant_(self.nb_y_layer.bias, 0)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = F.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        if self.net_stride == 128:
+            x = F.relu(self.bn5(self.layer5(x)))
+            x = F.relu(self.bn6(self.layer6(x)))
+        elif self.net_stride == 64:
+            x = F.relu(self.bn5(self.layer5(x)))
+        else:
+            pass
+        x1 = self.cls_layer(x)
+        x2 = self.x_layer(x)
+        x3 = self.y_layer(x)
+        x4 = self.nb_x_layer(x)
+        x5 = self.nb_y_layer(x)
+        return x1, x2, x3, x4, x5
+
+# net_stride output_size
+# 128        2x2
+# 64         4x4
+# 32         8x8
+# pip regression, resnet50
+class Pip_resnet50(nn.Module):
+    def __init__(self, resnet, num_nb, num_lms=68, input_size=256, net_stride=32):
+        super(Pip_resnet50, self).__init__()
+        self.num_nb = num_nb
+        self.num_lms = num_lms
+        self.input_size = input_size
+        self.net_stride = net_stride
+        self.conv1 = resnet.conv1
+        self.bn1 = resnet.bn1
+        self.maxpool = resnet.maxpool
+        self.sigmoid = nn.Sigmoid()
+        self.layer1 = resnet.layer1
+        self.layer2 = resnet.layer2
+        self.layer3 = resnet.layer3
+        self.layer4 = resnet.layer4
+        if self.net_stride == 128:
+            self.layer5 = nn.Conv2d(2048, 512, kernel_size=3, stride=2, padding=1)
+            self.bn5 = nn.BatchNorm2d(512)
+            self.layer6 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)
+            self.bn6 = nn.BatchNorm2d(512)
+            # init
+            nn.init.normal_(self.layer5.weight, std=0.001)
+            if self.layer5.bias is not None:
+                nn.init.constant_(self.layer5.bias, 0)
+            nn.init.constant_(self.bn5.weight, 1)
+            nn.init.constant_(self.bn5.bias, 0)
+
+            nn.init.normal_(self.layer6.weight, std=0.001)
+            if self.layer6.bias is not None:
+                nn.init.constant_(self.layer6.bias, 0)
+            nn.init.constant_(self.bn6.weight, 1)
+            nn.init.constant_(self.bn6.bias, 0)
+        elif self.net_stride == 64:
+            self.layer5 = nn.Conv2d(2048, 512, kernel_size=3, stride=2, padding=1)
+            self.bn5 = nn.BatchNorm2d(512)
+            # init
+            nn.init.normal_(self.layer5.weight, std=0.001)
+            if self.layer5.bias is not None:
+                nn.init.constant_(self.layer5.bias, 0)
+            nn.init.constant_(self.bn5.weight, 1)
+            nn.init.constant_(self.bn5.bias, 0)
+        elif self.net_stride == 32:
+            pass
+        else:
+            print('No such net_stride!')
+            exit(0)
+
+        self.cls_layer = nn.Conv2d(2048, num_lms, kernel_size=1, stride=1, padding=0)
+        self.x_layer = nn.Conv2d(2048, num_lms, kernel_size=1, stride=1, padding=0)
+        self.y_layer = nn.Conv2d(2048, num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_x_layer = nn.Conv2d(2048, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_y_layer = nn.Conv2d(2048, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+
+        nn.init.normal_(self.cls_layer.weight, std=0.001)
+        if self.cls_layer.bias is not None:
+            nn.init.constant_(self.cls_layer.bias, 0)
+
+        nn.init.normal_(self.x_layer.weight, std=0.001)
+        if self.x_layer.bias is not None:
+            nn.init.constant_(self.x_layer.bias, 0)
+
+        nn.init.normal_(self.y_layer.weight, std=0.001)
+        if self.y_layer.bias is not None:
+            nn.init.constant_(self.y_layer.bias, 0)
+
+        nn.init.normal_(self.nb_x_layer.weight, std=0.001)
+        if self.nb_x_layer.bias is not None:
+            nn.init.constant_(self.nb_x_layer.bias, 0)
+
+        nn.init.normal_(self.nb_y_layer.weight, std=0.001)
+        if self.nb_y_layer.bias is not None:
+            nn.init.constant_(self.nb_y_layer.bias, 0)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = F.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        if self.net_stride == 128:
+            x = F.relu(self.bn5(self.layer5(x)))
+            x = F.relu(self.bn6(self.layer6(x)))
+        elif self.net_stride == 64:
+            x = F.relu(self.bn5(self.layer5(x)))
+        else:
+            pass
+        x1 = self.cls_layer(x)
+        x2 = self.x_layer(x)
+        x3 = self.y_layer(x)
+        x4 = self.nb_x_layer(x)
+        x5 = self.nb_y_layer(x)
+        return x1, x2, x3, x4, x5
+
+# net_stride output_size
+# 128        2x2
+# 64         4x4
+# 32         8x8
+# pip regression, resnet18
+class Pip_resnet18(nn.Module):
+    def __init__(self, resnet, num_nb, num_lms=68, input_size=256, net_stride=32):
+        super(Pip_resnet18, self).__init__()
+        self.num_nb = num_nb
+        self.num_lms = num_lms
+        self.input_size = input_size
+        self.net_stride = net_stride
+        self.conv1 = resnet.conv1
+        self.bn1 = resnet.bn1
+        self.maxpool = resnet.maxpool
+        self.sigmoid = nn.Sigmoid()
+        self.layer1 = resnet.layer1
+        self.layer2 = resnet.layer2
+        self.layer3 = resnet.layer3
+        self.layer4 = resnet.layer4
+        if self.net_stride == 128:
+            self.layer5 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)
+            self.bn5 = nn.BatchNorm2d(512)
+            self.layer6 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)
+            self.bn6 = nn.BatchNorm2d(512)
+            # init
+            nn.init.normal_(self.layer5.weight, std=0.001)
+            if self.layer5.bias is not None:
+                nn.init.constant_(self.layer5.bias, 0)
+            nn.init.constant_(self.bn5.weight, 1)
+            nn.init.constant_(self.bn5.bias, 0)
+
+            nn.init.normal_(self.layer6.weight, std=0.001)
+            if self.layer6.bias is not None:
+                nn.init.constant_(self.layer6.bias, 0)
+            nn.init.constant_(self.bn6.weight, 1)
+            nn.init.constant_(self.bn6.bias, 0)
+        elif self.net_stride == 64:
+            self.layer5 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)
+            self.bn5 = nn.BatchNorm2d(512)
+            # init
+            nn.init.normal_(self.layer5.weight, std=0.001)
+            if self.layer5.bias is not None:
+                nn.init.constant_(self.layer5.bias, 0)
+            nn.init.constant_(self.bn5.weight, 1)
+            nn.init.constant_(self.bn5.bias, 0)
+        elif self.net_stride == 32:
+            pass
+        elif self.net_stride == 16:
+            self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=4, stride=2, padding=1, bias=False)
+            self.bn_deconv1 = nn.BatchNorm2d(512)
+            nn.init.normal_(self.deconv1.weight, std=0.001)
+            if self.deconv1.bias is not None:
+                nn.init.constant_(self.deconv1.bias, 0)
+            nn.init.constant_(self.bn_deconv1.weight, 1)
+            nn.init.constant_(self.bn_deconv1.bias, 0)
+        else:
+            print('No such net_stride!')
+            exit(0)
+
+        self.cls_layer = nn.Conv2d(512, num_lms, kernel_size=1, stride=1, padding=0)
+        self.x_layer = nn.Conv2d(512, num_lms, kernel_size=1, stride=1, padding=0)
+        self.y_layer = nn.Conv2d(512, num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_x_layer = nn.Conv2d(512, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_y_layer = nn.Conv2d(512, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+
+        nn.init.normal_(self.cls_layer.weight, std=0.001)
+        if self.cls_layer.bias is not None:
+            nn.init.constant_(self.cls_layer.bias, 0)
+
+        nn.init.normal_(self.x_layer.weight, std=0.001)
+        if self.x_layer.bias is not None:
+            nn.init.constant_(self.x_layer.bias, 0)
+
+        nn.init.normal_(self.y_layer.weight, std=0.001)
+        if self.y_layer.bias is not None:
+            nn.init.constant_(self.y_layer.bias, 0)
+
+        nn.init.normal_(self.nb_x_layer.weight, std=0.001)
+        if self.nb_x_layer.bias is not None:
+            nn.init.constant_(self.nb_x_layer.bias, 0)
+
+        nn.init.normal_(self.nb_y_layer.weight, std=0.001)
+        if self.nb_y_layer.bias is not None:
+            nn.init.constant_(self.nb_y_layer.bias, 0)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = F.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        if self.net_stride == 128:
+            x = F.relu(self.bn5(self.layer5(x)))
+            x = F.relu(self.bn6(self.layer6(x)))
+        elif self.net_stride == 64:
+            x = F.relu(self.bn5(self.layer5(x)))
+        elif self.net_stride == 16:
+            x = F.relu(self.bn_deconv1(self.deconv1(x)))
+        else:
+            pass
+        x1 = self.cls_layer(x)
+        x2 = self.x_layer(x)
+        x3 = self.y_layer(x)
+        x4 = self.nb_x_layer(x)
+        x5 = self.nb_y_layer(x)
+        return x1, x2, x3, x4, x5
+
+class Pip_mbnetv2(nn.Module):
+    def __init__(self, mbnet, num_nb, num_lms=68, input_size=256, net_stride=32):
+        super(Pip_mbnetv2, self).__init__()
+        self.num_nb = num_nb
+        self.num_lms = num_lms
+        self.input_size = input_size
+        self.net_stride = net_stride
+        self.features = mbnet.features
+        self.sigmoid = nn.Sigmoid()
+
+        self.cls_layer = nn.Conv2d(1280, num_lms, kernel_size=1, stride=1, padding=0)
+        self.x_layer = nn.Conv2d(1280, num_lms, kernel_size=1, stride=1, padding=0)
+        self.y_layer = nn.Conv2d(1280, num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_x_layer = nn.Conv2d(1280, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_y_layer = nn.Conv2d(1280, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+
+        nn.init.normal_(self.cls_layer.weight, std=0.001)
+        if self.cls_layer.bias is not None:
+            nn.init.constant_(self.cls_layer.bias, 0)
+
+        nn.init.normal_(self.x_layer.weight, std=0.001)
+        if self.x_layer.bias is not None:
+            nn.init.constant_(self.x_layer.bias, 0)
+
+        nn.init.normal_(self.y_layer.weight, std=0.001)
+        if self.y_layer.bias is not None:
+            nn.init.constant_(self.y_layer.bias, 0)
+
+        nn.init.normal_(self.nb_x_layer.weight, std=0.001)
+        if self.nb_x_layer.bias is not None:
+            nn.init.constant_(self.nb_x_layer.bias, 0)
+
+        nn.init.normal_(self.nb_y_layer.weight, std=0.001)
+        if self.nb_y_layer.bias is not None:
+            nn.init.constant_(self.nb_y_layer.bias, 0)
+
+    def forward(self, x):
+        x = self.features(x)
+        x1 = self.cls_layer(x)
+        x2 = self.x_layer(x)
+        x3 = self.y_layer(x)
+        x4 = self.nb_x_layer(x)
+        x5 = self.nb_y_layer(x)
+        return x1, x2, x3, x4, x5
+
+class Pip_mbnetv3(nn.Module):
+    def __init__(self, mbnet, num_nb, num_lms=68, input_size=256, net_stride=32):
+        super(Pip_mbnetv3, self).__init__()
+        self.num_nb = num_nb
+        self.num_lms = num_lms
+        self.input_size = input_size
+        self.net_stride = net_stride
+        self.features = mbnet.features
+        self.conv = mbnet.conv
+        self.sigmoid = nn.Sigmoid()
+
+        self.cls_layer = nn.Conv2d(960, num_lms, kernel_size=1, stride=1, padding=0)
+        self.x_layer = nn.Conv2d(960, num_lms, kernel_size=1, stride=1, padding=0)
+        self.y_layer = nn.Conv2d(960, num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_x_layer = nn.Conv2d(960, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_y_layer = nn.Conv2d(960, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+
+        nn.init.normal_(self.cls_layer.weight, std=0.001)
+        if self.cls_layer.bias is not None:
+            nn.init.constant_(self.cls_layer.bias, 0)
+
+        nn.init.normal_(self.x_layer.weight, std=0.001)
+        if self.x_layer.bias is not None:
+            nn.init.constant_(self.x_layer.bias, 0)
+
+        nn.init.normal_(self.y_layer.weight, std=0.001)
+        if self.y_layer.bias is not None:
+            nn.init.constant_(self.y_layer.bias, 0)
+
+        nn.init.normal_(self.nb_x_layer.weight, std=0.001)
+        if self.nb_x_layer.bias is not None:
+            nn.init.constant_(self.nb_x_layer.bias, 0)
+
+        nn.init.normal_(self.nb_y_layer.weight, std=0.001)
+        if self.nb_y_layer.bias is not None:
+            nn.init.constant_(self.nb_y_layer.bias, 0)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.conv(x)
+        x1 = self.cls_layer(x)
+        x2 = self.x_layer(x)
+        x3 = self.y_layer(x)
+        x4 = self.nb_x_layer(x)
+        x5 = self.nb_y_layer(x)
+        return x1, x2, x3, x4, x5
+    
+    
+if __name__ == '__main__':
+    pass
+

third_party/PIPNet/lib/networks_gssl.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import numpy as np
+import time
+
+# net_stride output_size
+# 128        2x2
+# 64         4x4
+# 32         8x8
+# pip regression, resnet18, for GSSL
+class Pip_resnet18(nn.Module):
+    def __init__(self, resnet, num_nb, num_lms=68, input_size=256, net_stride=32):
+        super(Pip_resnet18, self).__init__()
+        self.num_nb = num_nb
+        self.num_lms = num_lms
+        self.input_size = input_size
+        self.net_stride = net_stride
+        self.conv1 = resnet.conv1
+        self.bn1 = resnet.bn1
+        self.maxpool = resnet.maxpool
+        self.sigmoid = nn.Sigmoid()
+        self.layer1 = resnet.layer1
+        self.layer2 = resnet.layer2
+        self.layer3 = resnet.layer3
+        self.layer4 = resnet.layer4
+
+        self.my_maxpool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        
+        self.cls_layer = nn.Conv2d(512, num_lms, kernel_size=1, stride=1, padding=0)
+        self.x_layer = nn.Conv2d(512, num_lms, kernel_size=1, stride=1, padding=0)
+        self.y_layer = nn.Conv2d(512, num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_x_layer = nn.Conv2d(512, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+        self.nb_y_layer = nn.Conv2d(512, num_nb*num_lms, kernel_size=1, stride=1, padding=0)
+
+        # init
+        nn.init.normal_(self.cls_layer.weight, std=0.001)
+        if self.cls_layer.bias is not None:
+            nn.init.constant_(self.cls_layer.bias, 0)
+
+        nn.init.normal_(self.x_layer.weight, std=0.001)
+        if self.x_layer.bias is not None:
+            nn.init.constant_(self.x_layer.bias, 0)
+
+        nn.init.normal_(self.y_layer.weight, std=0.001)
+        if self.y_layer.bias is not None:
+            nn.init.constant_(self.y_layer.bias, 0)
+
+        nn.init.normal_(self.nb_x_layer.weight, std=0.001)
+        if self.nb_x_layer.bias is not None:
+            nn.init.constant_(self.nb_x_layer.bias, 0)
+
+        nn.init.normal_(self.nb_y_layer.weight, std=0.001)
+        if self.nb_y_layer.bias is not None:
+            nn.init.constant_(self.nb_y_layer.bias, 0)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = F.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        cls1 = self.cls_layer(x)
+        offset_x = self.x_layer(x)
+        offset_y = self.y_layer(x)
+        nb_x = self.nb_x_layer(x)
+        nb_y = self.nb_y_layer(x)
+        x = self.my_maxpool(x)
+        cls2 = self.cls_layer(x)
+        x = self.my_maxpool(x)
+        cls3 = self.cls_layer(x)
+        return cls1, cls2, cls3, offset_x, offset_y, nb_x, nb_y
+
+if __name__ == '__main__':
+    pass
+    

third_party/PIPNet/lib/preprocess.py

@@ -0,0 +1,554 @@
+import os, cv2
+import hdf5storage
+import numpy as np
+import sys
+
+def process_300w(root_folder, folder_name, image_name, label_name, target_size):
+    image_path = os.path.join(root_folder, folder_name, image_name)
+    label_path = os.path.join(root_folder, folder_name, label_name)
+
+    with open(label_path, 'r') as ff:
+        anno = ff.readlines()[3:-1]
+        anno = [x.strip().split() for x in anno]
+        anno = [[int(float(x[0])), int(float(x[1]))] for x in anno]
+        image = cv2.imread(image_path)
+        image_height, image_width, _ = image.shape
+        anno_x = [x[0] for x in anno]
+        anno_y = [x[1] for x in anno]
+        bbox_xmin = min(anno_x)
+        bbox_ymin = min(anno_y)
+        bbox_xmax = max(anno_x)
+        bbox_ymax = max(anno_y)
+        bbox_width = bbox_xmax - bbox_xmin
+        bbox_height = bbox_ymax - bbox_ymin
+        scale = 1.1 
+        bbox_xmin -= int((scale-1)/2*bbox_width)
+        bbox_ymin -= int((scale-1)/2*bbox_height)
+        bbox_width *= scale
+        bbox_height *= scale
+        bbox_width = int(bbox_width)
+        bbox_height = int(bbox_height)
+        bbox_xmin = max(bbox_xmin, 0)
+        bbox_ymin = max(bbox_ymin, 0)
+        bbox_width = min(bbox_width, image_width-bbox_xmin-1)
+        bbox_height = min(bbox_height, image_height-bbox_ymin-1)
+        anno = [[(x-bbox_xmin)/bbox_width, (y-bbox_ymin)/bbox_height] for x,y in anno]
+
+        bbox_xmax = bbox_xmin + bbox_width
+        bbox_ymax = bbox_ymin + bbox_height
+        image_crop = image[bbox_ymin:bbox_ymax, bbox_xmin:bbox_xmax, :]
+        image_crop = cv2.resize(image_crop, (target_size, target_size))
+        return image_crop, anno
+        
+def process_cofw(image, bbox, anno, target_size):
+    image_height, image_width, _ = image.shape
+    anno_x = anno[:29]
+    anno_y = anno[29:58]
+    ################################
+    xmin, ymin, width, height = bbox
+    xmax = xmin + width -1
+    ymax = ymin + height -1
+    ################################
+    xmin = max(xmin, 0)
+    ymin = max(ymin, 0)
+    xmax = min(xmax, image_width-1)
+    ymax = min(ymax, image_height-1)
+    anno_x = (anno_x - xmin) / (xmax - xmin)
+    anno_y = (anno_y - ymin) / (ymax - ymin)
+    anno = np.concatenate([anno_x.reshape(-1,1), anno_y.reshape(-1,1)], axis=1)
+    anno = list(anno)
+    anno = [list(x) for x in anno]
+    image_crop = image[int(ymin):int(ymax), int(xmin):int(xmax), :]
+    image_crop = cv2.resize(image_crop, (target_size, target_size))
+    return image_crop, anno
+
+def process_wflw(anno, target_size):
+    image_name = anno[-1]
+    image_path = os.path.join('..', 'data', 'WFLW', 'WFLW_images', image_name)
+    image = cv2.imread(image_path)
+    image_height, image_width, _ = image.shape
+    lms = anno[:196]
+    lms = [float(x) for x in lms]
+    lms_x = lms[0::2]
+    lms_y = lms[1::2]
+    lms_x = [x if x >=0 else 0 for x in lms_x] 
+    lms_x = [x if x <=image_width else image_width for x in lms_x] 
+    lms_y = [y if y >=0 else 0 for y in lms_y] 
+    lms_y = [y if y <=image_height else image_height for y in lms_y] 
+    lms = [[x,y] for x,y in zip(lms_x, lms_y)]
+    lms = [x for z in lms for x in z]
+    bbox = anno[196:200]
+    bbox = [float(x) for x in bbox]
+    attrs = anno[200:206]
+    attrs = np.array([int(x) for x in attrs])
+    bbox_xmin, bbox_ymin, bbox_xmax, bbox_ymax = bbox
+
+    width = bbox_xmax - bbox_xmin
+    height = bbox_ymax - bbox_ymin
+    scale = 1.2  
+    bbox_xmin -= width * (scale-1)/2
+    bbox_ymin -= height * (scale-1)/2
+    bbox_xmax += width * (scale-1)/2
+    bbox_ymax += height * (scale-1)/2
+    bbox_xmin = max(bbox_xmin, 0)
+    bbox_ymin = max(bbox_ymin, 0)
+    bbox_xmax = min(bbox_xmax, image_width-1)
+    bbox_ymax = min(bbox_ymax, image_height-1)
+    width = bbox_xmax - bbox_xmin
+    height = bbox_ymax - bbox_ymin
+    image_crop = image[int(bbox_ymin):int(bbox_ymax), int(bbox_xmin):int(bbox_xmax), :]
+    image_crop = cv2.resize(image_crop, (target_size, target_size))
+
+    tmp1 = [bbox_xmin, bbox_ymin]*98
+    tmp1 = np.array(tmp1)
+    tmp2 = [width, height]*98
+    tmp2 = np.array(tmp2)
+    lms = np.array(lms) - tmp1
+    lms = lms / tmp2
+    lms = lms.tolist()
+    lms = zip(lms[0::2], lms[1::2])
+    return image_crop, list(lms) 
+
+def process_aflw(root_folder, image_name, bbox, anno, target_size):
+    image = cv2.imread(os.path.join(root_folder, 'AFLW', 'flickr', image_name))
+    image_height, image_width, _ = image.shape
+    anno_x = anno[:19]
+    anno_y = anno[19:]
+    anno_x = [x if x >=0 else 0 for x in anno_x] 
+    anno_x = [x if x <=image_width else image_width for x in anno_x] 
+    anno_y = [y if y >=0 else 0 for y in anno_y] 
+    anno_y = [y if y <=image_height else image_height for y in anno_y] 
+    anno_x_min = min(anno_x)
+    anno_x_max = max(anno_x)
+    anno_y_min = min(anno_y)
+    anno_y_max = max(anno_y)
+    xmin, xmax, ymin, ymax = bbox
+    
+    xmin = max(xmin, 0)
+    ymin = max(ymin, 0)
+    xmax = min(xmax, image_width-1)
+    ymax = min(ymax, image_height-1)
+
+    image_crop = image[int(ymin):int(ymax), int(xmin):int(xmax), :]
+    image_crop = cv2.resize(image_crop, (target_size, target_size))
+
+    anno_x = (np.array(anno_x) - xmin) / (xmax - xmin)
+    anno_y = (np.array(anno_y) - ymin) / (ymax - ymin)
+
+    anno = np.concatenate([anno_x.reshape(-1,1), anno_y.reshape(-1,1)], axis=1).flatten()
+    anno = zip(anno[0::2], anno[1::2])
+    return image_crop, anno
+
+def gen_meanface(root_folder, data_name):
+    with open(os.path.join(root_folder, data_name, 'train.txt'), 'r') as f:
+        annos = f.readlines()
+    annos = [x.strip().split()[1:] for x in annos]
+    annos = [[float(x) for x in anno] for anno in annos]
+    annos = np.array(annos)
+    meanface = np.mean(annos, axis=0)
+    meanface = meanface.tolist()
+    meanface = [str(x) for x in meanface]
+    
+    with open(os.path.join(root_folder, data_name, 'meanface.txt'), 'w') as f:
+        f.write(' '.join(meanface))
+
+def convert_wflw(root_folder, data_name):
+    with open(os.path.join('../data/WFLW/test.txt'), 'r') as f:
+        annos = f.readlines()
+    annos = [x.strip().split() for x in annos]
+    annos_new = []
+    for anno in annos:
+        annos_new.append([])
+        # name
+        annos_new[-1].append(anno[0])
+        anno = anno[1:]
+        # jaw
+        for i in range(17):
+            annos_new[-1].append(anno[i*2*2])
+            annos_new[-1].append(anno[i*2*2+1])
+        # left eyebrow
+        annos_new[-1].append(anno[33*2])
+        annos_new[-1].append(anno[33*2+1])
+        annos_new[-1].append(anno[34*2])
+        annos_new[-1].append(str((float(anno[34*2+1])+float(anno[41*2+1]))/2))
+        annos_new[-1].append(anno[35*2])
+        annos_new[-1].append(str((float(anno[35*2+1])+float(anno[40*2+1]))/2))
+        annos_new[-1].append(anno[36*2])
+        annos_new[-1].append(str((float(anno[36*2+1])+float(anno[39*2+1]))/2))
+        annos_new[-1].append(anno[37*2])
+        annos_new[-1].append(str((float(anno[37*2+1])+float(anno[38*2+1]))/2))
+        # right eyebrow
+        annos_new[-1].append(anno[42*2])
+        annos_new[-1].append(str((float(anno[42*2+1])+float(anno[50*2+1]))/2))
+        annos_new[-1].append(anno[43*2])
+        annos_new[-1].append(str((float(anno[43*2+1])+float(anno[49*2+1]))/2))
+        annos_new[-1].append(anno[44*2])
+        annos_new[-1].append(str((float(anno[44*2+1])+float(anno[48*2+1]))/2))
+        annos_new[-1].append(anno[45*2])
+        annos_new[-1].append(str((float(anno[45*2+1])+float(anno[47*2+1]))/2))
+        annos_new[-1].append(anno[46*2])
+        annos_new[-1].append(anno[46*2+1])
+        # nose
+        for i in range(51, 60):
+            annos_new[-1].append(anno[i*2])
+            annos_new[-1].append(anno[i*2+1])
+        # left eye
+        annos_new[-1].append(anno[60*2])
+        annos_new[-1].append(anno[60*2+1])
+        annos_new[-1].append(str(0.666*float(anno[61*2])+0.333*float(anno[62*2])))
+        annos_new[-1].append(str(0.666*float(anno[61*2+1])+0.333*float(anno[62*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[63*2])+0.333*float(anno[62*2])))
+        annos_new[-1].append(str(0.666*float(anno[63*2+1])+0.333*float(anno[62*2+1])))
+        annos_new[-1].append(anno[64*2])
+        annos_new[-1].append(anno[64*2+1])
+        annos_new[-1].append(str(0.666*float(anno[65*2])+0.333*float(anno[66*2])))
+        annos_new[-1].append(str(0.666*float(anno[65*2+1])+0.333*float(anno[66*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[67*2])+0.333*float(anno[66*2])))
+        annos_new[-1].append(str(0.666*float(anno[67*2+1])+0.333*float(anno[66*2+1])))
+        # right eye
+        annos_new[-1].append(anno[68*2])
+        annos_new[-1].append(anno[68*2+1])
+        annos_new[-1].append(str(0.666*float(anno[69*2])+0.333*float(anno[70*2])))
+        annos_new[-1].append(str(0.666*float(anno[69*2+1])+0.333*float(anno[70*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[71*2])+0.333*float(anno[70*2])))
+        annos_new[-1].append(str(0.666*float(anno[71*2+1])+0.333*float(anno[70*2+1])))
+        annos_new[-1].append(anno[72*2])
+        annos_new[-1].append(anno[72*2+1])
+        annos_new[-1].append(str(0.666*float(anno[73*2])+0.333*float(anno[74*2])))
+        annos_new[-1].append(str(0.666*float(anno[73*2+1])+0.333*float(anno[74*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[75*2])+0.333*float(anno[74*2])))
+        annos_new[-1].append(str(0.666*float(anno[75*2+1])+0.333*float(anno[74*2+1])))
+        # mouth
+        for i in range(76, 96):
+            annos_new[-1].append(anno[i*2])
+            annos_new[-1].append(anno[i*2+1])
+
+    with open(os.path.join(root_folder, data_name, 'test.txt'), 'w') as f:
+        for anno in annos_new:
+            f.write(' '.join(anno)+'\n')
+
+
+def gen_data(root_folder, data_name, target_size):
+    if not os.path.exists(os.path.join(root_folder, data_name, 'images_train')):
+        os.mkdir(os.path.join(root_folder, data_name, 'images_train'))
+    if not os.path.exists(os.path.join(root_folder, data_name, 'images_test')):
+        os.mkdir(os.path.join(root_folder, data_name, 'images_test'))
+
+    ################################################################################################################
+    if data_name == 'data_300W':
+        folders_train = ['afw', 'helen/trainset', 'lfpw/trainset']
+        annos_train = {}
+        for folder_train in folders_train:
+            all_files = sorted(os.listdir(os.path.join(root_folder, data_name, folder_train)))
+            image_files = [x for x in all_files if '.pts' not in x]
+            label_files = [x for x in all_files if '.pts' in x]
+            assert len(image_files) == len(label_files)
+            for image_name, label_name in zip(image_files, label_files):
+                print(image_name)
+                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_train, image_name, label_name, target_size)
+                image_crop_name = folder_train.replace('/', '_')+'_'+image_name
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_train', image_crop_name), image_crop)
+                annos_train[image_crop_name] = anno
+        with open(os.path.join(root_folder, data_name, 'train.txt'), 'w') as f:
+            for image_crop_name, anno in annos_train.items():
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+        
+
+        folders_test = ['helen/testset', 'lfpw/testset', 'ibug']
+        annos_test = {}
+        for folder_test in folders_test:
+            all_files = sorted(os.listdir(os.path.join(root_folder, data_name, folder_test)))
+            image_files = [x for x in all_files if '.pts' not in x]
+            label_files = [x for x in all_files if '.pts' in x]
+            assert len(image_files) == len(label_files)
+            for image_name, label_name in zip(image_files, label_files):
+                print(image_name)
+                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_test, image_name, label_name, target_size)
+                image_crop_name = folder_test.replace('/', '_')+'_'+image_name
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_test', image_crop_name), image_crop)
+                annos_test[image_crop_name] = anno
+        with open(os.path.join(root_folder, data_name, 'test.txt'), 'w') as f:
+            for image_crop_name, anno in annos_test.items():
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        annos = None
+        with open(os.path.join(root_folder, data_name, 'test.txt'), 'r') as f:
+            annos = f.readlines()
+        with open(os.path.join(root_folder, data_name, 'test_common.txt'), 'w') as f:
+            for anno in annos:
+                if not 'ibug' in anno:
+                    f.write(anno)
+        with open(os.path.join(root_folder, data_name, 'test_challenge.txt'), 'w') as f:
+            for anno in annos:
+                if 'ibug' in anno:
+                    f.write(anno)
+
+        gen_meanface(root_folder, data_name)
+    ################################################################################################################
+    elif data_name == 'COFW':
+        train_file = 'COFW_train_color.mat'
+        train_mat = hdf5storage.loadmat(os.path.join(root_folder, 'COFW', train_file))
+        images = train_mat['IsTr']
+        bboxes = train_mat['bboxesTr']
+        annos = train_mat['phisTr']
+
+        count = 1
+        with open(os.path.join(root_folder, 'COFW', 'train.txt'), 'w') as f:
+            for i in range(images.shape[0]):
+                image = images[i, 0]
+                # grayscale
+                if len(image.shape) == 2:
+                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+                # swap rgb channel to bgr
+                else:
+                    image = image[:,:,::-1]
+                bbox = bboxes[i, :]
+                anno = annos[i, :]
+                image_crop, anno = process_cofw(image, bbox, anno, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'cofw_train_' + '0' * pad_num + str(count) + '.jpg'
+                print(image_crop_name)
+                cv2.imwrite(os.path.join(root_folder, 'COFW', 'images_train', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+                count += 1
+
+        test_file = 'COFW_test_color.mat'
+        test_mat = hdf5storage.loadmat(os.path.join(root_folder, 'COFW', test_file))
+        images = test_mat['IsT']
+        bboxes = test_mat['bboxesT']
+        annos = test_mat['phisT']
+
+        count = 1
+        with open(os.path.join(root_folder, 'COFW', 'test.txt'), 'w') as f:
+            for i in range(images.shape[0]):
+                image = images[i, 0]
+                # grayscale
+                if len(image.shape) == 2:
+                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+                # swap rgb channel to bgr
+                else:
+                    image = image[:,:,::-1]
+                bbox = bboxes[i, :]
+                anno = annos[i, :]
+                image_crop, anno = process_cofw(image, bbox, anno, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'cofw_test_' + '0' * pad_num + str(count) + '.jpg'
+                print(image_crop_name)
+                cv2.imwrite(os.path.join(root_folder, 'COFW', 'images_test', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+                count += 1
+        gen_meanface(root_folder, data_name)
+    ################################################################################################################
+    elif data_name == 'WFLW':
+        train_file = 'list_98pt_rect_attr_train.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_rect_attr_train_test', train_file), 'r') as f:
+            annos_train = f.readlines()
+        annos_train = [x.strip().split() for x in annos_train]
+        count = 1
+        with open(os.path.join(root_folder, 'WFLW', 'train.txt'), 'w') as f:
+            for anno_train in annos_train:
+                image_crop, anno = process_wflw(anno_train, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'wflw_train_' + '0' * pad_num + str(count) + '.jpg'
+                print(image_crop_name)
+                cv2.imwrite(os.path.join(root_folder, 'WFLW', 'images_train', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+                count += 1
+
+        test_file = 'list_98pt_rect_attr_test.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_rect_attr_train_test', test_file), 'r') as f:
+            annos_test = f.readlines()
+        annos_test = [x.strip().split() for x in annos_test]
+        names_mapping = {}
+        count = 1
+        with open(os.path.join(root_folder, 'WFLW', 'test.txt'), 'w') as f:
+            for anno_test in annos_test:
+                image_crop, anno = process_wflw(anno_test, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'wflw_test_' + '0' * pad_num + str(count) + '.jpg'
+                print(image_crop_name)
+                names_mapping[anno_test[0]+'_'+anno_test[-1]] = [image_crop_name, anno]
+                cv2.imwrite(os.path.join(root_folder, 'WFLW', 'images_test', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in list(anno):
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+                count += 1
+
+        test_pose_file = 'list_98pt_test_largepose.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_pose_file), 'r') as f:
+            annos_pose_test = f.readlines()
+        names_pose = [x.strip().split() for x in annos_pose_test]
+        names_pose = [x[0]+'_'+x[-1] for x in names_pose]
+        with open(os.path.join(root_folder, 'WFLW', 'test_pose.txt'), 'w') as f:
+            for name_pose in names_pose:
+                if name_pose in names_mapping:
+                    image_crop_name, anno = names_mapping[name_pose]
+                    f.write(image_crop_name+' ')
+                    for x,y in anno:
+                        f.write(str(x)+' '+str(y)+' ')
+                    f.write('\n')
+                else:
+                    print('error!')
+                    exit(0)
+
+        test_expr_file = 'list_98pt_test_expression.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_expr_file), 'r') as f:
+            annos_expr_test = f.readlines()
+        names_expr = [x.strip().split() for x in annos_expr_test]
+        names_expr = [x[0]+'_'+x[-1] for x in names_expr]
+        with open(os.path.join(root_folder, 'WFLW', 'test_expr.txt'), 'w') as f:
+            for name_expr in names_expr:
+                if name_expr in names_mapping:
+                    image_crop_name, anno = names_mapping[name_expr]
+                    f.write(image_crop_name+' ')
+                    for x,y in anno:
+                        f.write(str(x)+' '+str(y)+' ')
+                    f.write('\n')
+                else:
+                    print('error!')
+                    exit(0)
+
+        test_illu_file = 'list_98pt_test_illumination.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_illu_file), 'r') as f:
+            annos_illu_test = f.readlines()
+        names_illu = [x.strip().split() for x in annos_illu_test]
+        names_illu = [x[0]+'_'+x[-1] for x in names_illu]
+        with open(os.path.join(root_folder, 'WFLW', 'test_illu.txt'), 'w') as f:
+            for name_illu in names_illu:
+                if name_illu in names_mapping:
+                    image_crop_name, anno = names_mapping[name_illu]
+                    f.write(image_crop_name+' ')
+                    for x,y in anno:
+                        f.write(str(x)+' '+str(y)+' ')
+                    f.write('\n')
+                else:
+                    print('error!')
+                    exit(0)
+
+        test_mu_file = 'list_98pt_test_makeup.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_mu_file), 'r') as f:
+            annos_mu_test = f.readlines()
+        names_mu = [x.strip().split() for x in annos_mu_test]
+        names_mu = [x[0]+'_'+x[-1] for x in names_mu]
+        with open(os.path.join(root_folder, 'WFLW', 'test_mu.txt'), 'w') as f:
+            for name_mu in names_mu:
+                if name_mu in names_mapping:
+                    image_crop_name, anno = names_mapping[name_mu]
+                    f.write(image_crop_name+' ')
+                    for x,y in anno:
+                        f.write(str(x)+' '+str(y)+' ')
+                    f.write('\n')
+                else:
+                    print('error!')
+                    exit(0)
+
+        test_occu_file = 'list_98pt_test_occlusion.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_occu_file), 'r') as f:
+            annos_occu_test = f.readlines()
+        names_occu = [x.strip().split() for x in annos_occu_test]
+        names_occu = [x[0]+'_'+x[-1] for x in names_occu]
+        with open(os.path.join(root_folder, 'WFLW', 'test_occu.txt'), 'w') as f:
+            for name_occu in names_occu:
+                if name_occu in names_mapping:
+                    image_crop_name, anno = names_mapping[name_occu]
+                    f.write(image_crop_name+' ')
+                    for x,y in anno:
+                        f.write(str(x)+' '+str(y)+' ')
+                    f.write('\n')
+                else:
+                    print('error!')
+                    exit(0)
+
+
+        test_blur_file = 'list_98pt_test_blur.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_blur_file), 'r') as f:
+            annos_blur_test = f.readlines()
+        names_blur = [x.strip().split() for x in annos_blur_test]
+        names_blur = [x[0]+'_'+x[-1] for x in names_blur]
+        with open(os.path.join(root_folder, 'WFLW', 'test_blur.txt'), 'w') as f:
+            for name_blur in names_blur:
+                if name_blur in names_mapping:
+                    image_crop_name, anno = names_mapping[name_blur]
+                    f.write(image_crop_name+' ')
+                    for x,y in anno:
+                        f.write(str(x)+' '+str(y)+' ')
+                    f.write('\n')
+                else:
+                    print('error!')
+                    exit(0)
+        gen_meanface(root_folder, data_name)
+    ################################################################################################################
+    elif data_name == 'AFLW':
+        mat = hdf5storage.loadmat('../data/AFLW/AFLWinfo_release.mat')
+        bboxes = mat['bbox']
+        annos = mat['data']
+        mask_new = mat['mask_new']
+        nameList = mat['nameList']
+        ra = mat['ra'][0]
+        train_indices = ra[:20000]
+        test_indices = ra[20000:]
+
+        with open(os.path.join(root_folder, 'AFLW', 'train.txt'), 'w') as f:
+            for index in train_indices:
+                # from matlab index
+                image_name = nameList[index-1][0][0]
+                bbox = bboxes[index-1]
+                anno = annos[index-1]
+                image_crop, anno = process_aflw(root_folder, image_name, bbox, anno, target_size)
+                pad_num = 5-len(str(index))
+                image_crop_name = 'aflw_train_' + '0' * pad_num + str(index) + '.jpg'
+                print(image_crop_name)
+                cv2.imwrite(os.path.join(root_folder, 'AFLW', 'images_train', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        with open(os.path.join(root_folder, 'AFLW', 'test.txt'), 'w') as f:
+            for index in test_indices:
+                # from matlab index
+                image_name = nameList[index-1][0][0]
+                bbox = bboxes[index-1]
+                anno = annos[index-1]
+                image_crop, anno = process_aflw(root_folder, image_name, bbox, anno, target_size)
+                pad_num = 5-len(str(index))
+                image_crop_name = 'aflw_test_' + '0' * pad_num + str(index) + '.jpg'
+                print(image_crop_name)
+                cv2.imwrite(os.path.join(root_folder, 'AFLW', 'images_test', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+        gen_meanface(root_folder, data_name)
+    else:
+        print('Wrong data!')
+
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        print('please input the data name.')
+        print('1. data_300W')
+        print('2. COFW')
+        print('3. WFLW')
+        print('4. AFLW')
+        exit(0)
+    else:
+        data_name = sys.argv[1]
+        gen_data('../data', data_name, 256)
+
+

third_party/PIPNet/lib/preprocess_gssl.py

@@ -0,0 +1,544 @@
+import os, cv2
+import hdf5storage
+import numpy as np
+import sys
+
+def process_300w(root_folder, folder_name, image_name, label_name, target_size):
+    image_path = os.path.join(root_folder, folder_name, image_name)
+    label_path = os.path.join(root_folder, folder_name, label_name)
+
+    with open(label_path, 'r') as ff:
+        anno = ff.readlines()[3:-1]
+        anno = [x.strip().split() for x in anno]
+        anno = [[int(float(x[0])), int(float(x[1]))] for x in anno]
+        image = cv2.imread(image_path)
+        image_height, image_width, _ = image.shape
+        anno_x = [x[0] for x in anno]
+        anno_y = [x[1] for x in anno]
+        bbox_xmin = min(anno_x)
+        bbox_ymin = min(anno_y)
+        bbox_xmax = max(anno_x)
+        bbox_ymax = max(anno_y)
+        bbox_width = bbox_xmax - bbox_xmin
+        bbox_height = bbox_ymax - bbox_ymin
+        scale = 1.3
+        bbox_xmin -= int((scale-1)/2*bbox_width)
+        bbox_ymin -= int((scale-1)/2*bbox_height)
+        bbox_width *= scale
+        bbox_height *= scale
+        bbox_width = int(bbox_width)
+        bbox_height = int(bbox_height)
+        bbox_xmin = max(bbox_xmin, 0)
+        bbox_ymin = max(bbox_ymin, 0)
+        bbox_width = min(bbox_width, image_width-bbox_xmin-1)
+        bbox_height = min(bbox_height, image_height-bbox_ymin-1)
+        anno = [[(x-bbox_xmin)/bbox_width, (y-bbox_ymin)/bbox_height] for x,y in anno]
+
+        bbox_xmax = bbox_xmin + bbox_width
+        bbox_ymax = bbox_ymin + bbox_height
+        image_crop = image[bbox_ymin:bbox_ymax, bbox_xmin:bbox_xmax, :]
+        image_crop = cv2.resize(image_crop, (target_size, target_size))
+        return image_crop, anno
+        
+def process_wflw(anno, target_size):
+    image_name = anno[-1]
+    image_path = os.path.join('..', 'data', 'WFLW', 'WFLW_images', image_name)
+    image = cv2.imread(image_path)
+    image_height, image_width, _ = image.shape
+    lms = anno[:196]
+    lms = [float(x) for x in lms]
+    lms_x = lms[0::2]
+    lms_y = lms[1::2]
+    lms_x = [x if x >=0 else 0 for x in lms_x] 
+    lms_x = [x if x <=image_width else image_width for x in lms_x] 
+    lms_y = [y if y >=0 else 0 for y in lms_y] 
+    lms_y = [y if y <=image_height else image_height for y in lms_y] 
+    lms = [[x,y] for x,y in zip(lms_x, lms_y)]
+    lms = [x for z in lms for x in z]
+    bbox = anno[196:200]
+    bbox = [float(x) for x in bbox]
+    attrs = anno[200:206]
+    attrs = np.array([int(x) for x in attrs])
+    bbox_xmin, bbox_ymin, bbox_xmax, bbox_ymax = bbox
+
+    width = bbox_xmax - bbox_xmin
+    height = bbox_ymax - bbox_ymin
+    scale = 1.2
+    bbox_xmin -= width * (scale-1)/2
+    # remove a part of top area for alignment, see details in paper
+    bbox_ymin += height * (scale-1)/2
+    bbox_xmax += width * (scale-1)/2
+    bbox_ymax += height * (scale-1)/2
+    bbox_xmin = max(bbox_xmin, 0)
+    bbox_ymin = max(bbox_ymin, 0)
+    bbox_xmax = min(bbox_xmax, image_width-1)
+    bbox_ymax = min(bbox_ymax, image_height-1)
+    width = bbox_xmax - bbox_xmin
+    height = bbox_ymax - bbox_ymin
+    image_crop = image[int(bbox_ymin):int(bbox_ymax), int(bbox_xmin):int(bbox_xmax), :]
+    image_crop = cv2.resize(image_crop, (target_size, target_size))
+
+    tmp1 = [bbox_xmin, bbox_ymin]*98
+    tmp1 = np.array(tmp1)
+    tmp2 = [width, height]*98
+    tmp2 = np.array(tmp2)
+    lms = np.array(lms) - tmp1
+    lms = lms / tmp2
+    lms = lms.tolist()
+    lms = zip(lms[0::2], lms[1::2])
+    return image_crop, list(lms) 
+
+def process_celeba(root_folder, image_name, bbox, target_size):
+    image = cv2.imread(os.path.join(root_folder, 'CELEBA', 'img_celeba', image_name))
+    image_height, image_width, _ = image.shape
+    xmin, ymin, xmax, ymax = bbox
+    width = xmax - xmin + 1
+    height = ymax - ymin + 1
+    scale = 1.2
+    xmin -= width * (scale-1)/2
+    # remove a part of top area for alignment, see details in paper
+    ymin += height * (scale+0.1-1)/2
+    xmax += width * (scale-1)/2
+    ymax += height * (scale-1)/2
+    xmin = max(xmin, 0)
+    ymin = max(ymin, 0)
+    xmax = min(xmax, image_width-1)
+    ymax = min(ymax, image_height-1)
+    image_crop = image[int(ymin):int(ymax), int(xmin):int(xmax), :]
+    image_crop = cv2.resize(image_crop, (target_size, target_size))
+    return image_crop
+
+def process_cofw_68_train(image, bbox, anno, target_size):
+    image_height, image_width, _ = image.shape
+    anno_x = anno[:29]
+    anno_y = anno[29:58]
+    xmin, ymin, width, height = bbox
+    xmax = xmin + width -1
+    ymax = ymin + height -1
+    scale = 1.3
+    xmin -= width * (scale-1)/2
+    ymin -= height * (scale-1)/2
+    xmax += width * (scale-1)/2
+    ymax += height * (scale-1)/2
+    xmin = max(xmin, 0)
+    ymin = max(ymin, 0)
+    xmax = min(xmax, image_width-1)
+    ymax = min(ymax, image_height-1)
+    anno_x = (anno_x - xmin) / (xmax - xmin)
+    anno_y = (anno_y - ymin) / (ymax - ymin)
+    anno = np.concatenate([anno_x.reshape(-1,1), anno_y.reshape(-1,1)], axis=1)
+    anno = list(anno)
+    anno = [list(x) for x in anno]
+    image_crop = image[int(ymin):int(ymax), int(xmin):int(xmax), :]
+    image_crop = cv2.resize(image_crop, (target_size, target_size))
+    return image_crop, anno
+
+def process_cofw_68_test(image, bbox, anno, target_size):
+    image_height, image_width, _ = image.shape
+    anno_x = anno[:,0].flatten()
+    anno_y = anno[:,1].flatten()
+
+    xmin, ymin, width, height = bbox
+    xmax = xmin + width -1
+    ymax = ymin + height -1
+
+    scale = 1.3
+    xmin -= width * (scale-1)/2
+    ymin -= height * (scale-1)/2
+    xmax += width * (scale-1)/2
+    ymax += height * (scale-1)/2
+    xmin = max(xmin, 0)
+    ymin = max(ymin, 0)
+    xmax = min(xmax, image_width-1)
+    ymax = min(ymax, image_height-1)
+    anno_x = (anno_x - xmin) / (xmax - xmin)
+    anno_y = (anno_y - ymin) / (ymax - ymin)
+    anno = np.concatenate([anno_x.reshape(-1,1), anno_y.reshape(-1,1)], axis=1)
+    anno = list(anno)
+    anno = [list(x) for x in anno]
+    image_crop = image[int(ymin):int(ymax), int(xmin):int(xmax), :]
+    image_crop = cv2.resize(image_crop, (target_size, target_size))
+    return image_crop, anno
+
+def gen_meanface(root_folder, data_name):
+    with open(os.path.join(root_folder, data_name, 'train_300W.txt'), 'r') as f:
+        annos = f.readlines()
+    annos = [x.strip().split()[1:] for x in annos]
+    annos = [[float(x) for x in anno] for anno in annos]
+    annos = np.array(annos)
+    meanface = np.mean(annos, axis=0)
+    meanface = meanface.tolist()
+    meanface = [str(x) for x in meanface]
+    
+    with open(os.path.join(root_folder, data_name, 'meanface.txt'), 'w') as f:
+        f.write(' '.join(meanface))
+
+def convert_wflw(root_folder, data_name):
+    with open(os.path.join(root_folder, data_name, 'test_WFLW_98.txt'), 'r') as f:
+        annos = f.readlines()
+    annos = [x.strip().split() for x in annos]
+    annos_new = []
+    for anno in annos:
+        annos_new.append([])
+        # name
+        annos_new[-1].append(anno[0])
+        anno = anno[1:]
+        # jaw
+        for i in range(17):
+            annos_new[-1].append(anno[i*2*2])
+            annos_new[-1].append(anno[i*2*2+1])
+        # left eyebrow
+        annos_new[-1].append(anno[33*2])
+        annos_new[-1].append(anno[33*2+1])
+        annos_new[-1].append(anno[34*2])
+        annos_new[-1].append(str((float(anno[34*2+1])+float(anno[41*2+1]))/2))
+        annos_new[-1].append(anno[35*2])
+        annos_new[-1].append(str((float(anno[35*2+1])+float(anno[40*2+1]))/2))
+        annos_new[-1].append(anno[36*2])
+        annos_new[-1].append(str((float(anno[36*2+1])+float(anno[39*2+1]))/2))
+        annos_new[-1].append(anno[37*2])
+        annos_new[-1].append(str((float(anno[37*2+1])+float(anno[38*2+1]))/2))
+        # right eyebrow
+        annos_new[-1].append(anno[42*2])
+        annos_new[-1].append(str((float(anno[42*2+1])+float(anno[50*2+1]))/2))
+        annos_new[-1].append(anno[43*2])
+        annos_new[-1].append(str((float(anno[43*2+1])+float(anno[49*2+1]))/2))
+        annos_new[-1].append(anno[44*2])
+        annos_new[-1].append(str((float(anno[44*2+1])+float(anno[48*2+1]))/2))
+        annos_new[-1].append(anno[45*2])
+        annos_new[-1].append(str((float(anno[45*2+1])+float(anno[47*2+1]))/2))
+        annos_new[-1].append(anno[46*2])
+        annos_new[-1].append(anno[46*2+1])
+        # nose
+        for i in range(51, 60):
+            annos_new[-1].append(anno[i*2])
+            annos_new[-1].append(anno[i*2+1])
+        # left eye
+        annos_new[-1].append(anno[60*2])
+        annos_new[-1].append(anno[60*2+1])
+        annos_new[-1].append(str(0.666*float(anno[61*2])+0.333*float(anno[62*2])))
+        annos_new[-1].append(str(0.666*float(anno[61*2+1])+0.333*float(anno[62*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[63*2])+0.333*float(anno[62*2])))
+        annos_new[-1].append(str(0.666*float(anno[63*2+1])+0.333*float(anno[62*2+1])))
+        annos_new[-1].append(anno[64*2])
+        annos_new[-1].append(anno[64*2+1])
+        annos_new[-1].append(str(0.666*float(anno[65*2])+0.333*float(anno[66*2])))
+        annos_new[-1].append(str(0.666*float(anno[65*2+1])+0.333*float(anno[66*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[67*2])+0.333*float(anno[66*2])))
+        annos_new[-1].append(str(0.666*float(anno[67*2+1])+0.333*float(anno[66*2+1])))
+        # right eye
+        annos_new[-1].append(anno[68*2])
+        annos_new[-1].append(anno[68*2+1])
+        annos_new[-1].append(str(0.666*float(anno[69*2])+0.333*float(anno[70*2])))
+        annos_new[-1].append(str(0.666*float(anno[69*2+1])+0.333*float(anno[70*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[71*2])+0.333*float(anno[70*2])))
+        annos_new[-1].append(str(0.666*float(anno[71*2+1])+0.333*float(anno[70*2+1])))
+        annos_new[-1].append(anno[72*2])
+        annos_new[-1].append(anno[72*2+1])
+        annos_new[-1].append(str(0.666*float(anno[73*2])+0.333*float(anno[74*2])))
+        annos_new[-1].append(str(0.666*float(anno[73*2+1])+0.333*float(anno[74*2+1])))
+        annos_new[-1].append(str(0.666*float(anno[75*2])+0.333*float(anno[74*2])))
+        annos_new[-1].append(str(0.666*float(anno[75*2+1])+0.333*float(anno[74*2+1])))
+        # mouth
+        for i in range(76, 96):
+            annos_new[-1].append(anno[i*2])
+            annos_new[-1].append(anno[i*2+1])
+
+    with open(os.path.join(root_folder, data_name, 'test_WFLW.txt'), 'w') as f:
+        for anno in annos_new:
+            f.write(' '.join(anno)+'\n')
+
+def gen_data(root_folder, data_name, target_size):
+    if not os.path.exists(os.path.join(root_folder, data_name, 'images_train')):
+        os.mkdir(os.path.join(root_folder, data_name, 'images_train'))
+    if not os.path.exists(os.path.join(root_folder, data_name, 'images_test')):
+        os.mkdir(os.path.join(root_folder, data_name, 'images_test'))
+    ################################################################################################################
+    if data_name == 'CELEBA':
+        os.system('rmdir ../data/CELEBA/images_test')
+        with open(os.path.join(root_folder, data_name, 'celeba_bboxes.txt'), 'r') as f:
+            bboxes = f.readlines()
+
+        bboxes = [x.strip().split() for x in bboxes]
+        with open(os.path.join(root_folder, data_name, 'train.txt'), 'w') as f:
+            for bbox in bboxes:
+                image_name = bbox[0]
+                print(image_name)
+                f.write(image_name+'\n')
+                bbox = bbox[1:]
+                bbox = [int(x) for x in bbox]
+                image_crop = process_celeba(root_folder, image_name, bbox, target_size)
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_train', image_name), image_crop)
+    ################################################################################################################
+    elif data_name == 'data_300W_CELEBA':
+        os.system('cp -r ../data/CELEBA/images_train ../data/data_300W_CELEBA/.')
+        os.system('cp ../data/CELEBA/train.txt ../data/data_300W_CELEBA/train_CELEBA.txt')
+
+        os.system('rmdir ../data/data_300W_CELEBA/images_test')
+        if not os.path.exists(os.path.join(root_folder, data_name, 'images_test_300W')):
+            os.mkdir(os.path.join(root_folder, data_name, 'images_test_300W'))
+        if not os.path.exists(os.path.join(root_folder, data_name, 'images_test_COFW')):
+            os.mkdir(os.path.join(root_folder, data_name, 'images_test_COFW'))
+        if not os.path.exists(os.path.join(root_folder, data_name, 'images_test_WFLW')):
+            os.mkdir(os.path.join(root_folder, data_name, 'images_test_WFLW'))
+
+        # train for data_300W
+        folders_train = ['afw', 'helen/trainset', 'lfpw/trainset']
+        annos_train = {}
+        for folder_train in folders_train:
+            all_files = sorted(os.listdir(os.path.join(root_folder, 'data_300W', folder_train)))
+            image_files = [x for x in all_files if '.pts' not in x]
+            label_files = [x for x in all_files if '.pts' in x]
+            assert len(image_files) == len(label_files)
+            for image_name, label_name in zip(image_files, label_files):
+                print(image_name)
+                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_train, image_name, label_name, target_size)
+                image_crop_name = folder_train.replace('/', '_')+'_'+image_name
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_train', image_crop_name), image_crop)
+                annos_train[image_crop_name] = anno
+        with open(os.path.join(root_folder, data_name, 'train_300W.txt'), 'w') as f:
+            for image_crop_name, anno in annos_train.items():
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        # test for data_300W
+        folders_test = ['helen/testset', 'lfpw/testset', 'ibug']
+        annos_test = {}
+        for folder_test in folders_test:
+            all_files = sorted(os.listdir(os.path.join(root_folder, 'data_300W', folder_test)))
+            image_files = [x for x in all_files if '.pts' not in x]
+            label_files = [x for x in all_files if '.pts' in x]
+            assert len(image_files) == len(label_files)
+            for image_name, label_name in zip(image_files, label_files):
+                print(image_name)
+                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_test, image_name, label_name, target_size)
+                image_crop_name = folder_test.replace('/', '_')+'_'+image_name
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_test_300W', image_crop_name), image_crop)
+                annos_test[image_crop_name] = anno
+        with open(os.path.join(root_folder, data_name, 'test_300W.txt'), 'w') as f:
+            for image_crop_name, anno in annos_test.items():
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        # test for COFW_68
+        test_mat = hdf5storage.loadmat(os.path.join('../data/COFW', 'COFW_test_color.mat'))
+        images = test_mat['IsT']
+        
+        bboxes_mat = hdf5storage.loadmat(os.path.join('../data/data_300W_CELEBA', 'cofw68_test_bboxes.mat'))
+        bboxes = bboxes_mat['bboxes']
+        image_num = images.shape[0]
+        with open('../data/data_300W_CELEBA/test_COFW.txt', 'w') as f:
+            for i in range(image_num):
+                image = images[i,0]
+                # grayscale
+                if len(image.shape) == 2:
+                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+                # swap rgb channel to bgr
+                else:
+                    image = image[:,:,::-1]
+            
+                bbox = bboxes[i,:]
+                anno_mat = hdf5storage.loadmat(os.path.join('../data/data_300W_CELEBA/cofw68_test_annotations', str(i+1)+'_points.mat'))
+                anno = anno_mat['Points']
+                image_crop, anno = process_cofw_68_test(image, bbox, anno, target_size)
+                pad_num = 4-len(str(i+1))
+                image_crop_name = 'cofw_test_' + '0' * pad_num + str(i+1) + '.jpg'
+                cv2.imwrite(os.path.join('../data/data_300W_CELEBA/images_test_COFW', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        # test for WFLW_68
+        test_file = 'list_98pt_rect_attr_test.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_rect_attr_train_test', test_file), 'r') as f:
+            annos_test = f.readlines()
+        annos_test = [x.strip().split() for x in annos_test]
+        names_mapping = {}
+        count = 1
+        with open(os.path.join(root_folder, 'data_300W_CELEBA', 'test_WFLW_98.txt'), 'w') as f:
+            for anno_test in annos_test:
+                image_crop, anno = process_wflw(anno_test, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'wflw_test_' + '0' * pad_num + str(count) + '.jpg'
+                print(image_crop_name)
+                names_mapping[anno_test[0]+'_'+anno_test[-1]] = [image_crop_name, anno]
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_test_WFLW', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in list(anno):
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+                count += 1
+
+        convert_wflw(root_folder, data_name)
+
+        gen_meanface(root_folder, data_name)
+    ################################################################################################################
+    elif data_name == 'data_300W_COFW_WFLW':
+
+        os.system('rmdir ../data/data_300W_COFW_WFLW/images_test')
+        if not os.path.exists(os.path.join(root_folder, data_name, 'images_test_300W')):
+            os.mkdir(os.path.join(root_folder, data_name, 'images_test_300W'))
+        if not os.path.exists(os.path.join(root_folder, data_name, 'images_test_COFW')):
+            os.mkdir(os.path.join(root_folder, data_name, 'images_test_COFW'))
+        if not os.path.exists(os.path.join(root_folder, data_name, 'images_test_WFLW')):
+            os.mkdir(os.path.join(root_folder, data_name, 'images_test_WFLW'))
+
+        # train for data_300W
+        folders_train = ['afw', 'helen/trainset', 'lfpw/trainset']
+        annos_train = {}
+        for folder_train in folders_train:
+            all_files = sorted(os.listdir(os.path.join(root_folder, 'data_300W', folder_train)))
+            image_files = [x for x in all_files if '.pts' not in x]
+            label_files = [x for x in all_files if '.pts' in x]
+            assert len(image_files) == len(label_files)
+            for image_name, label_name in zip(image_files, label_files):
+                print(image_name)
+                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_train, image_name, label_name, target_size)
+                image_crop_name = folder_train.replace('/', '_')+'_'+image_name
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_train', image_crop_name), image_crop)
+                annos_train[image_crop_name] = anno
+        with open(os.path.join(root_folder, data_name, 'train_300W.txt'), 'w') as f:
+            for image_crop_name, anno in annos_train.items():
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        # test for data_300W
+        folders_test = ['helen/testset', 'lfpw/testset', 'ibug']
+        annos_test = {}
+        for folder_test in folders_test:
+            all_files = sorted(os.listdir(os.path.join(root_folder, 'data_300W', folder_test)))
+            image_files = [x for x in all_files if '.pts' not in x]
+            label_files = [x for x in all_files if '.pts' in x]
+            assert len(image_files) == len(label_files)
+            for image_name, label_name in zip(image_files, label_files):
+                print(image_name)
+                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_test, image_name, label_name, target_size)
+                image_crop_name = folder_test.replace('/', '_')+'_'+image_name
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_test_300W', image_crop_name), image_crop)
+                annos_test[image_crop_name] = anno
+        with open(os.path.join(root_folder, data_name, 'test_300W.txt'), 'w') as f:
+            for image_crop_name, anno in annos_test.items():
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        # train for COFW_68
+        ###################
+        train_file = 'COFW_train_color.mat'
+        train_mat = hdf5storage.loadmat(os.path.join(root_folder, 'COFW', train_file))
+        images = train_mat['IsTr']
+        bboxes = train_mat['bboxesTr']
+        annos = train_mat['phisTr']
+
+        count = 1
+        with open('../data/data_300W_COFW_WFLW/train_COFW.txt', 'w') as f:
+            for i in range(images.shape[0]):
+                image = images[i, 0]
+                # grayscale
+                if len(image.shape) == 2:
+                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+                # swap rgb channel to bgr
+                else:
+                    image = image[:,:,::-1]
+                bbox = bboxes[i, :]
+                anno = annos[i, :]
+                image_crop, anno = process_cofw_68_train(image, bbox, anno, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'cofw_train_' + '0' * pad_num + str(count) + '.jpg'
+                f.write(image_crop_name+'\n')
+                cv2.imwrite(os.path.join(root_folder, 'data_300W_COFW_WFLW', 'images_train', image_crop_name), image_crop)
+                count += 1
+        ###################
+
+        # test for COFW_68
+        test_mat = hdf5storage.loadmat(os.path.join('../data/COFW', 'COFW_test_color.mat'))
+        images = test_mat['IsT']
+        
+        bboxes_mat = hdf5storage.loadmat(os.path.join('../data/data_300W_COFW_WFLW', 'cofw68_test_bboxes.mat'))
+        bboxes = bboxes_mat['bboxes']
+        image_num = images.shape[0]
+        with open('../data/data_300W_COFW_WFLW/test_COFW.txt', 'w') as f:
+            for i in range(image_num):
+                image = images[i,0]
+                # grayscale
+                if len(image.shape) == 2:
+                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+                # swap rgb channel to bgr
+                else:
+                    image = image[:,:,::-1]
+            
+                bbox = bboxes[i,:]
+                anno_mat = hdf5storage.loadmat(os.path.join('../data/data_300W_COFW_WFLW/cofw68_test_annotations', str(i+1)+'_points.mat'))
+                anno = anno_mat['Points']
+                image_crop, anno = process_cofw_68_test(image, bbox, anno, target_size)
+                pad_num = 4-len(str(i+1))
+                image_crop_name = 'cofw_test_' + '0' * pad_num + str(i+1) + '.jpg'
+                cv2.imwrite(os.path.join('../data/data_300W_COFW_WFLW/images_test_COFW', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in anno:
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+
+        # train for WFLW_68
+        train_file = 'list_98pt_rect_attr_train.txt'
+        with open(os.path.join('../data', 'WFLW', 'WFLW_annotations', 'list_98pt_rect_attr_train_test', train_file), 'r') as f:
+            annos_train = f.readlines()
+        annos_train = [x.strip().split() for x in annos_train]
+        count = 1
+        with open('../data/data_300W_COFW_WFLW/train_WFLW.txt', 'w') as f:
+            for anno_train in annos_train:
+                image_crop, anno = process_wflw(anno_train, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'wflw_train_' + '0' * pad_num + str(count) + '.jpg'
+                print(image_crop_name)
+                f.write(image_crop_name+'\n')
+                cv2.imwrite(os.path.join(root_folder, 'data_300W_COFW_WFLW', 'images_train', image_crop_name), image_crop)
+                count += 1
+
+        # test for WFLW_68
+        test_file = 'list_98pt_rect_attr_test.txt'
+        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_rect_attr_train_test', test_file), 'r') as f:
+            annos_test = f.readlines()
+        annos_test = [x.strip().split() for x in annos_test]
+        names_mapping = {}
+        count = 1
+        with open(os.path.join(root_folder, 'data_300W_COFW_WFLW', 'test_WFLW_98.txt'), 'w') as f:
+            for anno_test in annos_test:
+                image_crop, anno = process_wflw(anno_test, target_size)
+                pad_num = 4-len(str(count))
+                image_crop_name = 'wflw_test_' + '0' * pad_num + str(count) + '.jpg'
+                print(image_crop_name)
+                names_mapping[anno_test[0]+'_'+anno_test[-1]] = [image_crop_name, anno]
+                cv2.imwrite(os.path.join(root_folder, data_name, 'images_test_WFLW', image_crop_name), image_crop)
+                f.write(image_crop_name+' ')
+                for x,y in list(anno):
+                    f.write(str(x)+' '+str(y)+' ')
+                f.write('\n')
+                count += 1
+
+        convert_wflw(root_folder, data_name)
+
+        gen_meanface(root_folder, data_name)
+    else:
+        print('Wrong data!')
+
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        print('please input the data name.')
+        print('1. CELEBA')
+        print('2. data_300W_CELEBA')
+        print('3. data_300W_COFW_WFLW')
+        exit(0)
+    else:
+        data_name = sys.argv[1]
+        gen_data('../data', data_name, 256)
+
+

third_party/PIPNet/lib/tools.py

@@ -0,0 +1,174 @@
+import cv2
+import sys
+
+
+from math import floor
+from third_party.PIPNet.FaceBoxesV2.faceboxes_detector import *
+
+import torch
+import torch.nn.parallel
+import torch.utils.data
+import torchvision.transforms as transforms
+import torchvision.models as models
+
+from third_party.PIPNet.lib.networks import *
+from third_party.PIPNet.lib.functions import *
+from third_party.PIPNet.reverse_index import ri1, ri2
+
+
+make_abs_path = lambda fn: os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), fn))
+
+
+class Config:
+    def __init__(self):
+        self.det_head = "pip"
+        self.net_stride = 32
+        self.batch_size = 16
+        self.init_lr = 0.0001
+        self.num_epochs = 60
+        self.decay_steps = [30, 50]
+        self.input_size = 256
+        self.backbone = "resnet101"
+        self.pretrained = True
+        self.criterion_cls = "l2"
+        self.criterion_reg = "l1"
+        self.cls_loss_weight = 10
+        self.reg_loss_weight = 1
+        self.num_lms = 98
+        self.save_interval = self.num_epochs
+        self.num_nb = 10
+        self.use_gpu = True
+        self.gpu_id = 3
+
+
+def get_lmk_model():
+
+    cfg = Config()
+
+    resnet101 = models.resnet101(pretrained=cfg.pretrained)
+    net = Pip_resnet101(
+        resnet101,
+        cfg.num_nb,
+        num_lms=cfg.num_lms,
+        input_size=cfg.input_size,
+        net_stride=cfg.net_stride,
+    )
+
+    if cfg.use_gpu:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device("cpu")
+    net = net.to(device)
+
+    weight_file = make_abs_path('../../../weights/PIPNet/epoch59.pth')
+    state_dict = torch.load(weight_file, map_location=device)
+    net.load_state_dict(state_dict)
+
+    detector = FaceBoxesDetector(
+        "FaceBoxes",
+        make_abs_path("./../../weights/PIPNet/FaceBoxesV2.pth"),
+        use_gpu=torch.cuda.is_available(),
+        device=device,
+    )
+    return net, detector
+
+
+def demo_image(
+    image_file,
+    net,
+    detector,
+    input_size=256,
+    net_stride=32,
+    num_nb=10,
+    use_gpu=True,
+    device="cuda:0",
+):
+
+    my_thresh = 0.6
+    det_box_scale = 1.2
+    net.eval()
+    preprocess = transforms.Compose(
+        [
+            transforms.Resize((256, 256)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+    reverse_index1, reverse_index2, max_len = ri1, ri2, 17
+    # image = cv2.imread(image_file)
+    image = image_file
+    image_height, image_width, _ = image.shape
+    detections, _ = detector.detect(image, my_thresh, 1)
+    lmks = []
+    for i in range(len(detections)):
+        det_xmin = detections[i][2]
+        det_ymin = detections[i][3]
+        det_width = detections[i][4]
+        det_height = detections[i][5]
+        det_xmax = det_xmin + det_width - 1
+        det_ymax = det_ymin + det_height - 1
+
+        det_xmin -= int(det_width * (det_box_scale - 1) / 2)
+        # remove a part of top area for alignment, see paper for details
+        det_ymin += int(det_height * (det_box_scale - 1) / 2)
+        det_xmax += int(det_width * (det_box_scale - 1) / 2)
+        det_ymax += int(det_height * (det_box_scale - 1) / 2)
+        det_xmin = max(det_xmin, 0)
+        det_ymin = max(det_ymin, 0)
+        det_xmax = min(det_xmax, image_width - 1)
+        det_ymax = min(det_ymax, image_height - 1)
+        det_width = det_xmax - det_xmin + 1
+        det_height = det_ymax - det_ymin + 1
+
+        # cv2.rectangle(image, (det_xmin, det_ymin), (det_xmax, det_ymax), (0, 0, 255), 2)
+
+        det_crop = image[det_ymin:det_ymax, det_xmin:det_xmax, :]
+        det_crop = cv2.resize(det_crop, (input_size, input_size))
+        inputs = Image.fromarray(det_crop[:, :, ::-1].astype("uint8"), "RGB")
+        inputs = preprocess(inputs).unsqueeze(0)
+        inputs = inputs.to(device)
+        (
+            lms_pred_x,
+            lms_pred_y,
+            lms_pred_nb_x,
+            lms_pred_nb_y,
+            outputs_cls,
+            max_cls,
+        ) = forward_pip(net, inputs, preprocess, input_size, net_stride, num_nb)
+        lms_pred = torch.cat((lms_pred_x, lms_pred_y), dim=1).flatten()
+        tmp_nb_x = lms_pred_nb_x[reverse_index1, reverse_index2].view(98, max_len)
+        tmp_nb_y = lms_pred_nb_y[reverse_index1, reverse_index2].view(98, max_len)
+        tmp_x = torch.mean(torch.cat((lms_pred_x, tmp_nb_x), dim=1), dim=1).view(-1, 1)
+        tmp_y = torch.mean(torch.cat((lms_pred_y, tmp_nb_y), dim=1), dim=1).view(-1, 1)
+        lms_pred_merge = torch.cat((tmp_x, tmp_y), dim=1).flatten()
+        lms_pred = lms_pred.cpu().numpy()
+        lms_pred_merge = lms_pred_merge.cpu().numpy()
+        lmk_ = []
+        for i in range(98):
+            x_pred = lms_pred_merge[i * 2] * det_width
+            y_pred = lms_pred_merge[i * 2 + 1] * det_height
+
+            # cv2.circle(
+            #     image,
+            #     (int(x_pred) + det_xmin, int(y_pred) + det_ymin),
+            #     1,
+            #     (0, 0, 255),
+            #     1,
+            # )
+
+            lmk_.append([int(x_pred) + det_xmin, int(y_pred) + det_ymin])
+        lmks.append(np.array(lmk_))
+
+    # image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    # cv2.imwrite("./1_out.jpg", image_bgr)
+
+    return lmks
+
+
+if __name__ == "__main__":
+    net, detector = get_lmk_model()
+    demo_image(
+        "/apdcephfs/private_ahbanliang/codes/Real-ESRGAN-master/tmp_frames/yanikefu/frame00000046.png",
+        net,
+        detector,
+    )

third_party/PIPNet/lib/train.py

@@ -0,0 +1,196 @@
+import cv2, os
+import sys
+sys.path.insert(0, '..')
+import numpy as np
+from PIL import Image
+import logging
+import copy
+import importlib
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.utils.data
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+import torchvision.models as models
+
+from networks import *
+import data_utils
+from functions import * 
+from mobilenetv3 import mobilenetv3_large
+
+if not len(sys.argv) == 2:
+    print('Format:')
+    print('python lib/train.py config_file')
+    exit(0)
+experiment_name = sys.argv[1].split('/')[-1][:-3]
+data_name = sys.argv[1].split('/')[-2]
+config_path = '.experiments.{}.{}'.format(data_name, experiment_name)
+
+my_config = importlib.import_module(config_path, package='PIPNet')
+Config = getattr(my_config, 'Config')
+cfg = Config()
+cfg.experiment_name = experiment_name
+cfg.data_name = data_name
+
+os.environ['CUDA_VISIBLE_DEVICES'] = str(cfg.gpu_id)
+
+if not os.path.exists(os.path.join('./snapshots', cfg.data_name)):
+    os.mkdir(os.path.join('./snapshots', cfg.data_name))
+save_dir = os.path.join('./snapshots', cfg.data_name, cfg.experiment_name)
+if not os.path.exists(save_dir):
+    os.mkdir(save_dir)
+
+if not os.path.exists(os.path.join('./logs', cfg.data_name)):
+    os.mkdir(os.path.join('./logs', cfg.data_name))
+log_dir = os.path.join('./logs', cfg.data_name, cfg.experiment_name)
+if not os.path.exists(log_dir):
+    os.mkdir(log_dir)
+
+logging.basicConfig(filename=os.path.join(log_dir, 'train.log'), level=logging.INFO)
+
+print('###########################################')
+print('experiment_name:', cfg.experiment_name)
+print('data_name:', cfg.data_name)
+print('det_head:', cfg.det_head)
+print('net_stride:', cfg.net_stride)
+print('batch_size:', cfg.batch_size)
+print('init_lr:', cfg.init_lr)
+print('num_epochs:', cfg.num_epochs)
+print('decay_steps:', cfg.decay_steps)
+print('input_size:', cfg.input_size)
+print('backbone:', cfg.backbone)
+print('pretrained:', cfg.pretrained)
+print('criterion_cls:', cfg.criterion_cls)
+print('criterion_reg:', cfg.criterion_reg)
+print('cls_loss_weight:', cfg.cls_loss_weight)
+print('reg_loss_weight:', cfg.reg_loss_weight)
+print('num_lms:', cfg.num_lms)
+print('save_interval:', cfg.save_interval)
+print('num_nb:', cfg.num_nb)
+print('use_gpu:', cfg.use_gpu)
+print('gpu_id:', cfg.gpu_id)
+print('###########################################')
+logging.info('###########################################')
+logging.info('experiment_name: {}'.format(cfg.experiment_name))
+logging.info('data_name: {}'.format(cfg.data_name))
+logging.info('det_head: {}'.format(cfg.det_head))
+logging.info('net_stride: {}'.format(cfg.net_stride))
+logging.info('batch_size: {}'.format(cfg.batch_size))
+logging.info('init_lr: {}'.format(cfg.init_lr))
+logging.info('num_epochs: {}'.format(cfg.num_epochs))
+logging.info('decay_steps: {}'.format(cfg.decay_steps))
+logging.info('input_size: {}'.format(cfg.input_size))
+logging.info('backbone: {}'.format(cfg.backbone))
+logging.info('pretrained: {}'.format(cfg.pretrained))
+logging.info('criterion_cls: {}'.format(cfg.criterion_cls))
+logging.info('criterion_reg: {}'.format(cfg.criterion_reg))
+logging.info('cls_loss_weight: {}'.format(cfg.cls_loss_weight))
+logging.info('reg_loss_weight: {}'.format(cfg.reg_loss_weight))
+logging.info('num_lms: {}'.format(cfg.num_lms))
+logging.info('save_interval: {}'.format(cfg.save_interval))
+logging.info('num_nb: {}'.format(cfg.num_nb))
+logging.info('use_gpu: {}'.format(cfg.use_gpu))
+logging.info('gpu_id: {}'.format(cfg.gpu_id))
+logging.info('###########################################')
+
+if cfg.det_head == 'pip':
+    meanface_indices, _, _, _ = get_meanface(os.path.join('data', cfg.data_name, 'meanface.txt'), cfg.num_nb)
+
+
+if cfg.det_head == 'pip':
+    if cfg.backbone == 'resnet18':
+        resnet18 = models.resnet18(pretrained=cfg.pretrained)
+        net = Pip_resnet18(resnet18, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+    elif cfg.backbone == 'resnet50':
+        resnet50 = models.resnet50(pretrained=cfg.pretrained)
+        net = Pip_resnet50(resnet50, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+    elif cfg.backbone == 'resnet101':
+        resnet101 = models.resnet101(pretrained=cfg.pretrained)
+        net = Pip_resnet101(resnet101, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+    elif cfg.backbone == 'mobilenet_v2':
+        mbnet = models.mobilenet_v2(pretrained=cfg.pretrained)
+        net = Pip_mbnetv2(mbnet, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+    elif cfg.backbone == 'mobilenet_v3':
+        mbnet = mobilenetv3_large()
+        if cfg.pretrained:
+            mbnet.load_state_dict(torch.load('lib/mobilenetv3-large-1cd25616.pth'))
+        net = Pip_mbnetv3(mbnet, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+    else:
+        print('No such backbone!')
+        exit(0)
+else:
+    print('No such head:', cfg.det_head)
+    exit(0)
+
+if cfg.use_gpu:
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+else:
+    device = torch.device("cpu")
+net = net.to(device)
+
+criterion_cls = None
+if cfg.criterion_cls == 'l2':
+    criterion_cls = nn.MSELoss()
+elif cfg.criterion_cls == 'l1':
+    criterion_cls = nn.L1Loss()
+else:
+    print('No such cls criterion:', cfg.criterion_cls)
+
+criterion_reg = None
+if cfg.criterion_reg == 'l1':
+    criterion_reg = nn.L1Loss()
+elif cfg.criterion_reg == 'l2':
+    criterion_reg = nn.MSELoss()
+else:
+    print('No such reg criterion:', cfg.criterion_reg)
+
+points_flip = None
+if cfg.data_name == 'data_300W':
+    points_flip = [17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 27, 26, 25, 24, 23, 22, 21, 20, 19, 18, 28, 29, 30, 31, 36, 35, 34, 33, 32, 46, 45, 44, 43, 48, 47, 40, 39, 38, 37, 42, 41, 55, 54, 53, 52, 51, 50, 49, 60, 59, 58, 57, 56, 65, 64, 63, 62, 61, 68, 67, 66]
+    points_flip = (np.array(points_flip)-1).tolist()
+    assert len(points_flip) == 68
+elif cfg.data_name == 'WFLW':
+    points_flip = [32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 46, 45, 44, 43, 42, 50, 49, 48, 47, 37, 36, 35, 34, 33, 41, 40, 39, 38, 51, 52, 53, 54, 59, 58, 57, 56, 55, 72, 71, 70, 69, 68, 75, 74, 73, 64, 63, 62, 61, 60, 67, 66, 65, 82, 81, 80, 79, 78, 77, 76, 87, 86, 85, 84, 83, 92, 91, 90, 89, 88, 95, 94, 93, 97, 96]
+    assert len(points_flip) == 98
+elif cfg.data_name == 'COFW':
+    points_flip = [2, 1, 4, 3, 7, 8, 5, 6, 10, 9, 12, 11, 15, 16, 13, 14, 18, 17, 20, 19, 21, 22, 24, 23, 25, 26, 27, 28, 29]
+    points_flip = (np.array(points_flip)-1).tolist()
+    assert len(points_flip) == 29
+elif cfg.data_name == 'AFLW':
+    points_flip = [6, 5, 4, 3, 2, 1, 12, 11, 10, 9, 8, 7, 15, 14, 13, 18, 17, 16, 19]
+    points_flip = (np.array(points_flip)-1).tolist()
+    assert len(points_flip) == 19
+else:
+    print('No such data!')
+    exit(0)
+
+normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+
+if cfg.pretrained:  
+    optimizer = optim.Adam(net.parameters(), lr=cfg.init_lr)
+else:
+    optimizer = optim.Adam(net.parameters(), lr=cfg.init_lr, weight_decay=5e-4)
+scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=cfg.decay_steps, gamma=0.1)
+
+labels = get_label(cfg.data_name, 'train.txt')
+
+if cfg.det_head == 'pip':
+    train_data = data_utils.ImageFolder_pip(os.path.join('data', cfg.data_name, 'images_train'), 
+                                              labels, cfg.input_size, cfg.num_lms, 
+                                              cfg.net_stride, points_flip, meanface_indices,
+                                              transforms.Compose([
+                                              transforms.RandomGrayscale(0.2),
+                                              transforms.ToTensor(),
+                                              normalize]))
+else:
+    print('No such head:', cfg.det_head)
+    exit(0)
+
+train_loader = torch.utils.data.DataLoader(train_data, batch_size=cfg.batch_size, shuffle=True, num_workers=8, pin_memory=True, drop_last=True)
+
+train_model(cfg.det_head, net, train_loader, criterion_cls, criterion_reg, cfg.cls_loss_weight, cfg.reg_loss_weight, cfg.num_nb, optimizer, cfg.num_epochs, scheduler, save_dir, cfg.save_interval, device)
+

third_party/PIPNet/lib/train_gssl.py

@@ -0,0 +1,303 @@
+import cv2, os
+import sys
+sys.path.insert(0, '..')
+import numpy as np
+from PIL import Image
+import logging
+import importlib
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.utils.data
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+import torchvision.models as models
+
+from networks_gssl import *
+import data_utils_gssl
+from functions_gssl import * 
+
+if not len(sys.argv) == 2:
+    print('Format:')
+    print('python lib/train_gssl.py config_file')
+    exit(0)
+experiment_name = sys.argv[1].split('/')[-1][:-3]
+data_name = sys.argv[1].split('/')[-2]
+config_path = '.experiments.{}.{}'.format(data_name, experiment_name)
+
+my_config = importlib.import_module(config_path, package='PIPNet')
+Config = getattr(my_config, 'Config')
+cfg = Config()
+cfg.experiment_name = experiment_name
+cfg.data_name = data_name
+
+os.environ['CUDA_VISIBLE_DEVICES'] = str(cfg.gpu_id)
+
+if not os.path.exists(os.path.join('./snapshots', cfg.data_name)):
+    os.mkdir(os.path.join('./snapshots', cfg.data_name))
+save_dir = os.path.join('./snapshots', cfg.data_name, cfg.experiment_name)
+if not os.path.exists(save_dir):
+    os.mkdir(save_dir)
+
+if not os.path.exists(os.path.join('./logs', cfg.data_name)):
+    os.mkdir(os.path.join('./logs', cfg.data_name))
+log_dir = os.path.join('./logs', cfg.data_name, cfg.experiment_name)
+if not os.path.exists(log_dir):
+    os.mkdir(log_dir)
+
+logging.basicConfig(filename=os.path.join(log_dir, 'train.log'), level=logging.INFO)
+
+print('###########################################')
+print('experiment_name:', cfg.experiment_name)
+print('data_name:', cfg.data_name)
+print('det_head:', cfg.det_head)
+print('net_stride:', cfg.net_stride)
+print('batch_size:', cfg.batch_size)
+print('init_lr:', cfg.init_lr)
+print('num_epochs:', cfg.num_epochs)
+print('decay_steps:', cfg.decay_steps)
+print('input_size:', cfg.input_size)
+print('backbone:', cfg.backbone)
+print('pretrained:', cfg.pretrained)
+print('criterion_cls:', cfg.criterion_cls)
+print('criterion_reg:', cfg.criterion_reg)
+print('cls_loss_weight:', cfg.cls_loss_weight)
+print('reg_loss_weight:', cfg.reg_loss_weight)
+print('num_lms:', cfg.num_lms)
+print('save_interval:', cfg.save_interval)
+print('num_nb:', cfg.num_nb)
+print('use_gpu:', cfg.use_gpu)
+print('gpu_id:', cfg.gpu_id)
+print('curriculum:', cfg.curriculum)
+print('###########################################')
+logging.info('###########################################')
+logging.info('experiment_name: {}'.format(cfg.experiment_name))
+logging.info('data_name: {}'.format(cfg.data_name))
+logging.info('det_head: {}'.format(cfg.det_head))
+logging.info('net_stride: {}'.format(cfg.net_stride))
+logging.info('batch_size: {}'.format(cfg.batch_size))
+logging.info('init_lr: {}'.format(cfg.init_lr))
+logging.info('num_epochs: {}'.format(cfg.num_epochs))
+logging.info('decay_steps: {}'.format(cfg.decay_steps))
+logging.info('input_size: {}'.format(cfg.input_size))
+logging.info('backbone: {}'.format(cfg.backbone))
+logging.info('pretrained: {}'.format(cfg.pretrained))
+logging.info('criterion_cls: {}'.format(cfg.criterion_cls))
+logging.info('criterion_reg: {}'.format(cfg.criterion_reg))
+logging.info('cls_loss_weight: {}'.format(cfg.cls_loss_weight))
+logging.info('reg_loss_weight: {}'.format(cfg.reg_loss_weight))
+logging.info('num_lms: {}'.format(cfg.num_lms))
+logging.info('save_interval: {}'.format(cfg.save_interval))
+logging.info('num_nb: {}'.format(cfg.num_nb))
+logging.info('use_gpu: {}'.format(cfg.use_gpu))
+logging.info('gpu_id: {}'.format(cfg.gpu_id))
+logging.info('###########################################')
+
+if cfg.curriculum:
+    # self-training with curriculum
+    task_type_list = ['cls3', 'cls2', 'std', 'std', 'std']
+else:
+    # standard self-training 
+    task_type_list = ['std']*3
+
+meanface_indices, reverse_index1, reverse_index2, max_len = get_meanface(os.path.join('data', cfg.data_name, 'meanface.txt'), cfg.num_nb)
+
+if cfg.det_head == 'pip':
+    if cfg.backbone == 'resnet18':
+        resnet18 = models.resnet18(pretrained=cfg.pretrained)
+        net = Pip_resnet18(resnet18, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+    else:
+        print('No such backbone!')
+        exit(0)
+else:
+    print('No such head:', cfg.det_head)
+    exit(0)
+
+if cfg.use_gpu:
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+else:
+    device = torch.device("cpu")
+net = net.to(device)
+
+criterion_cls = None
+if cfg.criterion_cls == 'l2':
+    criterion_cls = nn.MSELoss(reduction='sum')
+elif cfg.criterion_cls == 'l1':
+    criterion_cls = nn.L1Loss()
+else:
+    print('No such cls criterion:', cfg.criterion_cls)
+
+criterion_reg = None
+if cfg.criterion_reg == 'l1':
+    criterion_reg = nn.L1Loss(reduction='sum')
+elif cfg.criterion_reg == 'l2':
+    criterion_reg = nn.MSELoss()
+else:
+    print('No such reg criterion:', cfg.criterion_reg)
+
+points_flip = [17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 27, 26, 25, 24, 23, 22, 21, 20, 19, 18, 28, 29, 30, 31, 36, 35, 34, 33, 32, 46, 45, 44, 43, 48, 47, 40, 39, 38, 37, 42, 41, 55, 54, 53, 52, 51, 50, 49, 60, 59, 58, 57, 56, 65, 64, 63, 62, 61, 68, 67, 66]
+points_flip = (np.array(points_flip)-1).tolist()
+assert len(points_flip) == 68
+
+normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+
+optimizer = optim.Adam(net.parameters(), lr=cfg.init_lr)
+scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=cfg.decay_steps, gamma=0.1)
+
+labels = get_label(cfg.data_name, 'train_300W.txt', 'std')
+
+train_data = data_utils_gssl.ImageFolder_pip(os.path.join('data', cfg.data_name, 'images_train'), 
+                                        labels, cfg.input_size, cfg.num_lms, 
+                                        cfg.net_stride, points_flip, meanface_indices,
+                                        transforms.Compose([
+                                        transforms.RandomGrayscale(0.2),
+                                        transforms.ToTensor(),
+                                        normalize]))
+
+train_loader = torch.utils.data.DataLoader(train_data, batch_size=cfg.batch_size, shuffle=True, num_workers=8, pin_memory=True, drop_last=True)
+
+train_model(cfg.det_head, net, train_loader, criterion_cls, criterion_reg, cfg.cls_loss_weight, cfg.reg_loss_weight, cfg.num_nb, optimizer, cfg.num_epochs, scheduler, save_dir, cfg.save_interval, device)
+
+###############
+# test
+norm_indices = [36, 45]
+
+preprocess = transforms.Compose([transforms.Resize((cfg.input_size, cfg.input_size)), transforms.ToTensor(), normalize])
+test_data_list = ['300W', 'COFW', 'WFLW']
+for test_data in test_data_list:
+    labels = get_label(cfg.data_name, 'test_'+test_data+'.txt')
+    nmes = []
+    norm = None
+    for label in labels:
+        image_name = label[0]
+        lms_gt = label[1]
+        image_path = os.path.join('data', cfg.data_name, 'images_test_'+test_data, image_name)
+        image = cv2.imread(image_path)
+        image = cv2.resize(image, (cfg.input_size, cfg.input_size))
+        inputs = Image.fromarray(image[:,:,::-1].astype('uint8'), 'RGB')
+        inputs = preprocess(inputs).unsqueeze(0)
+        inputs = inputs.to(device)
+        lms_pred_x, lms_pred_y, lms_pred_nb_x, lms_pred_nb_y, outputs_cls, max_cls = forward_pip(net, inputs, preprocess, cfg.input_size, cfg.net_stride, cfg.num_nb)
+        # inter-ocular
+        norm = np.linalg.norm(lms_gt.reshape(-1, 2)[norm_indices[0]] - lms_gt.reshape(-1, 2)[norm_indices[1]])
+        #############################
+        # merge neighbor predictions
+        lms_pred = torch.cat((lms_pred_x, lms_pred_y), dim=1).flatten().cpu().numpy()
+        tmp_nb_x = lms_pred_nb_x[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+        tmp_nb_y = lms_pred_nb_y[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+        tmp_x = torch.mean(torch.cat((lms_pred_x, tmp_nb_x), dim=1), dim=1).view(-1,1)
+        tmp_y = torch.mean(torch.cat((lms_pred_y, tmp_nb_y), dim=1), dim=1).view(-1,1)
+        lms_pred_merge = torch.cat((tmp_x, tmp_y), dim=1).flatten().cpu().numpy()
+        #############################
+        nme = compute_nme(lms_pred_merge, lms_gt, norm)
+        nmes.append(nme)
+    
+    print('{} nme: {}'.format(test_data, np.mean(nmes)))
+    logging.info('{} nme: {}'.format(test_data, np.mean(nmes)))
+
+for ti, task_type in enumerate(task_type_list):
+    print('###################################################')
+    print('Iter:', ti, 'task_type:', task_type)
+    ###############
+    # estimate
+    if cfg.data_name == 'data_300W_COFW_WFLW':
+        est_data_list = ['COFW', 'WFLW']
+    elif cfg.data_name == 'data_300W_CELEBA': 
+        est_data_list = ['CELEBA']
+    else:
+        print('No such data!')
+        exit(0)
+    est_preds = []
+    for est_data in est_data_list:
+        labels = get_label(cfg.data_name, 'train_'+est_data+'.txt')
+        for label in labels:
+            image_name = label[0]
+            #print(image_name)
+            image_path = os.path.join('data', cfg.data_name, 'images_train', image_name)
+            image = cv2.imread(image_path)
+            image = cv2.resize(image, (cfg.input_size, cfg.input_size))
+            inputs = Image.fromarray(image[:,:,::-1].astype('uint8'), 'RGB')
+            inputs = preprocess(inputs).unsqueeze(0)
+            inputs = inputs.to(device)
+            lms_pred_x, lms_pred_y, lms_pred_nb_x, lms_pred_nb_y, outputs_cls, max_cls = forward_pip(net, inputs, preprocess, cfg.input_size, cfg.net_stride, cfg.num_nb)
+            #############################
+            # merge neighbor predictions
+            lms_pred = torch.cat((lms_pred_x, lms_pred_y), dim=1).flatten().cpu().numpy()
+            tmp_nb_x = lms_pred_nb_x[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+            tmp_nb_y = lms_pred_nb_y[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+            tmp_x = torch.mean(torch.cat((lms_pred_x, tmp_nb_x), dim=1), dim=1).view(-1,1)
+            tmp_y = torch.mean(torch.cat((lms_pred_y, tmp_nb_y), dim=1), dim=1).view(-1,1)
+            lms_pred_merge = torch.cat((tmp_x, tmp_y), dim=1).flatten().cpu().numpy()
+            #############################
+            est_preds.append([image_name, task_type, lms_pred_merge])
+    
+    ################
+    # GSSL 
+    if cfg.det_head == 'pip':
+        if cfg.backbone == 'resnet18':
+            resnet18 = models.resnet18(pretrained=cfg.pretrained)
+            net = Pip_resnet18(resnet18, cfg.num_nb, num_lms=cfg.num_lms, input_size=cfg.input_size, net_stride=cfg.net_stride)
+        else:
+            print('No such backbone!')
+            exit(0)
+    else:
+        print('No such head:', cfg.det_head)
+        exit(0)
+
+    net = net.to(device)
+    optimizer = optim.Adam(net.parameters(), lr=cfg.init_lr)
+    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=cfg.decay_steps, gamma=0.1)
+    labels = get_label(cfg.data_name, 'train_300W.txt', 'std')
+    labels += est_preds
+    
+    train_data = data_utils_gssl.ImageFolder_pip(os.path.join('data', cfg.data_name, 'images_train'), 
+                                            labels, cfg.input_size, cfg.num_lms, 
+                                            cfg.net_stride, points_flip, meanface_indices,
+                                            transforms.Compose([
+                                            transforms.RandomGrayscale(0.2),
+                                            transforms.ToTensor(),
+                                            normalize]))
+    
+    train_loader = torch.utils.data.DataLoader(train_data, batch_size=cfg.batch_size, shuffle=True, num_workers=8, pin_memory=True, drop_last=True)
+    
+    train_model(cfg.det_head, net, train_loader, criterion_cls, criterion_reg, cfg.cls_loss_weight, cfg.reg_loss_weight, cfg.num_nb, optimizer, cfg.num_epochs, scheduler, save_dir, cfg.save_interval, device)
+    
+    ###############
+    # test
+    preprocess = transforms.Compose([transforms.Resize((cfg.input_size, cfg.input_size)), transforms.ToTensor(), normalize])
+    test_data_list = ['300W', 'COFW', 'WFLW']
+    for test_data in test_data_list:
+        labels = get_label(cfg.data_name, 'test_'+test_data+'.txt')
+        nmes = []
+        norm = None
+        for label in labels:
+            image_name = label[0]
+            lms_gt = label[1]
+            image_path = os.path.join('data', cfg.data_name, 'images_test_'+test_data, image_name)
+            image = cv2.imread(image_path)
+            image = cv2.resize(image, (cfg.input_size, cfg.input_size))
+            inputs = Image.fromarray(image[:,:,::-1].astype('uint8'), 'RGB')
+            inputs = preprocess(inputs).unsqueeze(0)
+            inputs = inputs.to(device)
+            lms_pred_x, lms_pred_y, lms_pred_nb_x, lms_pred_nb_y, outputs_cls, max_cls = forward_pip(net, inputs, preprocess, cfg.input_size, cfg.net_stride, cfg.num_nb)
+            # inter-ocular
+            norm = np.linalg.norm(lms_gt.reshape(-1, 2)[norm_indices[0]] - lms_gt.reshape(-1, 2)[norm_indices[1]])
+            #############################
+            # merge neighbor predictions
+            lms_pred = torch.cat((lms_pred_x, lms_pred_y), dim=1).flatten().cpu().numpy()
+            tmp_nb_x = lms_pred_nb_x[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+            tmp_nb_y = lms_pred_nb_y[reverse_index1, reverse_index2].view(cfg.num_lms, max_len)
+            tmp_x = torch.mean(torch.cat((lms_pred_x, tmp_nb_x), dim=1), dim=1).view(-1,1)
+            tmp_y = torch.mean(torch.cat((lms_pred_y, tmp_nb_y), dim=1), dim=1).view(-1,1)
+            lms_pred_merge = torch.cat((tmp_x, tmp_y), dim=1).flatten().cpu().numpy()
+            #############################
+            nme = compute_nme(lms_pred_merge, lms_gt, norm)
+            nmes.append(nme)
+        
+        print('{} nme: {}'.format(test_data, np.mean(nmes)))
+        logging.info('{} nme: {}'.format(test_data, np.mean(nmes)))
+
+

third_party/PIPNet/requirements.txt

@@ -0,0 +1,3 @@
+opencv-python
+scipy
+Cython

third_party/PIPNet/reverse_index.py

@@ -0,0 +1,3338 @@
+ri1 = [
+    1,
+    2,
+    3,
+    4,
+    5,
+    33,
+    1,
+    2,
+    3,
+    4,
+    5,
+    33,
+    1,
+    2,
+    3,
+    4,
+    5,
+    0,
+    2,
+    3,
+    4,
+    5,
+    6,
+    33,
+    0,
+    2,
+    3,
+    4,
+    5,
+    6,
+    33,
+    0,
+    2,
+    3,
+    0,
+    1,
+    3,
+    4,
+    5,
+    6,
+    0,
+    1,
+    3,
+    4,
+    5,
+    6,
+    0,
+    1,
+    3,
+    4,
+    5,
+    0,
+    1,
+    2,
+    4,
+    5,
+    6,
+    7,
+    0,
+    1,
+    2,
+    4,
+    5,
+    6,
+    7,
+    0,
+    1,
+    2,
+    0,
+    1,
+    2,
+    3,
+    5,
+    6,
+    7,
+    8,
+    0,
+    1,
+    2,
+    3,
+    5,
+    6,
+    7,
+    8,
+    0,
+    1,
+    2,
+    3,
+    4,
+    6,
+    7,
+    8,
+    9,
+    1,
+    2,
+    3,
+    4,
+    6,
+    7,
+    8,
+    9,
+    1,
+    2,
+    3,
+    4,
+    5,
+    7,
+    8,
+    9,
+    10,
+    2,
+    3,
+    4,
+    5,
+    7,
+    8,
+    9,
+    10,
+    2,
+    3,
+    4,
+    5,
+    6,
+    8,
+    9,
+    10,
+    3,
+    4,
+    5,
+    6,
+    8,
+    9,
+    10,
+    3,
+    4,
+    5,
+    4,
+    5,
+    6,
+    7,
+    9,
+    10,
+    11,
+    4,
+    5,
+    6,
+    7,
+    9,
+    10,
+    11,
+    4,
+    5,
+    6,
+    4,
+    5,
+    6,
+    7,
+    8,
+    10,
+    11,
+    12,
+    4,
+    5,
+    6,
+    7,
+    8,
+    10,
+    11,
+    12,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    11,
+    12,
+    13,
+    76,
+    5,
+    6,
+    7,
+    8,
+    9,
+    11,
+    12,
+    13,
+    7,
+    8,
+    9,
+    10,
+    12,
+    13,
+    14,
+    76,
+    88,
+    7,
+    8,
+    9,
+    10,
+    12,
+    13,
+    14,
+    76,
+    8,
+    9,
+    10,
+    11,
+    13,
+    14,
+    15,
+    8,
+    9,
+    10,
+    11,
+    13,
+    14,
+    15,
+    8,
+    9,
+    10,
+    10,
+    11,
+    12,
+    14,
+    15,
+    16,
+    10,
+    11,
+    12,
+    14,
+    15,
+    16,
+    10,
+    11,
+    12,
+    14,
+    15,
+    11,
+    12,
+    13,
+    15,
+    16,
+    17,
+    11,
+    12,
+    13,
+    15,
+    16,
+    17,
+    11,
+    12,
+    13,
+    15,
+    16,
+    12,
+    13,
+    14,
+    16,
+    17,
+    18,
+    12,
+    13,
+    14,
+    16,
+    17,
+    18,
+    12,
+    13,
+    14,
+    16,
+    17,
+    13,
+    14,
+    15,
+    17,
+    18,
+    19,
+    13,
+    14,
+    15,
+    17,
+    18,
+    19,
+    13,
+    14,
+    15,
+    17,
+    18,
+    14,
+    15,
+    16,
+    18,
+    19,
+    20,
+    14,
+    15,
+    16,
+    18,
+    19,
+    20,
+    14,
+    15,
+    16,
+    18,
+    19,
+    15,
+    16,
+    17,
+    19,
+    20,
+    21,
+    15,
+    16,
+    17,
+    19,
+    20,
+    21,
+    15,
+    16,
+    17,
+    19,
+    20,
+    16,
+    17,
+    18,
+    20,
+    21,
+    22,
+    16,
+    17,
+    18,
+    20,
+    21,
+    22,
+    16,
+    17,
+    18,
+    20,
+    21,
+    17,
+    18,
+    19,
+    21,
+    22,
+    23,
+    24,
+    17,
+    18,
+    19,
+    21,
+    22,
+    23,
+    24,
+    17,
+    18,
+    19,
+    18,
+    19,
+    20,
+    22,
+    23,
+    24,
+    25,
+    82,
+    18,
+    19,
+    20,
+    22,
+    23,
+    24,
+    25,
+    82,
+    18,
+    19,
+    20,
+    21,
+    23,
+    24,
+    25,
+    26,
+    27,
+    19,
+    20,
+    21,
+    23,
+    24,
+    25,
+    26,
+    27,
+    19,
+    20,
+    21,
+    22,
+    24,
+    25,
+    26,
+    27,
+    28,
+    20,
+    21,
+    22,
+    24,
+    25,
+    26,
+    27,
+    28,
+    20,
+    21,
+    22,
+    23,
+    25,
+    26,
+    27,
+    28,
+    21,
+    22,
+    23,
+    25,
+    26,
+    27,
+    28,
+    21,
+    22,
+    23,
+    21,
+    22,
+    23,
+    24,
+    26,
+    27,
+    28,
+    29,
+    21,
+    22,
+    23,
+    24,
+    26,
+    27,
+    28,
+    29,
+    21,
+    22,
+    23,
+    24,
+    25,
+    27,
+    28,
+    29,
+    30,
+    22,
+    23,
+    24,
+    25,
+    27,
+    28,
+    29,
+    30,
+    22,
+    23,
+    24,
+    25,
+    26,
+    28,
+    29,
+    30,
+    31,
+    23,
+    24,
+    25,
+    26,
+    28,
+    29,
+    30,
+    31,
+    23,
+    24,
+    25,
+    26,
+    27,
+    29,
+    30,
+    31,
+    32,
+    24,
+    25,
+    26,
+    27,
+    29,
+    30,
+    31,
+    32,
+    24,
+    25,
+    26,
+    27,
+    28,
+    30,
+    31,
+    32,
+    25,
+    26,
+    27,
+    28,
+    30,
+    31,
+    32,
+    25,
+    26,
+    27,
+    26,
+    27,
+    28,
+    29,
+    31,
+    32,
+    26,
+    27,
+    28,
+    29,
+    31,
+    32,
+    26,
+    27,
+    28,
+    29,
+    31,
+    26,
+    27,
+    28,
+    29,
+    30,
+    32,
+    46,
+    26,
+    27,
+    28,
+    29,
+    30,
+    32,
+    46,
+    26,
+    27,
+    28,
+    27,
+    28,
+    29,
+    30,
+    31,
+    46,
+    27,
+    28,
+    29,
+    30,
+    31,
+    46,
+    27,
+    28,
+    29,
+    30,
+    31,
+    0,
+    1,
+    2,
+    3,
+    34,
+    41,
+    60,
+    0,
+    1,
+    2,
+    3,
+    34,
+    41,
+    60,
+    0,
+    1,
+    2,
+    0,
+    33,
+    35,
+    40,
+    41,
+    60,
+    0,
+    33,
+    35,
+    40,
+    41,
+    60,
+    0,
+    33,
+    35,
+    40,
+    41,
+    33,
+    34,
+    36,
+    37,
+    39,
+    40,
+    41,
+    60,
+    61,
+    62,
+    33,
+    34,
+    36,
+    37,
+    39,
+    40,
+    41,
+    34,
+    35,
+    37,
+    38,
+    39,
+    40,
+    63,
+    64,
+    34,
+    35,
+    37,
+    38,
+    39,
+    40,
+    63,
+    64,
+    34,
+    36,
+    38,
+    39,
+    51,
+    64,
+    36,
+    38,
+    39,
+    51,
+    64,
+    36,
+    38,
+    39,
+    51,
+    64,
+    36,
+    38,
+    36,
+    37,
+    39,
+    51,
+    52,
+    63,
+    64,
+    65,
+    36,
+    37,
+    39,
+    51,
+    52,
+    63,
+    64,
+    65,
+    36,
+    35,
+    36,
+    37,
+    38,
+    40,
+    62,
+    63,
+    64,
+    65,
+    66,
+    67,
+    96,
+    35,
+    36,
+    37,
+    38,
+    40,
+    33,
+    34,
+    35,
+    36,
+    37,
+    38,
+    39,
+    41,
+    60,
+    61,
+    62,
+    63,
+    65,
+    66,
+    67,
+    96,
+    33,
+    0,
+    1,
+    2,
+    33,
+    34,
+    35,
+    40,
+    60,
+    61,
+    67,
+    0,
+    1,
+    2,
+    33,
+    34,
+    35,
+    40,
+    43,
+    49,
+    50,
+    51,
+    68,
+    43,
+    49,
+    50,
+    51,
+    68,
+    43,
+    49,
+    50,
+    51,
+    68,
+    43,
+    49,
+    42,
+    44,
+    45,
+    48,
+    49,
+    50,
+    68,
+    69,
+    42,
+    44,
+    45,
+    48,
+    49,
+    50,
+    68,
+    69,
+    42,
+    42,
+    43,
+    45,
+    46,
+    47,
+    48,
+    49,
+    70,
+    42,
+    43,
+    45,
+    46,
+    47,
+    48,
+    49,
+    70,
+    42,
+    32,
+    44,
+    46,
+    47,
+    48,
+    71,
+    72,
+    73,
+    32,
+    44,
+    46,
+    47,
+    48,
+    71,
+    72,
+    73,
+    32,
+    29,
+    30,
+    31,
+    32,
+    45,
+    47,
+    72,
+    29,
+    30,
+    31,
+    32,
+    45,
+    47,
+    72,
+    29,
+    30,
+    31,
+    30,
+    31,
+    32,
+    44,
+    45,
+    46,
+    48,
+    71,
+    72,
+    73,
+    30,
+    31,
+    32,
+    44,
+    45,
+    46,
+    48,
+    42,
+    43,
+    44,
+    45,
+    46,
+    47,
+    49,
+    50,
+    69,
+    70,
+    71,
+    72,
+    73,
+    74,
+    75,
+    97,
+    42,
+    42,
+    43,
+    44,
+    48,
+    50,
+    68,
+    69,
+    70,
+    74,
+    75,
+    97,
+    42,
+    43,
+    44,
+    48,
+    50,
+    68,
+    42,
+    43,
+    49,
+    51,
+    52,
+    68,
+    69,
+    75,
+    42,
+    43,
+    49,
+    51,
+    52,
+    68,
+    69,
+    75,
+    42,
+    37,
+    38,
+    42,
+    50,
+    52,
+    53,
+    64,
+    68,
+    37,
+    38,
+    42,
+    50,
+    52,
+    53,
+    64,
+    68,
+    37,
+    51,
+    53,
+    54,
+    51,
+    53,
+    54,
+    51,
+    53,
+    54,
+    51,
+    53,
+    54,
+    51,
+    53,
+    54,
+    51,
+    53,
+    51,
+    52,
+    54,
+    55,
+    56,
+    57,
+    59,
+    51,
+    52,
+    54,
+    55,
+    56,
+    57,
+    59,
+    51,
+    52,
+    54,
+    52,
+    53,
+    55,
+    56,
+    57,
+    58,
+    59,
+    52,
+    53,
+    55,
+    56,
+    57,
+    58,
+    59,
+    52,
+    53,
+    55,
+    53,
+    54,
+    56,
+    57,
+    76,
+    77,
+    78,
+    88,
+    53,
+    54,
+    56,
+    57,
+    76,
+    77,
+    78,
+    88,
+    53,
+    53,
+    54,
+    55,
+    57,
+    58,
+    77,
+    78,
+    79,
+    88,
+    53,
+    54,
+    55,
+    57,
+    58,
+    77,
+    78,
+    79,
+    53,
+    54,
+    55,
+    56,
+    58,
+    59,
+    78,
+    79,
+    80,
+    90,
+    53,
+    54,
+    55,
+    56,
+    58,
+    59,
+    78,
+    53,
+    54,
+    56,
+    57,
+    59,
+    79,
+    80,
+    81,
+    82,
+    92,
+    53,
+    54,
+    56,
+    57,
+    59,
+    79,
+    80,
+    53,
+    54,
+    57,
+    58,
+    80,
+    81,
+    82,
+    92,
+    53,
+    54,
+    57,
+    58,
+    80,
+    81,
+    82,
+    92,
+    53,
+    0,
+    1,
+    2,
+    3,
+    4,
+    33,
+    34,
+    41,
+    61,
+    62,
+    66,
+    67,
+    96,
+    0,
+    1,
+    2,
+    3,
+    0,
+    1,
+    33,
+    34,
+    35,
+    40,
+    41,
+    60,
+    62,
+    63,
+    65,
+    66,
+    67,
+    96,
+    0,
+    1,
+    33,
+    33,
+    34,
+    35,
+    36,
+    37,
+    38,
+    39,
+    40,
+    41,
+    60,
+    61,
+    63,
+    64,
+    65,
+    66,
+    67,
+    96,
+    35,
+    36,
+    37,
+    38,
+    39,
+    40,
+    51,
+    52,
+    61,
+    62,
+    64,
+    65,
+    66,
+    67,
+    96,
+    35,
+    36,
+    36,
+    37,
+    38,
+    39,
+    51,
+    52,
+    53,
+    63,
+    65,
+    66,
+    96,
+    36,
+    37,
+    38,
+    39,
+    51,
+    52,
+    36,
+    37,
+    38,
+    39,
+    52,
+    61,
+    62,
+    63,
+    64,
+    66,
+    67,
+    96,
+    36,
+    37,
+    38,
+    39,
+    52,
+    41,
+    60,
+    61,
+    62,
+    63,
+    64,
+    65,
+    67,
+    96,
+    41,
+    60,
+    61,
+    62,
+    63,
+    64,
+    65,
+    67,
+    0,
+    1,
+    2,
+    3,
+    33,
+    34,
+    35,
+    40,
+    41,
+    60,
+    61,
+    62,
+    65,
+    66,
+    96,
+    0,
+    1,
+    42,
+    43,
+    49,
+    50,
+    51,
+    52,
+    53,
+    69,
+    74,
+    75,
+    97,
+    42,
+    43,
+    49,
+    50,
+    51,
+    52,
+    42,
+    43,
+    44,
+    48,
+    49,
+    50,
+    51,
+    68,
+    70,
+    71,
+    73,
+    74,
+    75,
+    97,
+    42,
+    43,
+    44,
+    42,
+    43,
+    44,
+    45,
+    46,
+    47,
+    48,
+    49,
+    50,
+    68,
+    69,
+    71,
+    72,
+    73,
+    74,
+    75,
+    97,
+    31,
+    32,
+    44,
+    45,
+    46,
+    47,
+    48,
+    69,
+    70,
+    72,
+    73,
+    74,
+    75,
+    97,
+    31,
+    32,
+    44,
+    28,
+    29,
+    30,
+    31,
+    32,
+    45,
+    46,
+    47,
+    70,
+    71,
+    73,
+    74,
+    97,
+    28,
+    29,
+    30,
+    31,
+    29,
+    30,
+    31,
+    32,
+    44,
+    45,
+    46,
+    47,
+    48,
+    70,
+    71,
+    72,
+    74,
+    75,
+    97,
+    29,
+    30,
+    47,
+    68,
+    69,
+    70,
+    71,
+    72,
+    73,
+    75,
+    97,
+    47,
+    68,
+    69,
+    70,
+    71,
+    72,
+    73,
+    75,
+    42,
+    43,
+    49,
+    50,
+    52,
+    68,
+    69,
+    70,
+    71,
+    72,
+    73,
+    74,
+    97,
+    42,
+    43,
+    49,
+    50,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    55,
+    77,
+    87,
+    88,
+    89,
+    95,
+    6,
+    7,
+    8,
+    9,
+    55,
+    56,
+    76,
+    78,
+    86,
+    87,
+    88,
+    89,
+    95,
+    55,
+    56,
+    76,
+    78,
+    86,
+    87,
+    88,
+    89,
+    54,
+    55,
+    56,
+    57,
+    58,
+    76,
+    77,
+    79,
+    80,
+    85,
+    86,
+    87,
+    88,
+    89,
+    90,
+    94,
+    95,
+    54,
+    55,
+    56,
+    57,
+    58,
+    59,
+    77,
+    78,
+    80,
+    81,
+    84,
+    85,
+    86,
+    89,
+    90,
+    91,
+    94,
+    54,
+    57,
+    58,
+    59,
+    78,
+    79,
+    81,
+    82,
+    83,
+    84,
+    85,
+    90,
+    91,
+    92,
+    93,
+    94,
+    54,
+    58,
+    59,
+    80,
+    82,
+    83,
+    84,
+    91,
+    92,
+    93,
+    58,
+    59,
+    80,
+    82,
+    83,
+    84,
+    91,
+    92,
+    20,
+    21,
+    22,
+    23,
+    24,
+    25,
+    26,
+    59,
+    81,
+    83,
+    91,
+    92,
+    93,
+    20,
+    21,
+    22,
+    23,
+    17,
+    18,
+    19,
+    20,
+    21,
+    22,
+    23,
+    81,
+    82,
+    84,
+    91,
+    92,
+    93,
+    17,
+    18,
+    19,
+    20,
+    16,
+    17,
+    18,
+    19,
+    20,
+    81,
+    82,
+    83,
+    85,
+    91,
+    92,
+    93,
+    94,
+    16,
+    17,
+    18,
+    19,
+    14,
+    15,
+    16,
+    17,
+    18,
+    83,
+    84,
+    86,
+    87,
+    90,
+    93,
+    94,
+    95,
+    14,
+    15,
+    16,
+    17,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    76,
+    77,
+    85,
+    87,
+    88,
+    89,
+    94,
+    95,
+    11,
+    12,
+    13,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    76,
+    77,
+    86,
+    88,
+    89,
+    95,
+    9,
+    10,
+    11,
+    12,
+    13,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    55,
+    76,
+    77,
+    86,
+    87,
+    89,
+    95,
+    7,
+    8,
+    9,
+    55,
+    56,
+    76,
+    77,
+    78,
+    79,
+    86,
+    87,
+    88,
+    90,
+    95,
+    55,
+    56,
+    76,
+    77,
+    78,
+    79,
+    56,
+    57,
+    58,
+    78,
+    79,
+    80,
+    83,
+    84,
+    85,
+    86,
+    87,
+    89,
+    91,
+    92,
+    93,
+    94,
+    95,
+    58,
+    59,
+    79,
+    80,
+    81,
+    82,
+    83,
+    84,
+    85,
+    90,
+    92,
+    93,
+    94,
+    58,
+    59,
+    79,
+    80,
+    19,
+    20,
+    21,
+    22,
+    23,
+    24,
+    25,
+    59,
+    81,
+    82,
+    83,
+    84,
+    91,
+    93,
+    19,
+    20,
+    21,
+    18,
+    19,
+    79,
+    80,
+    81,
+    82,
+    83,
+    84,
+    85,
+    90,
+    91,
+    92,
+    94,
+    18,
+    19,
+    79,
+    80,
+    15,
+    16,
+    17,
+    78,
+    79,
+    80,
+    83,
+    84,
+    85,
+    86,
+    87,
+    89,
+    90,
+    91,
+    93,
+    95,
+    15,
+    13,
+    14,
+    15,
+    76,
+    77,
+    78,
+    85,
+    86,
+    87,
+    88,
+    89,
+    90,
+    94,
+    13,
+    14,
+    15,
+    76,
+    34,
+    35,
+    36,
+    38,
+    39,
+    40,
+    41,
+    60,
+    61,
+    62,
+    63,
+    64,
+    65,
+    66,
+    67,
+    34,
+    35,
+    43,
+    44,
+    45,
+    47,
+    48,
+    49,
+    50,
+    68,
+    69,
+    70,
+    71,
+    72,
+    73,
+    74,
+    75,
+    43,
+    44,
+]
+
+
+ri2 = [
+    0,
+    2,
+    4,
+    6,
+    8,
+    4,
+    0,
+    2,
+    4,
+    6,
+    8,
+    4,
+    0,
+    2,
+    4,
+    6,
+    8,
+    0,
+    0,
+    2,
+    4,
+    6,
+    8,
+    8,
+    0,
+    0,
+    2,
+    4,
+    6,
+    8,
+    8,
+    0,
+    0,
+    2,
+    1,
+    1,
+    0,
+    2,
+    4,
+    6,
+    1,
+    1,
+    0,
+    2,
+    4,
+    6,
+    1,
+    1,
+    0,
+    2,
+    4,
+    3,
+    2,
+    1,
+    0,
+    2,
+    4,
+    6,
+    3,
+    2,
+    1,
+    0,
+    2,
+    4,
+    6,
+    3,
+    2,
+    1,
+    6,
+    3,
+    3,
+    1,
+    0,
+    2,
+    4,
+    7,
+    6,
+    3,
+    3,
+    1,
+    0,
+    2,
+    4,
+    7,
+    6,
+    6,
+    4,
+    3,
+    1,
+    0,
+    2,
+    4,
+    8,
+    6,
+    4,
+    3,
+    1,
+    0,
+    2,
+    4,
+    8,
+    6,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    9,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    9,
+    7,
+    6,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    6,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    6,
+    5,
+    3,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    7,
+    5,
+    3,
+    9,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    5,
+    9,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    5,
+    9,
+    9,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    5,
+    8,
+    9,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    5,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    5,
+    9,
+    9,
+    7,
+    5,
+    3,
+    1,
+    0,
+    2,
+    5,
+    9,
+    9,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    9,
+    5,
+    3,
+    1,
+    0,
+    2,
+    4,
+    9,
+    5,
+    3,
+    6,
+    3,
+    1,
+    0,
+    2,
+    6,
+    6,
+    3,
+    1,
+    0,
+    2,
+    6,
+    6,
+    3,
+    1,
+    0,
+    2,
+    7,
+    3,
+    1,
+    0,
+    3,
+    7,
+    7,
+    3,
+    1,
+    0,
+    3,
+    7,
+    7,
+    3,
+    1,
+    0,
+    3,
+    6,
+    3,
+    1,
+    1,
+    3,
+    6,
+    6,
+    3,
+    1,
+    1,
+    3,
+    6,
+    6,
+    3,
+    1,
+    1,
+    3,
+    7,
+    3,
+    1,
+    1,
+    3,
+    7,
+    7,
+    3,
+    1,
+    1,
+    3,
+    7,
+    7,
+    3,
+    1,
+    1,
+    3,
+    6,
+    3,
+    0,
+    1,
+    3,
+    6,
+    6,
+    3,
+    0,
+    1,
+    3,
+    6,
+    6,
+    3,
+    0,
+    1,
+    3,
+    7,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    2,
+    0,
+    1,
+    3,
+    5,
+    2,
+    0,
+    1,
+    3,
+    5,
+    5,
+    2,
+    0,
+    1,
+    3,
+    5,
+    5,
+    2,
+    0,
+    1,
+    3,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    8,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    8,
+    4,
+    2,
+    0,
+    5,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    9,
+    5,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    9,
+    5,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    9,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    9,
+    4,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    9,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    9,
+    4,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    7,
+    4,
+    2,
+    0,
+    9,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    6,
+    9,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    6,
+    9,
+    9,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    6,
+    9,
+    4,
+    2,
+    0,
+    1,
+    3,
+    5,
+    6,
+    9,
+    8,
+    4,
+    2,
+    0,
+    1,
+    3,
+    4,
+    6,
+    8,
+    4,
+    2,
+    0,
+    1,
+    3,
+    4,
+    6,
+    8,
+    6,
+    4,
+    2,
+    0,
+    1,
+    3,
+    3,
+    5,
+    6,
+    4,
+    2,
+    0,
+    1,
+    3,
+    3,
+    5,
+    6,
+    6,
+    4,
+    2,
+    0,
+    1,
+    2,
+    3,
+    6,
+    4,
+    2,
+    0,
+    1,
+    2,
+    3,
+    6,
+    4,
+    2,
+    6,
+    4,
+    2,
+    0,
+    1,
+    1,
+    6,
+    4,
+    2,
+    0,
+    1,
+    1,
+    6,
+    4,
+    2,
+    0,
+    1,
+    8,
+    6,
+    4,
+    2,
+    0,
+    0,
+    9,
+    8,
+    6,
+    4,
+    2,
+    0,
+    0,
+    9,
+    8,
+    6,
+    4,
+    8,
+    6,
+    4,
+    2,
+    0,
+    6,
+    8,
+    6,
+    4,
+    2,
+    0,
+    6,
+    8,
+    6,
+    4,
+    2,
+    0,
+    2,
+    4,
+    5,
+    8,
+    3,
+    1,
+    6,
+    2,
+    4,
+    5,
+    8,
+    3,
+    1,
+    6,
+    2,
+    4,
+    5,
+    7,
+    1,
+    1,
+    5,
+    0,
+    8,
+    7,
+    1,
+    1,
+    5,
+    0,
+    8,
+    7,
+    1,
+    1,
+    5,
+    0,
+    7,
+    1,
+    2,
+    8,
+    6,
+    0,
+    5,
+    9,
+    8,
+    8,
+    7,
+    1,
+    2,
+    8,
+    6,
+    0,
+    5,
+    8,
+    2,
+    1,
+    4,
+    0,
+    6,
+    7,
+    9,
+    8,
+    2,
+    1,
+    4,
+    0,
+    6,
+    7,
+    9,
+    8,
+    1,
+    0,
+    5,
+    5,
+    7,
+    1,
+    0,
+    5,
+    5,
+    7,
+    1,
+    0,
+    5,
+    5,
+    7,
+    1,
+    0,
+    4,
+    0,
+    2,
+    2,
+    6,
+    6,
+    2,
+    8,
+    4,
+    0,
+    2,
+    2,
+    6,
+    6,
+    2,
+    8,
+    4,
+    4,
+    0,
+    2,
+    1,
+    4,
+    7,
+    4,
+    4,
+    5,
+    9,
+    9,
+    7,
+    4,
+    0,
+    2,
+    1,
+    4,
+    5,
+    2,
+    0,
+    3,
+    9,
+    9,
+    4,
+    2,
+    7,
+    5,
+    4,
+    8,
+    9,
+    8,
+    6,
+    6,
+    5,
+    5,
+    7,
+    9,
+    0,
+    0,
+    3,
+    3,
+    2,
+    6,
+    7,
+    5,
+    7,
+    9,
+    0,
+    0,
+    3,
+    3,
+    2,
+    5,
+    0,
+    6,
+    7,
+    2,
+    5,
+    0,
+    6,
+    7,
+    2,
+    5,
+    0,
+    6,
+    7,
+    2,
+    5,
+    1,
+    1,
+    8,
+    5,
+    0,
+    4,
+    9,
+    7,
+    1,
+    1,
+    8,
+    5,
+    0,
+    4,
+    9,
+    7,
+    1,
+    8,
+    1,
+    1,
+    7,
+    4,
+    0,
+    6,
+    9,
+    8,
+    1,
+    1,
+    7,
+    4,
+    0,
+    6,
+    9,
+    8,
+    7,
+    2,
+    1,
+    0,
+    6,
+    9,
+    8,
+    9,
+    7,
+    2,
+    1,
+    0,
+    6,
+    9,
+    8,
+    9,
+    7,
+    8,
+    5,
+    4,
+    2,
+    2,
+    1,
+    6,
+    8,
+    5,
+    4,
+    2,
+    2,
+    1,
+    6,
+    8,
+    5,
+    4,
+    9,
+    7,
+    6,
+    3,
+    0,
+    0,
+    3,
+    6,
+    2,
+    7,
+    9,
+    7,
+    6,
+    3,
+    0,
+    0,
+    3,
+    7,
+    3,
+    0,
+    3,
+    5,
+    2,
+    2,
+    9,
+    8,
+    4,
+    5,
+    7,
+    6,
+    7,
+    9,
+    6,
+    7,
+    2,
+    0,
+    4,
+    2,
+    1,
+    3,
+    2,
+    7,
+    9,
+    5,
+    8,
+    2,
+    0,
+    4,
+    2,
+    1,
+    3,
+    0,
+    4,
+    3,
+    1,
+    5,
+    2,
+    6,
+    8,
+    0,
+    4,
+    3,
+    1,
+    5,
+    2,
+    6,
+    8,
+    0,
+    5,
+    6,
+    5,
+    5,
+    1,
+    5,
+    8,
+    8,
+    5,
+    6,
+    5,
+    5,
+    1,
+    5,
+    8,
+    8,
+    5,
+    0,
+    1,
+    9,
+    0,
+    1,
+    9,
+    0,
+    1,
+    9,
+    0,
+    1,
+    9,
+    0,
+    1,
+    9,
+    0,
+    1,
+    7,
+    0,
+    1,
+    9,
+    9,
+    9,
+    9,
+    7,
+    0,
+    1,
+    9,
+    9,
+    9,
+    9,
+    7,
+    0,
+    1,
+    4,
+    0,
+    5,
+    2,
+    0,
+    2,
+    4,
+    4,
+    0,
+    5,
+    2,
+    0,
+    2,
+    4,
+    4,
+    0,
+    5,
+    6,
+    5,
+    0,
+    8,
+    6,
+    6,
+    9,
+    6,
+    6,
+    5,
+    0,
+    8,
+    6,
+    6,
+    9,
+    6,
+    6,
+    3,
+    2,
+    0,
+    2,
+    7,
+    7,
+    5,
+    7,
+    8,
+    3,
+    2,
+    0,
+    2,
+    7,
+    7,
+    5,
+    7,
+    2,
+    0,
+    2,
+    1,
+    1,
+    2,
+    4,
+    3,
+    5,
+    7,
+    2,
+    0,
+    2,
+    1,
+    1,
+    2,
+    4,
+    4,
+    3,
+    7,
+    1,
+    0,
+    5,
+    4,
+    8,
+    8,
+    8,
+    4,
+    3,
+    7,
+    1,
+    0,
+    5,
+    4,
+    7,
+    4,
+    7,
+    0,
+    9,
+    6,
+    6,
+    6,
+    7,
+    4,
+    7,
+    0,
+    9,
+    6,
+    6,
+    6,
+    7,
+    4,
+    5,
+    6,
+    7,
+    8,
+    2,
+    5,
+    4,
+    1,
+    9,
+    6,
+    1,
+    9,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    3,
+    4,
+    6,
+    2,
+    3,
+    1,
+    2,
+    9,
+    7,
+    4,
+    0,
+    5,
+    8,
+    9,
+    3,
+    9,
+    6,
+    5,
+    6,
+    7,
+    7,
+    3,
+    1,
+    7,
+    4,
+    2,
+    3,
+    6,
+    4,
+    1,
+    4,
+    0,
+    8,
+    5,
+    3,
+    3,
+    1,
+    8,
+    8,
+    9,
+    7,
+    3,
+    1,
+    0,
+    5,
+    8,
+    3,
+    8,
+    5,
+    8,
+    4,
+    2,
+    8,
+    4,
+    3,
+    9,
+    1,
+    1,
+    7,
+    8,
+    8,
+    4,
+    2,
+    8,
+    4,
+    3,
+    9,
+    6,
+    5,
+    9,
+    7,
+    9,
+    6,
+    0,
+    0,
+    3,
+    5,
+    2,
+    9,
+    6,
+    5,
+    9,
+    7,
+    9,
+    3,
+    4,
+    1,
+    5,
+    5,
+    3,
+    2,
+    1,
+    9,
+    3,
+    4,
+    1,
+    5,
+    5,
+    3,
+    2,
+    9,
+    8,
+    8,
+    9,
+    6,
+    7,
+    9,
+    9,
+    6,
+    0,
+    0,
+    5,
+    6,
+    2,
+    4,
+    9,
+    8,
+    4,
+    8,
+    8,
+    2,
+    3,
+    2,
+    8,
+    1,
+    8,
+    1,
+    9,
+    4,
+    8,
+    8,
+    2,
+    3,
+    2,
+    3,
+    5,
+    8,
+    8,
+    1,
+    3,
+    9,
+    0,
+    3,
+    7,
+    8,
+    5,
+    0,
+    5,
+    3,
+    5,
+    8,
+    9,
+    6,
+    5,
+    6,
+    8,
+    6,
+    1,
+    4,
+    7,
+    6,
+    4,
+    2,
+    5,
+    4,
+    2,
+    4,
+    0,
+    9,
+    8,
+    6,
+    4,
+    3,
+    3,
+    4,
+    9,
+    1,
+    1,
+    0,
+    4,
+    7,
+    2,
+    9,
+    8,
+    6,
+    8,
+    7,
+    7,
+    5,
+    4,
+    5,
+    2,
+    5,
+    8,
+    1,
+    1,
+    6,
+    7,
+    8,
+    7,
+    7,
+    5,
+    9,
+    8,
+    8,
+    9,
+    9,
+    7,
+    4,
+    7,
+    9,
+    5,
+    0,
+    0,
+    1,
+    6,
+    3,
+    9,
+    8,
+    9,
+    5,
+    5,
+    2,
+    4,
+    3,
+    2,
+    3,
+    1,
+    9,
+    5,
+    5,
+    2,
+    4,
+    3,
+    2,
+    3,
+    6,
+    9,
+    9,
+    6,
+    8,
+    1,
+    0,
+    6,
+    8,
+    9,
+    5,
+    3,
+    4,
+    6,
+    9,
+    9,
+    6,
+    9,
+    8,
+    6,
+    6,
+    5,
+    6,
+    7,
+    8,
+    4,
+    2,
+    0,
+    8,
+    7,
+    9,
+    8,
+    6,
+    6,
+    1,
+    5,
+    2,
+    7,
+    5,
+    3,
+    2,
+    0,
+    3,
+    1,
+    5,
+    2,
+    7,
+    5,
+    3,
+    2,
+    0,
+    7,
+    4,
+    3,
+    4,
+    9,
+    7,
+    5,
+    1,
+    3,
+    7,
+    7,
+    6,
+    7,
+    2,
+    2,
+    3,
+    4,
+    6,
+    7,
+    4,
+    3,
+    4,
+    6,
+    9,
+    0,
+    0,
+    9,
+    9,
+    6,
+    9,
+    7,
+    0,
+    7,
+    2,
+    8,
+    5,
+    3,
+    3,
+    3,
+    2,
+    5,
+    7,
+    6,
+    7,
+    8,
+    3,
+    2,
+    7,
+    4,
+    4,
+    8,
+    5,
+    1,
+    6,
+    2,
+    3,
+    5,
+    0,
+    2,
+    3,
+    5,
+    1,
+    6,
+    2,
+    3,
+    5,
+    0,
+    2,
+    7,
+    6,
+    6,
+    6,
+    7,
+    8,
+    9,
+    8,
+    4,
+    2,
+    8,
+    0,
+    8,
+    7,
+    6,
+    6,
+    6,
+    8,
+    7,
+    6,
+    5,
+    7,
+    8,
+    9,
+    3,
+    1,
+    1,
+    3,
+    1,
+    2,
+    8,
+    7,
+    6,
+    5,
+    7,
+    5,
+    4,
+    5,
+    9,
+    7,
+    5,
+    5,
+    1,
+    4,
+    5,
+    1,
+    5,
+    7,
+    5,
+    4,
+    5,
+    8,
+    5,
+    4,
+    6,
+    8,
+    8,
+    2,
+    2,
+    8,
+    4,
+    9,
+    0,
+    9,
+    8,
+    5,
+    4,
+    6,
+    9,
+    8,
+    4,
+    4,
+    6,
+    8,
+    5,
+    8,
+    2,
+    5,
+    5,
+    4,
+    6,
+    1,
+    9,
+    8,
+    4,
+    9,
+    8,
+    5,
+    4,
+    6,
+    7,
+    1,
+    3,
+    1,
+    1,
+    3,
+    2,
+    9,
+    8,
+    5,
+    4,
+    6,
+    9,
+    8,
+    7,
+    7,
+    8,
+    9,
+    9,
+    6,
+    0,
+    2,
+    8,
+    1,
+    5,
+    5,
+    9,
+    8,
+    7,
+    3,
+    6,
+    3,
+    0,
+    2,
+    8,
+    3,
+    4,
+    3,
+    6,
+    0,
+    3,
+    6,
+    3,
+    0,
+    2,
+    8,
+    8,
+    6,
+    8,
+    1,
+    0,
+    1,
+    9,
+    6,
+    3,
+    6,
+    9,
+    6,
+    6,
+    9,
+    7,
+    1,
+    8,
+    6,
+    5,
+    6,
+    2,
+    0,
+    3,
+    4,
+    3,
+    9,
+    5,
+    3,
+    0,
+    9,
+    6,
+    5,
+    6,
+    2,
+    9,
+    8,
+    8,
+    7,
+    7,
+    9,
+    9,
+    7,
+    2,
+    0,
+    1,
+    8,
+    5,
+    5,
+    9,
+    8,
+    8,
+    9,
+    8,
+    9,
+    8,
+    1,
+    4,
+    0,
+    0,
+    4,
+    8,
+    1,
+    4,
+    7,
+    9,
+    8,
+    9,
+    8,
+    8,
+    9,
+    9,
+    6,
+    4,
+    7,
+    7,
+    4,
+    0,
+    4,
+    7,
+    9,
+    1,
+    9,
+    6,
+    6,
+    8,
+    8,
+    9,
+    9,
+    4,
+    1,
+    8,
+    5,
+    0,
+    0,
+    4,
+    1,
+    9,
+    8,
+    8,
+    9,
+    9,
+    4,
+    9,
+    7,
+    7,
+    8,
+    7,
+    7,
+    8,
+    5,
+    3,
+    0,
+    2,
+    3,
+    2,
+    0,
+    3,
+    9,
+    7,
+    7,
+    7,
+    9,
+    8,
+    7,
+    7,
+    8,
+    4,
+    3,
+    0,
+    3,
+    4,
+    3,
+    0,
+    2,
+    7,
+    7,
+]

third_party/PIPNet/run_demo.sh

@@ -0,0 +1,11 @@
+# image
+python lib/demo.py experiments/WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10.py images/1.jpg
+#python lib/demo.py experiments/data_300W_CELEBA/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py images/2.jpg
+
+# video 
+#python lib/demo_video.py experiments/WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10.py videos/002.avi
+#python lib/demo_video.py experiments/data_300W_CELEBA/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py videos/007.avi
+
+# camera 
+#python lib/demo_video.py experiments/WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10.py camera
+

third_party/PIPNet/run_test.sh

@@ -0,0 +1,34 @@
+# supervised learning
+
+# 300W, resnet18
+#python lib/test.py experiments/data_300W/pip_32_16_60_r18_l2_l1_10_1_nb10.py test.txt images_test
+# 300W, resnet101
+#python lib/test.py experiments/data_300W/pip_32_16_60_r101_l2_l1_10_1_nb10.py test.txt images_test
+
+# COFW, resnet18
+#python lib/test.py experiments/COFW/pip_32_16_60_r18_l2_l1_10_1_nb10.py test.txt images_test
+# COFW, resnet101
+#python lib/test.py experiments/COFW/pip_32_16_60_r101_l2_l1_10_1_nb10.py test.txt images_test
+
+# WFLW, resnet18
+#python lib/test.py experiments/WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10.py test.txt images_test
+# WFLW, resnet101
+#python lib/test.py experiments/WFLW/pip_32_16_60_r101_l2_l1_10_1_nb10.py test.txt images_test
+
+# AFLW, resnet18
+#python lib/test.py experiments/AFLW/pip_32_16_60_r18_l2_l1_10_1_nb10.py test.txt images_test
+# AFLW, resnet101
+#python lib/test.py experiments/AFLW/pip_32_16_60_r101_l2_l1_10_1_nb10.py test.txt images_test
+
+######################################################################################
+# GSSL
+
+# 300W + COFW_68 (unlabeled) + WFLW_68 (unlabeled), resnet18, with curriculum
+#python lib/test.py experiments/data_300W_COFW_WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py test_300W.txt images_test_300W
+#python lib/test.py experiments/data_300W_COFW_WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py test_COFW.txt images_test_COFW
+#python lib/test.py experiments/data_300W_COFW_WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py test_WFLW.txt images_test_WFLW
+
+# 300W + CelebA (unlabeled), resnet18, with curriculum
+#python lib/test.py experiments/data_300W_CELEBA/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py test_300W.txt images_test_300W
+#python lib/test.py experiments/data_300W_CELEBA/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py test_COFW.txt images_test_COFW
+#python lib/test.py experiments/data_300W_CELEBA/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py test_WFLW.txt images_test_WFLW

third_party/PIPNet/run_train.sh

@@ -0,0 +1,33 @@
+######################################################################################
+# supervised learning
+
+# 300W, resnet18
+#python lib/train.py experiments/data_300W/pip_32_16_60_r18_l2_l1_10_1_nb10.py
+# 300W, resnet101
+#python lib/train.py experiments/data_300W/pip_32_16_60_r101_l2_l1_10_1_nb10.py
+
+# COFW, resnet18
+#python lib/train.py experiments/COFW/pip_32_16_60_r18_l2_l1_10_1_nb10.py
+# COFW, resnet101
+#python lib/train.py experiments/COFW/pip_32_16_60_r101_l2_l1_10_1_nb10.py
+
+# WFLW, resnet18
+#python lib/train.py experiments/WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10.py
+# WFLW, resnet101
+#python lib/train.py experiments/WFLW/pip_32_16_60_r101_l2_l1_10_1_nb10.py
+
+# AFLW, resnet18
+#python lib/train.py experiments/AFLW/pip_32_16_60_r18_l2_l1_10_1_nb10.py
+# AFLW, resnet101
+#python lib/train.py experiments/AFLW/pip_32_16_60_r101_l2_l1_10_1_nb10.py
+
+######################################################################################
+# GSSL
+
+# 300W + COFW_68 (unlabeled) + WFLW_68 (unlabeled), resnet18, with curriculum
+#python lib/train_gssl.py experiments/data_300W_COFW_WFLW/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py 
+
+# 300W + CelebA (unlabeled), resnet18, with curriculum
+#nohup python lib/train_gssl.py experiments/data_300W_CELEBA/pip_32_16_60_r18_l2_l1_10_1_nb10_wcc.py &
+
+

weights/PIPNet/FaceBoxesV2.pth

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

weights/PIPNet/epoch59.pth

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

weights/arcface/mouth_net_28_56_84_112.pth

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