1MR's picture
Update app.py
2d6f3da verified
Raw History Blame Contribute Delete
2.24 kB
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import tensorflow as tf
import numpy as np
import shutil
import os
# Initialize FastAPI app
app = FastAPI()
# Class labels
class_labels = {
0: 'Baked Potato', 1: 'Burger', 2: 'Crispy Chicken', 3: 'Donut', 4: 'Fries',
5: 'Hot Dog', 6: 'Jalapeno', 7: 'Kiwi', 8: 'Lemon', 9: 'Lettuce',
10: 'Mango', 11: 'Onion', 12: 'Orange', 13: 'Pizza', 14: 'Taquito',
15: 'Apple', 16: 'Banana', 17: 'Beetroot', 18: 'Bell Pepper', 19: 'Bread',
20: 'Cabbage', 21: 'Carrot', 22: 'Cauliflower', 23: 'Cheese',
24: 'Chilli Pepper', 25: 'Corn', 26: 'Crab', 27: 'Cucumber',
28: 'Eggplant', 29: 'Eggs', 30: 'Garlic', 31: 'Ginger', 32: 'Grapes',
33: 'Milk', 34: 'Salmon', 35: 'Yogurt'
}
# Load the trained model
model = tf.keras.models.load_model("model_unfreezeNewCorrectpredict.keras")
# Image preprocessing function
def load_and_prep_image(file_path, img_shape=224):
img = tf.io.read_file(file_path)
img = tf.image.decode_image(img, channels=3)
img = tf.image.resize(img, size=[img_shape, img_shape])
img = tf.expand_dims(img, axis=0)
return img
# Predict label function
def predict_label(model, image_path, class_names):
img = load_and_prep_image(image_path, img_shape=224)
pred = model.predict(img)
pred_class_index = np.argmax(pred, axis=1)[0]
pred_class_name = class_names[pred_class_index]
return pred_class_name
@app.get("/")
def read_root():
return {"message": "This is My Nutrionguid App"}
# API endpoint for prediction
@app.post("/predict")
async def predict_image(file: UploadFile = File(...)):
try:
# Save the uploaded file
file_location = f"./temp_{file.filename}"
with open(file_location, "wb") as f:
shutil.copyfileobj(file.file, f)
# Predict the label
prediction = predict_label(model, file_location, class_labels)
# Remove the temporary file
os.remove(file_location)
return {"predicted_label": prediction}
except Exception as e:
return JSONResponse(
status_code=500,
content={"error": f"An error occurred: {str(e)}"}
)