image_classifier/image_classifier/image_classifier.py

#!/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):
        ''' 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

        if os.access(directory, os.R_OK):
            with os.scandir(directory) as directory_item:
                for entry in directory_item:
                    if not entry.name.startswith('.') and entry.is_file():
                        self.process_file(os.path.join(directory, entry.name))

    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')}"
            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}'...")
                    shutil.copy(file, person_path)
            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}'...")
                shutil.move(file, new_path)
            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_config_file.configuration_option()
def __main__(debug_level, log_file, faces_directory, directory, no_move, people_folder):
    return image_classifier(debug_level, log_file, faces_directory, directory, no_move,
                            people_folder)


if __name__ == "__main__":
    __main__()