313 lines
14 KiB
Python
Executable file
313 lines
14 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
# -*- encoding: utf-8 -*-
|
|
#
|
|
# This script is licensed under GNU GPL version 2.0 or above
|
|
# (c) 2021 Antonio J. Delgado
|
|
#
|
|
|
|
import sys
|
|
import os
|
|
import logging
|
|
import json
|
|
import shutil
|
|
import re
|
|
import click
|
|
import click_config_file
|
|
from logging.handlers import SysLogHandler
|
|
import face_recognition
|
|
import pyexiv2
|
|
import PIL
|
|
|
|
|
|
class CustomFormatter(logging.Formatter):
|
|
"""Logging colored formatter, adapted from https://stackoverflow.com/a/56944256/3638629"""
|
|
|
|
grey = '\x1b[38;21m'
|
|
blue = '\x1b[38;5;39m'
|
|
yellow = '\x1b[38;5;226m'
|
|
red = '\x1b[38;5;196m'
|
|
bold_red = '\x1b[31;1m'
|
|
reset = '\x1b[0m'
|
|
|
|
def __init__(self, fmt):
|
|
super().__init__()
|
|
self.fmt = fmt
|
|
self.FORMATS = {
|
|
logging.DEBUG: self.grey + self.fmt + self.reset,
|
|
logging.INFO: self.blue + self.fmt + self.reset,
|
|
logging.WARNING: self.yellow + self.fmt + self.reset,
|
|
logging.ERROR: self.red + self.fmt + self.reset,
|
|
logging.CRITICAL: self.bold_red + self.fmt + self.reset
|
|
}
|
|
|
|
def format(self, record):
|
|
log_fmt = self.FORMATS.get(record.levelno)
|
|
formatter = logging.Formatter(log_fmt)
|
|
return formatter.format(record)
|
|
|
|
|
|
class image_classifier:
|
|
|
|
def __init__(self, debug_level, log_file, faces_directory, directory, no_move,
|
|
people_folder, recursive):
|
|
''' Initial function called when object is created '''
|
|
self.debug_level = debug_level
|
|
if log_file is None:
|
|
home_path = os.environ.get('HOME', os.environ.get('USERPROFILE', os.getcwd()))
|
|
log_file = os.path.join(home_path, 'log', 'image_classifier.log')
|
|
self.log_file = log_file
|
|
self._init_log()
|
|
self.faces_directory = faces_directory
|
|
self.directory = directory
|
|
self.known_people = self.load_known_people()
|
|
self.no_move = no_move
|
|
self.people_folder = people_folder
|
|
self.recursive = recursive
|
|
|
|
if os.access(directory, os.R_OK):
|
|
if self.recursive:
|
|
entries = self.recursive_scandir(directory)
|
|
else:
|
|
entries = os.scandir(directory)
|
|
self._log.debug(f"Processing {len(entries)} files...")
|
|
for entry in entries:
|
|
if not entry.name.startswith('.') and entry.is_file():
|
|
self.process_file(entry.path)
|
|
|
|
def recursive_scandir(path, ignore_hidden_files=True):
|
|
files = []
|
|
try:
|
|
for file in os.scandir(path):
|
|
if not file.name.startswith('.'):
|
|
if file.is_file():
|
|
files.append(file)
|
|
elif file.is_dir(follow_symlinks=False):
|
|
more_files = self.recursive_scandir(
|
|
file.path,
|
|
ignore_hidden_files=ignore_hidden_files
|
|
)
|
|
if more_files:
|
|
files = files + more_files
|
|
except PermissionError as error:
|
|
self._log.warning(f"Permission denied accessing folder '{path}'")
|
|
return files
|
|
|
|
|
|
def process_metadata(self, file):
|
|
self.metadata = pyexiv2.ImageMetadata(file)
|
|
self.metadata.read()
|
|
if 'Xmp.iptcExt.PersonInImage' in self.metadata.xmp_keys:
|
|
self._log.debug(f"People (before): \
|
|
{self.metadata['Xmp.iptcExt.PersonInImage'].raw_value} \
|
|
(type: {type(self.metadata['Xmp.iptcExt.PersonInImage'].raw_value)})")
|
|
|
|
def process_file(self, file):
|
|
''' Process a file, find faces, add EXIF information and
|
|
move it to the folder of the day'''
|
|
self._log.debug(f"Looking for faces in file '{file}'...")
|
|
folder_date = 'unknown-time'
|
|
dirname = os.path.dirname(os.path.realpath(file))
|
|
filename = os.path.basename(file)
|
|
people = list()
|
|
if not os.access(file, os.R_OK):
|
|
self._log.error(f"The file '{file}' is not readable.")
|
|
else:
|
|
if self.is_image(file):
|
|
self.process_metadata(file)
|
|
people = self.find_faces(file)
|
|
if people:
|
|
self._log.debug(f"Found {len(people)} known people in the image.")
|
|
self._log.debug(json.dumps(people, indent=2))
|
|
self.append_people(file, people)
|
|
if 'Exif.Photo.DateTimeOriginal' in self.metadata.exif_keys:
|
|
original_date = self.metadata['Exif.Photo.DateTimeOriginal'].value
|
|
self._log.debug(f"File creation time: {original_date} \
|
|
(type: {type(original_date)})")
|
|
folder_date = original_date.strftime('%Y.%m.%d')
|
|
if folder_date == 'unknown-time':
|
|
match = re.search(r'(?P<year>20[0-9]{2})[\-/\._]?\
|
|
(?P<month>[0-1]?[0-9])[\-/\._]?(?P<day>[0-3]?[0-9])', filename)
|
|
if match:
|
|
folder_date = f"{match.group('year')}.{match.group('month')}.\
|
|
{match.group('day')}"
|
|
else:
|
|
match = re.search(r'(?P<day>[0-3]?[0-9])[\-/\._]?\
|
|
(?P<month>[0-1]?[0-9])[\-/\._]?(?P<year>20[0-9]{2})', filename)
|
|
if match:
|
|
folder_date = f"{match.group('year')}.{match.group('month')}.\
|
|
{match.group('day')}"
|
|
else:
|
|
self._log.debug(f"The file '{file}' doesn't seem to be an image for PIL.")
|
|
folder = os.path.join(dirname, folder_date, filename)
|
|
self._log.debug(f"Time based folder name section '{folder_date}'")
|
|
new_path = os.path.dirname(folder)
|
|
if not os.path.exists(new_path):
|
|
os.makedirs(new_path)
|
|
if self.people_folder:
|
|
for person in people:
|
|
person_path = os.path.join(self.people_folder,
|
|
person.replace(' ', '.'),
|
|
folder_date)
|
|
if not os.path.exists(person_path):
|
|
os.makedirs(person_path)
|
|
self._log.info(f"Copying file '{file}' to person '{person}' folder,\
|
|
'{person_path}'...")
|
|
try:
|
|
shutil.copy(file, person_path)
|
|
except FileNotFoundError as error:
|
|
self._log.error(f"Error copying. File not found. {error}")
|
|
if not self.no_move:
|
|
if os.path.exists(os.path.join(new_path, filename)):
|
|
self._log.debug(f"Destination '{new_path}/{filename}' exists, removing it...")
|
|
os.remove(os.path.join(new_path, filename))
|
|
self._log.info(f"Moving file '{file}' to '{new_path}'...")
|
|
try:
|
|
shutil.move(file, new_path)
|
|
except FileNotFoundError as error:
|
|
self._log.error(f"Error copying. File not found. {error}")
|
|
else:
|
|
self._log.info(f"NOT moving file '{file}' to '{new_path}' because of --no-move")
|
|
|
|
def print_metadata(self):
|
|
print("IPTC keys:")
|
|
for key in self.metadata.iptc_keys:
|
|
print(f" {key}: '{self.metadata[key].raw_value}'")
|
|
print("EXIF keys:")
|
|
for key in self.metadata.exif_keys:
|
|
print(f" {key}: '{self.metadata[key].raw_value}'")
|
|
print("XMP keys:")
|
|
for key in self.metadata.xmp_keys:
|
|
print(f" {key}: '{self.metadata[key].raw_value}'")
|
|
|
|
def append_people(self, file, people):
|
|
new_list = list()
|
|
if 'Xmp.iptcExt.PersonInImage' in self.metadata.xmp_keys:
|
|
for person in self.metadata['Xmp.iptcExt.PersonInImage'].raw_value:
|
|
new_list.append(person)
|
|
for person in people:
|
|
if person not in new_list:
|
|
self._log.debug(f"Adding person '{person}'...")
|
|
new_list.append(person)
|
|
if 'Xmp.iptcExt.PersonInImage' in self.metadata.xmp_keys:
|
|
self.metadata['Xmp.iptcExt.PersonInImage'].value = new_list
|
|
else:
|
|
self.metadata['Xmp.iptcExt.PersonInImage'] = pyexiv2.XmpTag('Xmp.iptcExt.PersonInImage',
|
|
new_list)
|
|
self._log.debug(f"People (after): \
|
|
{self.metadata['Xmp.iptcExt.PersonInImage'].raw_value} \
|
|
(type: {type(self.metadata['Xmp.iptcExt.PersonInImage'].raw_value)})")
|
|
try:
|
|
self.metadata.write()
|
|
self._log.debug(f"Updated file '{file}'.")
|
|
except OSError as error:
|
|
self._log.error(f"Error writing metadata to picture file. {error}")
|
|
|
|
def is_json(self, data):
|
|
try:
|
|
json.loads(data)
|
|
except TypeError:
|
|
return False
|
|
return True
|
|
|
|
def load_known_people(self):
|
|
known_people = list()
|
|
self._log.debug(f"Looking for known faces in directory '{self.faces_directory}'...")
|
|
if os.access(self.faces_directory, os.R_OK):
|
|
with os.scandir(self.faces_directory) as faces_items:
|
|
for entry in faces_items:
|
|
if (not entry.name.startswith('.') and entry.is_file() and
|
|
self.is_image(self.faces_directory + os.sep + entry.name)):
|
|
self._log.debug(f"Identifying face in file '{entry.name}'...")
|
|
person = dict()
|
|
person['filename'] = face_recognition.load_image_file(self.faces_directory +
|
|
os.sep + entry.name)
|
|
person['name'] = os.path.basename(os.path.splitext(self.faces_directory +
|
|
os.sep + entry.name)[0])
|
|
encodings = face_recognition.face_encodings(person['filename'])
|
|
if len(encodings) > 0:
|
|
person['encoding'] = encodings[0]
|
|
known_people.append(person)
|
|
else:
|
|
self._log.info(f"No faces found in file '{entry.name}'.")
|
|
return known_people
|
|
|
|
def find_faces(self, file):
|
|
''' Find faces in an image/video file '''
|
|
people = list()
|
|
if self.is_image(file):
|
|
image = face_recognition.load_image_file(file)
|
|
encodings = face_recognition.face_encodings(image)
|
|
self._log.debug(f"Found {len(encodings)} faces.")
|
|
for known_person in self.known_people:
|
|
for encoding in encodings:
|
|
if face_recognition.compare_faces([known_person['encoding']], encoding)[0]:
|
|
if known_person['name'] not in people:
|
|
people.append(known_person['name'])
|
|
else:
|
|
return False
|
|
return people
|
|
|
|
def is_image(self, file):
|
|
try:
|
|
PIL.Image.open(file)
|
|
except OSError as error:
|
|
self._log.debug(f"File '{file}' is not readable by PIL. {error}")
|
|
return False
|
|
except PIL.UnidentifiedImageError as error:
|
|
self._log.debug(f"File '{file}' is not an image recognizable by PIL. {error}")
|
|
return False
|
|
return True
|
|
|
|
def _init_log(self):
|
|
''' Initialize log object '''
|
|
self._log = logging.getLogger("image_classifier")
|
|
self._log.setLevel(logging.DEBUG)
|
|
|
|
sysloghandler = SysLogHandler()
|
|
sysloghandler.setLevel(logging.DEBUG)
|
|
self._log.addHandler(sysloghandler)
|
|
|
|
streamhandler = logging.StreamHandler(sys.stdout)
|
|
streamhandler.setLevel(logging.getLevelName(self.debug_level))
|
|
# formatter = '%(asctime)s | %(levelname)8s | %(message)s'
|
|
formatter = '[%(levelname)s] %(message)s'
|
|
streamhandler.setFormatter(CustomFormatter(formatter))
|
|
self._log.addHandler(streamhandler)
|
|
|
|
if not os.path.exists(os.path.dirname(self.log_file)):
|
|
os.mkdir(os.path.dirname(self.log_file))
|
|
|
|
filehandler = logging.handlers.RotatingFileHandler(self.log_file, maxBytes=102400000)
|
|
# create formatter
|
|
formatter = logging.Formatter('%(asctime)s %(name)-12s %(levelname)-8s %(message)s')
|
|
filehandler.setFormatter(formatter)
|
|
filehandler.setLevel(logging.DEBUG)
|
|
self._log.addHandler(filehandler)
|
|
return True
|
|
|
|
|
|
@click.command()
|
|
@click.option("--debug-level", "-d", default="INFO",
|
|
type=click.Choice(
|
|
["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "NOTSET"],
|
|
case_sensitive=False,
|
|
),
|
|
help='Set the debug level for the standard output.')
|
|
@click.option('--log-file', '-l', help="File to store all debug messages.")
|
|
@click.option("--faces-directory", "-f", required=True, help="Folder containing the pictures that \
|
|
identify people. Filename would be used as the name for the person. Just one person per picture.")
|
|
@click.option("--directory", "-d", required=True, help="Folder with the pictures to classify.")
|
|
@click.option("--no-move", "-n", is_flag=True, help="Don't move files, just add people's tag.")
|
|
@click.option('--people-folder', '-p', help="Define a folder for people's folders and copy \
|
|
pictures to each person's folder. Be sure to have deduplication in the filesystem to avoid using \
|
|
too much storage.")
|
|
@click.option('--recursive', '-r', is_flag=True, help='Recursively search for files in the provided --directory')
|
|
@click_config_file.configuration_option()
|
|
def __main__(debug_level, log_file, faces_directory, directory, no_move,
|
|
people_folder, recursive):
|
|
return image_classifier(debug_level, log_file, faces_directory, directory, no_move,
|
|
people_folder, recursive)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
__main__()
|