import argparse
import glob
import os
import pickle
from pathlib import Path

import numpy as np
from PIL import Image
from tqdm import tqdm

from src.align.align_trans import get_reference_facial_points, warp_and_crop_face

# sys.path.append("../../")
from src.align.detector import detect_faces

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="face alignment")
    parser.add_argument(
        "-source_root",
        "--source_root",
        help="specify your source dir",
        default="../../data/fiw-videos/new-processed/",
        type=str,
    )
    parser.add_argument(
        "-dest_root",
        "--dest_root",
        help="specify your destination dir",
        default="../../data/fiw-videos/new-processed/",
        type=str,
    )
    parser.add_argument(
        "-crop_size",
        "--crop_size",
        help="specify size of aligned faces, align and crop with padding",
        default=112,
        type=int,
    )
    args = parser.parse_args()

    source_root = args.source_root  # specify your source dir
    dest_root = args.dest_root  # specify your destination dir
    crop_size = (
        args.crop_size
    )  # specify size of aligned faces, align and crop with padding
    scale = crop_size / 112.0
    reference = get_reference_facial_points(default_square=True) * scale

    cwd = os.getcwd()  # delete '.DS_Store' existed in the source_root
    os.chdir(source_root)
    os.system("find . -name '*.DS_Store' -type f -delete")
    os.chdir(cwd)

    imfiles = [
        f
        for f in glob.glob(f"{source_root}F????/MID*/faces/msceleb*")
        if Path(f).is_file()
    ]
    # images = {imfile.replace(source_root, ''): Image.open(imfile) for imfile in imfiles}
    meta = {}
    # for subfolder in tqdm(os.listdir(source_root)):

    for imfile in tqdm(imfiles):
        ref = imfile.replace(source_root, "")
        print("Processing\t{}".format(imfile))
        img = Image.open(imfile)
        try:  # Handle exception
            bbs, landmarks = detect_faces(img)
        except Exception:
            print("{} is discarded due to exception!".format(imfile))
            continue
        ref = imfile.replace(source_root, "")
        ndetections = len(landmarks)
        if (
            ndetections == 0
        ):  # If the landmarks cannot be detected, the img will be discarded
            print("{} is discarded due to non-detected landmarks!".format(imfile))
            meta[ref] = []
            continue

        li_meta = []
        for i in range(ndetections):
            im_meta = {}
            im_meta["face"] = i
            im_meta["landmarks"] = landmarks[i]
            im_meta["bb"] = bbs[i]

            facial5points = [[landmarks[i][j], landmarks[i][j + 5]] for j in range(5)]
            warped_face = warp_and_crop_face(
                np.array(img),
                facial5points,
                reference,
                crop_size=(crop_size, crop_size),
            )
            img_warped = Image.fromarray(warped_face)
            image_name = imfile.replace("images", "cropped").replace(
                ".jpg", "-{:02d}.jpg".format(i)
            )
            # im_meta['ref'] = "/".join(image_name.split('/')[-5:])
            img_warped.save(image_name)
            li_meta.append(im_meta)
        meta[ref] = li_meta
    with open(source_root + "cropped-meta.pkl", "wb") as f:
        pickle.dump(meta, f)

[((17, 13, 17, 66), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((49, 10, 49, 21), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((50, 4, 50, 25), 'os.chdir', 'os.chdir', ({(50, 13, 50, 24): 'source_root'}, {}), '(source_root)', False, 'import os\n'), ((51, 4, 51, 58), 'os.system', 'os.system', ({(51, 14, 51, 57): '"""find . -name \'*.DS_Store\' -type f -delete"""'}, {}), '("find . -name \'*.DS_Store\' -type f -delete")', False, 'import os\n'), ((52, 4, 52, 17), 'os.chdir', 'os.chdir', ({(52, 13, 52, 16): 'cwd'}, {}), '(cwd)', False, 'import os\n'), ((63, 18, 63, 31), 'tqdm.tqdm', 'tqdm', ({(63, 23, 63, 30): 'imfiles'}, {}), '(imfiles)', False, 'from tqdm import tqdm\n'), ((47, 16, 47, 64), 'src.align.align_trans.get_reference_facial_points', 'get_reference_facial_points', (), '', False, 'from src.align.align_trans import get_reference_facial_points, warp_and_crop_face\n'), ((66, 14, 66, 32), 'PIL.Image.open', 'Image.open', ({(66, 25, 66, 31): 'imfile'}, {}), '(imfile)', False, 'from PIL import Image\n'), ((104, 8, 104, 28), 'pickle.dump', 'pickle.dump', ({(104, 20, 104, 24): 'meta', (104, 26, 104, 27): 'f'}, {}), '(meta, f)', False, 'import pickle\n'), ((56, 17, 56, 69), 'glob.glob', 'glob.glob', ({(56, 27, 56, 68): 'f"""{source_root}F????/MID*/faces/msceleb*"""'}, {}), "(f'{source_root}F????/MID*/faces/msceleb*')", False, 'import glob\n'), ((68, 29, 68, 46), 'src.align.detector.detect_faces', 'detect_faces', ({(68, 42, 68, 45): 'img'}, {}), '(img)', False, 'from src.align.detector import detect_faces\n'), ((95, 25, 95, 53), 'PIL.Image.fromarray', 'Image.fromarray', ({(95, 41, 95, 52): 'warped_face'}, {}), '(warped_face)', False, 'from PIL import Image\n'), ((90, 16, 90, 29), 'numpy.array', 'np.array', ({(90, 25, 90, 28): 'img'}, {}), '(img)', True, 'import numpy as np\n'), ((57, 11, 57, 18), 'pathlib.Path', 'Path', ({(57, 16, 57, 17): 'f'}, {}), '(f)', False, 'from pathlib import Path\n')]