# 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()