lua-scripts-w-api/contrib/face_recognition.lua

--[[
  Face recognition for darktable

  Copyright (c) 2017  Sebastian Witt
   
  darktable is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  darktable is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with darktable.  If not, see <http://www.gnu.org/licenses/>.
]]

--[[
face_recognition
Add a new storage option to send images to face_recognition.
Images are exported to darktable tmp dir first.
A directory with known faces must exist, the image name are the
tag names which will be used.
Multiple images for one face can exist, add a number to it, the
number will be removed from the tag, for example:
People|IknowYou1.jpg
People|IknowYou2.jpg
People|Another.jpg
People|Youtoo.jpg

ADDITIONAL SOFTWARE NEEDED FOR THIS SCRIPT
* https://github.com/ageitgey/face_recognition
* https://github.com/darktable-org/lua-scripts/tree/master/lib

USAGE
* require this file from your main luarc config file.

This plugin will add a new storage option and calls face_recognition after export.
]]

local dt = require "darktable"
local du = require "lib/dtutils"
local df = require "lib/dtutils.file"
local dtsys = require "lib/dtutils.system"
local gettext = dt.gettext

-- constants

local MODULE = "face_recognition"
local PS = dt.configuration.running_os == "windows" and '\\' or '/'
local OUTPUT = dt.configuration.tmp_dir .. PS .. "facerecognition.txt"

-- namespace

local fc = {}

-- ensure we meet the minimum api
du.check_min_api_version("5.0.0", "face_recognition")

-- Tell gettext where to find the .mo file translating messages for a particular domain
gettext.bindtextdomain("face_recognition", dt.configuration.config_dir.."/lua/locale/")

local function _(msgid)
    return gettext.dgettext("face_recognition", msgid)
end

-- preferences

if not dt.preferences.read(MODULE, "initialized", "bool") then
  dt.preferences.write(MODULE, "unknown_tag", "string", "unknown_person")
  dt.preferences.write(MODULE, "ignore_tags", "string", "")
  dt.preferences.write(MODULE, "tolerance", "float", 0.6)
  dt.preferences.write(MODULE, "num_cores", "integer", 0)
  dt.preferences.write(MODULE, "known_image_path", "directory", dt.configuration.config_dir .. "/face_recognition")
  dt.preferences.write(MODULE, "export_format", "integer", 1)
  dt.preferences.write(MODULE, "max_width", "integer", 1000)
  dt.preferences.write(MODULE, "max_height", "integer", 1000)
  dt.preferences.write(MODULE, "initialized", "bool", true)
  dt.preferences.write(MODULE, "no_persons_found_tag", "string", "no_persons_found")
end

local function build_image_table(images)
  local image_table = {}
  local file_extension = ""
  local tmp_dir = dt.configuration.tmp_dir .. PS
  local ff = fc.export_format.value
  local cnt = 0

  -- check for plugin-data and direct_edit and build image table accordingly

  if string.match(ff, "JPEG") then
    file_extension = ".jpg"
  elseif string.match(ff, "PNG") then
    file_extension = ".png"
  elseif string.match(ff, "TIFF") then
    file_extension = ".tif"
  end

  for _,img in ipairs(images) do
    image_table[img] = tmp_dir .. df.get_basename(img.filename) .. file_extension
    cnt = cnt + 1
  end

  return image_table, cnt
end

local function stop_job(job)
  job.valid = false
end

local function do_export(img_tbl)
  local exporter = nil
  local upsize = false
  local upscale = false
  local ff = fc.export_format.value
  local height = dt.preferences.read(MODULE, "max_height", "integer")
  local width = dt.preferences.read(MODULE, "max_width", "integer")
  local images = 0
  for k,v in pairs(img_tbl) do
    images = images + 1
  end

  -- get the export format parameters
  if string.match(ff, "JPEG") then
    exporter = dt.new_format("jpeg")
    exporter.quality = 80
  elseif string.match(ff, "PNG") then
    exporter = dt.new_format("png")
    exporter.bpp = 8
  elseif string.match(ff, "TIFF") then
    exporter = dt.new_format("tiff")
    exporter.bpp = 8
  end
  exporter.max_height = height
  exporter.max_width = width

  -- export the images
  local job = dt.gui.create_job(_("export images"), true, stop_job)
  local exp_cnt = 0
  local percent_step = 1.0 / images
  job.percent = 0.0
  for img,export in pairs(img_tbl) do
    exp_cnt = exp_cnt + 1
    dt.print(string.format(_("Exporting image %i of %i images"), exp_cnt, images))
    exporter:write_image(img, export, upsize)
    job.percent = job.percent + percent_step
  end
  job.valid = false

  -- return success, or not
  return true
end

local function save_preferences()
  dt.preferences.write(MODULE, "unknown_tag", "string", fc.unknown_tag.text)
  dt.preferences.write(MODULE, "ignore_tags", "string", fc.ignore_tags.text)
  dt.preferences.write(MODULE, "max_width", "integer", tonumber(fc.width.text))
  dt.preferences.write(MODULE, "max_height", "integer", tonumber(fc.height.text))
  dt.preferences.write(MODULE, "num_cores", "integer", fc.num_cores.value)
  dt.preferences.write(MODULE, "no_persons_found_tag", "string", fc.no_persons_found_tag.text)
  local val = fc.tolerance.value
  val = string.gsub(tostring(val), ",", ".")
  dt.preferences.write(MODULE, "tolerance", "float", tonumber(val))
end

-- Check if image has ignored tag attached
local function ignoreByTag (image, ignoreTags)
  local tags = image:get_tags ()
  local ignoreImage = false
  -- For each image tag
  for _,t in ipairs (tags) do
    -- Check if it contains a ignore tag
    for _,it in ipairs (ignoreTags) do
      if string.find (t.name, it, 1, true) then
        -- The image has ignored tag attached
        ignoreImage = true
        dt.print_log ("Face recognition: Ignored tag: " .. it .. " found in " .. image.id .. ":" .. t.name)
      end
    end
  end
  
  return ignoreImage
end

local function cleanup(img_list)
  for _, img in ipairs(img_list) do
    os.remove(img)
  end
  os.remove(OUTPUT)
end

local function face_recognition ()

  local bin_path = df.check_if_bin_exists("face_recognition")

  if not bin_path then
    dt.print(_("Face recognition not found"))
    return
  end

  save_preferences()

  -- Get preferences
  local knownPath = dt.preferences.read(MODULE, "known_image_path", "directory")
  local nrCores = dt.preferences.read(MODULE, "num_cores", "integer")
  local ignoreTagString = dt.preferences.read(MODULE, "ignore_tags", "string")
  local unknownTag = dt.preferences.read(MODULE, "unknown_tag", "string")
  local nonpersonsfoundTag = dt.preferences.read(MODULE, "no_persons_found_tag", "string")

  -- face_recognition uses -1 for all cores, we use 0 in preferences
  if nrCores < 1 then
    nrCores = -1
  end
  
  -- Split ignore tags (if any)
  ignoreTags = {}
  for tag in string.gmatch(ignoreTagString, '([^,]+)') do
    table.insert (ignoreTags, tag)
    dt.print_log ("Face recognition: Ignore tag: " .. tag)
  end
  
  -- list of exported images

  local image_table, cnt = build_image_table(dt.gui.action_images)

  if cnt > 0 then
    local success = do_export(image_table)
    if success then
      -- do the face recognition
      local img_list = {}

      for img,v in pairs(image_table) do
        table.insert (img_list, v)
      end

      -- Get path of exported images
      local path = df.get_path (img_list[1])
      dt.print_log ("Face recognition: Path to unknown images: " .. path)
      os.setlocale("C")
      local tolerance = dt.preferences.read(MODULE, "tolerance", "float")
      
      local command = bin_path ..  " --cpus " .. nrCores .. " --tolerance " .. tolerance .. " " .. knownPath .. " " .. path .. " > " .. OUTPUT
      os.setlocale()
      dt.print_log("Face recognition: Running command: " .. command)
      dt.print(_("Starting face recognition..."))

      dtsys.external_command(command)

      -- Open output file
      local f = io.open(OUTPUT, "rb")
      
      if not f then
        dt.print(_("Face recognition failed"))
      else
        dt.print(_("Face recognition finished"))
        f:close ()
      end
      
      -- Read output
      dt.print(_("processing results..."))
      local result = {}
      for line in io.lines(OUTPUT) do 
        if not string.match(line, "^WARNING:") then
          local file, tag = string.match (line, "(.*),(.*)$")
          tag = string.gsub (tag, "%d*$", "")
          dt.print_log ("File:"..file .." Tag:".. tag)
          if result[file] ~= nil then
            table.insert (result[file], tag)
          else
            result[file] = {tag}
          end
        end
      end
      
      -- Attach tags
      for file,tags in pairs(result) do
        -- Find image in table
        for img,file2 in pairs(image_table) do
          if file == file2 then
            for _,t in ipairs (tags) do
              -- Check if image is ignored
              if ignoreByTag (img, ignoreTags) then
                dt.print_log("Face recognition: Ignoring image with ID " .. img.id)
              else
                -- Check of unrecognized unknown_person
                if t == "unknown_person" then
                  t = unknownTag
                end
                -- Check of unrecognized no_persons_found
                if t == "no_persons_found" then
                  t = nonpersonsfoundTag
                end
                if t ~= "" and t ~= nil then
                  dt.print_log ("ImgId:" .. img.id .. " Tag:".. t)
                  -- Create tag if it does not exists
                  local tag = dt.tags.create (t)
                  img:attach_tag (tag)
                end
              end
            end
          end
        end
      end
      dt.print(_("face recognition complete"))
      cleanup(img_list)
    else
      dt.print(_("image export failed"))
      return
    end
  else
    dt.print(_("no images selected"))
    return
  end


end

-- build the interface

fc.unknown_tag = dt.new_widget("entry"){
  text = dt.preferences.read(MODULE, "unknown_tag", "string"),
  tooltip = _("tag to be used for unknown person"),
  editable = true,
}

fc.no_persons_found_tag = dt.new_widget("entry"){
  text = dt.preferences.read(MODULE, "no_persons_found_tag", "string"),
  tooltip = _("tag to be used when no persons are found"),
  editable = true,
}

fc.ignore_tags = dt.new_widget("entry"){
  text = dt.preferences.read(MODULE, "ignore_tags", "string"),
  tooltip = _("tags of images to ignore"),
  editable = true,
}

fc.tolerance = dt.new_widget("slider"){
  label = _("tolerance"),
  tooltip = ("detection tolerance - 0.6 default - lower if too many faces detected"),
  soft_min = 0.0,
  hard_min = 0.0,
  soft_max = 1.0,
  soft_min = 1.0,
  step = 0.1,
  digits = 1,
  value = 0.0,
}

fc.num_cores = dt.new_widget("slider"){
  label = _("processor cores"),
  tooltip = _("number of processor cores to use, 0 for all"),
  soft_min = 0,
  soft_max = 16,
  hard_min = 0,
  hard_max = 64,
  step = 1,
  digits = 0,
  value = dt.preferences.read(MODULE, "num_cores", "integer"),
}

fc.known_image_path = dt.new_widget("file_chooser_button"){
  title = _("known image directory"),
  tooltip = _("face data directory"),
  value = dt.preferences.read(MODULE, "known_image_path", "directory"),
  is_directory = true,
  changed_callback = function(this)
    dt.preferences.write(MODULE, "known_image_path", "directory", this.value)
  end
}

fc.export_format = dt.new_widget("combobox"){
  label = _("export image format"),
  tooltip = _("format for exported images"),
  selected = dt.preferences.read(MODULE, "export_format", "integer"),
  changed_callback = function(this)
    dt.preferences.write(MODULE, "export_format", "integer", this.selected)
  end,
  "JPEG", "PNG", "TIFF",
}

fc.width = dt.new_widget("entry"){
  text = tostring(dt.preferences.read(MODULE, "max_width", "integer")),
  tooltip = _("maximum exported image width"),
  editable = true,
}

fc.height = dt.new_widget("entry"){
  text = tostring(dt.preferences.read(MODULE, "max_height", "integer")),
  tooltip = _("maximum exported image height"),
  editable = true,
}

fc.execute = dt.new_widget("button"){
  label = "detect faces",
  clicked_callback = function(this) 
    face_recognition()
  end
}

local widgets = {
  dt.new_widget("label"){ label = _("unknown person tag")},
  fc.unknown_tag,
  dt.new_widget("label"){ label = _("no persons found tag")},
  fc.no_persons_found_tag,
  dt.new_widget("label"){ label = _("tags of images to ignore")},
  fc.ignore_tags,
  dt.new_widget("label"){ label = _("face data directory")},
  fc.known_image_path,
}

if dt.configuration.running_os == "windows" or dt.configuration.running_os == "macos" then
  table.insert(widgets, df.executable_path_widget({"face_recognition"}))
end
table.insert(widgets, dt.new_widget("section_label"){ label = _("processing options")})
table.insert(widgets, fc.tolerance)
table.insert(widgets, fc.num_cores)
table.insert(widgets, fc.export_format)
table.insert(widgets, dt.new_widget("box"){
  orientation = "horizontal",
  dt.new_widget("label"){ label = _("width  ")},
  fc.width,
})
table.insert(widgets, dt.new_widget("box"){
  orientation = "horizontal",
  dt.new_widget("label"){ label = _("height ")},
  fc.height,
})
table.insert(widgets, fc.execute)

fc.widget = dt.new_widget("box"){
  orientation = vertical,
  table.unpack(widgets)
}

--fc.tolerance.value = dt.preferences.read(MODULE, "tolerance", "float")

-- Register
--dt.register_storage("module_face_recognition", _("Face recognition"), show_status, face_recognition)

dt.register_lib(
  "face_recognition",     -- Module name
  _("face recognition"),     -- Visible name
  true,                -- expandable
  false,               -- resetable
  {[dt.gui.views.lighttable] = {"DT_UI_CONTAINER_PANEL_RIGHT_CENTER", 300}},   -- containers
  fc.widget,
  nil,-- view_enter
  nil -- view_leave
)

fc.tolerance.value = dt.preferences.read(MODULE, "tolerance", "float")

--
-- vim: shiftwidth=2 expandtab tabstop=2 cindent syntax=lua