sulthonkaf's picture
Update app.py
f2f1372 verified
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from PIL import Image
import numpy as np
import tensorflow as tf
import pickle
import io
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.mount("/", StaticFiles(directory=".", html=True), name="static")
model = tf.keras.models.load_model("model_2.h5")
with open("label_encoder.pkl", "rb") as f:
label_encoder = pickle.load(f)
class_names = label_encoder.inverse_transform(np.arange(len(label_encoder.classes_)))
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
contents = await file.read()
image = Image.open(io.BytesIO(contents)).convert("RGB")
image = image.resize((224, 224))
img_array = np.array(image) / 255.0
img_array = np.expand_dims(img_array, axis=0)
prediction = model.predict(img_array)[0]
predicted_index = np.argmax(prediction)
predicted_label = class_names[predicted_index]
confidence = float(prediction[predicted_index]) * 100
all_probs = {
class_names[i]: float(prob)
for i, prob in enumerate(prediction)
}
return JSONResponse(content={
"predicted_label": predicted_label,
"confidence": round(confidence, 2),
"all_probabilities": all_probs
})
@app.get("/")
def redirect_to_ui():
return RedirectResponse("/index.html")