Update app.py

app.py

@@ -1,32 +1,39 @@
-from fastapi import FastAPI
+import fastapi
 from transformers import pipeline
-import numpy as np
-from pydantic import BaseModel, Field
+import pickle
 
-app = FastAPI(docs_url="/")
+# Load the model from the pickle file
+with open("model.pkl", "rb") as f:
+    model = pickle.load(f)
 
-def preprocess_image(image: np.ndarray):
+# 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
 
-class Image(BaseModel):
-    image: np.ndarray = Field(...)
-
+# Define an endpoint to predict the output
 @app.post("/predict")
-def predict_endpoint(image: Image):
+async def predict_endpoint(image: fastapi.File):
     # Preprocess the image
-    image = preprocess_image(image.image)
+    image = preprocess_image(image)
 
-    # Create a pipeline for image classification
-    classifier = pipeline("image-classification")
+    # Make a prediction
+    prediction = model(image)
 
-    # Classify the image
-    result = classifier(image)
+    # Return the prediction
+    return {"prediction": prediction}
 
-    # Return the classification result
-    return result
+# Start the FastAPI app
+if _name_ == "_main_":
+    import uvicorn
 
+    uvicorn.run(app, host="0.0.0.0", port=8000)