File size: 3,219 Bytes
cd9ae77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | 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
|