pawanperera5/fastAPI-ImageClassification

1from fastapi import FastAPI, File, UploadFile
2from tensorflow.keras.models import load_model
3from PIL import Image
4import numpy as np
5import os
6os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
7
8app = FastAPI()
9
10# Load the pre-trained Keras model
11model = load_model("new_model.keras")
12
13# Define the list of celebrities
14celebrities = ['Angelina Jolie', 'Brad Pitt', 'Hugh Jackman', 'Johnny Depp', 'Leonardo DiCaprio']
15
16# Preprocess the uploaded image
17def preprocess_image(image: Image.Image) -> np.array:
18    # Resize the image to the input size expected by your model (e.g., 150x150)
19    image = image.resize((100, 100))
20    # Convert the image to a numpy array and normalize it
21    image = np.array(image) / 255.0
22    # Add batch dimension
23    image = np.expand_dims(image, axis=0)
24    return image
25
26@app.get("/")
27async def read_root():
28    return {
29        "message": "Welcome to the Celebrity Prediction API!",
30        "available_celebrities": celebrities
31    }
32
33@app.post("/predict/")
34async def predict(file: UploadFile = File(...)):
35    # Read the uploaded file and convert it to an image
36    img = Image.open(file.file)
37    # Preprocess the image for model input
38    img_array = preprocess_image(img)
39    # Predict the celebrity
40    predictions = model.predict(img_array)
41    # Get the index of the highest probability
42    predicted_index = np.argmax(predictions)
43    # Get the celebrity name
44    predicted_celebrity = celebrities[predicted_index]
45    return {"celebrity": predicted_celebrity}