Spaces:
Sleeping
Sleeping
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 | |
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) |