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

# Import custom image filter library
import filters as flt

class apply_filters:
    def __init__(self):

        if len(sys.argv) < 4:
            # Parse configuration from json file
            self.conf_fn = sys.argv[1]
            self.preset_name = sys.argv[2]
            try:
                self.conf_file = open(self.conf_fn)
            except(FileNotFoundError):
                print("Creating new file")
                self.conf_file = open(self.conf_fn, "x")
            self.conf = json.load(self.conf_file)
            
            self.parse_conf()
        else:
            # Parse arguments from command line
            self.parse_arguments()
            self.input_file = self.args.input_file
            self.output_file = self.args.output_file
            self.dpi = self.args.dpi
            self.filters = self.args.filters
            self.flip = self.args.flip
            self.width = self.args.width if self.args.width else 0

            # Save preset to config file
            if self.args.config_file:

                self.conf_file = self.args.config_file

                if self.args.preset:
                    self.preset_name = self.args.preset
                else:
                    # If no preset name given, create one from time
                    self.preset_name = "preset_" + datetime.now().strftime("%d_%m_%Y_%H_%M_%S")
            self.write_conf()

        # Convert dimensions
        self.img = Image.open(self.input_file)
        self.convert_dpi()
        self.resize_image()
        self.flip_image()

        # Apply all filters
        self.apply_filter()


    def parse_conf(self):

        # Parse configuration file if given.
        preset = self.conf[self.preset_name]

        self.input_file = preset['inputFile']
        self.output_file = preset['outputFile']
        self.dpi = preset['dpi']
        self.filters = preset['filters']
        self.flip = preset['flip']

        # Width in config file should be set to 0 even if not used, 
        # checking nonetheless
        try:
            self.width = preset['width']
        except(KeyError):
            self.width = 0


    def write_conf(self):

        temp = {}
        temp_fp = open(self.conf_file, 'a+')
        # If config file new or empty, loading fails
        try:
            temp = json.load(temp_fp)
        except(json.decoder.JSONDecodeError):
            temp = {}
        
        # Create new preset
        temp[self.preset_name] = ({
            "inputFile": self.input_file,
            "outputFile": self.output_file,
            "dpi": self.dpi,
            "flip": self.flip,
            "width": self.width,
            "filters": self.filters
            })
        
        # Update json file with new preset
        with open(self.conf_file, 'w') as temp_fp:
            json.dump(temp, temp_fp, indent=4)


    def parse_arguments(self):

        # Parse arguments
        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")

        # predefined width
        parser.add_argument('-w', "--width", type= int, help = "Option to input predefined width")

        # boolean switch
        parser.add_argument('-f', "--flip", help="Flip input image",
                            type=bool, action=ap.BooleanOptionalAction)

        # file with configuration containing presets
        parser.add_argument('-c', "--config_file", help = "Config file to save the preset into")

        # new preset name
        parser.add_argument('-p', "--preset", help = "Name of the newly created preset")

        # array of unknown length, all filter names saved inside
        parser.add_argument('filters', type = str, nargs = '*', help = "List of filter names")

        self.args = parser.parse_args()


    def filter_factory(self, filter_name):
        # selects filter method of filters library
        # better this than a 100 if/else
        return getattr(flt, filter_name)


    def convert_dpi(self):

        # conversion from inches to milimeters
        self.size = np.empty(2)
        self.size[0] = self.img.size[0]  # / self.dpi * 25.4 # width
        self.size[1] = self.img.size[1]  # / self.dpi * 25.4 # height


    def resize_image(self):

        # open image as python image object
        print("Resize image", file = sys.stderr)

        # resize if target width passed as an argument
        if self.width:
            self.size[0] = int(
                self.width * self.img.size[0] / self.img.size[1])
            self.size[1] = self.width
            print(self.size)
            self.img = self.img.resize(
                (np.array(self.size).astype(int)))
        else:
            self.convert_dpi()
            #self.img = self.img.resize((np.array(self.size)).astype(int))

        #convert to numpy array for further processing
        self.img = np.array(self.img)


    def flip_image(self):

        # flip image when mirroring is needed
        # should be used only if we want a positive form
        
        if self.flip:
            print("Flipping image", file = sys.stderr)
            self.img = cv.flip(self.img, 1) # 1 for vertical flip


    def apply_filter(self):
        
        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_size(self, size):
        print("Width: " + str(size[0]), file = sys.stderr)
        print("Height: " + str(size[1]), file = sys.stderr)


    def save_image(self):
        # Save processed image.
        plt.xticks([]), plt.yticks([])
        plt.axis('off')
        print("Saving image", file = sys.stderr)

        # TODO idk what dpi means, and if it should be put in here
        plt.savefig(self.output_file)#, dpi=self.dpi)  


app = apply_filters()