import fastapi
from transformers import pipeline
import h5py

# Load the model from the h5 file
model = h5py.File("model_finetuned.h5", "r")["model"]

# Define a function to preprocess the image
def preprocess_image(image):
    # Resize the image to a fixed size
    image = image.resize((224, 224))
    # Convert the image to a NumPy array
    image = np.array(image)
    # Normalize the image
    image = image / 255.0
    # Return the image
    return image

# Define an endpoint to predict the output
@app.post("/predict")
async def predict_endpoint(image: fastapi.File):
    # Preprocess the image
    image = preprocess_image(image)
    # Make a prediction
    prediction = model.predict(image)
    # Return the prediction
    return {"prediction": prediction}

# Start the FastAPI app
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)