python sortify_main.py

# Date: 03/09/2025
#Author: Rosey

#Aim: To use my custom dataset and my trained model to test and perform object detection
# A smart sorting bin which actually sorts out bins into various classes be it nylon,bottles or cans
# Data was collected: About 200 samples of plastic bottles, nylon and cans were snapped
#  
#Job Status: COMPLETED!

#import necessary libraries
import cv2
import serial
import time
from ultralytics import YOLO

# Load model
model = YOLO("best.pt")

# Connect to Arduino
arduino = serial.Serial('COM10', 9600)  # change COM10 to your correct port
time.sleep(2)  # wait for Arduino reset

# Open webcam
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # Run detection
    results = model.predict(frame, conf=0.6, verbose=False)

    for r in results:
        for box in r.boxes:
            cls = int(box.cls)   # fixed indexing
            label = model.names[cls].lower()

            command = None
            if label == "can":
                command = b'C'
            elif label == "nylon":
                command = b'N'
            elif label == "bottle":
                command = b'B'

            if command:
                print(f"Sending to Arduino: {command.decode()}  (Detected: {label})")
                arduino.write(command)

            # Draw on frame
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0,255,0), 2)
            cv2.putText(frame, label, (x1, y1-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)

    cv2.imshow("Sorter", frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
arduino.close()