Face recognition with DB

Mix.install(
  [
    {:ecto, "~> 3.11"},
    {:ecto_sql, "~> 3.11"},
    {:jason, "~> 1.4"},
    {:kino, "~> 0.14"},
    {:postgrex, "~> 0.17.3"},
    {:pgvector, "~> 0.2.0"},
    {:evision, github: "cocoa-xu/evision", branch: "main"}
  ],
  system_env: [
    {"EVISION_PREFER_PRECOMPILED", "false"}
  ]
)

画像の準備

image_files = Path.wildcard("/home/livebook/evision/test-images/*.{jpg,png}")

images =
  image_files
  |> Enum.map(fn image_file ->
    Evision.imread(image_file)
  end)

Kino.Layout.grid(images, columns: 4)

顔特徴量の取得

recognizer =
  Evision.Zoo.FaceRecognition.SFace.init(:default_model,
    backend: Evision.Constant.cv_DNN_BACKEND_OPENCV(),
    target: Evision.Constant.cv_DNN_TARGET_CPU(),
    distance_type: :cosine_similarity,
    cosine_threshold: 0.363,
    l2_norm_threshold: 1.128
  )

detector =
  Evision.Zoo.FaceDetection.YuNet.init(:default_model,
    backend: Evision.Constant.cv_DNN_BACKEND_OPENCV(),
    target: Evision.Constant.cv_DNN_TARGET_CPU(),
    nms_threshold: 0.3,
    conf_threshold: 0.8,
    top_k: 5
  )

[feature_list, visualized_list] =
  images
  |> Enum.reduce([[], []], fn image, [feature_acc, visualized_acc] ->
    results = Evision.Zoo.FaceDetection.YuNet.infer(detector, image)

    bbox = Evision.Mat.to_nx(results, Nx.BinaryBackend)[0][0..-2//1]

    feature =
      recognizer
      |> Evision.Zoo.FaceRecognition.SFace.infer(image, bbox)
      |> Evision.Mat.to_nx()
      |> Evision.Mat.from_nx()

    visualized = Evision.Zoo.FaceDetection.YuNet.visualize(image, results[0])

    [[feature | feature_acc], [visualized | visualized_acc]]
  end)
  |> Enum.map(&Enum.reverse/1)

Kino.Layout.grid(visualized_list, columns: 4)

feature_list
|> hd()
|> Evision.Mat.to_nx()

Evision.Zoo.FaceRecognition.SFace.match_feature(
  recognizer,
  Enum.at(feature_list, 8),
  Enum.at(feature_list, 15)
)

Evision.Zoo.FaceRecognition.SFace.match_feature(
  recognizer,
  Enum.at(feature_list, 8),
  Enum.at(feature_list, 0)
)

vectors =
  feature_list
  |> Enum.zip(image_files)
  |> Enum.map(fn {feature, image_file} ->
    values =
      feature
      |> Evision.Mat.to_nx(Nx.BinaryBackend)
      |> Nx.flatten()
      |> Nx.to_list()

    %{
      embedding: values,
      file_path: image_file
    }
  end)

DB 接続

Postgrex.Types.define(
  ExtendedTypes,
  [Pgvector.Extensions.Vector] ++ Ecto.Adapters.Postgres.extensions(),
  []
)

defmodule Repo do
  use Ecto.Repo,
    otp_app: :my_notebook,
    adapter: Ecto.Adapters.Postgres
end

opts = [
  hostname: "postgres_for_livebook",
  port: 5432,
  username: "postgres",
  password: System.fetch_env!("LB_DB_PASSWORD"),
  database: "postgres",
  types: ExtendedTypes
]

Kino.start_child({Repo, opts})

拡張機能の追加

defmodule Migrations.CreateVectorExtension do
  use Ecto.Migration

  def up do
    execute("CREATE EXTENSION IF NOT EXISTS vector")
  end

  def down do
    execute("DROP EXTENSION vector")
  end
end

Ecto.Migrator.up(Repo, 21, Migrations.CreateVectorExtension)

テーブル作成

defmodule Migrations.CreateFaceTable do
  use Ecto.Migration

  def change do
    create table(:face) do
      add(:file_path, :string)
      add(:embedding, :vector, size: 128)
    end
  end
end

Ecto.Migrator.up(Repo, 31, Migrations.CreateFaceTable)

DB への登録

defmodule Face do
  use Ecto.Schema

  schema "face" do
    field(:file_path, :string)
    field(:embedding, Pgvector.Ecto.Vector)
  end
end

Repo.insert_all(Face, vectors)

Repo.all(Face)

似ている顔の検索

query =
  """
  SELECT
    src.file_path as src_file_path,
    dst.file_path as dst_file_path,
    src.embedding <=> dst.embedding as cos_distance
  FROM
    face AS src
  INNER JOIN
    face AS dst
  ON
    src.id < dst.id
  AND
    src.embedding <=> dst.embedding < 0.5
  ORDER BY
    cos_distance ASC
  """

{:ok, result} = Ecto.Adapters.SQL.query(Repo, query, [])

result.rows
|> Enum.map(fn [src_file_path, dst_file_path, distance] ->
  src_img = Evision.imread(src_file_path)
  dst_img = Evision.imread(dst_file_path)

  [
    distance,
    Kino.Layout.grid([src_img, dst_img], columns: 2)
  ]
  |> Kino.Layout.grid(columns: 1)
end)
|> Kino.Layout.grid(columns: 4)