|
import cv2 |
|
import tensorflow as tf |
|
import numpy as np |
|
from fastapi import FastAPI, UploadFile, File |
|
from pydantic import BaseModel |
|
from huggingface_hub import from_pretrained_keras |
|
|
|
app = FastAPI() |
|
|
|
def load_ben_color(image, sigmaX=10): |
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
|
image = cv2.resize(image, (224, 224)) |
|
image = cv2.addWeighted(image, 4, cv2.GaussianBlur(image, (0, 0), sigmaX), -4, 128) |
|
return image |
|
|
|
def clahe(image): |
|
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) |
|
r, g, b = cv2.split(image) |
|
r = clahe.apply(r) |
|
g = clahe.apply(g) |
|
b = clahe.apply(b) |
|
result = cv2.merge((r, g, b)) |
|
return result |
|
|
|
def filter1(image): |
|
image = load_ben_color(image) |
|
return image |
|
|
|
def filter2(image): |
|
image = clahe(image) |
|
image = cv2.resize(image, (224, 224)) |
|
return image |
|
|
|
def predict(image, model, filter_func): |
|
model_image = filter_func(image) |
|
model_image = np.array([model_image], dtype=np.float32) / 255.0 |
|
predictions = model(tf.constant(model_image)) |
|
return predictions.numpy() |
|
|
|
def result(predictions): |
|
class_labels = ["Age related Macular Degeneration", "Cataract", "Diabetes", "Glaucoma", "Hypertension", "Normal", "Others", "Pathological Myopia"] |
|
predictions = np.array(predictions) |
|
predictions = predictions.tolist()[0] |
|
predictions_index = np.argmax(predictions) |
|
return class_labels[predictions_index] |
|
|
|
|
|
models_names = ["ODIR-B-2K-5Class-LastTrain-Xception", "ODIR-B-2K-6Class-LastTrain-Xception"] |
|
model_paths = ["ODIR-B-2K-5Class-LastTrain-Xception", "ODIR-B-2K-6Class-LastTrain-Xception"] |
|
|
|
filters = [filter1, filter2] |
|
|
|
models = [] |
|
for model_path in model_paths: |
|
model = tf.saved_model.load(model_path) |
|
models.append(model) |
|
|
|
class PredictionResponse(BaseModel): |
|
predictions: dict |
|
|
|
@app.post("/predict", response_model=PredictionResponse) |
|
async def predict_endpoint(file: UploadFile = File(...)): |
|
contents = await file.read() |
|
nparr = np.frombuffer(contents, np.uint8) |
|
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) |
|
|
|
result_json = {} |
|
for i in range(len(models)): |
|
model = models[i] |
|
prediction = predict(image, model, filters[i]) |
|
result_json[models_names[i]] = result(prediction) |
|
|
|
return {"predictions": result_json} |
|
|