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

# Libraries for image processing
import numpy as np
import matplotlib.pyplot as plt
#from PIL import Image
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()
        self.params = {}
        self.filters = []
        
        # Parse configuration from json file
        if self.args.config:
            self.config_file = self.args.config[0]
            self.preset_name = self.args.config[1]
            self.config = json.load(open(self.config_file))
            self.parse_conf()

        # If no config file given, expect filters in command line
        else:
            if not self.args.filters:
                print("No filters given, saving original image")

            print("No config file given, using command line arguments")
            i = 0

            for filter in self.args.filters:
                if filter.find('=') == -1:
                    # if no '=' in filter, it is a new filter
                    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])

        if self.args.stl_file and len(self.args.stl_file) == 3:
            self.stl_file = self.args.stl_file[0]
            self.height_line = float(self.args.stl_file[1])
            self.height_base = float(self.args.stl_file[2])
            self.mode = "2d"

        elif self.args.stl_file and len(self.args.stl_file) == 2:
            self.stl_file = self.args.stl_file[0]
            self.height_line = float(self.args.stl_file[1])
            self.mode = "3d"
    
        else:
            print("No STL file given, saving image only")
            exit(1)

        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()
        else:
            print("Input file does not exist", file=sys.stderr)
            exit(1)

    def run(self):
        # read as numpy.array
        self.img = cv.imread(
            self.input_file, cv.IMREAD_GRAYSCALE).astype(np.uint8)

        self.width = self.img.shape[1]
        self.height = self.img.shape[0]
        self.print_size(self.img.shape)
        fig = plt.figure(figsize=(self.width/self.dpi, self.height/self.dpi),
                         frameon=False, dpi=self.dpi)

        ax = plt.Axes(fig, [0., 0., 1., 1.])
        ax.set_axis_off()
        fig.add_axes(ax)

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

        # Apply all filters and save image
        self.apply_filter()
        self.save_image(fig, ax)
        plt.close()
        if self.args.stl_file:
            self.make_model()

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

        possible_params = {"h", "searchWindowSize", "templateWindowSize",
                           "ksize", "kernel", "sigmaX", "sigmaY",
                           "sigmaColor", "sigmaSpace", "d", "anchor", "iterations",
                           "op", "strength", "amount", "radius", "weight", "channelAxis"}
        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 and store filters with their parameters
        '''

        if self.preset_name in self.config:
            filter_array = self.config[self.preset_name]
            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():
                    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:
            print("Preset not found", file=sys.stderr)

    def parse_arguments(self):
        ''' Parse arguments from command line
        '''

        parser = ap.ArgumentParser(prog='main.py',
                                   description='Program for processing a 2D image into 3D fingerprint.',
                                   usage='%(prog)s [-h] [-m | --mirror | --no-mirror] input_file output_file dpi \
                                   ([-c config_file preset | --config config_file preset] | [filters ...]) \
                                   [-s stl_file | --stl_file stl_file depth_total depth_line]')

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

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

        # 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="make planar model from processed image", required=False)

        # file with configuration containing presets, new 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=value1 param2=value2 filter_name2 param1=value1...]")

        self.args = parser.parse_args()

    def filter_factory(self, filter_name):
        ''' Selects filter method of filters library.
        '''

        print("Applying " + filter_name + " filter ", end='')
        return getattr(flt, filter_name)

    def mirror_image(self):
        ''' Mirror image when mirroring is needed,
        should be used only 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_filter(self):
        ''' Apply filters to image.

        Apply the filters one by one, if none were given, just save original image output.
        '''

        if len(self.filters) == 0:
            # No filter given, just save the image
            pass
        else:
            # Apply all filters
            for i, filter_name in enumerate(self.filters):
                filter = self.filter_factory(filter_name)
                filter.apply(self, self.params[i+1])

    def save_image(self, fig, ax):
        ''' Save processed image.
        Colormap set to grayscale to avoid color mismatch.
        '''

        print("Saving image to ", self.output_file, file=sys.stderr)
        ax.imshow(self.img, cmap="gray")
        fig.savefig(fname=self.output_file, dpi='figure')

    def print_size(self, size):
        print("Image of height: " + str(size[0]) + 
            " px and width: " + str(size[1]) + " px", file=sys.stderr)

    def make_model(self):
        '''After processing image, make a lithophane from it.
        '''

        print("Making heighthmap", file=sys.stderr)
        self.prepare_heightmap()

        if self.mode == "2d":
            print("Converting to stl format", file=sys.stderr)
            self.make_stl_planar()
            plt.show()
            print(f"Saving lithophane to ", self.stl_file, file=sys.stderr)
            self.save_stl_2d()
        elif self.mode == "3d":
            self.map_image_to_3d()
            plt.show()
            self.save_stl_3d()
        else:
            print("Mode not supported", file=sys.stderr)
            exit(1)      

    def prepare_heightmap(self):
        ''' Create numpy meshgrid.
        Modify image values to get usable depth values.
        '''

        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)

        if self.mode == "2d":
            if self.height_base <= 0:
                print("Depth of plate height must be positive", file=sys.stderr)
                exit(1)

            if self.height_line + self.height_base <= 0:
                print("Line depth must be less than plate thickness", file=sys.stderr)
                exit(1)
        
            print("Base height:", self.height_base,
                "mm, lines depth/height:", self.height_line, "mm")

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

        if self.mode == "3d":
            #TODO add some checks and print info
            pass

    def add_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 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
        '''

        # 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_2d = np.meshgrid(x, y)

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

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

        count = 0
        vertices = []
        faces = []

        # Iterate over all vertices, create faces
        for i in range(self.height - 1):
            for j in range(self.width - 1):

                vertices.append([vertex_arr[i][j]])
                vertices.append([vertex_arr[i][j+1]])
                vertices.append([vertex_arr[i+1][j]])
                vertices.append([vertex_arr[i+1][j+1]])

                count = self.add_faces(faces, count)

        # Add faces for the backside of the lithophane
        null_arr = np.copy(vertex_arr)
        for i in range(self.height):
            for j in range(self.width):
                null_arr[i][j][2] = 0
    
        # Back side faces
        for i in range(self.height - 1):
            for j in range(self.width - 1): 

                vertices.append([null_arr[i][j]])
                vertices.append([null_arr[i+1][j]])
                vertices.append([null_arr[i][j+1]])
                vertices.append([null_arr[i+1][j+1]])

                count = self.add_faces(faces, count)

        # Horizontal side faces
        for j in range(self.height - 1):
            vertices.append([vertex_arr[j][0]])
            vertices.append([vertex_arr[j+1][0]])
            vertices.append([null_arr[j][0]])
            vertices.append([null_arr[j+1][0]])

            count = self.add_faces(faces, count)

            max = self.width - 1

            vertices.append([vertex_arr[j+1][max]])
            vertices.append([vertex_arr[j][max]])
            vertices.append([null_arr[j+1][max]])
            vertices.append([null_arr[j][max]])

            count = self.add_faces(faces, count)

        # Vertical side faces
        for j in range(self.width - 1):
            vertices.append([vertex_arr[0][j+1]])
            vertices.append([vertex_arr[0][j]])
            vertices.append([null_arr[0][j+1]])
            vertices.append([null_arr[0][j]])

            count = self.add_faces(faces, count)

            max = self.height - 1

            vertices.append([vertex_arr[max][j]])
            vertices.append([vertex_arr[max][j+1]])
            vertices.append([null_arr[max][j]])
            vertices.append([null_arr[max][j+1]])

            count = self.add_faces(faces, count)

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

        # Create the mesh - vertices.shape (no_faces, 3, 3)
        self.stl_mesh_2d = mesh.Mesh(np.zeros(faces.shape[0], dtype=mesh.Mesh.dtype))
        for i, face in enumerate(faces):
            for j in range(3):
                self.stl_mesh_2d.vectors[i][j] = vertices[face[j], :]

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

        self.stl_mesh_2d.save(self.stl_file)

    def map_image_to_3d(self):
        ''' Map fingerprint to finger model.
        '''

        x = np.linspace(0, self.width * 25.4 / self.dpi, self.width)
        y = np.linspace(0, self.height * 25.4 / self.dpi, self.height)

        z1 = np.logspace(0, 10, int(np.ceil(self.width / 2)), base=0.7)
        z2 = np.logspace(10, 0, int(np.floor(self.width / 2)), base=0.7)
        ztemp = 5*np.concatenate((z1, z2))

        z = np.array([])
        for i in range(self.height):
            z = np.concatenate((z, ztemp * pow(np.log(i+2), -1)))
        z = z.reshape(-1, 1)

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

        vertex_arr = np.vstack(list(map(np.ravel, self.meshgrid_3d))).T
        vertex_arr = np.concatenate((vertex_arr, z), axis=1)
        vertex_arr = vertex_arr.reshape(self.height, self.width, 3)

        count = 0
        vertices = []
        faces = []

        # Iterate over all vertices, create faces
        for i in range(self.height - 1):
            for j in range(self.width - 1):

                vertices.append([vertex_arr[i][j]])
                vertices.append([vertex_arr[i][j+1]])
                vertices.append([vertex_arr[i+1][j]])
                vertices.append([vertex_arr[i+1][j+1]])

                count = self.add_faces(faces, count)

        #self.finger_base = mesh.Mesh(np.zeros(, dtype=mesh.Mesh.dtype))

        # linear projection
            # extrude lines in 1 direction 
        # cylinder / circular projection
            # extrude lines in direction of a suitable cylinder
        # normal projection
            # extrude lines in the direction of normals of given finger model

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

        # Create the mesh - vertices.shape (no_faces, 3, 3)
        self.mesh_finger = mesh.Mesh(np.zeros(faces.shape[0], dtype=mesh.Mesh.dtype))
        for i, face in enumerate(faces):
            for j in range(3):
                self.mesh_finger.vectors[i][j] = vertices[face[j], :]

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

        self.mesh_finger.save(self.stl_file)

image = app()