FingerprintSOM / app.py
Overglitch's picture
Update app.py
dcb2259 verified
raw
history blame
3.6 kB
from fastapi import FastAPI, File, UploadFile, HTTPException
from PIL import Image
import numpy as np
import pickle
from io import BytesIO
import math
app = FastAPI()
# Cargar el modelo SOM previamente entrenado
with open("som.pkl", "rb") as tf:
som = pickle.load(tf)
M = np.array([
[ 0., -1., -1., -1., -1., 2., -1., -1., -1., 3.],
[-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
[-1., -1., -1., 1., -1., -1., -1., -1., -1., -1.],
[ 1., -1., -1., -1., -1., -1., -1., -1., -1., 0.],
[-1., -1., -1., -1., 1., -1., -1., -1., -1., -1.],
[-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
[ 3., -1., -1., -1., -1., -1., -1., -1., -1., 3.],
[-1., -1., -1., 0., -1., -1., 3., -1., -1., -1.],
[-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
[ 2., -1., -1., -1., 1., -1., -1., -1., -1., 2.]
])
def predict_fingerprint(image):
try:
processed_image = preprocess_image(image)
winner = som.winner(processed_image)
fingerprint_type = get_fingerprint_type(winner)
return fingerprint_type
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
def preprocess_image(image):
# Guardar la imagen en formato TIFF
processed_image = representativo("temp_image.tif")
processed_image_resized = processed_image.reshape(1, -1)
return processed_image_resized
def get_fingerprint_type(winner):
labels = {0: "LL", 1: "RL", 2: "WH", 3: "AR"}
fingerprint_type = labels[int(M[winner[0], winner[1]])]
return fingerprint_type
@app.post("/predict/")
async def predict_fingerprint_api(file: UploadFile = File(...)):
try:
contents = await file.read()
image = Image.open(BytesIO(contents)).convert('L')
image.resize((256,256))
image.save("temp_image.tif")
image = Image.open("temp_image.tiff")
print(f"\n\n\n\nSIZE:{image.size}")
fingerprint_type = predict_fingerprint(image)
return {"prediction": fingerprint_type}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
def sobel(I):
m, n = I.shape
Gx = np.zeros([m-2, n-2], np.float32)
Gy = np.zeros([m-2, n-2], np.float32)
gx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
gy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
for j in range(1, m-2):
for i in range(1, n-2):
Gx[j-1, i-1] = np.sum(I[j-1:j+2, i-1:i+2] * gx)
Gy[j-1, i-1] = np.sum(I[j-1:j+2, i-1:i+2] * gy)
return Gx, Gy
def medfilt2(G, d=3):
m, n = G.shape
temp = np.zeros([m+2*(d//2), n+2*(d//2)], np.float32)
salida = np.zeros([m, n], np.float32)
temp[1:m+1, 1:n+1] = G
for i in range(1, m):
for j in range(1, n):
A = np.sort(temp[i-1:i+2, j-1:j+2].reshape(-1))
salida[i-1, j-1] = A[d+1]
return salida
def orientacion(patron, w):
Gx, Gy = sobel(patron)
Gx = medfilt2(Gx)
Gy = medfilt2(Gy)
m, n = Gx.shape
mOrientaciones = np.zeros([m//w, n//w], np.float32)
for i in range(m//w):
for j in range(n//w):
YY = np.sum(2 * Gx[i*w:(i+1)*w, j:j+1] * Gy[i*w:(i+1)*w, j:j+1])
XX = np.sum(Gx[i*w:(i+1)*w, j:j+1]**2 - Gy[i*w:(i+1)*w, j:j+1]**2)
mOrientaciones[i, j] = (0.5 * math.atan2(YY, XX) + math.pi/2.0) * (180.0 / math.pi)
return mOrientaciones
def representativo(archivo):
im = Image.open(archivo)
m, n = im.size
imarray = np.array(im, np.float32)
patron = imarray[1:m-1, 1:n-1]
EE = orientacion(patron, 14)
return np.asarray(EE).reshape(-1)