First code generated byt chat GPT
import serial
import serial.tools.list_ports
import time
def determine_color(r, g, b):
return 'Red' #return Red by default as a placeholder
# TODO:
# write a function to discriminate between red, blue, yellow, and white marbles, based on real RBG readings
# Return a string 'Red', 'Blue', 'Yellow', or 'White'.
def find_arduino_serial_port():
ports = serial.tools.list_ports.comports()
for port in ports:
if 'Arduino' in port.description or 'usbmodem' in port.device:
return port.device
return None
def initialize_serial_connection(arduino_port):
try:
ser = serial.Serial(arduino_port, 9600, timeout=1)
time.sleep(1) # Wait for connection to establish
return ser
except serial.SerialException as e:
print(f"Error opening serial connection: {e}")
return None
def process_marble(ser):
# command hardware to input a marble, read the sensor and then receive the RGB:
r, g, b = read_sensor(ser)
if r is not None and g is not None and b is not None: #sanity check
print(f"RGB Values: R: {r}, G: {g}, B: {b}") #print to terminal for debugging and sampling of data
color = determine_color(r, g, b) #decide which color marble it is. This function should be tuned based on real world data samples
print(f"Detected color: {color}") #debug output
chute_number = get_chute_number(color)
print(f"Outputting marble to chute {chute_number}...")
send_command(ser, f"O{chute_number}")
time.sleep(0.5) #the chute servo takes some time to move, wait so as not to crowd the serial connection
print("Ready for next marble.")
else:
print("Failed to read sensor data.")
time.sleep(0.5)
def send_command(ser, command):
print(f"[Sent] {command}")
ser.write(f"{command}\n".encode())
def read_sensor(ser):
ser.reset_input_buffer()
send_command(ser, "I")
return wait_for_sensor_data(ser)
def wait_for_sensor_data(ser): #hardware servo takes some time to move, and sensor reading takes almost a second. Wait for a reply so as not to crown the serial.
start_time = time.time()
timeout = 5 # Timeout in seconds
while (time.time() - start_time) < timeout:
if ser.in_waiting:
line = ser.readline().decode().strip()
print(f"[Received] {line}")
if line.startswith("R:"): #the RGB data comes as a line like R: 300 G: 812 B: 642
parts = line.split(", ")
if len(parts) == 3:
r = int(parts[0].split(": ")[1])
g = int(parts[1].split(": ")[1])
b = int(parts[2].split(": ")[1])
return r, g, b
time.sleep(0.1)
return None, None, None
def get_chute_number(color):
return {
'Red': 1,
'Blue': 2,
'Yellow': 3,
'White': 4
}.get(color, 1)
def main():
arduino_port = find_arduino_serial_port()
if arduino_port is None:
print("Hardware not found. Please check the connection.")
return
ser = initialize_serial_connection(arduino_port)
if ser is None:
return
try:
while True:
process_marble(ser)
except KeyboardInterrupt:
print("Program interrupted by user")
finally:
if ser:
ser.close()
print("Serial connection closed.")
if __name__ == "__main__":
main()