| import os |
|
|
|
|
| DAMAGE_CHECK_PARTS = { |
| 'driver_front_side', |
| 'driver_rear_side', |
| 'passenger_front_side', |
| 'passenger_rear_side', |
| } |
|
|
|
|
| def _class_name(model, class_index): |
| """Read a class label from YOLO names whether it is a dict or a list.""" |
| names = model.names |
| if isinstance(names, dict): |
| return names.get(class_index, 'unknown') |
| if 0 <= class_index < len(names): |
| return names[class_index] |
| return 'unknown' |
|
|
|
|
| def run_single_inference(model, filepath): |
| """Run one YOLO model on one image and normalize the top prediction.""" |
| if model is None: |
| raise RuntimeError('Inference model is not loaded.') |
|
|
| results = model(filepath, verbose=False) |
| if not results: |
| return {'class': 'unknown', 'confidence': 0.0} |
|
|
| result = results[0] |
|
|
| if result.probs is not None: |
| class_index = int(result.probs.top1) |
| confidence = float(result.probs.top1conf) |
| class_name = _class_name(model, class_index) |
| elif result.boxes is not None and len(result.boxes) > 0: |
| class_index = int(result.boxes.cls[0].item()) |
| confidence = float(result.boxes.conf[0].item()) |
| class_name = _class_name(model, class_index) |
| else: |
| class_name = 'unknown' |
| confidence = 0.0 |
|
|
| return { |
| 'class': class_name, |
| 'confidence': round(confidence, 4), |
| } |
|
|
|
|
| def process_images(parts_model, damage_model, image_inputs): |
| """ |
| Process uploaded images while preserving each browser-visible filename/index. |
| |
| image_inputs accepts dictionaries with path, filename, and index. Plain path |
| strings are also accepted for compatibility with older callers. |
| """ |
| 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 fallback_index, image_input in enumerate(image_inputs): |
| if isinstance(image_input, dict): |
| filepath = image_input['path'] |
| filename = image_input.get('filename') or os.path.basename(filepath) |
| client_index = image_input.get('index', fallback_index) |
| else: |
| filepath = image_input |
| filename = os.path.basename(filepath) |
| client_index = fallback_index |
|
|
| print(f'Processing {filename}...') |
|
|
| part_prediction = run_single_inference(parts_model, filepath) |
| predicted_part = part_prediction['class'] |
|
|
| 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'.") |
| damage_prediction = { |
| 'class': 'correct', |
| 'confidence': 1.0, |
| 'note': 'Result by default, not by model inference.', |
| } |
|
|
| final_results.append({ |
| 'index': client_index, |
| 'filename': filename, |
| 'part_prediction': part_prediction, |
| 'damage_prediction': damage_prediction, |
| }) |
|
|
| return final_results |
|
|
