MichalMlodawski/

open-closed-eye-detection

import os
from pathlib import Path
from ultralytics import YOLO
import cv2
import logging
import argparse

def setup_logging():
    logging.basicConfig(level=logging.INFO, 
                        format='%(asctime)s - %(levelname)s - %(message)s')

def process_images(model_path, test_images_path):
    try:
        # Path to the results directory
        results_path = os.path.join(test_images_path, 'result')
        
        # Create the results folder
        os.makedirs(results_path, exist_ok=True)
        logging.info(f'Created results directory: {results_path}')
        
        # Load the model
        model = YOLO(model_path)
        logging.info(f'Loaded model from: {model_path}')
        
        # Process images
        for img_file in Path(test_images_path).glob('*.*'):
            if img_file.suffix.lower() in ['.jpg', '.jpeg', '.png']:  # Supports JPG, JPEG, and PNG formats
                logging.info(f'Processing file: {img_file}')
                # Detect objects in the image
                results = model(img_file)
                
                for result in results:
                    # Get the result image with detections drawn
                    result_img = result.plot()
                    
                    # Save the result image to the results_path folder
                    result_image_path = os.path.join(results_path, img_file.name)
                    cv2.imwrite(result_image_path, result_img)
                    logging.info(f'Saved result image to: {result_image_path}')
        
        logging.info("Image processing completed.")
    except Exception as e:
        logging.error(f'An error occurred: {e}')

def main():
    parser = argparse.ArgumentParser(description='Process images using YOLO model.')
    parser.add_argument('model_path', type=str, help='Path to the YOLO model.')
    parser.add_argument('test_images_path', type=str, help='Path to the directory containing test images.')
    
    args = parser.parse_args()
    setup_logging()
    
    process_images(args.model_path, args.test_images_path)

if __name__ == "__main__":
    main()