Update app.py

app.py

@@ -1,13 +1,11 @@
 # app.py
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
-from ultralytics import YOLO
-import uvicorn
+from ultralyticsplus import YOLO
 import numpy as np
 from PIL import Image
 import io
-import requests
-import os
+import uvicorn
 
 app = FastAPI()
 
@@ -20,62 +18,60 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# -----------------------------
-# 다운로드할 모델 경로 (HF Repo)
-# -----------------------------
-MODEL_URL = "https://huggingface.co/OpenSistemas/YOLOv8-crack-seg/resolve/main/yolov8n/weights/best.pt"
-MODEL_PATH = "best.pt"
+# ===========================================
+#  YOLOv8 segmentation model (from keremberke)
+# ===========================================
+MODEL_ID = "keremberke/yolov8n-building-segmentation"
 
-# -----------------------------
-# 최초 실행 시 모델 다운로드
-# -----------------------------
-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.")
+print(f"🔵 Loading model: {MODEL_ID}")
+model = YOLO(MODEL_ID)
+print("✅ Model loaded")
 
+@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)
 
-# -----------------------------
-# YOLOv8 모델 로드
-# -----------------------------
-print("🔵 Loading YOLOv8 crack segmentation model...")
-model = YOLO(MODEL_PATH)
-print("✅ Model loaded!")
+        # YOLO inference
+        results = model(np_img)
+        result = results[0]
 
+        # segmentation mask 여부로 crack 유무 판단
+        has_mask = result.masks is not None
 
-@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)
+        if not has_mask:
+            return {
+                "data": [
+                    {"label": "normal", "confidence": 1.0}
+                ]
+            }
 
-    # YOLO inference
-    result = model(np_img)[0]
+        # 가장 높은 confidence 추출
+        if result.boxes is not None and len(result.boxes) > 0:
+            conf = float(result.boxes.conf.max().item())
+        else:
+            conf = 0.85  # 기본값
 
-    # segmentation mask confidence 추출
-    if result.masks is None:
-        # 균열 없음
         return {
             "data": [
-                {"label": "normal", "confidence": 1.0}
+                {
+                    "label": "crack",
+                    "confidence": conf
+                }
             ]
         }
 
-    # segmentation mask가 있을 경우 → crack 감지됨
-    # mask의 confidence 평균값 사용
-    conf = float(result.boxes.conf.cpu().numpy().max()) if result.boxes else 0.9
-
-    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__":