Skip to content

Commit cb6199c

Browse files
supertobisupertobi
authored
Merge pull request darktable-org#100 from GHswitt/master
Storage plugin for face_recognition
2 parents c9eb087 + 8e6552e commit cb6199c

File tree

1 file changed

+215
-0
lines changed

1 file changed

+215
-0
lines changed

‎contrib/face_recognition.lua

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
--[[
2+
Face recognition for darktable
3+
4+
Copyright (c) 2017 Sebastian Witt
5+
6+
darktable is free software: you can redistribute it and/or modify
7+
it under the terms of the GNU General Public License as published by
8+
the Free Software Foundation, either version 3 of the License, or
9+
(at your option) any later version.
10+
11+
darktable is distributed in the hope that it will be useful,
12+
but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
GNU General Public License for more details.
15+
16+
You should have received a copy of the GNU General Public License
17+
along with darktable. If not, see <http://www.gnu.org/licenses/>.
18+
]]
19+
20+
--[[
21+
face_recognition
22+
Add a new storage option to send images to face_recognition.
23+
Images are exported to darktable tmp dir first.
24+
A directory with known faces must exist, the image name are the
25+
tag names which will be used.
26+
Multiple images for one face can exist, add a number to it, the
27+
number will be removed from the tag, for example:
28+
People|IknowYou1.jpg
29+
People|IknowYou2.jpg
30+
People|Another.jpg
31+
People|Youtoo.jpg
32+
33+
ADDITIONAL SOFTWARE NEEDED FOR THIS SCRIPT
34+
* https://github.com/ageitgey/face_recognition
35+
* https://github.com/darktable-org/lua-scripts/tree/master/lib
36+
37+
USAGE
38+
* require this file from your main luarc config file.
39+
40+
This plugin will add a new storage option and calls face_recognition after export.
41+
]]
42+
43+
local dt = require "darktable"
44+
local df = require "lib/dtutils.file"
45+
local gettext = dt.gettext
46+
47+
-- works with darktable API version from 2.0.0 to 5.0.0
48+
dt.configuration.check_version(...,{2,0,0},{3,0,0},{4,0,0},{5,0,0})
49+
50+
-- Tell gettext where to find the .mo file translating messages for a particular domain
51+
gettext.bindtextdomain("face_recognition", dt.configuration.config_dir.."/lua/locale/")
52+
53+
local function _(msgid)
54+
return gettext.dgettext("face_recognition", msgid)
55+
end
56+
57+
-- Preference: Tag for unknown_person
58+
dt.preferences.register("FaceRecognition",
59+
"unknownTag",
60+
"string", -- type
61+
_("Face recognition: Unknown tag"), -- label
62+
_("Tag for faces that are not recognized"), -- tooltip
63+
"unknown_person")
64+
-- Preference: Images with this substring in tags are ignored
65+
dt.preferences.register("FaceRecognition",
66+
"ignoreTags",
67+
"string", -- type
68+
_("Face recognition: Ignore tag"), -- label
69+
_("Images with this substring in tags are ignored, separate multiple strings with ,"), -- tooltip
70+
"")
71+
-- Preference: Number of CPU cores to use
72+
dt.preferences.register("FaceRecognition",
73+
"nrCores",
74+
"integer", -- type
75+
_("Face recognition: Nr of CPU cores"), -- label
76+
_("Number of CPU cores to use, 0 for all"), -- tooltip
77+
0, -- default
78+
0, -- min
79+
64) -- max
80+
-- Preference: Known faces path
81+
dt.preferences.register("FaceRecognition",
82+
"knownImagePath",
83+
"directory", -- type
84+
_("Face recognition: Known images"), -- label
85+
_("Path to images with known faces, files named after tag to apply"), -- tooltip
86+
"~/.config/darktable/face_recognition") -- default -- default
87+
88+
local function show_status (storage, image, format, filename, number, total, high_quality, extra_data)
89+
dt.print("Export to Face recognition "..tostring(number).."/"..tostring(total))
90+
end
91+
92+
-- Check if image has ignored tag attached
93+
local function ignoreByTag (image, ignoreTags)
94+
local tags = image:get_tags ()
95+
local ignoreImage = false
96+
-- For each image tag
97+
for _,t in ipairs (tags) do
98+
-- Check if it contains a ignore tag
99+
for _,it in ipairs (ignoreTags) do
100+
if string.find (t.name, it, 1, true) then
101+
-- The image has ignored tag attached
102+
ignoreImage = true
103+
dt.print_error ("Face recognition: Ignored tag: " .. it .. " found in " .. image.id .. ":" .. t.name)
104+
end
105+
end
106+
end
107+
108+
return ignoreImage
109+
end
110+
111+
local function face_recognition (storage, image_table, extra_data) --finalize
112+
if not df.check_if_bin_exists("face_recognition") then
113+
dt.print(_("Face recognition not found"))
114+
return
115+
end
116+
117+
-- Get preferences
118+
local knownPath = dt.preferences.read("FaceRecognition", "knownImagePath", "directory")
119+
local nrCores = dt.preferences.read("FaceRecognition", "nrCores", "integer")
120+
local ignoreTagString = dt.preferences.read("FaceRecognition", "ignoreTags", "string")
121+
local unknownTag = dt.preferences.read("FaceRecognition", "unknownTag", "string")
122+
123+
-- face_recognition uses -1 for all cores, we use 0 in preferences
124+
if nrCores < 1 then
125+
nrCores = -1
126+
end
127+
128+
-- Split ignore tags (if any)
129+
ignoreTags = {}
130+
for tag in string.gmatch(ignoreTagString, '([^,]+)') do
131+
table.insert (ignoreTags, tag)
132+
dt.print_error ("Face recognition: Ignore tag: " .. tag)
133+
end
134+
135+
-- list of exported images
136+
local img_list = {}
137+
138+
for img,v in pairs(image_table) do
139+
table.insert (img_list, v)
140+
end
141+
142+
-- Get path of exported images
143+
local path = df.get_path (img_list[1])
144+
dt.print_error ("Face recognition: Path to unknown images: " .. path)
145+
146+
-- Output file
147+
local output = path .. "facerecognition.txt"
148+
149+
local command = "face_recognition --cpus " .. nrCores .. " " .. knownPath .. " " .. path .. " > " .. output
150+
dt.print_error("Face recognition: Running command: " .. command)
151+
dt.print(_("Starting face recognition..."))
152+
153+
dt.control.execute(command)
154+
155+
-- Remove exported images
156+
for _,v in ipairs(img_list) do
157+
os.remove (v)
158+
end
159+
160+
-- Open output file
161+
local f = io.open(output, "rb")
162+
163+
if not f then
164+
dt.print(_("Face recognition failed"))
165+
else
166+
dt.print(_("Face recognition finished"))
167+
f:close ()
168+
end
169+
170+
-- Read output
171+
local result = {}
172+
for line in io.lines(output) do
173+
local file, tag = string.match (line, "(.*),(.*)$")
174+
tag = string.gsub (tag, "%d*$", "")
175+
dt.print_error ("File:"..file .." Tag:".. tag)
176+
if result[file] ~= nil then
177+
table.insert (result[file], tag)
178+
else
179+
result[file] = {tag}
180+
end
181+
end
182+
183+
-- Attach tags
184+
for file,tags in pairs(result) do
185+
-- Find image in table
186+
for img,file2 in pairs(image_table) do
187+
if file == file2 then
188+
for _,t in ipairs (tags) do
189+
-- Check if image is ignored
190+
if ignoreByTag (img, ignoreTags) then
191+
dt.print_error("Face recognition: Ignoring image with ID " .. img.id)
192+
else
193+
-- Check of unrecognized unknown_person
194+
if t == "unknown_person" then
195+
t = unknownTag
196+
end
197+
dt.print_error ("ImgId:" .. img.id .. " Tag:".. t)
198+
-- Create tag if it does not exists
199+
local tag = dt.tags.create (t)
200+
img:attach_tag (tag)
201+
end
202+
end
203+
end
204+
end
205+
end
206+
207+
--os.remove (output)
208+
209+
end
210+
211+
-- Register
212+
dt.register_storage("module_face_recognition", _("Face recognition"), show_status, face_recognition)
213+
214+
--
215+
-- vim: shiftwidth=2 expandtab tabstop=2 cindent syntax=lua

0 commit comments

Comments
 (0)