Color Sorter - Code

Return back to project page: Color Sorter - Jakub Vojtek

Python code for the Color Sorter project:

 
import cv2
from util import get_limits
import math
from buildhat import Motor
import time

colors = {'green': [0, 255, 0], 'yellow': [0, 255, 255], 'red': [0, 0, 255], 'blue': [255, 0, 0]}
current_color = 'None'

left_box = Motor('A')
right_box = Motor('B')
separator = Motor('C')

left_side = ('red', 'green')
right_side = ('blue', 'yellow')


def separate(direction):
    global separator
    if direction == 'left':
        separator.run_for_degrees(-360)
        print("separating left")
    else:
        separator.run_for_degrees(360)
        print("separating right")


def rotate(side):
    global left_box, right_box
    if side == 'left':
        left_box.run_for_degrees(180)
        print("turning left box")
    else:
        right_box.run_for_degrees(180)
        print("turning right box")


def set_current_color():
    cap = cv2.VideoCapture(0)
    print("generating")
    global current_color
    ret, frame = cap.read()
    if not ret:
        return None

    height, width, _ = frame.shape

    upper_half = height // 2
    third_width = width // 3
    middle_third_start = width // 3
    middle_third_end = 2 * width // 3

    frame = frame[upper_half:, middle_third_start:middle_third_end]

    hsvImage = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    closest_contour_color = None
    closest_distance = float('inf')

    for color_name, color_value in colors.items():
        lowerLimit, upperLimit = get_limits(color=color_value)

        mask = cv2.inRange(hsvImage, lowerLimit, upperLimit)

        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 10000]

        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            cx, cy = x + w // 2, y + h // 2

            distance = math.sqrt((cx - third_width // 2) ** 2 + (cy - (height - upper_half) // 2) ** 2)
            if distance < closest_distance:
                closest_distance = distance
                closest_contour_color = color_name

    if closest_contour_color:
        current_color = closest_contour_color
    else:
        current_color = None
    print("last color: ", current_color)
    cap.release()
    cv2.destroyAllWindows()


def sort_logic():
    print("sorting")
    global current_color, left_side, right_side
    if current_color in left_side:
        if current_color == left_side[0]:
            separate('left')
        else:
            rotate('left')
            separate('left')
            left_side = left_side[::-1]
    elif current_color in right_side:
        if current_color == right_side[0]:
            separate('right')
        else:
            rotate('right')
            separate('right')
            right_side = right_side[::-1]
    time.sleep(1)


if __name__ == '__main__':
    try:
        while True:
            set_current_color()
            sort_logic()
            time.sleep(1)  # Add a delay to avoid excessive looping
    except KeyboardInterrupt:
        print("Sorting interrupted by user.")
    finally:
        left_box.stop()
        right_box.stop()
        separator.stop()
        left_box = None
        right_box = None
        separator = None