cdpr.py

import argparse
import os
import subprocess
import tempfile

import PIL.Image
import imagehash

DATA_OFFSET = 270, 0
DATA_PADDING_RIGHT = 211
FONT_PADDING = 1
FONT_WIDTH = 19
FONT_HEIGHT = 39
CHARACTERS_PATH = 'characters'
HASH_SIZE = 8

argument_parser = argparse.ArgumentParser()
argument_parser.add_argument('path')

_character_hashes = {}


def recognize_character(character_image):
    hash = imagehash.average_hash(character_image, HASH_SIZE)
    result = list(sorted(_character_hashes.items(), key=lambda item: item[1] - hash))
    return result[0][0]


def extract_frames(path):
    destination = tempfile.mkdtemp()
    subprocess.run(['ffmpeg', '-i', path, '-r', '1/1', os.path.join(destination, '%05d.png')], stdout=subprocess.DEVNULL,
                   stderr=subprocess.DEVNULL)
    return destination


def extract_frame_characters(frame):
    x, y = DATA_OFFSET
    max_width, max_height = frame.size
    max_width -= DATA_PADDING_RIGHT
    index = 0

    characters = []
    while y <= max_height:
        character_image = frame.crop((x, y, x + FONT_WIDTH, y + FONT_HEIGHT))
        character = recognize_character(character_image)
        characters.append(character)

        index += 1
        x += FONT_WIDTH + FONT_PADDING

        if x >= max_width:
            x = DATA_OFFSET[0]
            y += FONT_HEIGHT + FONT_PADDING

    return ''.join(characters).replace(' ', '')


def initialize_data():
    for filename in os.listdir(CHARACTERS_PATH):
        path = os.path.join(CHARACTERS_PATH, filename)
        character = os.path.splitext(filename)[0]
        image = PIL.Image.open(path)
        _character_hashes[character] = imagehash.average_hash(image, HASH_SIZE)

    _character_hashes['/'] = _character_hashes['slash']
    del _character_hashes['slash']


def main(arguments):
    initialize_data()

    frame_destination = extract_frames(arguments.path)
    for filename in sorted(os.listdir(frame_destination)):
        path = os.path.join(frame_destination, filename)
        image = PIL.Image.open(path)
        characters = extract_frame_characters(image)
        print(filename, len(characters))


if __name__ == '__main__':
    arguments = argument_parser.parse_args()
    argument_parser.exit(main(arguments))