Update app.py

app.py

@@ -1,6 +1,8 @@
 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
@@ -8,6 +10,7 @@ import uvicorn
 
 app = FastAPI()
 
+# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,10 +19,23 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-print("🔵 Loading model: best.pt")
-model = YOLO("best.pt")
-print("✅ Model loaded")
+# ===========================================
+#       DOWNLOAD CRACK MODEL FROM HF
+# ===========================================
+print("🔵 Downloading crack model from Hugging Face...")
+model_path = hf_hub_download(
+    repo_id="cazzz307/yolov8-crack-detection",
+    filename="best.pt"
+)
+print("📥 Model downloaded:", model_path)
+
+print("🔵 Loading YOLO model...")
+model = YOLO(model_path)
+print("✅ Crack Model Loaded Successfully")
 
+# ===========================================
+#               PREDICTION API
+# ===========================================
 @app.post("/predict")
 async def predict(img: UploadFile = File(...)):
     try:
@@ -30,20 +46,24 @@ async def predict(img: UploadFile = File(...)):
         results = model(np_img)
         result = results[0]
 
-        has_mask = result.masks is not None
-
-        if not has_mask:
+        # crack detection: check boxes
+        if result.boxes is None or len(result.boxes) == 0:
             return {
-                "data": [{"label": "normal", "confidence": 1.0}]
+                "data": [
+                    {"label": "normal", "confidence": 1.0}
+                ]
             }
 
-        if result.boxes is not None and len(result.boxes) > 0:
-            conf = float(result.boxes.conf.max().item())
-        else:
-            conf = 0.85  
+        # There are crack boxes
+        conf = float(result.boxes.conf.max().item())
 
         return {
-            "data": [{"label": "crack", "confidence": conf}]
+            "data": [
+                {
+                    "label": "crack",
+                    "confidence": conf
+                }
+            ]
         }
 
     except Exception as e:
@@ -53,5 +73,6 @@ async def predict(img: UploadFile = File(...)):
             "error": str(e)
         }
 
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)