Update app.py

app.py

@@ -1,25 +1,38 @@
 from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import JSONResponse
-from tensorflow.keras.models import load_model
-from PIL import Image
+from fastapi.middleware.cors import CORSMiddleware
 import numpy as np
-import io
+from PIL import Image
+import tensorflow as tf
+
+# Load model and classes
+model = tf.keras.models.load_model("hf_keras_model.keras")
+class_names = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']
 
+# Initialize app
 app = FastAPI()
 
-# Load model and class names
-model = load_model("hf_keras_model.keras")
-class_names = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']
+# Allow all CORS (for frontend/test requests)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/")
+def root():
+    return {"message": "API is working!"}
 
-@app.post("/predict/")
+@app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-    image = image.resize((150, 150))
+    # Load image
+    image = Image.open(file.file).convert("RGB").resize((150, 150))
     img_array = np.array(image) / 255.0
     img_array = np.expand_dims(img_array, axis=0)
 
-    preds = model.predict(img_array)[0]
-    confidences = {class_names[i]: float(preds[i]) for i in range(len(class_names))}
+    # Predict
+    predictions = model.predict(img_array)[0]
+    results = {class_names[i]: float(predictions[i]) for i in range(len(class_names))}
+    top_class = class_names[np.argmax(predictions)]
 
-    return JSONResponse(content=confidences)
+    return {"top_prediction": top_class, "all_predictions": results}