BP_DP-xlanro00/src/main.py

"""! @file main.py
    @brief Main file for the application
    @author xlanro00
"""

# Import basic libraries
import argparse as ap
from os.path import exists
import hashlib

# Libraries for image processing
import numpy as np
import matplotlib.pyplot as plt
import cv2 as cv
from stl import mesh
import trimesh
import trimesh.transformations as tmtra
import trimesh.remesh as tmrem

# Import custom image filter library
import filters as flt
import config_parser as cp
import log
import math


class fingerprint_app:
    '''Main class for the application.
    '''

    def __init__(self):
        # Parse arguments from command line
        self.parse_arguments()

        # List and dict for filters and corresponding parameters
        self.filters = []
        self.params = {}

        # Parse configuration from json config file
        if self.args.config:
            self.config_file, self.preset_name = self.args.config
            cp.parse_conf(self.preset_name, self.filters,
                          self.params, self.config_file)

        elif self.args.filters:
            filter_index = 0
            log.print_message(
                "No config file given, using command line arguments")

            # Otherwise expect filters from command line
            for filter_part in self.args.filters:
                # If no '=' char in filter, it is a new filter name
                if filter_part.find('=') == -1:
                    self.filters.append(filter_part)
                    filter_index += 1
                    # create empty dict for params
                    self.params[filter_index] = {}
                # Otherwise it's a parameter for current filter
                else:
                    key, value = filter_part.split('=')
                    self.params[filter_index][key] = value

                cp.parse_params(self.params[filter_index])

            # If database flag is set, save filters to database as a new preset
            if self.args.database:
                cp.save_preset(self.filters, self.params,
                               self.args.database[0])

        else:
            log.print_message("No filters given, saving original image")

        # Set input and output file paths, dpi and mirror flag for easier readability
        self.input_file = self.args.input_file
        self.output_file = self.args.output_file
        self.dpi = self.args.dpi
        self.mirror = True if self.args.mirror else False

        if exists(self.input_file):
            self.run_filtering()
        else:
            log.error_exit("Input file " + self.input_file +
                           " does not exist")

        if self.args.stl:
            self.parse_stl()

    def parse_arguments(self):
        '''Parse arguments from command line using argparse library.
        '''

        parser = ap.ArgumentParser(prog='main.py',
                                   description='Program for transforming a 2D image into 3D fingerprint.',
                                   usage='%(prog)s [-h] [-m | --mirror | --no-mirror] input_file output_file dpi ([-c | --config config_file preset] | [filters ...]) [-s | --stl p height_line height_base | --stl c height_line curv_rate_x curv_rate_y | --stl m height_line iter finger_x finger_y finger_z] [-d | --database database_filename]')

        # positional arguments
        parser.add_argument("input_file", type=str, help="input file path")
        parser.add_argument("output_file", type=str, help="output file path")
        parser.add_argument("dpi", type=int, help="dpi of used scanner")

        # boolean switch argument
        parser.add_argument('-m', '--mirror', type=bool, action=ap.BooleanOptionalAction,
                            help="switch to mirror input image")

        # another boolean switch argument, this time with value, name of the new file and dimensions
        parser.add_argument('-s', "--stl", type=str, nargs='*',
                            help="create stl model from processed image")

        # configuration file containing presets, preset name
        # pair argument - give both or none
        parser.add_argument('-c', '--config', nargs=2,
                            help='pair: name of the config file with presets, name of the preset')

        # array of unknown length, all filter names saved inside
        parser.add_argument('filters', type=str, nargs='*',
                            help="list of filter names and their parameters in form [filter_name1 param1=value param2=value filter_name2 param1=value...]")

        parser.add_argument('-d', '--database', nargs=1,
                            help='switch to store presets in config database')

        self.args = parser.parse_args()

    def parse_stl(self):
        # Get stl filename
        self.stl_path = self.output_file.rsplit('/', 1)[0] + '/'
        self.mode = self.args.stl[0]
        log.print_message("Stl generation in", self.mode, "mode")

        # Default values for stl generation parameters
        def_val = {"hl": 2, "hb": 10, "crx": 2, "cry": 2, "it": 2, "fx": 0, "fy": 0, "fz": 0}

        # Get mode and model parameters
        if self.mode == 'p':
            
            self.height_line = float(self.args.stl[1]) if len(
                self.args.stl) > 1 else def_val.get("hl")
            self.height_base = float(self.args.stl[2]) if len(
                self.args.stl) > 2 else def_val.get("hb")
            
            if len(self.args.stl) < 3:
                log.print_message(
                    "Warning: Too few arguments, using some default values")
            log.print_message("Base height:", self.height_base,
                                "mm, lines depth/height:", self.height_line, "mm")

        elif self.mode == 'c':

            self.height_line = float(self.args.stl[1]) if len(
                self.args.stl) > 1 else def_val.get("hl")
            self.height_base = float(self.args.stl[2]) if len(
                self.args.stl) > 2 else def_val.get("hb")
            self.curv_rate_x = float(self.args.stl[3]) if len(
                self.args.stl) > 3 else def_val.get("crx")
            self.curv_rate_y = float(self.args.stl[4]) if len(
                self.args.stl) > 4 else def_val.get("cry")
            
            if len(self.args.stl) < 5:
                log.print_message(
                    "Warning: Too few arguments, using some default values")
            log.print_message("Line height:", self.height_line, "mm, base height:", self.height_base,
                                "mm, x axis curvature:", self.curv_rate_x, ", y axis curvature:", self.curv_rate_y)
            
        elif self.mode == 'm':

            self.height_line = float(self.args.stl[1]) if len(
                self.args.stl) > 1 else def_val.get("hl")/10
            self.iter = int(self.args.stl[2]) if len(
                self.args.stl) > 2 else def_val.get("it")
            self.finger_x = float(self.args.stl[3]) if len(
                self.args.stl) > 3 else def_val.get("fx")
            self.finger_y = float(self.args.stl[4]) if len(
                self.args.stl) > 4 else def_val.get("fy")
            self.finger_z = float(self.args.stl[5]) if len(
                self.args.stl) > 5 else def_val.get("fz")
            
            if len(self.args.stl) < 6:
                log.print_message(
                    "Warning: Too few arguments, using some default values")
            log.print_message("Line height:", self.height_line, "mm, iterations:", self.iter,
                                ", finger position:", self.finger_x, self.finger_y, self.finger_z)

        else:
            log.error_exit("Unrecognized generation mode")

        self.run_stl()

# ------------------------- FILTERING -------------------------#

    def run_filtering(self):
        '''Read input file, store as numpy.array, uint8, grayscale.
        Call function to apply the filters and a function to save it to output file.
        '''

        self.img = cv.imread(
            self.input_file, cv.IMREAD_GRAYSCALE).astype(np.uint8)

        # Gets empty figure and ax with dimensions of input image
        self.height, self.width = self.img.shape
        self.fig, ax = self.get_empty_figure()

        if self.mirror is True:
            self.mirror_image()

        # Apply all filters and save image
        self.apply_filters()
        self.save_image(self.fig, ax)
        plt.close()

    def get_empty_figure(self):
        '''Return empty figure with one ax, which has dimensions of the input image.
        '''

        size = (self.width/self.dpi, self.height/self.dpi)
        fig = plt.figure(figsize=size, frameon=False, dpi=self.dpi)
        ax = plt.Axes(fig, [0., 0., 1., 1.])
        ax.set_axis_off()
        fig.add_axes(ax)
        return fig, ax

    def mirror_image(self):
        '''Mirror image using opencv, should be used if we want a positive model.
        '''

        log.print_message("Mirroring image")
        self.img = cv.flip(self.img, 1)  # 1 for vertical mirror

    def apply_filters(self):
        '''Apply filters to image one by one.
        In case none were given, pass and save original image to the output file.
        '''

        if len(self.filters) != 0:
            for i, filter_name in enumerate(self.filters):
                # Get filter class from filter.py, use the apply method
                try:
                    filter = getattr(flt, filter_name)
                except AttributeError:
                    log.error_exit("Filter " + filter_name + " not found")
                log.print_message("Applying filter:", filter_name)
                for param in self.params[i+1]:
                    if self.params[i+1][param] is not None:
                        log.print_message("\twith parameter", param,
                                          "=", str(self.params[i+1][param]))

                filter.apply(self, self.params[i+1])
        else:
            pass

    def save_image(self, fig, ax):
        '''Save processed image to the output file.
        '''

        log.print_message("Saving image to", self.output_file)

        # Colormap must be set to grayscale to avoid color mismatch.
        ax.imshow(self.img, cmap="gray")
        fig.savefig(fname=self.output_file)

# ------------------------- STL GENERATION -------------------------#

    def run_stl(self):
        '''Make heightmap, create mesh and save as stl file.
        '''

        self.prepare_heightmap()
        # create ID for the model from all its parameters
        self.get_ID()

        # Create a mesh using one of two modes
        if self.mode == "p":
            self.make_stl_planar()

        elif self.mode == "c":
            self.make_stl_curved()

        elif self.mode == "m":
            self.make_stl_map()

        plt.show()
        self.save_stl()

    def prepare_heightmap(self):
        '''Scale image values to get values from 0 to 255. 
        Then compute base and papilar lines height.
        Check validity of dimension parameters.
        Prepare meshgrid, array which later serves to store point coordinates.
        '''

        if self.img.dtype != np.uint8:
            log.print_message("Converting heightmap to uint8")
            self.img = self.img / np.max(self.img) * 255
            self.img = self.img.astype(np.uint8)

        if self.mode == "p":
            if self.height_base <= 0:
                log.error_exit("Depth of plate height must be positive")

            if self.height_line + self.height_base <= 0:
                log.error_exit("Line depth must be less than plate thickness")

        if self.mode == "c":
            # Don't need to check curvature, check only heights
            if self.height_base <= 0 or self.height_line <= 0:
                log.error_exit("Base and line height must both be positive")

        if self.mode == "m":
            if self.height_line <= 0:
                log.error_exit("Line height must be positive")
            if self.iter < 0:
                log.error_exit("Number of iterations must be positive orr zero")
            self.height_base = 0

        # TODO: curved height base could be done here?
        # Transform image values to get a heightmap
        self.img = (self.height_base + (1 - self.img/255)
                    * self.height_line)

        # This sets the size of stl model and number of subdivisions / triangles
        x = np.linspace(0, self.width * 25.4 / self.dpi, self.width)
        y = np.linspace(0, self.height * 25.4 / self.dpi, self.height)

        self.meshgrid = np.meshgrid(x, y)

    def write_stl_header(self):
        '''Write parameter string to stl header.
        This header is 80 bytes long, so the data needs to be shortened to fit.
        If the parameter string is too long, a warning is printed and the data is truncated.
        '''

        # Truncate if necessary
        if (len(self.param_string) >= 80):
            self.param_string = self.param_string[:80]
            log.print_message("Warning: Parameter string too long, truncating")

        # Overwrite stl header (which is only 80 bytes)
        log.print_message("Writing info to stl header")
        with open(self.stl_filename, "r+") as f:
            f.write(self.param_string)

    def get_ID(self):
        '''Get a unique ID for the model, which is used in filename and on the model backside.
        Also create parameter string for stl header, which is used to create ID using hash function SHA512.
        '''

        # these are the same for all types of models
        param_list = [self.input_file, str(self.dpi)]

        # add parameters specific to the model creation process
        if self.args.config:
            param_list.append(self.config_file)
            param_list.append(self.preset_name)
        else:
            # add filters with their params
            filter_list = []
            for i in range(len(self.filters)):
                tmp_params = []
                for j in self.params[i+1]:
                    if self.params[i+1][j] != None:
                        tmp_params.append(
                            str(j[:1] + ":" + str(self.params[i+1][j])))
                tmp_params = ",".join(tmp_params)
                tmp = str(self.filters[i][0:1] + self.filters[i][-1:])
                if tmp_params != "":
                    tmp = tmp + ";" + str(tmp_params)
                filter_list.append(tmp)
            filter_string = ">".join(filter_list)
            param_list.append(filter_string)

        # these are the same for all types of models
        param_list.append(str(self.height_line))
        param_list.append(str(self.height_base))

        # add parameters specific to the model type
        if self.mode == "c":
            param_list.append(str(self.curv_rate_x))
            param_list.append(str(self.curv_rate_y))

        if self.mode == "m":
            param_list.append(str(self.height_line))
            param_list.append(str(self.iter))
            param_list.append(str(self.finger_x))
            param_list.append(str(self.finger_y))
            param_list.append(str(self.finger_z))

        if self.mode == "p":
            param_list.append("P")
        elif self.mode == "c":
            param_list.append("C")
        elif self.mode == "m":
            param_list.append("M")

        if self.args.mirror:
            param_list.append("F")

        # string that will later be put inside the header of an stl file
        # fill the rest with the ending char to rewrite any leftover header
        # this is done for easier parsing of the header
        self.param_string = "\\".join(param_list)
        self.param_string = self.param_string + \
            "\n" * (80 - len(self.param_string))

        # hash the param string to get unique ID, this will be put in filename and on the back of the model
        # not using built-in hash function because it's seed cannot be set to constant number
        # don't need to worry about collisions and security, just need a relatively unique ID
        self.id = str(hashlib.md5(
            self.param_string.encode('utf-8')).hexdigest())[:10]

    def append_faces(self, faces, c):
        ''' Function to add faces to the  list of faces.
        '''

        faces.append([c, c + 1, c + 2])
        faces.append([c + 1, c + 3, c + 2])
        return c + 4

    def engrave_text(self, bottom_vert_arr, top_vert_arr):
        '''Engrave text on the back of the model.
        Create an empty image, fill it with color and draw text on it.
        '''

        fig, ax = self.get_empty_figure()

        # paint the background black
        ax.plot([0, 1], [0, 1], c="black", lw=self.width)

        # extract filename
        text = self.stl_path.split("/")[-1].split(".")[0] + self.id
        fontsize = 28

        # create text object, paint it white
        t = ax.text(0.5, 0.5, text, ha="center", va="center",
                    fontsize=fontsize, c="white", rotation=90, wrap=True, clip_on=True)

        # adjust fontsize to fit text in the image
        # matplotlib does not support multiline text, wrapping is broken
        rend = fig.canvas.get_renderer()
        while (t.get_window_extent(rend).width > self.width or t.get_window_extent(rend).height > self.height):
            fontsize -= 0.3
            t.set_fontsize(fontsize)

        # update figure, read pixels and reshape to 3d array
        fig.canvas.draw()
        data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))

        # scale inscription layer to suitable height
        data = (data/255)/10

        plt.close()

        # TODO: this is very badly written, fix it
        # TODO: this does not always work, fix it
        # add the bottom array
        OFFSET = 0.01
        
        for i in range(self.height):
            if self.mode == "p":
                for j in range(self.width):
                    bottom_vert_arr[i][j][2] = data[i][j][0]
            elif self.mode == "c":
                for j in range(self.width):
                    bottom_vert_arr[i][j][2] += data[i][j][0]
                    if (bottom_vert_arr[i][j][2] < (top_vert_arr[i][0][2])-OFFSET):
                        bottom_vert_arr[i][j][2] = top_vert_arr[i][0][2]-OFFSET
                    if (bottom_vert_arr[i][j][2] < (top_vert_arr[0][j][2])-OFFSET):
                        bottom_vert_arr[i][j][2] = top_vert_arr[0][j][2]-OFFSET

        return bottom_vert_arr

    def create_stl_mesh(self, faces, vertices):
        '''Create mesh from faces and vertices.
        '''

        # Convert lists to numpy arrays
        faces = np.array(faces)
        vertices = np.array(vertices)
        c = 0

        # Create the mesh - vertices.shape (no_faces, 3, 3)
        self.stl_model = mesh.Mesh(
            np.zeros(faces.shape[0], dtype=mesh.Mesh.dtype))
        for i, face in enumerate(faces):
            for j in range(3):
                self.stl_model.vectors[i][j] = vertices[face[j], :]
            
            # Prints out generation progress
            if i % 100 == 0:
                percentage = round(i/len(faces) * 100, 2)
                if percentage > c:
                    log.print_message("Creating model " + str(c) + "%")
                    c += 10
        
        log.print_message("Model creation finished")

    def create_faces(self, top_vert_arr, bottom_vert_arr):
        '''Create faces for the model.

        Iterate over all vertices, append to vector and create faces from indices
        '''

        count = 0
        vertices = []
        faces = []

        # Front side faces
        for i in range(self.height - 1):
            for j in range(self.width - 1):

                vertices.append([top_vert_arr[i][j]])
                vertices.append([top_vert_arr[i][j+1]])
                vertices.append([top_vert_arr[i+1][j]])
                vertices.append([top_vert_arr[i+1][j+1]])
                count = self.append_faces(faces, count)

        # Back side faces
        for i in range(self.height - 1):
            for j in range(self.width - 1):

                vertices.append([bottom_vert_arr[i][j]])
                vertices.append([bottom_vert_arr[i+1][j]])
                vertices.append([bottom_vert_arr[i][j+1]])
                vertices.append([bottom_vert_arr[i+1][j+1]])
                count = self.append_faces(faces, count)

        # Horizontal side faces
        for i in range(self.height - 1):

            vertices.append([top_vert_arr[i][0]])
            vertices.append([top_vert_arr[i+1][0]])
            vertices.append([bottom_vert_arr[i][0]])
            vertices.append([bottom_vert_arr[i+1][0]])
            count = self.append_faces(faces, count)
            max = self.width - 1

            vertices.append([top_vert_arr[i+1][max]])
            vertices.append([top_vert_arr[i][max]])
            vertices.append([bottom_vert_arr[i+1][max]])
            vertices.append([bottom_vert_arr[i][max]])
            count = self.append_faces(faces, count)

        # Vertical side faces
        for j in range(self.width - 1):

            vertices.append([top_vert_arr[0][j+1]])
            vertices.append([top_vert_arr[0][j]])
            vertices.append([bottom_vert_arr[0][j+1]])
            vertices.append([bottom_vert_arr[0][j]])
            count = self.append_faces(faces, count)
            max = self.height - 1

            vertices.append([top_vert_arr[max][j]])
            vertices.append([top_vert_arr[max][j+1]])
            vertices.append([bottom_vert_arr[max][j]])
            vertices.append([bottom_vert_arr[max][j+1]])
            count = self.append_faces(faces, count)

        return faces, vertices

    def make_stl_planar(self):
        '''
        Create vertices from meshgrid, add depth values from image.
        Create faces from vertices. Add vectors and faces to the model.
        '''

        # Add the image matrix to the 2D meshgrid and create 1D array of 3D points
        tmp_vert_arr = np.vstack(list(map(np.ravel, self.meshgrid))).T
        top_vert_arr = np.concatenate(
            (tmp_vert_arr, (self.img / 10).reshape(-1, 1)), axis=1)

        # Convert 1D array back to matrix of 3D points
        top_vert_arr = top_vert_arr.reshape(self.height, self.width, 3)

        # Prepare image with plotted text for the backside of the lithophane
        bottom_vert_arr = np.copy(top_vert_arr)

        # Engrave text on the back of the model
        self.engrave_text(bottom_vert_arr, top_vert_arr)

        # Create all vertices, faces
        faces, vertices = self.create_faces(top_vert_arr, bottom_vert_arr)

        # Add the created vertices and faces to a mesh
        self.create_stl_mesh(faces, vertices)

    def make_stl_curved(self):
        '''Compute curved surface.
        Create mesh from meshgrid.
        Create vertices from meshgrid, add depth values from image.
        Create faces from vertices. Add vectors and faces to the model.
        '''

        # Calculate the curved surface values
        x = np.arange(self.width)
        y = np.arange(self.height)[:, np.newaxis]
        x = (2*x / self.width) - 1
        z = np.sqrt(1 - x**2) * self.curv_rate_x**2
        z = np.tile(z, (self.height, 1))
        z *= np.sqrt((1 - ((self.height - y) / self.height)**2)
                    * self.curv_rate_y**2)
        z = z.reshape(-1, 1)

        # Make a copy of z for the bottom side
        z_cpy = z.copy()

        # Reshape img and add it to height values
        self.img = (self.img / 10).reshape(-1, 1)
        z += self.img

        vert_arr_tmp = np.vstack(list(map(np.ravel, self.meshgrid))).T

        # for top side
        top_vert_arr = np.concatenate((vert_arr_tmp, z), axis=1)
        top_vert_arr = top_vert_arr.reshape(self.height, self.width, 3)

        # for bottom side
        bottom_vert_arr = np.concatenate((vert_arr_tmp, z_cpy), axis=1)
        bottom_vert_arr = bottom_vert_arr.reshape(self.height, self.width, 3)

        # Engrave text on the back of the model
        self.engrave_text(bottom_vert_arr, top_vert_arr)

        # Create all vertices, faces
        faces, vertices = self.create_faces(top_vert_arr, bottom_vert_arr)

        # Add the created vertices and faces to a mesh
        self.create_stl_mesh(faces, vertices)

    def make_stl_map(self):
        '''Map fingerprint to a given finger model.
        '''

        # Conversion constants for mm and pixels
        mm2px = self.dpi/25.4
        px2mm = 25.4/self.dpi

        # Finds the image pixel closest to finger vertice in 2D plane
        def find_nearest(ver1, ver2, img):
            searched_point = np.array([ver1, ver2])

            min1 = math.floor(ver1*mm2px)
            max1 = math.ceil(ver1*mm2px)
            min2 = math.floor(ver2*mm2px)
            max2 = math.ceil(ver2*mm2px)
            min_dist_point = img[min2][min1]

            for i in range(min2, max2 - 1):
                for j in range(min1, max1 - 1):
                    if np.linalg.norm(img[i][j] - searched_point) < min_dist_point:
                        min_dist_point = img[i][j]
            return min_dist_point

        # Load the finger model
        finger = trimesh.load("res/finger-mod.stl")

        # Implicitly translate it to match middle of the fingerprint
        # Later this can be modified
        x = (self.width * px2mm / 2) + self.finger_x
        y = (self.height * px2mm / 2) + self.finger_y
        z = self.finger_z
        matrix = tmtra.translation_matrix([x, y, z])
        finger.apply_transform(matrix)

        # Subdivide the finger mesh to allow for more precision
        vertices, faces = tmrem.subdivide_loop(
            finger.vertices, finger.faces, iterations=self.iter)

        # For logging progress
        c = 0
        for k, vertice in enumerate(vertices):
            # Skip vertices under plane xy
            # also skip vertices under the fingerprint image, 
            # they are all unused
            if vertice[2] < 0 or vertice[1] > self.height * px2mm:
                continue

            # This is the easiest way to avoid indexing errors
            # Those errors are caused by vertices outside of the image
            # When this occurs, input parameters need to be adjusted
            try:
                # Find the closest point in the image
                # To the 2D image projection of vertice, add its value
                point = find_nearest(vertice[0], vertice[1], self.img)
            except IndexError:
                log.error_exit("Fingerprint image is outside of the finger model")

            vertices[k][2] += point

            # Prints out generation progress
            if k % 1000 == 0:
                percentage = round(k/len(vertices) * 100, 2)
                if percentage > c:
                    log.print_message("Carving finger: " + str(c) + "%")
                    c += 10

        self.stl_model = trimesh.Trimesh(vertices, faces)
        log.print_message("Carving finger finished")

    def save_stl(self):
        '''Save final mesh to stl file.
        '''

        # Create output file name, save it and write header with file info
        self.stl_filename = self.output_file.split(
            ".")[0] + "_" + self.id + ".stl"
        if (self.mode == "m"):
            self.stl_model.export(file_obj=self.stl_filename)
        else:
            self.stl_model.save(self.stl_filename)

        log.print_message("Saving model to", self.stl_filename)
        self.write_stl_header()


fingerprint_app()