Update app.py

app.py

@@ -1,39 +1,38 @@
 import fastapi
-from transformers import pipeline
-import pickle
+import torch
+from transformers import AutoModelForImageClassification
 
-# Load the model from the pickle file
-with open("model_finetuned.pkl", "rb") as f:
-    model = pickle.load(f)
+# Load the model from the local file
+model = AutoModelForImageClassification.from_pretrained("./model.ckpt")
 
 # Define a function to preprocess the image
 def preprocess_image(image):
-    # Resize the image to a fixed size
-    image = image.resize((224, 224))
+  # Resize the image to a fixed size
+  image = image.resize((224, 224))
 
-    # Convert the image to a NumPy array
-    image = np.array(image)
+  # Convert the image to a NumPy array
+  image = np.array(image)
 
-    # Normalize the image
-    image = image / 255.0
+  # Normalize the image
+  image = image / 255.0
 
-    # Return the image
-    return image
+  # 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)
+  # Preprocess the image
+  image = preprocess_image(image)
 
-    # Make a prediction
-    prediction = model(image)
+  # Make a prediction
+  prediction = model(image)
 
-    # Return the prediction
-    return {"prediction": prediction}
+  # Return the prediction
+  return {"prediction": prediction}
 
 # Start the FastAPI app
 if _name_ == "_main_":
-    import uvicorn
+  import uvicorn
 
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+  uvicorn.run(app, host="0.0.0.0", port=8000)