Spaces:
Sleeping
Sleeping
File size: 3,119 Bytes
5c8508a |
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 |
# 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 |