Update app.py

app.py

@@ -1,16 +1,13 @@
-from fastapi import FastAPI, Request, File, UploadFile
-from fastapi.responses import HTMLResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.templating import Jinja2Templates
-import gradio as gr
+from fastapi import FastAPI, Request, File, UploadFile, HTTPException
+from fastapi.responses import HTMLResponse, JSONResponse
+from starlette.staticfiles import StaticFiles
 import tensorflow as tf
 import numpy as np
 from PIL import Image
+import io
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.templating import Jinja2Templates
 import uvicorn
-import os
-import shutil
-import threading
 
 # Load mô hình TFLite
 interpreter = tf.lite.Interpreter(model_path="model_wcpj_pro.tflite")
@@ -19,9 +16,9 @@ interpreter.allocate_tensors()
 input_details = interpreter.get_input_details()
 output_details = interpreter.get_output_details()
 
-class_names = ['wc_clean', 'wc_moderately_dirty', 'wc_slightly_dirty', 'wc_very_dirty']
+class_names = ['wc_clean', 'wc_moderately_dirty', 'wc_slightly_dirty', 'wc_very_dirty']  # Thay bằng nhãn thực tế của bạn
 
-# Khởi tạo FastAPI
+# Khởi tạo ứng dụng FastAPI
 app = FastAPI()
 
 # Cho phép CORS
@@ -37,56 +34,30 @@ app.add_middleware(
 templates = Jinja2Templates(directory="templates")
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Tạo thư mục lưu ảnh nếu chưa có
-UPLOAD_FOLDER = "uploads"
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-
 # Route trang chính
 @app.get("/", response_class=HTMLResponse)
 async def read_root(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 
-# API tải ảnh lên Hugging Face
-@app.post("/upload/")
-async def upload_file(file: UploadFile = File(...)):
-    file_location = f"{UPLOAD_FOLDER}/{file.filename}"
-    with open(file_location, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-    return {"message": "Tải ảnh lên thành công!", "file_name": file.filename}
-
-# Hàm xử lý ảnh & dự đoán
-def predict(image):
-    image = image.resize((224, 224))  # Resize ảnh về kích thước chuẩn của mô hình
-    image = np.array(image, dtype=np.float32) / 255.0  # Chuẩn hóa ảnh về [0,1]
-    image = np.expand_dims(image, axis=0)  # Thêm batch dimension
-
-    interpreter.set_tensor(input_details[0]['index'], image)
-    interpreter.invoke()
-    output_data = interpreter.get_tensor(output_details[0]['index'])
-
-    predicted_class = np.argmax(output_data)
-    return {class_names[i]: float(output_data[0][i]) for i in range(len(class_names))}
-
-# Tạo giao diện Gradio
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(),
-    title="WCPJ Floor Classification",
-    description="Tải ảnh sàn nhà vệ sinh lên để phân loại"
-)
-
-# Chạy Gradio trên một thread riêng
-def run_gradio():
-    interface.launch(server_name="0.0.0.0", server_port=7861, share=False)
-
-thread = threading.Thread(target=run_gradio)
-thread.start()
-
-# Route nhúng Gradio vào iframe
-@app.get("/gradio", response_class=HTMLResponse)
-async def gradio_page(request: Request):
-    return templates.TemplateResponse("gradio.html", {"request": request})
+# API tải ảnh lên và phân loại
+@app.post("/predict/")
+async def predict(file: UploadFile = File(...)):
+    try:
+        image = Image.open(io.BytesIO(await file.read()))
+        image = image.resize((224, 224))  # Resize ảnh về đúng kích thước mô hình yêu cầu
+        image = np.array(image, dtype=np.float32) / 255.0  # Chuẩn hóa về [0,1]
+        image = np.expand_dims(image, axis=0)  # Thêm batch dimension
+
+        interpreter.set_tensor(input_details[0]['index'], image)
+        interpreter.invoke()
+        output_data = interpreter.get_tensor(output_details[0]['index'])
+
+        predicted_class = np.argmax(output_data)
+        result = {class_names[i]: float(output_data[0][i]) for i in range(len(class_names))}
+        return JSONResponse(content={"prediction": result, "class": class_names[predicted_class]})
+
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Lỗi xử lý ảnh: {str(e)}")
 
 # Chạy FastAPI
 if __name__ == "__main__":