Added experimental 3d model mapping, added more docstrings, separated stl parsing from init method.

src/main.py

@@ -9,17 +9,24 @@ import sys
 import json
 from os.path import exists
 import hashlib
+import math
 
 # 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 custom image filter library
 import filters as flt
 
+
 class app:
+    '''Main class for the application.
+    '''
+
     def __init__(self):
         # Parse arguments from command line
         self.parse_arguments()
@@ -63,54 +70,11 @@ class app:
             self.run_filtering()
 
         else:
-
             self.error_exit("Input file " + self.input_file +
                             " does not exist")
 
-        if self.args.stl_file:
+        if self.args.stl:
-
+            self.parse_stl()
-            # 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.
@@ -118,7 +82,7 @@ class app:
 
         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]')
+                                   usage='%(prog)s [-h] [-m | --mirror | --no-mirror] [-p] 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]')
 
         # positional arguments
         parser.add_argument("input_file", type=str, help="input file path")
@@ -130,12 +94,13 @@ class app:
                             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='*',
+        # TODO: behaves absolutely randomly for some reason 
+        parser.add_argument('-s', "--stl", 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,
+        #parser.add_argument('-p', '--planar', type=bool, action=ap.BooleanOptionalAction,
-                            help="make stl shape planar instead of curved one")
+        #                    help="make stl shape planar instead of curved one")
 
         # configuration file containing presets, preset name
         # pair argument - give both or none
@@ -149,17 +114,19 @@ class app:
         self.args = parser.parse_args()
 
     def parse_params(self, params):
-        '''Parse parameters of filters. Set to None if not given.
+        '''Parse parameters of filters. Set to None if parameter is not given.
         They are later set to default values in the filter method apply.
+        :param params: dictionary of filter parameters
         '''
 
         # TODO: possibly too bloated, sending all possible params to each filter
+        # TODO: remove unnecessary params
         possible_params = {"h", "searchWindowSize", "templateWindowSize",
-                           "ksize", "kernel", "sigmaX", "sigmaY",
+                           "ksize", "kernel",
-                           "sigmaColor", "sigmaSpace", "d", "anchor", "iterations",
+                           "sigmaColor", "sigmaSpace", "diameter", "anchor", "iterations",
                            "op", "strength", "amount", "radius", "weight", "channelAxis",
-                           "theta", "sigma", "lambda", "gamma", "psi", "shape", "percent", 
+                           "theta", "sigma", "lambd", "gamma", "psi", "shape", "percent",
-                           "threshold", "maxval", "type", "margin", "color"}
+                           "threshold", "maxval", "type", "margin", "color", "truncate", "patch_size", "patch_distance"}
 
         for key in possible_params:
             if params.get(key) is None:
@@ -189,24 +156,80 @@ class app:
         else:
             self.error_exit("Preset not found")
 
+    def parse_stl(self):
+        # Get stl filename
+        self.stl_path = self.output_file.rsplit('/', 1)[0] + '/'
+
+        # Get mode and model parameters
+        if self.args.stl[0] == 'p':
+            self.mode = "planar"
+
+            # TODO: add default values for planar mode, not like this
+            if len(self.args.stl) < 3:
+                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[1])
+                self.height_base = float(self.args.stl[2])
+                print("Base height:", self.height_base,
+                    "mm, lines depth/height:", self.height_line, "mm")
+
+        elif self.args.stl[0] == 'c':
+            self.mode = "curved"
+
+            # TODO: add default values for curved mode, not like this
+            if len(self.args.stl) < 5:
+                self.height_line = 2
+                self.height_base = 10
+                self.curv_rate_x = 2
+                self.curv_rate_y = 6
+                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[1])
+                self.height_base = float(self.args.stl[2])
+                self.curv_rate_x = float(self.args.stl[3])
+                self.curv_rate_y = float(self.args.stl[4])
+                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)
+        elif self.args.stl[0] == 'm':
+            self.mode = "mapped"
+
+            # TODO: add default values for mapped mode, add finger model?
+            if len(self.args.stl) < 2:
+                print(
+                    "Warning: Too few arguments, using default values")
+            else:
+                self.height_line = float(self.args.stl[1])
+                self.finger_model = self.args.stl[2]
+
+        else:
+            self.error_exit("Unrecognized generation mode")
+
+        print("Stl generation in ", self.mode)
+        self.run_stl()
+
     def error_exit(self, message):
-        '''Print error message and exit the application.
+        '''Print given error message and exit the application.
+        :param message: error message to be printed
         '''
 
         print("ERROR:", message, file=sys.stderr)
         exit(1)
 
-#------------------------- FILTERING -------------------------#
+# ------------------------- FILTERING -------------------------#
 
     def run_filtering(self):
-        '''Load from input file, store as numpy.array, 
+        '''Read input file, store as numpy.array, uint8, grayscale.
-        process image using filters and save to output file.
+        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
+        # Gets empty figure and ax with dimensions of input image
         self.height, self.width = self.img.shape
         self.fig, ax = self.get_empty_figure()
 
@@ -219,8 +242,9 @@ class app:
         plt.close()
 
     def get_empty_figure(self):
-        '''Return empty figure with one ax of dimensions of input image.
+        '''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.])
@@ -241,10 +265,10 @@ class app:
         '''
 
         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)
+                print("Applying filter:", filter_name, file=sys.stderr)
                 filter.apply(self, self.params[i+1])
         else:
             pass
@@ -259,13 +283,14 @@ class app:
         ax.imshow(self.img, cmap="gray")
         fig.savefig(fname=self.output_file)
 
-#------------------------- STL GENERATION -------------------------#
+# ------------------------- 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()
 
         print("Creating mesh", file=sys.stderr)
@@ -278,8 +303,8 @@ class app:
             self.make_stl_curved()
 
         elif self.mode == "mapped":
-            # TODO: find a suitable finger model, try to map the fingerprint onto it
+            # TODO: find a more suitable finger model
-            pass
+            self.make_stl_map()
 
         else:
             self.error_exit("Mode not supported")
@@ -289,9 +314,10 @@ class app:
         print(f"Saving model to ", self.stl_filename, file=sys.stderr)
 
     def prepare_heightmap(self):
-        '''Modify image values to get usable height/depth values.
+        '''Scale image values to get values from 0 to 255. 
+        Then compute base and papilar lines height.
         Check validity of dimension parameters.
-        Prepare meshgrid.
+        Prepare meshgrid, array which later serves to store point coordinates.
         '''
 
         if self.img.dtype != np.uint8:
@@ -300,9 +326,6 @@ class app:
             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")
 
@@ -310,15 +333,20 @@ class app:
                 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")
 
+        if self.mode == "mapped":
+            if self.height_line <= 0:
+                self.error_exit("Line height must be positive")
+            if not exists(self.finger_model):
+                self.error_exit("Finger model file does not exist")
+            self.height_base = 0
+
+        # TODO: curved height base could be done here?
         # Transform image values to get a heightmap
-        self.img = (height_base + (1 - self.img/255)
+        self.img = (self.height_base + (1 - self.img/255)
                     * self.height_line)
 
         # This sets the size of stl model and number of subdivisions / triangles
@@ -328,11 +356,13 @@ class app:
         self.meshgrid = np.meshgrid(x, y)
 
     def write_stl_header(self):
-        '''Write stl header.
+        '''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):
+        if (len(self.param_string) >= 80):
             self.param_string = self.param_string[:80]
             print("Warning: Parameter string too long, truncating", file=sys.stderr)
 
@@ -342,7 +372,7 @@ class app:
             f.write(self.param_string)
 
     def get_ID(self):
-        '''Get unique ID for the model, used in filename and on the model backside.
+        '''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.
         '''
 
@@ -360,9 +390,10 @@ class app:
                 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.append(
+                            str(j[:1] + ":" + str(self.params[i+1][j])))
                 tmp_params = ",".join(tmp_params)
-                tmp = str(self.filters[i][0:3])
+                tmp = str(self.filters[i][0:1] + self.filters[i][-1:])
                 if tmp_params != "":
                     tmp = tmp + ";" + str(tmp_params)
                 filter_list.append(tmp)
@@ -378,27 +409,41 @@ class app:
             param_list.append(str(self.curv_rate_x))
             param_list.append(str(self.curv_rate_y))
 
+        if self.mode == "mapped":
+            #TODO
+            pass
+
         if self.mode == "planar":
             param_list.append("P")
+        elif self.mode == "curved":
+            param_list.append("C")
+        elif self.mode == "mapped":
+            param_list.append("M")
 
         if self.args.mirror:
-            param_list.append("M")
+            param_list.append("F")
 
         # string that will later be put inside the header of an stl file
-        self.param_string = "\\".join(param_list) + "\n"
+        # 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.sha512(self.param_string.encode('utf-8')).hexdigest())[:10]
+        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
+        ''' 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):
+    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.
         '''
@@ -410,11 +455,11 @@ class app:
 
         # extract filename
         text = self.stl_path.split("/")[-1].split(".")[0] + self.id
-        fontsize = 20
+        fontsize = 28
 
         # 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)
+                    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
@@ -434,14 +479,21 @@ class app:
         plt.close()
 
         # TODO: maybe don't use nested for loops, use numpy?
-        if self.mode == "planar":
+        # 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 == "planar":
                 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
+                    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
 
@@ -503,7 +555,7 @@ class app:
 
         # 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)
+        self.engrave_text(bottom_vert_arr, top_vert_arr)
 
         # Back side faces
         for i in range(self.height - 1):
@@ -568,10 +620,12 @@ class app:
         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[y] = ((z[y] * new) + (z[y] + new))/2
 
         z = z.reshape(-1, 1)
+        # make a copy of z for the bottom side
         z_cpy = np.copy(z)
+        # reshape img and add it to z
         self.img = (self.img / 10).reshape(-1, 1)
         z += self.img
 
@@ -589,7 +643,7 @@ class app:
         vertices = []
         faces = []
 
-        self.engrave_text(bottom_vert_arr)
+        self.engrave_text(bottom_vert_arr, top_vert_arr)
 
         # TODO: code bellow is duplicate of the code in planar generation
         # if not changed move to a separate function and simplify
@@ -660,15 +714,61 @@ class app:
 
         self.create_stl_mesh(faces, vertices)
 
+    def make_stl_map(self):
+        '''Map fingerprint to a given finger model.
+
+        Experimental, does not work very well...
+        '''
+
+        # TODO: maybe use trimesh.update_vertices
+
+
+        print("Mapping to finger")
+        # TODO: try to merge meshes? or stl files?
+        # trimesh library?
+        finger = trimesh.load(self.finger_model)
+        # TODO: connect with curved generation
+ 
+        # manually tried to allign two models and concatenated
+        fingerprint = trimesh.load('res/0-norm1_e5f52c0fe1.stl')
+
+        angle = math.pi
+        dir = [0, 0, 1]
+        center = [0, 0, 0]
+        mat = tmtra.rotation_matrix(angle, dir, center)
+        finger.apply_transform(mat)
+
+        angle = -3 / 4 * math.pi
+        dir = [1, 0, 0]
+        center = [0, 0, 0]
+        mat = tmtra.rotation_matrix(angle, dir, center)
+        finger.apply_transform(mat)
+
+
+        # TODO: random values that works for one finger model...
+        # TODO: this can later be modified to map finger to the core of the finger.
+        x = 2 + (self.width * 25.4 / self.dpi / 2)
+        y = 5 + (self.height * 25.4 / self.dpi / 2)
+        z = 20
+        mat = tmtra.translation_matrix(
+            [x, y, z])
+        finger.apply_transform(mat)
+
+        self.stl_model = trimesh.util.concatenate([finger, fingerprint])
+
+
     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_filename = self.output_file.split(
+            ".")[0] + "_" + self.id + ".stl"
+        if (self.mode == "mapped"):
+            self.stl_model.export(file_obj=self.stl_filename)
+        else:
             self.stl_model.save(self.stl_filename)
         self.write_stl_header()
 
 
-# run the application
+app()
-image = app()