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:
            with open(file, 'rb') as image_file:
                exif_info = exif.Image(image_file)
            if exif_info.has_exif:
                print(json.dumps(exif_info, indent=2))
            # get date
            # move to destination

    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.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 '''
        names = list()
        try:
            image = face_recognition.load_image_file(file)
            encodings = face_recognition.face_encodings(image)
            self._log.debug(f"{len(encodings)} found")
            for known_person in self.known_people:
                if known_person['encoding'] in encodings:
                    names.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 names

    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__()