Team_odyssey / app.py
surbi karki
Update app.py
c4ee05d verified
raw
history blame contribute delete
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])}