Source code for image_processing.sample_extract

"""
This class wraps all the functionality required to extract text from an image.
"""

import os
import cv2
import pytesseract
from PIL import Image
from image_preprocessing.build_director import BuildDirector
from image_processing.text_cleaner import TextCleaner
from image_processing.simplification_manager import SimplificationManager
from image_processing.barcode_manager import BarCodeManager
from image_preprocessing.template_matching import TemplateMatching
from image_processing.context_manager import ContextManager
from hutts_utils.hutts_logger import logger, prettify_json_message

__author__ = "Nicolai van Niekerk"
__copyright__ = "Copyright 2017, Java the Hutts"
__license__ = "BSD"
__maintainer__ = "Nicolai van Niekerk"
__email__ = "nicvaniek@gmail.com"
__status__ = "Development"

DESKTOP = os.path.join(os.path.join(os.path.expanduser('~')), 'Desktop')


class TextExtractor:
    """
    The TextExtractor extracts text from the ID image.
    """

    def __init__(self, preferences):
        """
        Initialise Text Extractor.

        Args:
            preferences (dict): User-specified CV techniques.

        Returns:
            None
        """
        self.preferences = preferences
        self.remove_face = 'false'
        self._context_manager = ContextManager()
        self._text_cleaner = TextCleaner()

    def extract(self, img):
        """
        This function is a sample that demonstrates how text would be extracted

        Args:
            img: The image of the ID that contains the text to be extracted

        Returns:
            id_details (obj): The extracted information
        """
        if 'remove_face' in self.preferences:
            self.remove_face = self.preferences['remove_face'] == 'true'
        logger.debug('self.remove_face: ' + str(self.remove_face))

        simplification_manager = SimplificationManager()
        barcode_manager = BarCodeManager()
        data = {}

        # Perform perspective transformation and read from barcode.
        logger.info('Performing perspective transformation...')
        image = simplification_manager.perspectiveTransformation(img, self.preferences['useIO'])
        if self.preferences['useIO']:
            cv2.imwrite(DESKTOP + "/output/3.png", image)
        barcode_data_found, barcode_scan_data, barcoded_image = barcode_manager.get_barcode_info(image)
        if barcode_data_found:
            logger.info('Barcode successfully scanned')
            data = {
                'identity_number': barcode_scan_data.decode('utf-8'),
            }

        # Process image
        if 'id_type' in self.preferences:
            identification_type = self.preferences['id_type']
            logger.info("No template matching required")
            logger.info("Identification type: " + identification_type)
        else:
            template_match = TemplateMatching()
            logger.info('Performing template matching...')
            identification_type = template_match.identify(barcoded_image)

        logger.info('Constructing text extraction pipeline')
        pipeline = BuildDirector.construct_text_extract_pipeline(self.preferences, identification_type)
        image = pipeline.process_text_extraction(self.preferences['useIO'], barcoded_image, self.remove_face)

        # Extract and return text
        filename = "{}.png".format(os.getpid())
        cv2.imwrite(filename, image)

        text = pytesseract.image_to_string(Image.open(filename))
        os.remove(filename)

        # Log the uncleaned string to terminal.
        # This is for demonstration purposes.
        logger.debug('-' * 50)
        logger.debug('String to clean:')
        logger.debug('-' * 50)
        [logger.debug(log_line) for log_line in text.split('\n')]
        logger.debug('-' * 50)
        logger.info('Cleaning up text...')
        # Clean the OCR output text.
        clean_text = self._text_cleaner.clean_up(text)
        # Log the cleaned string to terminal.
        # This is for demonstration purposes.
        logger.debug('-' * 50)
        logger.debug('Cleaned text:')
        logger.debug('-' * 50)
        [logger.debug(log_line) for log_line in clean_text.split('\n')]
        logger.debug('-' * 50)
        # Get ID information from cleaned text.
        logger.info('Placing extracted text in a dictionary...')
        id_context = self._context_manager.get_id_context(identification_type)
        id_details = id_context.get_id_info(clean_text, barcode_data=data)
        # Log the retrieved ID information extracted text to terminal.
        # This is for demonstration purposes.
        logger.debug('-' * 50)
        logger.debug('Extracted ID details:')
        logger.debug('-' * 50)
        [logger.debug(id_details_line) for id_details_line in prettify_json_message(id_details).split('\n')]
        logger.debug('-' * 50)
        # Return the extracted ID information.
        return id_details


class FaceExtractor:
    """
    The FaceExtractor extracts the face region for the image passed.
    """
    def extract(self, img, use_io):
        """
        This function is a sample that demonstrates how the face would be extracted.

        Args:
            img: The image of the ID that contains the face that must be extracted.
            use_io (boolean): Whether or not images should be written to disk

        Returns:
            image: The extracted and aligned facial image.
        """
        simplification_manager = SimplificationManager()

        # Perform perspective transformation
        logger.info('Performing perspective transformation...')
        perspective_image = simplification_manager.perspectiveTransformation(img, use_io)
        if use_io:
            cv2.imwrite(DESKTOP + "/output/10.png", perspective_image)

        # Process image
        logger.info('Constructing facial extraction pipeline...')
        pipeline = BuildDirector.construct_face_extract_pipeline()
        image = pipeline.process_face_extraction(perspective_image)

        return image