Team_odyssey / app.py
surbi karki
Update app.py
c4ee05d verified
raw
history blame
877 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import cv2
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import to_categorical
# FastAPI app setup
app = FastAPI()
# Load the pre-trained model
model = load_model('sample.h5')
# Predefined image size for prediction
size = 100
# FastAPI Model for input data
class ImageData(BaseModel):
image_path: str # Path to image file
# Prediction Endpoint
@app.post("/predict/")
async def predict(data: ImageData):
# Load the image and preprocess
image = cv2.imread(data.image_path, 0)
image = cv2.resize(image, (size, size))
image = np.asarray(image).reshape(1, size, size, 1)
# Make prediction
prediction = model.predict(image)
predicted_class = np.argmax(prediction, axis=1)
return {"predicted_class": int(predicted_class[0])}