Spaces:
Sleeping
Sleeping
| # processing.py | |
| import os | |
| from ultralytics import YOLO | |
| # --- Configuration --- | |
| # These are the specific parts that require a subsequent damage check. | |
| DAMAGE_CHECK_PARTS = { | |
| 'driver_front_side', | |
| 'driver_rear_side', | |
| 'passenger_front_side', | |
| 'passenger_rear_side', | |
| } | |
| def run_single_inference(model, filepath): | |
| """ | |
| Helper function to run inference for a single model and format the result. | |
| """ | |
| if model is None: | |
| return None # Return None if the model isn't loaded | |
| results = model(filepath, verbose=False) # verbose=False to keep logs clean | |
| result = results[0] | |
| # Check if it's a classification model with probabilities | |
| if result.probs is not None: | |
| probs = result.probs | |
| top1_index = probs.top1 | |
| top1_confidence = float(probs.top1conf) | |
| class_name = model.names[top1_index] | |
| else: # Fallback for detection models or if probs are not available | |
| # Assuming the top prediction is what we need | |
| top1_index = result.boxes.cls[0].int() if len(result.boxes) > 0 else 0 | |
| top1_confidence = float(result.boxes.conf[0]) if len(result.boxes) > 0 else 0.0 | |
| class_name = model.names[top1_index] if len(result.boxes) > 0 else "unknown" | |
| return { | |
| "class": class_name, | |
| "confidence": round(top1_confidence, 4) | |
| } | |
| def process_images(parts_model, damage_model, image_paths): | |
| """ | |
| Processes a list of images. | |
| 1. Runs the 'parts_model' on every image. | |
| 2. If the detected part is in DAMAGE_CHECK_PARTS, it then runs the 'damage_model'. | |
| 3. Otherwise, the damage status defaults to 'correct'. | |
| """ | |
| if parts_model is None or damage_model is None: | |
| raise RuntimeError("One or more models are not loaded. Check server logs.") | |
| final_results = [] | |
| for filepath in image_paths: | |
| filename = os.path.basename(filepath) | |
| print(f"Processing {filename}...") | |
| # 1. First, predict the part | |
| part_prediction = run_single_inference(parts_model, filepath) | |
| predicted_part = part_prediction.get("class") if part_prediction else "unknown" | |
| damage_prediction = None | |
| # 2. Conditionally predict the damage | |
| if predicted_part in DAMAGE_CHECK_PARTS: | |
| print(f" -> Part '{predicted_part}' requires damage check. Running damage model...") | |
| damage_prediction = run_single_inference(damage_model, filepath) | |
| else: | |
| print(f" -> Part '{predicted_part}' does not require damage check. Defaulting to 'correct'.") | |
| # 3. For other parts, default to 'correct' | |
| damage_prediction = { | |
| "class": "correct", | |
| "confidence": 1.0, | |
| "note": "Result by default, not by model inference." | |
| } | |
| # Assemble the final result for this image | |
| final_results.append({ | |
| "filename": filename, | |
| "part_prediction": part_prediction, | |
| "damage_prediction": damage_prediction | |
| }) | |
| return final_results |