Spaces:
Sleeping
Sleeping
| 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_))) | |
| 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 | |
| }) | |
| def redirect_to_ui(): | |
| return RedirectResponse("/index.html") |