はじめに

自分の顔がイケメンかどうかが気になりすぎるのでイケメン判定AIを作ります。

ソースコードはこちら↓
https://github.com/Kentea-Watanabe/Handsome_judgement

構成

Yolov5を用いて顔部分を切り取る
imgsimで画像の類似度計算

1. Yolov5を用いて顔部分を切り取る

# Yoloの顔検出器を使用する
yolo_weight = "/content/Handsome_judgement/yolov5_face_detection/weights/best.pt"
yolo_dir = "/content/Handsome_judgement/yolov5_face_detection/yolov5"
d_model = torch.hub.load(yolo_dir, 'custom', path=yolo_weight, source='local')

input_im = cv2.imread(input_im_path)
input_pred_im = d_model(input_im) 
input_det_im = input_pred_im.crop(save=False)
input_cropped_im = input_det_im[0]["im"]

2. imgsimで画像の類似度計算

イケメンを以下のように定義します。

イケメンだけのデータベースを作成し、イケメンの顔と入力画像の類似度が閾値を超えていたらイケメンであると判定する

  # {画像パス: 画像の距離}の辞書を作成
  im_dist_dic = {}

  # 画像同士の距離が近いほうが似ている画像となる。
  min_dist = 1000 # 初期化する。
  # 類似度判定のためのベクトル化
  db_vec = vtr.vectorize(input_cropped_im)
  input_vec = vtr.vectorize(db_cropped_im)
  dist = imgsim.distance(db_vec, input_vec)
  im_dist_dic[db_img_path] = dist

  # 類似度が高いものを更新する。
  if min_dist > dist:
    min_dist = dist

# イケメン判定
if min_dist < 20:
  print("イケメン判定結果 : True")
else:
  print("イケメン判定結果 : False")

最後に

イケメンだけのデータベースに、こっそり自分の顔画像をいれておこうと思います。
それでは、良いイケメンライフを。