import fastapi
from fastapi import FastAPI,UploadFile
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):
  # 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):
  # 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)