fastapicode / main.py
Zain's picture
Upload main.py
8898712 verified
raw
history blame contribute delete
No virus
2.45 kB
import os
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import img_to_array
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
# from tensorflow.keras.preprocessing import image
import numpy as np
import uvicorn
from PIL import Image
import io
import logging
# Suppress TensorFlow warnings
#tf.get_logger().setLevel(logging.ERROR)
app = FastAPI()
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Adjust as needed
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Replace with your actual model path
model_path = 'face_expression_detection_model3.h5' # Update with your new model's path
model = load_model(model_path)
# def preprocess_image(img: Image.Image, target_size=(224, 224)):
# img = img.resize(target_size)
# x = image.img_to_array(img)
# x = np.expand_dims(x, axis=0)
# x = x / 255.0 # Assuming normalization was done during training
# return x
def preprocess_image(img: Image.Image, target_size=(224, 224)):
img = img.resize(target_size)
x = img_to_array(img)
x = np.expand_dims(x, axis=0)
x = x / 255.0 # Assuming normalization was done during training
return x
@app.get("/")
async def read_root():
return {"message": "Welcome to the Emotion Classification API"}
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
contents = await file.read()
img = Image.open(io.BytesIO(contents))
x = preprocess_image(img)
prediction = model.predict(x)
predicted_class = np.argmax(prediction[0])
class_names = ['angry', 'happy', 'sad', 'surprised', 'neutral', 'disgusted', 'fearful']
stress_levels = {
'angry': 'high_stress',
'happy': 'neutral',
'sad': 'low_stress',
'surprised': 'low_stress',
'neutral': 'neutral',
'disgusted': 'low_stress',
'fearful': 'low_stress'
}
emotion = class_names[predicted_class]
stress_level = stress_levels[emotion]
result = {
#"predicted_class": emotion,
"predicted_class": stress_level
}
return JSONResponse(content=result)