detect_nuts_interactive.py

import os
from ultralytics import YOLO

def interactive_detection():
    # Path to the best trained model
    model_path = '/Users/riwasthapa/Desktop/LABELS/runs/detect/nut_detection2/weights/best.pt'
    
    if not os.path.exists(model_path):
        print(f"Error: Model weights not found at {model_path}")
        return

    # Load the model
    print("Loading model...")
    model = YOLO(model_path)

    while True:
        img_path = input("\nEnter the path to an image (or 'q' to quit): ").strip()
        
        if img_path.lower() == 'q':
            break
            
        if not os.path.exists(img_path):
            print(f"Error: File '{img_path}' does not exist.")
            continue

        print(f"Detecting nuts in {img_path}...")
        
        # Run inference
        results = model.predict(source=img_path, save=True, project='runs/inference', name='detections', exist_ok=True)
        
        # Display results summary
        for result in results:
            count = len(result.boxes)
            print(f"Successfully detected {count} nuts!")
            print(f"Saved result to: {result.save_dir}")

if __name__ == "__main__":
    interactive_detection()