Refactored app.py for single damage model analysis and removed unused car_part_detector_model.pt

app.py

@@ -13,27 +13,24 @@ import os
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="Car Parts & Damage Detection API")
+app = FastAPI(title="Damage Detection API")
 
 # Log model file presence
-model_files = ["car_part_detector_model.pt", "damage_general_model.pt"]
-for model_file in model_files:
-    if os.path.exists(model_file):
-        logger.info(f"Model file found: {model_file}")
-    else:
-        logger.error(f"Model file missing: {model_file}")
+model_file = "/home/mussie/Documents/damage_detection_1/damage_general_model.pt"
+if os.path.exists(model_file):
+    logger.info(f"Model file found: {model_file}")
+else:
+    logger.error(f"Model file missing: {model_file}")
+    raise RuntimeError(f"Model file missing: {model_file}")
 
-# Load YOLO models
+# Load YOLO model
 try:
-    logger.info("Loading car part model...")
-    car_part_model = YOLO("car_part_detector_model.pt")
-    logger.info("Car part model loaded successfully")
     logger.info("Loading damage model...")
-    damage_model = YOLO("damage_general_model.pt")
+    damage_model = YOLO(model_file)
     logger.info("Damage model loaded successfully")
 except Exception as e:
-    logger.error(f"Failed to load models: {str(e)}")
-    raise RuntimeError(f"Failed to load models: {str(e)}")
+    logger.error(f"Failed to load model: {str(e)}")
+    raise RuntimeError(f"Failed to load model: {str(e)}")
 
 def image_to_base64(img: np.ndarray) -> str:
     """Convert numpy image to base64 string."""
@@ -44,9 +41,9 @@ def image_to_base64(img: np.ndarray) -> str:
         logger.error(f"Error encoding image to base64: {str(e)}")
         raise
 
-@app.post("/predict", summary="Run inference on an image for car parts and damage")
+@app.post("/predict", summary="Run inference on an image for damage detection")
 async def predict(file: UploadFile = File(...)):
-    """Upload an image and get car parts and damage detection results."""
+    """Upload an image and get damage detection results."""
     logger.info("Received image upload")
     try:
         contents = await file.read()
@@ -55,25 +52,9 @@ async def predict(file: UploadFile = File(...)):
         logger.info(f"Image loaded: shape={img.shape}")
 
         blank_img = np.full((img.shape[0], img.shape[1], 3), 128, dtype=np.uint8)
-        car_part_img = blank_img.copy()
         damage_img = blank_img.copy()
-        car_part_text = "Car Parts: No detections"
         damage_text = "Damage: No detections"
 
-        try:
-            logger.info("Running car part detection...")
-            car_part_results = car_part_model(img)[0]
-            if car_part_results.boxes:
-                car_part_img = car_part_results.plot()[..., ::-1]
-                car_part_text = "Car Parts:\n" + "\n".join(
-                    f"- {car_part_results.names[int(cls)]} ({conf:.2f})"
-                    for conf, cls in zip(car_part_results.boxes.conf, car_part_results.boxes.cls)
-                )
-            logger.info("Car part detection completed")
-        except Exception as e:
-            car_part_text = f"Car Parts: Error: {str(e)}"
-            logger.error(f"Car part detection error: {str(e)}")
-
         try:
             logger.info("Running damage detection...")
             damage_results = damage_model(img)[0]
@@ -88,12 +69,9 @@ async def predict(file: UploadFile = File(...)):
             damage_text = f"Damage: Error: {str(e)}"
             logger.error(f"Damage detection error: {str(e)}")
 
-        car_part_img_base64 = image_to_base64(car_part_img)
         damage_img_base64 = image_to_base64(damage_img)
         logger.info("Returning prediction results")
         return JSONResponse({
-            "car_part_image": car_part_img_base64,
-            "car_part_text": car_part_text,
             "damage_image": damage_img_base64,
             "damage_text": damage_text
         })
@@ -105,4 +83,4 @@ async def predict(file: UploadFile = File(...)):
 async def root():
     """Check if the API is running."""
     logger.info("Health check accessed")
-    return {"message": "Car Parts & Damage Detection API is running"}
+    return {"message": "Damage Detection API is running"}

car_part_detector_model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:919ed7512ffbee9013a1384c467c605e1b5beaa21319db6a306b5d0aa4180e9e
-size 6010902