BP_DP-xlanro00/src/main.py

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

# Import libraries
import argparse as ap
import json
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import cv2 as cv

# Import custom image filter library
import filters as flt

class apply_filters:
    def __init__(self):
        # Parse and save arguments.
        self.parse_arguments()
        self.input_file = self.args.input_file
        self.output_file = self.args.output_file
        self.dpi = self.args.dpi  # should be around 500-1000dpi
        self.filters = self.args.filters
        self.flip = self.args.flip

        # Read input image, save its dimensions
        self.img = plt.imread(self.input_file)
        self.dimensions = self.img.shape
        self.print_debug(self.dimensions)
        
        # Convert dimensions
        self.convert_dpi()
        
        # Apply all filters
        self.apply_filter()

    def parse_arguments(self):
        parser = ap.ArgumentParser(prog = 'main.py', description = 'Program for processing a 2D image into 3D fingerprint.')

        # 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")
        parser.add_argument('-f', "--flip", help="Flip input image",
                            type = bool, action = ap.BooleanOptionalAction) # boolean switch

        # file with default presets
        parser.add_argument('-pf', "--preset_file", help = "File with presets")
        parser.add_argument('filters', type = str, nargs = '*', help = "List of filter names")
        self.args = parser.parse_args()

    def convert_dpi(self):
        height = self.dimensions[0]
        width = self.dimensions[1]
        self.height = height / self.dpi * 25.4 # conversion to milimeters
        self.width = width / self.dpi * 25.4
        self.print_debug([self.height, self.width, self.dimensions[2]])

    def filter_factory(self, filter_name):
        if filter_name == "average":
            return flt.filter_average
        elif filter_name == "blur":
            return flt.filter_blur
        elif filter_name == "gaussian":
            return flt.filter_gaussian
        elif filter_name == "median":
            return flt.filter_median
        elif filter_name == "erode":
            return flt.filter_erode
        elif filter_name == "dilate":
            return flt.filter_dilate
        elif filter_name == "opening":
            return flt.filter_opening
        elif filter_name == "closing":
            return flt.filter_closing
        else:
            raise ValueError("Invalid filter name")

    def process_image(self):
        # resize img to the new width
        # flip image when mirroring is needed
        if self.flip:
            print("Flipping image")
            self.img = cv.flip(self.img, 1)

    def apply_filter(self):
        self.process_image()
        if len(self.filters) == 0:
            # save original image
            filter = flt.filter_none
            filter.apply(self)
        else:
            for filter_name in self.filters:
                filter = self.filter_factory(filter_name)
                filter.apply(self)
        self.save_image()

    def print_debug(self, dimensions):
        print("Height: " + str(dimensions[0]))
        print("Width: " + str(dimensions[1]))

    def save_image(self):
        ''' Save processed image. '''
        plt.xticks([]), plt.yticks([])
        plt.axis('off')
        plt.savefig(self.output_file)

app = apply_filters()