megaface_utils.py

import argparse
import json
import os
import struct

import cv2 as cv
import numpy as np
import torch
import tqdm
from PIL import Image, ImageOps
from tqdm import tqdm

from config import device
from data_gen import data_transforms
from utils import align_face, get_central_face_attributes


def walkdir(folder, ext):
    # Walk through each files in a directory
    for dirpath, dirs, files in os.walk(folder):
        for filename in [f for f in files if f.lower().endswith(ext)]:
            yield os.path.abspath(os.path.join(dirpath, filename))


def crop_one_image(filepath, oldkey, newkey):
    new_fn = filepath.replace(oldkey, newkey)
    tardir = os.path.dirname(new_fn)
    if not os.path.isdir(tardir):
        os.makedirs(tardir)

    if not os.path.exists(new_fn):
        is_valid, bounding_boxes, landmarks = get_central_face_attributes(filepath)
        if is_valid:
            img = align_face(filepath, landmarks)
            cv.imwrite(new_fn, img)


def crop(path, oldkey, newkey):
    print('Counting images under {}...'.format(path))
    # Preprocess the total files count
    filecounter = 0
    for filepath in walkdir(path, '.jpg'):
        filecounter += 1

    for filepath in tqdm(walkdir(path, '.jpg'), total=filecounter, unit="files"):
        crop_one_image(filepath, oldkey, newkey)

    print('{} images were cropped successfully.'.format(filecounter))


def get_image(transformer, filepath, flip=False):
    img = Image.open(filepath)
    if flip:
        img = ImageOps.mirror(img)
    img = transformer(img)
    return img.to(device)


def gen_feature(path, model):
    model.eval()

    print('gen features {}...'.format(path))
    # Preprocess the total files count
    files = []
    for filepath in walkdir(path, ('.jpg', '.png')):
        files.append(filepath)
    file_count = len(files)

    transformer = data_transforms['val']

    batch_size = 128

    with torch.no_grad():
        for start_idx in tqdm(range(0, file_count, batch_size)):
            end_idx = min(file_count, start_idx + batch_size)
            length = end_idx - start_idx

            imgs_0 = torch.zeros([length, 3, 112, 112], dtype=torch.float, device=device)
            for idx in range(0, length):
                i = start_idx + idx
                filepath = files[i]
                imgs_0[idx] = get_image(transformer, filepath, flip=False)

            features_0 = model(imgs_0.to(device))
            features_0 = features_0.cpu().numpy()

            imgs_1 = torch.zeros([length, 3, 112, 112], dtype=torch.float, device=device)
            for idx in range(0, length):
                i = start_idx + idx
                filepath = files[i]
                imgs_1[idx] = get_image(transformer, filepath, flip=True)

            features_1 = model(imgs_1.to(device))
            features_1 = features_1.cpu().numpy()

            for idx in range(0, length):
                i = start_idx + idx
                filepath = files[i]
                filepath = filepath.replace(' ', '_')
                tarfile = filepath + '_0.bin'
                feature = features_0[idx] + features_1[idx]
                write_feature(tarfile, feature / np.linalg.norm(feature))


def read_feature(filename):
    f = open(filename, 'rb')
    rows, cols, stride, type_ = struct.unpack('iiii', f.read(4 * 4))
    mat = np.fromstring(f.read(rows * 4), dtype=np.dtype('float32'))
    return mat.reshape(rows, 1)


def write_feature(filename, m):
    header = struct.pack('iiii', m.shape[0], 1, 4, 5)
    f = open(filename, 'wb')
    f.write(header)
    f.write(m.data)


def remove_noise():
    megaface_count = 0
    for line in open('megaface/megaface_noises.txt', 'r'):
        filename = 'megaface/MegaFace_aligned/FlickrFinal2/' + line.strip() + '_0.bin'
        if os.path.exists(filename):
            # print(filename)
            os.remove(filename)
            megaface_count += 1

    print('remove noise - megaface: ' + str(megaface_count))

    facescrub_count = 0
    noise = set()
    for line in open('megaface/facescrub_noises.txt', 'r'):
        noise.add((line.strip().replace('.png', '.jpg') + '_0.bin'))

    for root, dirs, files in os.walk('megaface/FaceScrub_aligned'):
        for f in files:
            # print(f)
            if f in noise:
                filename = os.path.join(root, f)
                if os.path.exists(filename):
                    # print(filename)
                    os.remove(filename)
                    facescrub_count += 1

    print('remove noise - facescrub: ' + str(facescrub_count))


def test():
    root1 = '/root/lin/data/FaceScrub_aligned/Benicio Del Toro'
    root2 = '/root/lin/data/FaceScrub_aligned/Ben Kingsley'
    for f1 in os.listdir(root1):
        for f2 in os.listdir(root2):
            if f1.lower().endswith('.bin') and f2.lower().endswith('.bin'):
                filename1 = os.path.join(root1, f1)
                filename2 = os.path.join(root2, f2)
                fea1 = read_feature(filename1)
                fea2 = read_feature(filename2)
                print(((fea1 - fea2) ** 2).sum() ** 0.5)


def match_result():
    with open('matches_facescrub_megaface_0_1000000_1.json', 'r') as load_f:
        load_dict = json.load(load_f)
        print(load_dict)
        for i in range(len(load_dict)):
            print(load_dict[i]['probes'])


def pngtojpg(path):
    for root, dirs, files in os.walk(path):
        for f in files:
            if os.path.splitext(f)[1] == '.png':
                img = cv.imread(os.path.join(root, f))
                newfilename = f.replace(".png", ".jpg")
                cv.imwrite(os.path.join(root, newfilename), img)


def parse_args():
    parser = argparse.ArgumentParser(description='Train face network')
    # general
    parser.add_argument('--action', default='crop_megaface', help='action')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_args()
    if args.action == 'crop_megaface':
        crop('megaface/MegaFace/FlickrFinal2', 'MegaFace', 'MegaFace_aligned')
    elif args.action == 'crop_facescrub':
        crop('megaface/facescrub_images', 'facescrub', 'facescrub_aligned')
    elif args.action == 'gen_features':
        gen_feature('megaface/facescrub_images')
        gen_feature('megaface/MegaFace_aligned/FlickrFinal2')
        remove_noise()
    elif args.action == 'pngtojpg':
        pngtojpg('megaface/facescrub_images')
    elif args.action == 'remove_noise':
        remove_noise()