mussie1212 commited on
Commit
0ebf72e
·
1 Parent(s): 3d0f64f

Fix: Add OpenCV system dependencies

Browse files
Files changed (1) hide show
  1. app.py +40 -18
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 for JSON response."""
21
- _, buffer = cv2.imencode(".png", img)
22
- return base64.b64encode(buffer).decode("utf-8")
 
 
 
 
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
- Upload an image and get car parts and damage detection results.
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] # BGR to RGB
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] # BGR to RGB
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"}