Storage plugin for face_recognition
This commit is contained in:
GHswitt
163
contrib/face_recognition.lua
Normal file
163
contrib/face_recognition.lua
Normal file
@@ -0,0 +1,163 @@
|
||||
--[[
|
||||
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 df = require "lib/dtutils.file"
|
||||
local gettext = dt.gettext
|
||||
|
||||
-- works with darktable API version from 2.0.0 to 5.0.0
|
||||
dt.configuration.check_version(...,{2,0,0},{3,0,0},{4,0,0},{5,0,0})
|
||||
|
||||
-- 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
|
||||
|
||||
-- Register preference for known faces path
|
||||
dt.preferences.register("FaceRecognition",
|
||||
"knownImagePath",
|
||||
"directory", -- type
|
||||
_("Face recognition: Known images"), -- label
|
||||
_("Path to images with known faces, files named after tag to apply"), -- tooltip
|
||||
"~/.config/darktable/face_recognition") -- default
|
||||
-- ...and number of CPU cores to use
|
||||
dt.preferences.register("FaceRecognition",
|
||||
"nrCores",
|
||||
"integer", -- type
|
||||
_("Face recognition: Nr of CPU cores"), -- label
|
||||
_("Number of CPU cores to use, 0 for all"), -- tooltip
|
||||
0, -- default
|
||||
0, -- min
|
||||
64) -- max
|
||||
|
||||
local function show_status (storage, image, format, filename, number, total, high_quality, extra_data)
|
||||
dt.print("Export to Face recognition "..tostring(number).."/"..tostring(total))
|
||||
end
|
||||
|
||||
local function face_recognition (storage, image_table, extra_data) --finalize
|
||||
if not df.check_if_bin_exists("face_recognition") then
|
||||
dt.print(_("Face recognition not found"))
|
||||
return
|
||||
end
|
||||
|
||||
-- list of exported images
|
||||
local img_list = {}
|
||||
|
||||
for _,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_error ("FR: Unknown path: " .. path)
|
||||
|
||||
-- Path to known images and number of CPU cores to use
|
||||
local knownPath = dt.preferences.read("FaceRecognition", "knownImagePath", "directory")
|
||||
local nrCores = dt.preferences.read("FaceRecognition", "nrCores", "integer")
|
||||
|
||||
if nrCores < 1 then
|
||||
nrCores = -1
|
||||
end
|
||||
|
||||
-- Output file
|
||||
local output = path .. "fr.txt"
|
||||
|
||||
local command = "face_recognition --cpus " .. nrCores .. " " .. knownPath .. " " .. path .. " > " .. output
|
||||
dt.print_error("FR: " .. command)
|
||||
dt.print(_("Starting face recognition..."))
|
||||
|
||||
dt.control.execute(command)
|
||||
|
||||
-- Remove exported images
|
||||
for _,v in ipairs(img_list) do
|
||||
os.remove (v)
|
||||
end
|
||||
|
||||
-- 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
|
||||
local result = {}
|
||||
for line in io.lines(output) do
|
||||
local file, tag = string.match (line, "(.*),(.*)$")
|
||||
tag = string.gsub (tag, "%d*$", "")
|
||||
dt.print_error ("F:"..file .." T:".. tag)
|
||||
if result[file] ~= nil then
|
||||
table.insert (result[file], tag)
|
||||
else
|
||||
result[file] = {tag}
|
||||
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
|
||||
dt.print_error ("I:" .. img.id .. " T: ".. t)
|
||||
-- Create tag if it does not exists
|
||||
local tag = dt.tags.create (t)
|
||||
img:attach_tag (tag)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
--os.remove (output)
|
||||
|
||||
end
|
||||
|
||||
-- Register
|
||||
dt.register_storage("module_face_recognition", _("Face recognition"), show_status, face_recognition)
|
||||
|
||||
--
|
||||
-- vim: shiftwidth=2 expandtab tabstop=2 cindent syntax=lua
|
||||
Reference in New Issue
Block a user