import fastapi

from fastapi import FastAPI, UploadFile, File
from transformers import AutoModelForImageClassification


app = FastAPI()


# Load the model from the local file
model = AutoModelForImageClassification.from_pretrained(
    "andupets/real-estate-image-classification-30classes"
)


# Define a function to preprocess the image
def preprocess_image(image: UploadFile):
    # 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: UploadFile = File(...)
):  # Preprocess the image
    image = preprocess_image(image)

    # Make a prediction
    prediction = model(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)