Update app.py

app.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
 from transformers import pipeline
 import numpy as np
+from pydantic import BaseModel, Field
 
 app = FastAPI(docs_url="/")
 
@@ -12,13 +13,16 @@ def preprocess_image(image: np.ndarray):
     # Return the image
     return image
 
+class Image(BaseModel):
+    image: np.ndarray = Field(...)
+
 @app.post("/predict")
-def predict_endpoint(image: np.ndarray):
+def predict_endpoint(image: Image):
     # Preprocess the image
-    image = preprocess_image(image)
+    image = preprocess_image(image.image)
 
     # Create a pipeline for image classification
-    classifier = pipeline("andupets/real-estate-image-classification-30classes")
+    classifier = pipeline("image-classification")
 
     # Classify the image
     result = classifier(image)