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 datetime import datetime

# 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 apply_filters:
    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))
            print("Config loaded")
            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] = {}
                else:
                    # else it's a parameter for current filter
                    key, value = filter.split('=')
                    self.params[i][key] = value
                self.parse_params(self.params[i])

        self.input_file = self.args.input_file
        self.output_file = self.args.output_file
        self.dpi = self.args.dpi
        self.run()

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

        self.width = self.img.shape[1]
        self.height = self.img.shape[0]
        print(self.width, self.height)

        fig = plt.figure(figsize=(self.width, self.height),
                         frameon=False, dpi=self.dpi / 100)  # dpi is in cm

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

        if self.args.mirror:
            self.mirror_image()

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


    def parse_params(self, params):
        possible_params = {"h", "searchWindowSize", "templateWindowSize",
                           "ksize", "kernel", "sigmaX", "sigmaY",
                           "sigmaColor", "sigmaSpace", "d", "anchor", "iterations", 
                           "op", "strength"}
        for key in possible_params:
            try:
                params[key] = params[key]
            except KeyError:
                params[key] = None

    def parse_conf(self):
        # Parse configuration file if given.
        try:
            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])

        except(KeyError):
            print("Preset not found", file=sys.stderr)

    def parse_arguments(self):

        # Parse arguments
        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 ...])')

        # 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
        parser.add_argument('-m', "--mirror", help="mirror input image",
                            type=bool, action=ap.BooleanOptionalAction)

        parser.add_argument('-s', '--stl', help="make stl model from processed image",
                            type=bool, action=ap.BooleanOptionalAction)

        # file with configuration containing presets, new preset name
        # pair argument - give both or none
        parser.add_argument('-c', '--config', nargs=2, metavar=('config_file', 'preset'),
                            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", file=sys.stderr)
        return getattr(flt, filter_name)

    def resize_image(self):

        print("Resize image", file=sys.stderr)
        self.img = self.img.resize(
            (np.array(self.width, self.height)).astype(int))

    def mirror_image(self):
        ''' Mirror image when mirroring is needed,
        should be used only if we want a positive model
         '''

        #TODO make this automatic for positive STL
        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.
        
        Applies the filters one by one, if no filters 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 print_size(self, size):
        print("Width: " + str(size[0]), file=sys.stderr)
        print("Height: " + str(size[1]), file=sys.stderr)

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

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

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

        print("Making meshgrid", file=sys.stderr)
        self.make_meshgrid()
        print("Converting to stl format", file=sys.stderr)
        self.make_mesh()
        plt.show()
        self.save_model()

    def make_meshgrid(self):

        # Modify image to make it more suitable depth
        # values1 = (1 + (1 - self.img/255)/6) * 255/10 # this works
        # values2 = (1 - (1 - self.img/255)/6) * 255/10 # TODO: i dont know how to make white surrounding be extruded
        values1better = 28.05 - 0.01*self.img
        #values2better = 22.95 - 0.01*self.img
        # (np.around(values2[::300],3))

        # Add zero padding to image to make sides of the plate
        self.height = self.img.shape[0] + 2
        self.width = self.img.shape[1] + 2
        self.img = np.zeros([self.height, self.width])
        self.img[1:-1:1, 1:-1:1] = values1better

        # Create meshgrid for 3D model
        verticesX = np.around(np.linspace(0, self.width / 10, self.width), 3)
        verticesY = np.around(np.linspace(0, self.height / 10, self.height), 3)
        self.meshgrid = np.meshgrid(verticesX, verticesY)

    def make_mesh(self):
        # Convert meshgrid and image matrix to array of 3D points
        vertice_arr = np.vstack(list(map(np.ravel, self.meshgrid))).T
        z = (self.img / 10).reshape(-1, 1)
        vertice_arr = np.concatenate((vertice_arr, z), axis=1)

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

        count = 0
        vertices = []
        faces = []

        # Function to add faces to the list
        def add_faces(c):
            faces.append([c, c + 1, c + 2])
            faces.append([c + 1, c + 3, c + 2])
            c += 4
            return c

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

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

                count = add_faces(count)

        # Add faces for the backside of the lithophane
        # This makes it closed, so it can be printed
        vertices.append([vertice_arr[0][0]])
        vertices.append([vertice_arr[0][self.width - 1]])
        vertices.append([vertice_arr[self.height - 1][0]])
        vertices.append([vertice_arr[self.height - 1][self.width - 1]])

        count = add_faces(count)

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

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

    def save_model(self):
        print("Saving stl model", file=sys.stderr)
        self.model.save('res/test.stl')


image = apply_filters()