Spaces:
Sleeping
Sleeping
| import os | |
| from fastapi import FastAPI, File, UploadFile | |
| from PIL import Image | |
| import io | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from ultralytics import YOLO | |
| # --- DEBUGGING STEP 1 --- | |
| print("Application script is starting...") | |
| os.environ['TORCH_HOME'] = '/tmp/torch_cache' | |
| app = FastAPI(title="YOLOv8 Car Damage Detection API") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| model = YOLO('best.pt') | |
| def read_root(): | |
| return {"message": "Welcome to the Car Damage Detection API!"} | |
| async def detect_damage(file: UploadFile = File(...)): | |
| # 1. Get image and run the model (same as before) | |
| contents = await file.read() | |
| image = Image.open(io.BytesIO(contents)).convert("RGB") | |
| results = model(image) | |
| # 2. Process the results to determine unique conditions | |
| # Use a set to automatically handle duplicates (e.g., "broken" is only added once) | |
| conditions = set() | |
| # Define which detections count as "broken" or "dirty" | |
| broken_classes = {'scratch', 'rust', 'dent'} | |
| dirty_classes = {'dirt'} | |
| # Loop through all detected objects | |
| for result in results: | |
| for box in result.boxes: | |
| # Get the name and confidence of the detected object | |
| class_name = model.names[int(box.cls[0])] | |
| confidence = float(box.conf[0]) | |
| # Only consider detections with high confidence | |
| if confidence > 0.5: | |
| if class_name in broken_classes: | |
| conditions.add("Broken") | |
| elif class_name in dirty_classes: | |
| conditions.add("Dirty") | |
| # 3. Prepare the final output | |
| # If the conditions set is empty after checking all boxes, the car is clean | |
| if not conditions: | |
| final_output = ["No damage detected."] | |
| else: | |
| # Convert the set to a list for the JSON response | |
| final_output = list(conditions) | |
| # Return the simple list of conditions | |
| return {"conditions": final_output} | |
| # --- DEBUGGING STEP 2 --- | |
| # This will print all routes that FastAPI has successfully registered. | |
| print("\n--- Registered Routes ---") | |
| for route in app.routes: | |
| print(f"Path: {route.path}, Methods: {', '.join(route.methods)}") | |
| print("-------------------------\n") |