Spaces:
Runtime error
Runtime error
Commit
·
0ebf72e
1
Parent(s):
3d0f64f
Fix: Add OpenCV system dependencies
Browse files
app.py
CHANGED
@@ -1,86 +1,108 @@
|
|
1 |
from fastapi import FastAPI, File, UploadFile, HTTPException
|
2 |
from fastapi.responses import JSONResponse
|
|
|
3 |
from ultralytics import YOLO
|
4 |
import numpy as np
|
5 |
import cv2
|
6 |
from io import BytesIO
|
7 |
from PIL import Image
|
8 |
import base64
|
|
|
|
|
|
|
|
|
|
|
9 |
|
10 |
app = FastAPI(title="Car Parts & Damage Detection API")
|
11 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
12 |
# Load YOLO models
|
13 |
try:
|
|
|
14 |
car_part_model = YOLO("car_part_detector_model.pt")
|
|
|
|
|
15 |
damage_model = YOLO("damage_general_model.pt")
|
|
|
16 |
except Exception as e:
|
|
|
17 |
raise RuntimeError(f"Failed to load models: {str(e)}")
|
18 |
|
19 |
def image_to_base64(img: np.ndarray) -> str:
|
20 |
-
"""Convert numpy image to base64 string
|
21 |
-
|
22 |
-
|
|
|
|
|
|
|
|
|
23 |
|
24 |
@app.post("/predict", summary="Run inference on an image for car parts and damage")
|
25 |
async def predict(file: UploadFile = File(...)):
|
26 |
-
"""
|
27 |
-
|
28 |
-
Returns annotated images as base64 strings and text descriptions.
|
29 |
-
"""
|
30 |
try:
|
31 |
-
# Read and process image
|
32 |
contents = await file.read()
|
33 |
image = Image.open(BytesIO(contents)).convert("RGB")
|
34 |
img = np.array(image)
|
|
|
35 |
|
36 |
-
# Initialize default blank images (gray placeholder)
|
37 |
blank_img = np.full((img.shape[0], img.shape[1], 3), 128, dtype=np.uint8)
|
38 |
car_part_img = blank_img.copy()
|
39 |
damage_img = blank_img.copy()
|
40 |
-
|
41 |
-
# Initialize text results
|
42 |
car_part_text = "Car Parts: No detections"
|
43 |
damage_text = "Damage: No detections"
|
44 |
|
45 |
-
# Process car parts detection
|
46 |
try:
|
|
|
47 |
car_part_results = car_part_model(img)[0]
|
48 |
if car_part_results.boxes:
|
49 |
-
car_part_img = car_part_results.plot()[..., ::-1]
|
50 |
car_part_text = "Car Parts:\n" + "\n".join(
|
51 |
f"- {car_part_results.names[int(cls)]} ({conf:.2f})"
|
52 |
for conf, cls in zip(car_part_results.boxes.conf, car_part_results.boxes.cls)
|
53 |
)
|
|
|
54 |
except Exception as e:
|
55 |
car_part_text = f"Car Parts: Error: {str(e)}"
|
|
|
56 |
|
57 |
-
# Process damage detection
|
58 |
try:
|
|
|
59 |
damage_results = damage_model(img)[0]
|
60 |
if damage_results.boxes:
|
61 |
-
damage_img = damage_results.plot()[..., ::-1]
|
62 |
damage_text = "Damage:\n" + "\n".join(
|
63 |
f"- {damage_results.names[int(cls)]} ({conf:.2f})"
|
64 |
for conf, cls in zip(damage_results.boxes.conf, damage_results.boxes.cls)
|
65 |
)
|
|
|
66 |
except Exception as e:
|
67 |
damage_text = f"Damage: Error: {str(e)}"
|
|
|
68 |
|
69 |
-
# Convert output images to base64
|
70 |
car_part_img_base64 = image_to_base64(car_part_img)
|
71 |
damage_img_base64 = image_to_base64(damage_img)
|
72 |
-
|
73 |
return JSONResponse({
|
74 |
"car_part_image": car_part_img_base64,
|
75 |
"car_part_text": car_part_text,
|
76 |
"damage_image": damage_img_base64,
|
77 |
"damage_text": damage_text
|
78 |
})
|
79 |
-
|
80 |
except Exception as e:
|
|
|
81 |
raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")
|
82 |
|
83 |
@app.get("/", summary="Health check")
|
84 |
async def root():
|
85 |
"""Check if the API is running."""
|
|
|
86 |
return {"message": "Car Parts & Damage Detection API is running"}
|
|
|
1 |
from fastapi import FastAPI, File, UploadFile, HTTPException
|
2 |
from fastapi.responses import JSONResponse
|
3 |
+
import logging
|
4 |
from ultralytics import YOLO
|
5 |
import numpy as np
|
6 |
import cv2
|
7 |
from io import BytesIO
|
8 |
from PIL import Image
|
9 |
import base64
|
10 |
+
import os
|
11 |
+
|
12 |
+
# Setup logging
|
13 |
+
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
14 |
+
logger = logging.getLogger(__name__)
|
15 |
|
16 |
app = FastAPI(title="Car Parts & Damage Detection API")
|
17 |
|
18 |
+
# Log model file presence
|
19 |
+
model_files = ["car_part_detector_model.pt", "damage_general_model.pt"]
|
20 |
+
for model_file in model_files:
|
21 |
+
if os.path.exists(model_file):
|
22 |
+
logger.info(f"Model file found: {model_file}")
|
23 |
+
else:
|
24 |
+
logger.error(f"Model file missing: {model_file}")
|
25 |
+
|
26 |
# Load YOLO models
|
27 |
try:
|
28 |
+
logger.info("Loading car part model...")
|
29 |
car_part_model = YOLO("car_part_detector_model.pt")
|
30 |
+
logger.info("Car part model loaded successfully")
|
31 |
+
logger.info("Loading damage model...")
|
32 |
damage_model = YOLO("damage_general_model.pt")
|
33 |
+
logger.info("Damage model loaded successfully")
|
34 |
except Exception as e:
|
35 |
+
logger.error(f"Failed to load models: {str(e)}")
|
36 |
raise RuntimeError(f"Failed to load models: {str(e)}")
|
37 |
|
38 |
def image_to_base64(img: np.ndarray) -> str:
|
39 |
+
"""Convert numpy image to base64 string."""
|
40 |
+
try:
|
41 |
+
_, buffer = cv2.imencode(".png", img)
|
42 |
+
return base64.b64encode(buffer).decode("utf-8")
|
43 |
+
except Exception as e:
|
44 |
+
logger.error(f"Error encoding image to base64: {str(e)}")
|
45 |
+
raise
|
46 |
|
47 |
@app.post("/predict", summary="Run inference on an image for car parts and damage")
|
48 |
async def predict(file: UploadFile = File(...)):
|
49 |
+
"""Upload an image and get car parts and damage detection results."""
|
50 |
+
logger.info("Received image upload")
|
|
|
|
|
51 |
try:
|
|
|
52 |
contents = await file.read()
|
53 |
image = Image.open(BytesIO(contents)).convert("RGB")
|
54 |
img = np.array(image)
|
55 |
+
logger.info(f"Image loaded: shape={img.shape}")
|
56 |
|
|
|
57 |
blank_img = np.full((img.shape[0], img.shape[1], 3), 128, dtype=np.uint8)
|
58 |
car_part_img = blank_img.copy()
|
59 |
damage_img = blank_img.copy()
|
|
|
|
|
60 |
car_part_text = "Car Parts: No detections"
|
61 |
damage_text = "Damage: No detections"
|
62 |
|
|
|
63 |
try:
|
64 |
+
logger.info("Running car part detection...")
|
65 |
car_part_results = car_part_model(img)[0]
|
66 |
if car_part_results.boxes:
|
67 |
+
car_part_img = car_part_results.plot()[..., ::-1]
|
68 |
car_part_text = "Car Parts:\n" + "\n".join(
|
69 |
f"- {car_part_results.names[int(cls)]} ({conf:.2f})"
|
70 |
for conf, cls in zip(car_part_results.boxes.conf, car_part_results.boxes.cls)
|
71 |
)
|
72 |
+
logger.info("Car part detection completed")
|
73 |
except Exception as e:
|
74 |
car_part_text = f"Car Parts: Error: {str(e)}"
|
75 |
+
logger.error(f"Car part detection error: {str(e)}")
|
76 |
|
|
|
77 |
try:
|
78 |
+
logger.info("Running damage detection...")
|
79 |
damage_results = damage_model(img)[0]
|
80 |
if damage_results.boxes:
|
81 |
+
damage_img = damage_results.plot()[..., ::-1]
|
82 |
damage_text = "Damage:\n" + "\n".join(
|
83 |
f"- {damage_results.names[int(cls)]} ({conf:.2f})"
|
84 |
for conf, cls in zip(damage_results.boxes.conf, damage_results.boxes.cls)
|
85 |
)
|
86 |
+
logger.info("Damage detection completed")
|
87 |
except Exception as e:
|
88 |
damage_text = f"Damage: Error: {str(e)}"
|
89 |
+
logger.error(f"Damage detection error: {str(e)}")
|
90 |
|
|
|
91 |
car_part_img_base64 = image_to_base64(car_part_img)
|
92 |
damage_img_base64 = image_to_base64(damage_img)
|
93 |
+
logger.info("Returning prediction results")
|
94 |
return JSONResponse({
|
95 |
"car_part_image": car_part_img_base64,
|
96 |
"car_part_text": car_part_text,
|
97 |
"damage_image": damage_img_base64,
|
98 |
"damage_text": damage_text
|
99 |
})
|
|
|
100 |
except Exception as e:
|
101 |
+
logger.error(f"Inference error: {str(e)}")
|
102 |
raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")
|
103 |
|
104 |
@app.get("/", summary="Health check")
|
105 |
async def root():
|
106 |
"""Check if the API is running."""
|
107 |
+
logger.info("Health check accessed")
|
108 |
return {"message": "Car Parts & Damage Detection API is running"}
|