pouring.py

''' 
Pouring over side data collection 
Alex Gillespie
08/03/2023
 '''

from xarm.wrapper import XArmAPI
import os
import serial
import numpy as np
import pickle
import time
import threading
import queue
arm = XArmAPI("192.168.1.199")
# TODO
# SPECIFY: the directory you're working in
current_dir = r'C:\Users\yexin\Desktop\liquid\data'

weight_data_queue = queue.Queue()
# sets up sensors


def linear_interpolation(full_time, time, weight):
    interpolated_weights = np.interp(full_time, time, weight)
    return interpolated_weights


def collect_weight_data(serial_weight, stop_event):
    while not stop_event.is_set():
        line = serial_weight.readline()
        try:
            decoded_line = line.decode('ascii').split(',')[0]
            if decoded_line.split(":")[0] == "grams":
                float_value = float(decoded_line.split(":")[1].replace("\r\n", ""))
                weight_data_queue.put(float_value)
        except UnicodeDecodeError:
            pass


def findFrame(ts_target, ts_set):
    idx = (np.abs(ts_target-ts_set)).argmin()
    return idx


def setup_sensors():
    # Serial initialization
    baudrate = 115200
    # TODO
    # SPECIFY: your serial ports for the Teensys
    port1 = 'COM5'
    port2 = 'COM3'
    port_weight = 'COM4'
    buffer1 =  [] # for values of teensy 1
    buffer2 =  [] # for values of teensy 2
    buffer_weight = [] # for values of weight from Arduino uno
    serialTeensy_1 = serial.Serial(port1, baudrate)
    serialTeensy_2 = serial.Serial(port2, baudrate)
    serial_weight = serial.Serial(port_weight, 57600) #57600
    # checking on Teensy connections and resetting input buffer
    if serialTeensy_1 is None:
        raise RuntimeError('Serial Port is not found!')
    serialTeensy_1.reset_input_buffer()
    if serialTeensy_2 is None:
        raise RuntimeError('Serial Port is not found!')
    serialTeensy_2.reset_input_buffer()
    if serial_weight is None:
        raise RuntimeError('Serial Port is not found!')
    serial_weight.reset_input_buffer()
    return buffer1, buffer2, buffer_weight, serialTeensy_1, serialTeensy_2, serial_weight

# data collection process for pouring
def collect_data_pour(container, content, labels, iteration, rotate_speed, rotate_angle, p_duration):
    buffer1, buffer2, buffer_weight, serialTeensy_1, serialTeensy_2, serial_weight = setup_sensors()
    weight_time = []
    
    time.sleep(5)
    #buffer1, buffer2, buffer_weight, serialTeensy_1, serialTeensy_2 = setup_sensors()
    count = 0
    flag_first_time = True
    stop_event = threading.Event()
    weight_thread = threading.Thread(target=collect_weight_data, args=(serial_weight, stop_event))
    weight_thread.start()
    
    arm.set_position(410, 16.7, 260, rotate_angle[0], rotate_angle[1], rotate_angle[2], speed=rotate_speed, is_radian=False, wait=False)
    start = time.time()
    count_1 = 0
    # while the arm is moving, collect data
    while time.time() - start <= p_duration:
        buffer1.append(serialTeensy_1.read())
        buffer2.append(serialTeensy_2.read())
        
        try:
            float_value = weight_data_queue.get_nowait()  # non-blocking get
            buffer_weight.append(float_value)

            weight_time.append(time.time() - start)
        except queue.Empty:
            pass

        if len(buffer1) >= 18 and len(buffer2) >= 18:
            # These bytes indicate a new series of sensor values 
            new_line1 = int.from_bytes(buffer1[0], byteorder='little')
            new_line2 = int.from_bytes(buffer1[1],byteorder='little')
            new_line3 = int.from_bytes(buffer1[2],byteorder='little')
            new_line4 = int.from_bytes(buffer2[0], byteorder='little')
            new_line5 = int.from_bytes(buffer2[1],byteorder='little')
            new_line6 = int.from_bytes(buffer2[2],byteorder='little')
            # If all the bytes are 255, the maximum value (3 max bytes in a row for each teensy), that's how we know we're reading the sensor values
            if new_line1 == 255 and new_line2 == 255 and new_line3 == 255 and new_line4 == 255 and new_line5 == 255 and new_line6 == 255:
                # These teensy values store the time at which the values were taken 
                teensy_1_time = int.from_bytes(buffer1[13]+buffer1[14]+buffer1[15]+buffer1[16],byteorder='little')
                teensy_2_time = int.from_bytes(buffer2[13]+buffer2[14]+buffer2[15]+buffer2[16],byteorder='little')
                duration = time.time()-start
                # headers: container, content, labels, iteration, observation (in iteration), python time passed, Teensy time for sensor 1, Teensy time for sensor 2, sensors 1-10
                sensor_vals = np.array([container, content, labels, iteration, count, duration, teensy_1_time, teensy_2_time,
                int.from_bytes(buffer1[3]+buffer1[4],byteorder='little'), 
                int.from_bytes(buffer1[5]+buffer1[6],byteorder='little'), 
                int.from_bytes(buffer1[7]+buffer1[8],byteorder='little'), 
                int.from_bytes(buffer1[9]+buffer1[10],byteorder='little'),
                int.from_bytes(buffer1[11]+buffer1[12],byteorder='little'),
                int.from_bytes(buffer2[3]+buffer2[4],byteorder='little'), 
                int.from_bytes(buffer2[5]+buffer2[6],byteorder='little'), 
                int.from_bytes(buffer2[7]+buffer2[8],byteorder='little'), 
                int.from_bytes(buffer2[9]+buffer2[10],byteorder='little'),
                int.from_bytes(buffer2[11]+buffer2[12],byteorder='little')])
                # if this is the first sensor reading of the iteration, it creates a new array
                if flag_first_time:
                    array_record = sensor_vals
                    flag_first_time = False
                # if this is not the first sensor reading of the iteration, it adds on the values to the existing array
                else:
                    array_record = np.vstack((array_record, sensor_vals))
                    '''
                    print('Frequency: ', count/(time.time()-start)) # you'll watch these in the terminal to make sure the times all look like they're consistent
                    print('Python time: ', (time.time()-duration))
                    print('Teensy duration 1: ', teensy_1_time-prev_teensy_1_time)
                    print('Teensy duration 2: ', teensy_2_time-prev_teensy_2_time)
                    '''
                prev_teensy_1_time = teensy_1_time
                prev_teensy_2_time = teensy_2_time
                count +=1
            buffer1.pop(0)
            buffer2.pop(0)

    stop_event.set()
    weight_thread.join()
    arm.set_state(4)
    arm.set_state(0)
    
    array_record = array_record[500:]
    
    idx = findFrame(float(array_record[0][5]), np.array(weight_time))
    
    weight_array = np.column_stack((weight_time, buffer_weight))
    weight_array = weight_array[idx:]
    weight_inter = linear_interpolation(array_record[:, 5].astype(np.float32), weight_array[:, 0], weight_array[:, 1])
    
    
    array_record_d = array_record[:, 8 : 18].astype(np.float32)
    if np.max(array_record_d) > 1200 or np.min(array_record_d) < 500:
        print("Wrong value was collected")
    return array_record, prev_teensy_1_time, prev_teensy_2_time, weight_inter

# collects data for holding portion 
def collect_data_hold(container, content, labels, iteration, hold_time, array_record, prev_teensy_1_time, prev_teensy_2_time):
    #buffer1, buffer2, buffer_weight, serialTeensy_1, serialTeensy_2, serial_weight = setup_sensors()
    buffer1, buffer2, buffer_weight, serialTeensy_1, serialTeensy_2 = setup_sensors()
    count = 0
    start = time.time()
    duration = 0
    while duration < hold_time:
        buffer1.append(serialTeensy_1.read())
        buffer2.append(serialTeensy_2.read())
        #weight_line = serial_weight.readline()
        # try:
        #     decoded_line = weight_line.decode('ascii').split(',')[0]
        #     if decoded_line.split(":")[0] == "grams":
        #         float_value = float(decoded_line.split(":")[1].replace("\r\n", ""))
        #         buffer_weight.append(float_value)
        #         #print(buffer_weight)
        # except UnicodeDecodeError:
        #     pass
        
        # each series of bytes is 17 bytes long, so both the left and right teensy need to be at least 18 bytes long to start reading it
        if len(buffer1) >= 18 and len(buffer2) >= 18:
            # These bytes indicate a new series of sensor values 
            new_line1 = int.from_bytes(buffer1[0], byteorder='little')
            new_line2 = int.from_bytes(buffer1[1],byteorder='little')
            new_line3 = int.from_bytes(buffer1[2],byteorder='little')
            new_line4 = int.from_bytes(buffer2[0], byteorder='little')
            new_line5 = int.from_bytes(buffer2[1],byteorder='little')
            new_line6 = int.from_bytes(buffer2[2],byteorder='little')
            # If all the bytes are 255, the maximum value (3 max bytes in a row for each teensy), that's how we know we're reading the sensor values
            if new_line1 == 255 and new_line2 == 255 and new_line3 == 255 and new_line4 == 255 and new_line5 == 255 and new_line6 == 255:
                # These teensy values store the time at which the values were taken 
                teensy_1_time = int.from_bytes(buffer1[13]+buffer1[14]+buffer1[15]+buffer1[16],byteorder='little')
                teensy_2_time = int.from_bytes(buffer2[13]+buffer2[14]+buffer2[15]+buffer2[16],byteorder='little')
                duration = time.time()-start
                # headers: container, content, labels, iteration, observation (in iteration), python time passed, Teensy time for sensor 1, Teensy time for sensor 2, sensors 1-10
                sensor_vals = np.array([container, content, labels, iteration, count, duration, teensy_1_time, teensy_2_time,
                int.from_bytes(buffer1[3]+buffer1[4],byteorder='little'), 
                int.from_bytes(buffer1[5]+buffer1[6],byteorder='little'), 
                int.from_bytes(buffer1[7]+buffer1[8],byteorder='little'), 
                int.from_bytes(buffer1[9]+buffer1[10],byteorder='little'),
                int.from_bytes(buffer1[11]+buffer1[12],byteorder='little'),
                int.from_bytes(buffer2[3]+buffer2[4],byteorder='little'), 
                int.from_bytes(buffer2[5]+buffer2[6],byteorder='little'), 
                int.from_bytes(buffer2[7]+buffer2[8],byteorder='little'), 
                int.from_bytes(buffer2[9]+buffer2[10],byteorder='little'),
                int.from_bytes(buffer2[11]+buffer2[12],byteorder='little')])
                array_record = np.vstack((array_record, sensor_vals))
                print('Frequency: ', count/(time.time()-start)) # you'll watch these in the terminal to make sure the times all look like they're consistent
                print('Python time: ', (time.time()-duration))
                print('Teensy duration 1: ', teensy_1_time-prev_teensy_1_time)
                print('Teensy duration 2: ', teensy_2_time-prev_teensy_2_time)
                print()
                prev_teensy_1_time = teensy_1_time
                prev_teensy_2_time = teensy_2_time
                count +=1
            buffer1.pop(0)
            buffer2.pop(0)
            # buffer_weight.pop(0)
    return array_record

# positions arm for pick up
def ready_for_pickup(arm):
    arm.set_gripper_position(840, wait = True)
    arm.set_position(494, 16.7, 215, 44.9, -87.8, -134.9, speed=150, is_radian=False, wait=True)
    arm.set_position(494, 29.1, 80.1, 44.9, -87.8, -134.9, speed=150, is_radian=False, wait=True)

# picks up and pours container
def pick_up_pour(arm, closed_pos, grip_height, iteration, container, content, labels, hold_time, rotate_speed, rotate_angle, grasp_height, p_duration):
    arm.set_position(494, 29.1, 40.1 + grip_height + grasp_height, 44.9, -87.8, -134.9, speed=50, is_radian=False, wait=True)
    arm.set_gripper_position(closed_pos, wait = True)
    input('press enter to confirm height')
    arm.set_position(494, 16.7, 260, 44.9, -87.8, -134.9, speed=50, is_radian=False, wait=True)
    arm.set_position(410, 16.7, 260, 44.9, -87.8, -134.9, speed=50, is_radian=False, wait=True)
    sensor_vals, prev_teensy_1_time, prev_teensy_2_time, weight_array = collect_data_pour(container, content, labels, iteration, rotate_speed, rotate_angle, p_duration)
    #array_record = collect_data_hold(container, content, labels, iteration, hold_time, sensor_vals, prev_teensy_1_time, prev_teensy_2_time)
    return sensor_vals, weight_array

# places container back
def place_back(arm, grip_height, rotate_speed, grasp_height):
    arm.set_position(410, 16.7, 260, 44.9, -87.8, -134.9, speed=rotate_speed, is_radian=False, wait=True)
    arm.set_position(494, 16.7, 260, 44.9, -87.8, -134.9, speed=150, is_radian=False, wait=True)
    arm.set_position(494, 29.1, 40.1+grip_height + grasp_height, 44.9, -87.8, -134.9, speed=150, is_radian=False, wait=True)
    arm.set_gripper_position(840, wait = True)
    arm.set_position(494, 16.7, 215, 44.9, -87.8, -134.9, speed=150, is_radian=False, wait=True)

# returns arm to resting position
def resting_pos(arm):
    arm.set_gripper_position(840, wait = True)
    arm.set_position(338.0, 16.7, 260, 44.9, -87.8, -134.9, speed=150, is_radian=False, wait=True)

# sets height and position for whatever container is being used
def set_height_and_pos(container):
    if container == 'ceramic':
        closed_pos = 660
        grip_height = 28
    elif container == 'plastic':
        closed_pos = 615
        grip_height = 25
    elif container == 'paper':
        closed_pos = 610
        grip_height = 18.2
    elif container == 'foam':
        closed_pos = 625
        grip_height = 0
    elif container == 'silicon':
        closed_pos = 650
        grip_height = 54.8
    elif container == 'glass':
        closed_pos = 730
        grip_height = 26.5
    elif container == 'wood':
        closed_pos = 630
        grip_height = 30.1
        
    return closed_pos, grip_height

# creates the directories
def create_dirs(current_dir):
    # All the data will go into a specified directory
    direct_label = input('what is the title for the final directory? Sugested: date_purpose, ex: MM.DD_test ')
    new_dir = 'data_collection_'+direct_label
    path = os.path.join(current_dir, new_dir)
    # TODO
    # OPTIONAL SPECIFY: If you're adding the files to a directory that already exists, comment this out
    os.mkdir(path)
    os.chdir(current_dir+'/'+new_dir)

# runs all of the code

def get_parameters(rotate_speed_list, rotate_angle_list, grasp_height_list):
    while True:        
        rotate_command = input("select a rotate speed level, please type 1, 2, 3 to specify the level ")
        if rotate_command == "1":
            rotate_speed = rotate_speed_list[0]
            break
        elif rotate_command == "2":
            rotate_speed = rotate_speed_list[1]
            break
        elif rotate_command == "3":
            rotate_speed = rotate_speed_list[2]
            break
        else:
            print("Wrong level, please select again")
    
    while True:
        rotate_angle_command = input("select a rotate angle level, please type high or low to sepcify the level ")
        if rotate_angle_command == "high":
            rotate_angle = rotate_angle_list[0]
            break
        elif rotate_angle_command == "low":
            rotate_angle = rotate_angle_list[1]
            break
        else:
            print("Wrong level, please select again")
    
    while True:
        grasp_height_command = input("select a grasp height level, please type high or low to sepcify the level ")
        if grasp_height_command == "high":
            grasp_height = grasp_height_list[0]
            break
        elif grasp_height_command == "low":
            grasp_height = grasp_height_list[1]
            break
        else:
            print("Wrong level, please select again")
            
    while True:
        fill_level = input("select a fill level of your current cup, please type high or low to specify the level ")
        if fill_level == "high":
            break
        elif fill_level == "low":
            break
        else:
            print("Wrong level, please select again")
            
    return rotate_speed, rotate_angle, grasp_height, rotate_command, rotate_angle_command, grasp_height_command, fill_level


def trial():
    # TODO 
    # SPECIFY: containers and contents 
    containers = ['plastic']
    contents = ['rice']
    rotate_speed_list = [50, 100, 150]
    rotate_angle_list = [[87.4, 52.5, 178], [88.2, 35, 179.2]]
    grasp_height_list = [-10, -15]
    num_its = 10
    hold_time = 3
    p_duration = 20
    labels = 0
    arm.set_gripper_speed(1000)
    create_dirs(current_dir)
    for container in containers:
        closed_pos, grip_height = set_height_and_pos(container)
        for content in contents:
            print('\ncontainer: '+container+' \ncontent: '+content+'\n')
            rotate_speed, rotate_angle, grasp_height, rotate_command, rotate_angle_command, grasp_height_command, fill_level = get_parameters(rotate_speed_list, rotate_angle_list, grasp_height_list)
            for iteration in range(num_its):     
                ready_for_pickup(arm)            
                input('\npress enter to begin pouring')
                data_1, data_2 = pick_up_pour(arm, closed_pos, grip_height, iteration, container, content, labels, hold_time, rotate_speed, rotate_angle, grasp_height, p_duration)
                data = [data_1, data_2]
                with open(content+'_'+container+'_'+rotate_command+'_'+rotate_angle_command+'_'+grasp_height_command+'_'+fill_level+'_'+str(iteration)+'.pkl','wb') as file:
                    pickle.dump(data, file)
                #input('press enter when done pouring\n')
                #time.sleep(pouring_duration)
                place_back(arm, grip_height, rotate_speed, grasp_height)
                if iteration != num_its-1:
                    input('please reset the containers and contents. \npress enter to continue')
                else:
                    resting_pos(arm)
        labels +=1
trial()