BLにおける受け攻めを判定する、機械学習用顔画像データセットのラベリング補助スクリプト

背景

https://qiita.com/ishiwara51/items/3979fbc1c69b4f8ee2d3
のスクリプトで集めた機械学習用データセットに、正解ラベルを付与していくための補助スクリプトです。作った当時、「特定の人物のBL嗜好を学習して、男性の顔画像から受けか攻めかを判定できるモデルを作ろう」と試みており、そのためのスクリプトです。人の画像を一つずつ表示し、u(=受け), s(=攻め), n(=判別不能、適切な画像でないなど) の三つのうちどれかを入力すると、csvに対応関係が記録され、顔部分だけが切り取られた画像が保存されます。ます。当然ですが、ラベル設定を変えれば、他の人顔画像データセットのラベリングもできます。

動作環境

2019.07.05 動作確認済み。
動作環境：macOS10.14.5, python3.7.3

Pythonコード

# coding: utf-8
# To be executed in command line

import glob
import csv
import numpy as np
import cv2
import argparse
import os
import sys
import shutil


parser = argparse.ArgumentParser(description='Show the images in directory you\
                                input as an argument, \
                                and export a csv; sets of image\'s name and\
                                u or s or n(uke, seme, neutral) you input.',)
parser.add_argument('input_directory', action='store', type=str,
                    help='path of the directory image files are in.')

def main():
    args = parser.parse_args()
    new_dir_path = args.input_directory + '/new_' + args.input_directory.split('/')[-1]
    imgData = [r for r in glob.glob(args.input_directory + '/*' + '.jpg')]

    list_imgpath = []
    list_answer = []

    for count, i in enumerate(imgData):
        if count==30:
            break
        img=cv2.imread(i,cv2.IMREAD_GRAYSCALE)
        cv2.imshow(i,img)

        while True:
            answer = chr(cv2.waitKey(0))
            if answer in ['s', 'u', 'n']:
                break
            else:
                print('Wrong input.')
                print('Input a character in [s, u, n].')

        if answer == 'n':
            continue

        try:
            list_imgpath.append(face_trim_save(i, new_dir_path))
            list_answer.append(answer)
        except:
            print('Error in trimming and saving.')
        cv2.destroyAllWindows()
    print(list_imgpath)

    export_dict = {'Img': list_imgpath[:len(list_answer)], 'Answer': list_answer}
    with open(new_dir_path + '/dataset_for_BLlearn.csv', 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=['Img', 'Answer'])
        writer.writeheader()
        for i in range(len(list_answer)):
            writer.writerow({'Img': export_dict['Img'][i],
                             'Answer': export_dict['Answer'][i]})


def face_trim_save(img_name, new_dir_path):
    """
    入力された画像から顔の部分を切り出して保存します。

    引数:
        img_name : 入力画像のファイル名です。
        new_dir_path :　実行時に生成される、トリミング後の画像を保存するディレクトリのパスです。

    返り値:
        save_path : 新しく保存される、トリミング後の画像のパスです。
    """
    args = parser.parse_args()
    #cascade_pathには、haarcascade_frontalface_default.xmlのパスを入れること。
    cascade_path = ""
    origin_image_path = img_name

    image = cv2.imread(origin_image_path)
    if image is None:
        print('Not open : ')
        quit()

    cascade = cv2.CascadeClassifier(cascade_path)
    facerect = cascade.detectMultiScale(image, scaleFactor=1.1, minNeighbors=1, minSize=(10, 10))

    # 顔だけ切り出して保存
    save_path = str(new_dir_path) + '/' + img_name.split('/')[-1]
    if not os.path.exists(str(new_dir_path)):
        os.mkdir(str(new_dir_path))
    print(save_path)
    x = facerect[0][0]
    y = facerect[0][1]
    width = facerect[0][2]
    height = facerect[0][3]
    dst = image[y:y + height, x:x + width]
    #認識結果の保存
    cv2.imwrite(save_path, dst)
    print("saved!")

    return save_path

if __name__ == '__main__':
    print('Focus on the opencv window, and input u or s or n with keyboard.')
    print('u = UKE')
    print('s = SEME')
    print('n = Not ONE male face, already checked one, you cannot tell UKE or SEME...')
    print('Only the images that you input u or s will be trimed and recorded on csv.')
    main()

    print('finished!')