BP_DP-xlanro00/src/main.py

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

# Import basic libraries
import argparse as ap
import sys
import json
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 custom image filter library
import filters as flt

class app:
    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
            self.config = json.load(open(self.config_file))
            self.parse_conf()

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

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

        else:
            print("No filters given, saving original image")

        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:

            self.error_exit("Input file " + self.input_file +
                            " does not exist")

        if self.args.stl_file:

            # Get stl filename
            self.stl_path = self.output_file.rsplit('/', 1)[0] + '/'

            # Get mode and model parameters
            if self.args.planar:
                self.mode = "planar"
                
                # TODO: add default values for planar mode, not like this
                if len(self.args.stl_file) < 2:
                    self.height_line = 2
                    self.height_base = 10
                    print(
                        "Warning: Too few arguments, using default values (10mm base, 2mm lines)")
                else:
                    self.height_line = float(self.args.stl_file[0])
                    self.height_base = float(self.args.stl_file[1])
                    print("Base height:", self.height_base,
                          "mm, lines depth/height:", self.height_line, "mm")

            else:
                self.mode = "curved"

                # TODO: add default values for curved mode, not like this
                if len(self.args.stl_file) < 4:
                    self.height_line = 2
                    self.height_base = 10
                    self.curv_rate_x = 0.5
                    self.curv_rate_y = 0.5
                    print(
                        "Warning: Too few arguments, using default values (2mm lines, curvature 0.5 on x, 0.5 on y)")
                else:
                    self.height_line = float(self.args.stl_file[0])
                    self.height_base = float(self.args.stl_file[1])
                    self.curv_rate_x = float(
                        self.args.stl_file[2])  # finger depth
                    self.curv_rate_y = float(
                        self.args.stl_file[3])  # finger depth
                    print("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)

            print("Stl generation in ", self.mode)
            self.run_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] [-p] input_file output_file dpi ([-c | --config config_file preset] | [filters ...]) [-s | --stl_file height_line height_base | --stl_file height_line curv_rate_x curv_rate_y]')

        # 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_file', type=str, nargs='*',
                            help="create stl model from processed image")

        # another boolean switch argument, this enables planar mode
        parser.add_argument('-p', '--planar', type=bool, action=ap.BooleanOptionalAction,
                            help="make stl shape planar instead of curved one")

        # 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...]")

        self.args = parser.parse_args()

    def parse_params(self, params):
        '''Parse parameters of filters. Set to None if not given.
        They are later set to default values in the filter method apply.
        '''

        # TODO: possibly too bloated, sending all possible params to each filter
        possible_params = {"h", "searchWindowSize", "templateWindowSize",
                           "ksize", "kernel", "sigmaX", "sigmaY",
                           "sigmaColor", "sigmaSpace", "d", "anchor", "iterations",
                           "op", "strength", "amount", "radius", "weight", "channelAxis",
                           "theta", "sigma", "lambda", "gamma", "psi", "shape", "percent", 
                           "threshold", "maxval", "type", "margin", "color"}

        for key in possible_params:
            if params.get(key) is None:
                params[key] = None
            else:
                params[key] = params[key]

    def parse_conf(self):
        '''Parse configuration file if one was given. 
        Store filters and their parameters.
        '''

        # Find preset in config file
        if self.preset_name in self.config:
            filter_array = self.config[self.preset_name]
            # Iterate over filters in preset, store them and their parameters
            for i, filter in enumerate(range(len(filter_array)), start=1):
                self.filters.append(filter_array[filter]["name"])
                self.params[i] = {}
                for attribute, value in filter_array[filter].items():
                    # Filter name isn't needed in here
                    if attribute != "name":
                        self.params[i][attribute] = value
                self.parse_params(self.params[i])
            print("Loaded preset: " + self.preset_name +
                  " from file: " + self.config_file)
        else:
            self.error_exit("Preset not found")

    def error_exit(self, message):
        '''Print error message and exit the application.
        '''

        print("ERROR:", message, file=sys.stderr)
        exit(1)

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

    def run_filtering(self):
        '''Load from input file, store as numpy.array, 
        process image using filters and save 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 of dimensions of 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.
        '''

        print("Mirroring image", file=sys.stderr)
        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:
            # Apply all filters
            for i, filter_name in enumerate(self.filters):
                # Get filter class from filter.py, use the apply method
                filter = getattr(flt, filter_name)
                filter.apply(self, self.params[i+1])
        else:
            pass

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

        print("Saving image to", self.output_file, file=sys.stderr)

        # 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()
        self.get_ID()

        print("Creating mesh", file=sys.stderr)

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

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

        elif self.mode == "mapped":
            # TODO: find a suitable finger model, try to map the fingerprint onto it
            pass

        else:
            self.error_exit("Mode not supported")

        plt.show()
        self.save_stl()
        print(f"Saving model to ", self.stl_filename, file=sys.stderr)

    def prepare_heightmap(self):
        '''Modify image values to get usable height/depth values.
        Check validity of dimension parameters.
        Prepare meshgrid.
        '''

        if self.img.dtype != np.uint8:
            print("Converting to uint8", file=sys.stderr)
            self.img = self.img / np.max(self.img) * 255
            self.img = self.img.astype(np.uint8)

        if self.mode == "planar":
            # just renamed it for easier use
            height_base = self.height_base

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

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

        if self.mode == "curved":
            # still need this value later
            height_base = 0

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

        # Transform image values to get a heightmap
        self.img = (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 stl header.
        '''
        
        # Truncate if necessary
        if (len(self.param_string) > 80):
            self.param_string = self.param_string[:80]
            print("Warning: Parameter string too long, truncating", file=sys.stderr)
            
        # Overwrite stl header (which is only 80 bytes)
        print("Writing info to stl header", file=sys.stderr)
        with open(self.stl_filename, "r+") as f:
            f.write(self.param_string)

    def get_ID(self):
        '''Get unique ID for the model, 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[:3] + ":" + str(self.params[i+1][j])))
                tmp_params = ",".join(tmp_params)
                tmp = str(self.filters[i][0:3])
                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 == "curved":
            param_list.append(str(self.curv_rate_x))
            param_list.append(str(self.curv_rate_y))

        if self.mode == "planar":
            param_list.append("P")
        
        if self.args.mirror:
            param_list.append("M")

        # string that will later be put inside the header of an stl file
        self.param_string = "\\".join(param_list) + "\n"
        
        # 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.sha512(self.param_string.encode('utf-8')).hexdigest())[:10]

    def append_faces(self, faces, c):
        # Function to add faces to the list
        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):
        '''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 = 20

        # create text object, paint it white
        t = ax.text(0.5, 0.5, text, ha="center", va="center",
                    fontsize=30, 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: maybe don't use nested for loops, use numpy?
        if self.mode == "planar":
            for i in range(self.height):
                for j in range(self.width):
                    bottom_vert_arr[i][j][2] = data[i][j][0]
        elif self.mode == "curved":
            for i in range(self.height):
                for j in range(self.width):
                    bottom_vert_arr[i][j][2] += data[i][j][0] - self.height_base/10
            
        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)

        # 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], :]

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

        From wikipedia.org/wiki/STL_(file_format):
        ascii stl format consists of repeating structures:

        facet normal ni nj nk # normal vector
            outer loop
                vertex v1x v1y v1z # vertex 1
                vertex v2x v2y v2z # vertex 2
                vertex v3x v3y v3z # vertex 3
            endloop
        endfacet
        '''
        
        # Add the image matrix to the 2D meshgrid and create 1D array of 3D points
        top_vert_arr = np.vstack(list(map(np.ravel, self.meshgrid))).T
        z = (self.img / 10).reshape(-1, 1)
        top_vert_arr = np.concatenate((top_vert_arr, z), axis=1)

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

        count = 0
        vertices = []
        faces = []

        # TODO: don't like this, could be done using numpy vectorisation?
        # Iterate over all vertices, create 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)

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

        # 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)

        self.create_stl_mesh(faces, vertices)
       
    def make_stl_curved(self):
        '''Map fingerprint to finger model.
        '''

        # TODO: this might be done in a better way
        # instead of summing up the values, use their product - 0 ?
        z = np.array([])
        for x in range(self.width):
            z = np.append(z, np.sqrt(1 - (2*x/self.width - 1)**2)
                            * (self.curv_rate_x**2))
        z = np.tile(z, (self.height, 1))
        for y in range(self.height):
            new = np.sqrt((1 - ((self.height - y)/self.height)**2)
                          * (self.curv_rate_y**2))
            z[y] = z[y] + new

        z = z.reshape(-1, 1)
        z_cpy = np.copy(z)
        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)

        count = 0
        vertices = []
        faces = []

        self.engrave_text(bottom_vert_arr)

        # TODO: code bellow is duplicate of the code in planar generation
        # if not changed move to a separate function and simplify

        # Iterate over all vertices, create faces
        for i in range(self.height - 1):
            for j in range(self.width - 1):
                if (top_vert_arr[i][j][2] <= bottom_vert_arr[i][j][2]
                    or top_vert_arr[i+1][j][2] <= bottom_vert_arr[i+1][j][2]
                    or top_vert_arr[i][j+1][2] <= bottom_vert_arr[i][j+1][2]
                    or top_vert_arr[i+1][j+1][2] <= bottom_vert_arr[i+1][j+1][2]):
                    continue
                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)

        # Rotated back side faces
        for i in range(self.height - 1):
            for j in range(self.width - 1):
                if (top_vert_arr[i][j][2] <= bottom_vert_arr[i][j][2]):
                    continue

                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): # right

            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)

        self.create_stl_mesh(faces, vertices)

    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"
        self.stl_model.save(self.stl_filename)
        self.write_stl_header()


# run the application
image = app()