Update app.py

app.py

@@ -1,74 +1,32 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from PIL import Image
-import numpy as np
 import pickle
-from io import BytesIO
+from minisom import MiniSom
+import numpy as np
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List
 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))
+class InputData(BaseModel):
+    array: List[List[int]]
 
-def preprocess_image(image):
-    # Guardar la imagen en formato TIFF
-    processed_image = representativo(image)
-    processed_image_resized = processed_image.reshape(1, -1)
-    print(f"\n\n\nXD:{processed_image.size}\t{processed_image_resized.size}")
-    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 = FastAPI()
 
-@app.post("/predict/")
-async def predict_fingerprint_api(file: UploadFile = File(...)):
-    try:
-        contents = await file.read()
-        image = Image.open(BytesIO(contents)).convert('L')
-        image.save("temp_image.tif")
-        image = Image.open("temp_image.tif")
-        image.resize((256,256))
-        print(f"\n\n\n\nSIZE:{image.size}\t{image}")
-        fingerprint_type = predict_fingerprint(image)
-        return {"prediction": fingerprint_type}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+# Cargar el modelo SOM
+def load_model():
+    with open('som.pkl', 'rb') as fid:
+        somecoli = pickle.load(fid)
+    return somecoli
 
 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]])
+    gx = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]
+    gy = [[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)
+            Gx[j-1, i-1] = sum(sum(I[j-1:j+2, i-1:i+2] * gx))
+            Gy[j-1, i-1] = sum(sum(I[j-1:j+2, i-1:i+2] * gy))
     return Gx, Gy
 
 def medfilt2(G, d=3):
@@ -78,8 +36,8 @@ def medfilt2(G, d=3):
     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]
+            A = np.asarray(temp[i-1:i+2, j-1:j+2]).reshape(-1)
+            salida[i-1, j-1] = np.sort(A)[d+1]
     return salida
 
 def orientacion(patron, w):
@@ -90,14 +48,43 @@ def orientacion(patron, w):
     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)
+            YY = sum(sum(2*Gx[i*w:(i+1)*w, j:j+1] * Gy[i*w:(i+1)*w, j:j+1]))
+            XX = sum(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(im):
-    m, n = im.size
-    imarray = np.array(im, np.float32)
+def representativo(imarray):
+    imarray = np.squeeze(imarray)
+    m, n = imarray.shape
     patron = imarray[1:m-1, 1:n-1]
     EE = orientacion(patron, 14)
-    return np.asarray(EE).reshape(-1)
+    return np.asarray(EE).reshape(-1)
+
+som = load_model()
+
+MM = 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.]
+])
+
+@app.post("/predict/")
+async def predict(data: InputData):
+    try:
+        input_data = np.array(data.array).reshape(256, 256, 1)
+        representative_data = representativo(input_data)
+        representative_data = representative_data.reshape(1, -1)
+        
+        w = som.winner(representative_data)
+        prediction = MM[w]
+        
+        return {"prediction": prediction}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))