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])}