xlanro00
/
BP_DP-xlanro00
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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

Browse Source
master
Rostislav Lán 2 years ago
parent a95e53a69f
commit 239c675524
1 changed files with 211 additions and 111 deletions
Show Stats Download Patch File Download Diff File

322
src/main.py
Unescape View File

@ -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,62 +70,19 @@ class app:
self.run_filtering()
else:
self.error_exit("Input file " + self.input_file +
" does not exist")
if self.args.stl_file:
# 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()
if self.args.stl:
self.parse_stl()
def parse_arguments(self):
'''Parse arguments from command line using argparse library.
'''
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,
help="make stl shape planar instead of curved one")
#parser.add_argument('-p', '--planar', type=bool, action=ap.BooleanOptionalAction,
# 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",
"sigmaColor", "sigmaSpace", "d", "anchor", "iterations",
"ksize", "kernel",
"sigmaColor", "sigmaSpace", "diameter", "anchor", "iterations",
"op", "strength", "amount", "radius", "weight", "channelAxis",
"theta", "sigma", "lambda", "gamma", "psi", "shape", "percent",
"threshold", "maxval", "type", "margin", "color"}
"theta", "sigma", "lambd", "gamma", "psi", "shape", "percent",
"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,
process image using filters and save to output file.
'''Read input file, store as numpy.array, uint8, grayscale.
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
pass
# TODO: find a more suitable finger model
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,21 +356,23 @@ 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)
# Overwrite stl header (which is only 80 bytes)
print("Writing info to stl header", file=sys.stderr)
with open(self.stl_filename, "r+") as f:
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,31 +409,45 @@ 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")
if self.args.mirror:
elif self.mode == "curved":
param_list.append("C")
elif self.mode == "mapped":
param_list.append("M")
if self.args.mirror:
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.
'''
fig, ax = self.get_empty_figure()
# paint the background black
@ -410,12 +455,12 @@ 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
rend = fig.canvas.get_renderer()
@ -434,17 +479,24 @@ class app:
plt.close()
# TODO: maybe don't use nested for loops, use numpy?
if self.mode == "planar":
for i in range(self.height):
# 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):
elif self.mode == "curved":
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
def create_stl_mesh(self, faces, vertices):
'''Create mesh from faces and vertices.
'''
@ -476,7 +528,7 @@ class app:
endloop
endfacet
'''
# Add the image matrix to the 2D meshgrid and create 1D array of 3D points
top_vert_arr = np.vstack(list(map(np.ravel, self.meshgrid))).T
z = (self.img / 10).reshape(-1, 1)
@ -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):
@ -553,7 +605,7 @@ class app:
count = self.append_faces(faces, count)
self.create_stl_mesh(faces, vertices)
def make_stl_curved(self):
'''Map fingerprint to finger model.
'''
@ -563,15 +615,17 @@ class app:
z = np.array([])
for x in range(self.width):
z = np.append(z, np.sqrt(1 - (2*x/self.width - 1)**2)
* (self.curv_rate_x**2))
* (self.curv_rate_x**2))
z = np.tile(z, (self.height, 1))
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
@ -600,7 +654,7 @@ class app:
if (top_vert_arr[i][j][2] <= bottom_vert_arr[i][j][2]
or top_vert_arr[i+1][j][2] <= bottom_vert_arr[i+1][j][2]
or top_vert_arr[i][j+1][2] <= bottom_vert_arr[i][j+1][2]
or top_vert_arr[i+1][j+1][2] <= bottom_vert_arr[i+1][j+1][2]):
or top_vert_arr[i+1][j+1][2] <= bottom_vert_arr[i+1][j+1][2]):
continue
vertices.append([top_vert_arr[i][j]])
vertices.append([top_vert_arr[i][j+1]])
@ -623,7 +677,7 @@ class app:
count = self.append_faces(faces, count)
# Horizontal side faces
for i in range(self.height - 1): # right
for i in range(self.height - 1): # right
vertices.append([top_vert_arr[i][0]])
vertices.append([top_vert_arr[i+1][0]])
@ -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_model.save(self.stl_filename)
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
image = app()
app()

Write Preview
Loading…
Cancel
Save