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 click
import click_config_file
from logging.handlers import SysLogHandler
import face_recognition
import exif
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):
        ''' Initial function called when object is created '''
        self.debug_level = debug_level
        if log_file is None:
            log_file = os.path.join(os.environ.get('HOME', os.environ.get('USERPROFILE', os.getcwd())), '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()

        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(directory + os.sep + entry.name)

    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"Processing file '{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))
            with open(file, 'rb') as image_file:
                self.exif_info = exif.Image(image_file)
            if not self.exif_info.has_exif:
                self._log.debug("No exif info in the image.")
            else:
                # for key in dir(self.exif_info):
                #     if not key.startswith("_"):
                #         sys.stdout.write(f"{key}: ")
                #         sys.stdout.write(f"{self.exif_info[key]}")
                self.append_people_to_exif(people)
                with open(file, 'wb') as new_image_file:
                    new_image_file.write(self.exif_info.get_file())
                self._log.debug(f"Updated file '{file}'.")
                # get date
                self._log.debug(f"File time stamp: {self.exif_info.get('Image timestamp')} (type: {type(self.exif_info.get('Image timestamp'))})")
                # move to destination
        else:
            self._log.debug("Doesn't seem to be an image.")

    def append_people_to_exif(self, people):
        if self.is_json(self.exif_info.get('user_comment')):
            data = json.loads(self.exif_info['user_comment'])
            if 'PeopleDetected' not in data:
                data['PeopleDetected'] = list()
        else:
            data = dict()
            if self.exif_info.get('user_comment'):
                data["previous_user_comment"]=self.exif_info.get('user_comment')
            data['PeopleDetected'] = list()
        for person in people:
            data['PeopleDetected'].append(person)
        self._log.debug(f"New 'user_comment': {json.dumps(data, indent=2)}")
        self.exif_info.set("user_comment", json.dumps(data))

    def is_json(self, data):
        try:
            result = 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():
                        self._log.debug(f"Detecting known person 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])
                        person['encoding'] = face_recognition.face_encodings(person['filename'])[0]
                        known_people.append(person)
        return known_people

    def find_faces(self, file):
        ''' Find faces in an image/video file '''
        people = list()
        try:
            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'])
        except PIL.UnidentifiedImageError as error:
            self._log.debug(f"File '{file}' don't seem to be an image.")
            return False
        return people

    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. The filename would be used as the name for the person. Just one person per picture.")
@click.option("--directory","-d", required=True, help="Folder containing the pictures to classify.")
@click_config_file.configuration_option()
def __main__(debug_level, log_file, faces_directory, directory):
    object = image_classifier(debug_level, log_file, faces_directory, directory)

if __name__ == "__main__":
    __main__()