File size: 2,625 Bytes
f01c86d |
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 |
# Inference on unseen images for YOLOv8 parts segmentation
from ultralytics import YOLO
import os
from glob import glob
def run_inference(): # Get absolute paths
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
model_path = os.path.join(base_dir, 'models', 'parts', 'weights', 'weights', 'best.pt')
img_dir = os.path.join(base_dir, 'damage_detection_dataset', 'img')
train_dir = os.path.join(base_dir, 'data', 'data_yolo_for_training', 'car_parts_damage_dataset', 'images', 'train')
out_dir = os.path.join(base_dir, 'inference_results', 'parts')
# Validate paths
if not os.path.exists(model_path):
print(f"Error: Model weights not found at {model_path}")
return
if not os.path.exists(img_dir):
print(f"Error: Image directory not found at {img_dir}")
return
if not os.path.exists(train_dir):
print(f"Warning: Training directory not found at {train_dir}")
print("Will run inference on all images instead of just unseen ones")
train_imgs = set()
else:
# Get all images used for training
train_imgs = set(os.listdir(train_dir))
# Create output directory
os.makedirs(out_dir, exist_ok=True)
# Get all images in original dataset
all_imgs = set(os.listdir(img_dir))
# Select images not used in training
unseen_imgs = sorted(list(all_imgs - train_imgs))
if not unseen_imgs:
print(f"No images found for inference in {img_dir}")
return
try:
# Load model
model = YOLO(model_path)
# Class names for visualization
class_names = ['headlamp', 'front_bumper', 'hood', 'door', 'rear_bumper']
# Run inference on each unseen image
for img_name in unseen_imgs:
try:
img_path = os.path.join(img_dir, img_name)
results = model.predict(
source=img_path,
save=True,
project=out_dir,
name='',
imgsz=640,
conf=0.25,
classes=list(range(len(class_names))) # All classes
)
print(f'Processed: {img_name}')
except Exception as e:
print(f"Error processing {img_name}: {str(e)}")
continue
print(f'Inference complete. Results saved to {out_dir}')
except Exception as e:
print(f"Error loading model: {str(e)}")
return
if __name__ == '__main__':
run_inference()
|