Spaces:
Running
Running
File size: 1,826 Bytes
9840ff2 69a5ea6 eb8ec37 d6d90b8 36cdccd f450f93 5f1301a 69a5ea6 7ff4f82 d6d90b8 69a5ea6 5f1301a d6d90b8 fe3ea5b d6d90b8 36cdccd d6d90b8 f450f93 36cdccd 304c149 f450f93 9840ff2 d6d90b8 f450f93 d6d90b8 f450f93 9840ff2 d6d90b8 9840ff2 36cdccd d6d90b8 36cdccd 9840ff2 f450f93 303f396 f450f93 77341f2 d6d90b8 4be526a 69a5ea6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image
import io
import uvicorn
app = FastAPI()
# CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ===========================================
# DOWNLOAD CRACK MODEL FROM HF
# ===========================================
print("🔵 Loading local crack model...")
model = YOLO("best.pt")
print("✅ Crack Model Loaded Successfully")
# ===========================================
# PREDICTION API
# ===========================================
@app.post("/predict")
async def predict(img: UploadFile = File(...)):
try:
bytes_data = await img.read()
image = Image.open(io.BytesIO(bytes_data)).convert("RGB")
np_img = np.array(image)
results = model(np_img, conf=0.40)
result = results[0]
# crack detection: check boxes
if result.boxes is None or len(result.boxes) == 0:
return {
"data": [
{"label": "normal", "confidence": 1.0}
]
}
# There are crack boxes
conf = float(result.boxes.conf.max().item())
return {
"data": [
{
"label": "crack",
"confidence": conf
}
]
}
except Exception as e:
print("❌ Prediction error:", e)
return {
"data": [{"label": "normal", "confidence": 1.0}],
"error": str(e)
}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)
|