imdb_wiki数据集的解析

代码及解析：

import os
import cv2
import re
import time
import random
import shutil
import glob
import scipy.io

import numpy as np

from shutil import copyfile
from datetime import datetime, timedelta
from multiprocessing import Pool

from config import config, parser
from align_faces import FaceAligner

parser = parser['DATA']

def parse_matlab_date(x):
  """
  :param x: date string in matlab format
  :return: int, year
  """
  x, date = int(x), -1
  try:
    date = (datetime.fromordinal(int(x))
          + timedelta(days=x % 1)
          - timedelta(days=366)).year
  except:
    print("[convertMatlabDate] Failed to parse string {}".format(x))
  return date


def clear_dir(path):
  """
  remove all files in this directionary
  :param path: path to dir
  :return: 
  """
  if os.path.exists(path):
    shutil.rmtree(path)
    os.mkdir(path)
  return


def addlabels(data = 'wiki', clean = False):
  """
  move pictures to labled dir and rename to [Age]_[Gender]_[Name].jpg format
  :param data: 'wiki' or 'imdb'
  :return: 
  """
  # 1, clean previous
  origin_dir = config.wiki_raw if data == 'wiki' else config.imdb_raw
  if clean: clear_dir(config.labeled)

  # 2, read meta data
  mat = scipy.io.loadmat(origin_dir + data + '.mat')[data][0][0]

  # records
  no_face_image = 0
  multiple_face_image = 0
  wrong_age = 0
  wrong_gender = 0
  successful = 0

  # dob: 出生日期
  # dop: 照片拍照日期
  # path: 文件路径
  # gender: 性别
  # name: 名人的名字
  # mat[5][0]： 脸的位置 IMG(face_location（2）：face_location（4），face_location（1）：face_location（3），:)
  # face_score: 人脸得分
  # face_score2: 第二个人脸得分
  for dob, dop, path, gender, name, face_score, face_score2 \
          in zip(mat[0][0], mat[1][0], mat[2][0], mat[3][0], mat[4][0], mat[6][0], mat[7][0]):
    if face_score < 0 or not np.isnan(face_score2):
      if face_score < 0: no_face_image += 1        #没有人脸
      if not np.isnan(face_score2): multiple_face_image += 1    #多个人脸的都跳过，不记录
      continue

    age = dop - parse_matlab_date(dob)
    if age < int(parser['age_lower']) or age > int(parser['age_upper']):     #age不在年龄范围的去掉
      wrong_age += 1
      continue

    if gender not in [1.0, 0.0]:    #性别错误的去掉
      wrong_gender += 1
      continue

    newName = "{}_{}_{}.jpg".format(age,
                                    int(gender),
                                    name[0]
                                    .replace(' ', '')
                                    .replace('/', '')
                                    .replace(':', ''))
    # 2.1 check duplicate
    # 2.1 if duplicate exist, append a random number to it name
    newNameNoDupli = newName
    while os.path.exists(config.labeled + newNameNoDupli):
      newNameNoDupli = "{}{}{}".format(newName[:-4], random.randint(1, 9999), newName[-4:])
    # 2.2 save as a new file
    copyfile(origin_dir + path[0], config.labeled + newNameNoDupli)
    successful += 1
  print("{} Successful, {} no_face_image, {} multiple_face_image, {} wrong_age, {} wrong_gender"
        .format(successful, no_face_image, multiple_face_image, wrong_age, wrong_gender))
  return


# sort photos by their names
def sort_out_by_name(clean = False):
  pwd = os.getcwd()
  if clean:
    clear_dir(config.named)
  os.chdir(config.aligned)
  for img in glob.glob("*.jpg"):
    name = re.findall(r'[^_]*_[^_]*_([\D]*)[0-9]*.jpg', img)
    if not len(name): continue
    name = name[0].lower()
    if not os.path.exists(config.named + name + '/'):
      os.mkdir(config.named + name + '/')
    copyfile(img, config.named + name + '/' + img)
  os.chdir(pwd)



# TODO: any other ways to get around this public variable?
FL = FaceAligner()
def sub_align_face(picname):
  """
  sub thread function to get and store aligned faces
  :param picname: pic names
  :return: 
  """
  aligned = FL.getAligns(picname)
  if len(aligned) == 0:
    return
    # copyfile(picname, config.aligned + picname)
  cv2.imwrite(config.aligned + picname, aligned[0])


def creat_fgnet_val(clean = False):
  if clean:
    clear_dir(config.val)
  pwd = os.getcwd()
  os.chdir(config.fgnet_raw)
  for pic in glob.glob("*"):
    name, age = re.findall(r'(\d)*A(\d*).*', pic)[0]
    newName = "{}_1_{}.jpg".format(age,
                                    name[0]
                                    .replace(' ', '')
                                    .replace('/', '')
                                    .replace(':', ''))
    # 2.1 check duplicate
    # 2.1 if duplicate exist, append a random number to it name
    newNameNoDupli = newName
    while os.path.exists(config.labeled + newNameNoDupli):
      newNameNoDupli = "{}{}{}".format(newName[:-4], random.randint(1, 9999), newName[-4:])
    # 2.2 save as a new file
    copyfile(config.fgnet_raw + pic, config.val + newNameNoDupli)
  os.chdir(pwd)


def align_faces(clean = False):
  """
  get aligned faces from labeled folder and store it in aligned folder for training
  :param data: 'wiki' or 'imdb'
  :param clean: if set, clean aligned folder, else append or rewrite to it
  :return: 
  """
  if clean: clear_dir(config.aligned)
  os.chdir(config.labeled)
  jobs = glob.glob("*.jpg")

  # un-parallel
  # for picname in jobs:
  #   aligned = FL.getAligns(picname)
  #   if len(aligned) != 1: return
  #   cv2.imwrite(config.aligned + picname, aligned[0])

  # parallel
  with Pool() as pool:
    try:
      pool.map(sub_align_face, jobs)
    finally:
      pool.close()
  return



def sub_divideTrainVal(img):
  """
  distribute images randomly to train or test foled by 95% train prob
  :param img: image path
  :return: 
  """
  if np.random.rand() < float(parser['train_test_div']):
    copyfile(config.aligned + img, config.train + img)
  else:
    copyfile(config.aligned + img, config.val + img)
  return

def divideTrainVal():
  """
  distribute images randomly to train or test foled by 95% train prob
  :return: 
  """
  pwt = os.getcwd()
  os.chdir(config.aligned)

  # clean
  clear_dir(config.train)
  clear_dir(config.val)

  # read into mem
  # train, val = [], []

  # parallel
  with Pool() as pool:
    try:
      pool.map(sub_divideTrainVal, glob.glob("*.jpg"))
    finally:
      pool.close()

  os.chdir(pwt)
  return


if __name__ == "__main__":
  print("labeling..")
  addlabels(data='wiki', clean=True)
  
  print("aligning..")
  align_faces(clean = True)
  
  print("dividing..")
  divideTrainVal()

#   creat_fgnet_val(clean=True)
  pass

参考：

https://github.com/jingkuan/Age-Gender-Pred

https://blog.csdn.net/qxqxqzzz/article/details/106145608