|
import os |
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' |
|
|
|
from fastapi import FastAPI, UploadFile, File |
|
import json |
|
from PIL import Image |
|
from io import BytesIO |
|
import numpy as np |
|
|
|
from model import get_model |
|
|
|
app = FastAPI() |
|
|
|
IMAGE_WIDTH = 224 |
|
IMAGE_HEIGHT = 224 |
|
|
|
MODEL_WEIGHT_PATH = 'vgg_face_weights2.h5' |
|
model = get_model( |
|
image_shape = (IMAGE_WIDTH, IMAGE_HEIGHT, 3), |
|
num_classes = 6, |
|
model_weights = MODEL_WEIGHT_PATH |
|
) |
|
print(model.summary()) |
|
print("Model Loaded Successfully") |
|
|
|
|
|
def load_image(image_data): |
|
image = Image.open(BytesIO(image_data)) |
|
return image |
|
|
|
def preprocess(image): |
|
image = image.resize((IMAGE_WIDTH, IMAGE_HEIGHT)) |
|
|
|
image = np.array(image) |
|
image = np.expand_dims(image, axis=0) / 255.0 |
|
|
|
return image |
|
|
|
def get_prediction(image): |
|
probs = model.predict(image)[0] |
|
label = np.argmax(probs) |
|
|
|
return { |
|
'pred_probs': probs.tolist(), |
|
'label': int(label) |
|
} |
|
|
|
@app.get("/") |
|
def foo(): |
|
return { |
|
"status": "Face Expression Classifier" |
|
} |
|
|
|
|
|
@app.post("/get_prediction") |
|
async def predict(face_img: UploadFile = File(...)): |
|
image = load_image(await face_img.read()) |
|
|
|
image = preprocess(image) |
|
result = get_prediction(image) |
|
|
|
return { |
|
"result": json.dumps(result) |
|
} |