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
import math
from os.path import exists

# 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_file = self.args.stl_file[0]

            # Get mode and model parameters
            if self.args.planar:
                self.mode = "2d"

                if len(self.args.stl_file) < 3:
                    self.height_base = 10
                    self.height_line = 2
                    print(
                        "Warning: Too few arguments, using default values (10mm base, 2mm lines)")
                else:
                    self.height_line = float(self.args.stl_file[1])
                    self.height_base = float(self.args.stl_file[2])
                    print("Base height:", self.height_base,
                          "mm, lines depth/height:", self.height_line, "mm")

            else:
                self.mode = "3d"

                if len(self.args.stl_file) < 5:
                    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[1])
                    self.height_base = float(self.args.stl_file[2])
                    self.curv_rate_x = float(
                        self.args.stl_file[3])  # finger depth
                    self.curv_rate_y = float(
                        self.args.stl_file[4])  # 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(self.mode, "mode selected")
            self.run_stl()

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

        parser = ap.ArgumentParser(prog='main.py',
                                   description='Program for processing a 2D image into 3D fingerprint.',
                                   usage='%(prog)s [-h] [-m | --mirror | --no-mirror] [-p] input_file output_file dpi ([-c config_file preset | --config config_file preset] | [filters ...]) [-s stl_file | --stl stl_file height_line height_base | --stl_file 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 2d 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"}

        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.
        '''

        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)

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

        print("Height: " + str(self.height) + " px and width: "
              + str(self.width) + " px", file=sys.stderr)
        
        if self.mirror is True:
            self.mirror_image()

        # Apply all filters and save image
        self.apply_filters()
        self.save_image(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, dpi=self.dpi)

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

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

        self.prepare_heightmap()
        self.get_ID() 

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

        elif self.mode == "3d":
            self.make_stl_curved()

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

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

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

        # TODO: redo, too complicated, add extra params, redo checks

        if self.img.dtype == np.float32 or self.img.dtype == np.float64:
            print("Converting to uint8", file=sys.stderr)
            self.img = self.img * 255
            self.img = self.img.astype(np.uint8)
            print("Creating mesh", file=sys.stderr)

        if self.mode == "2d":
            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")

            # Transform image values to get a heightmap
            self.img = (self.height_base + (1 - self.img/255)
                        * self.height_line)

        if self.mode == "3d":
            # TODO check curvature values and print info
            # TODO: copy pasta code, remove
            # Transform image values to get a heightmap
            self.img = (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 get_ID(self):
        '''Get unique ID for the model.
        Consists of pair input_file + preset_name.
        '''
        # TODO: somehow compress this to fit it onto the model, maybe zlib
        self.id = self.input_file.split(
            "/")[-1].split(".")[0] + "_" + self.preset_name
        print(self.id)

    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_file.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 == "2d":
            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 == "3d":
            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
            
        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.
        '''

        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.
        '''
        # TODO: add a hash function to create filename specific to input image and preset
        self.stl_file = self.stl_file.split(".")[0] + "_" + self.id + "." + self.stl_file.split(".")[1]
        self.stl_model.save(self.stl_file)


# run the application
image = app()