Update app.py

app.py

@@ -1,13 +1,17 @@
 # app.py
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
-import requests
+from ultralytics import YOLO
 import uvicorn
+import numpy as np
+from PIL import Image
+import io
+import requests
 import os
 
 app = FastAPI()
 
-# CORS 설정
+# CORS 허용
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,41 +20,63 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# HuggingFace 모델 엔드포인트
-HF_API_URL = "https://api-inference.huggingface.co/models/keremberke/yolov8n-concrete-crack"
+# -----------------------------
+# 다운로드할 모델 경로 (HF Repo)
+# -----------------------------
+MODEL_URL = "https://huggingface.co/OpenSistemas/YOLOv8-crack-seg/resolve/main/yolov8n/weights/best.pt"
+MODEL_PATH = "best.pt"
+
+# -----------------------------
+# 최초 실행 시 모델 다운로드
+# -----------------------------
+if not os.path.exists(MODEL_PATH):
+    print("🔵 Downloading YOLOv8 crack model...")
+    r = requests.get(MODEL_URL)
+    with open(MODEL_PATH, "wb") as f:
+        f.write(r.content)
+    print("✅ YOLOv8 model downloaded.")
+
+
+# -----------------------------
+# YOLOv8 모델 로드
+# -----------------------------
+print("🔵 Loading YOLOv8 crack segmentation model...")
+model = YOLO(MODEL_PATH)
+print("✅ Model loaded!")
 
-# ❗ 토큰은 절대로 문자열로 넣으면 안 됨. 환경변수로만!
-HF_TOKEN = os.getenv("HF_TOKEN")
-headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
 
 @app.post("/predict")
 async def predict(img: UploadFile = File(...)):
+    # 이미지 읽기
     bytes_data = await img.read()
+    image = Image.open(io.BytesIO(bytes_data)).convert("RGB")
+    np_img = np.array(image)
 
-    response = requests.post(
-        HF_API_URL,
-        headers=headers,
-        data=bytes_data
-    )
-
-    try:
-        results = response.json()
-    except:
-        return {"data": [{"label": "normal", "confidence": 1.0}]}
+    # YOLO inference
+    result = model(np_img)[0]
 
-    if not isinstance(results, list) or len(results) == 0:
-        return {"data": [{"label": "normal", "confidence": 1.0}]}
+    # segmentation mask confidence 추출
+    if result.masks is None:
+        # 균열 없음
+        return {
+            "data": [
+                {"label": "normal", "confidence": 1.0}
+            ]
+        }
 
-    max_conf = max(item.get("score", 0) for item in results)
+    # segmentation mask가 있을 경우 → crack 감지됨
+    # mask의 confidence 평균값 사용
+    conf = float(result.boxes.conf.cpu().numpy().max()) if result.boxes else 0.9
 
     return {
         "data": [
             {
                 "label": "crack",
-                "confidence": float(max_conf)
+                "confidence": conf
             }
         ]
     }
 
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)