Roboteksperimentarium/robot/arlo.py

from time import sleep

import cv2

from .robot import Robot

START_VALUES = [60, 60]
THRESHOLD = 1.05
SLEEP_TIME = 2

FOCAL_LENGTH = 0

def test_forward(arlo, l_power, r_power):
    arlo.reset_encoder_counts()
    arlo.go_diff(l_power, r_power, 1, 1)
    sleep(SLEEP_TIME)
    arlo.stop()
    return (
        abs(int(arlo.read_left_wheel_encoder())),
        abs(int(arlo.read_right_wheel_encoder()))
    )

def test_back(arlo, l_power, r_power):
    arlo.reset_encoder_counts()
    arlo.go_diff(l_power, r_power, 0, 0)
    sleep(SLEEP_TIME)
    arlo.stop()
    return (
        abs(int(arlo.read_left_wheel_encoder())),
        abs(int(arlo.read_right_wheel_encoder()))
    )

def test_clockwise(arlo, l_power, r_power):
    arlo.reset_encoder_counts()
    arlo.go_diff(l_power, r_power, 1, 0)
    sleep(SLEEP_TIME)
    arlo.stop()
    return (
        abs(int(arlo.read_left_wheel_encoder())),
        abs(int(arlo.read_right_wheel_encoder()))
    )

def test_anticlockwise(arlo, l_power, r_power):
    arlo.reset_encoder_counts()
    arlo.go_diff(l_power, r_power, 0, 1)
    sleep(SLEEP_TIME)
    arlo.stop()
    return (
        abs(int(arlo.read_left_wheel_encoder())),
        abs(int(arlo.read_right_wheel_encoder()))
    )

class Arlo():
    def __init__(self, calibration_code = None) -> None:
        self.robot = Robot()
        if calibration_code is None:
            self._calibrate()
        else:
            self.calibration_code = calibration_code
            self._decode_calibration_code()

    def forward(self, distance) -> None:
        left, right = tuple(self.wheel_values[0])
        self.robot.go_diff(left, right, 1, 1)
        sleep((distance * self.timing[0])/1000)
        self.robot.stop()

    def backwards(self, distance) -> None:
        left, right = tuple(self.wheel_values[1])
        self.robot.go_diff(left, right, 0, 0)
        sleep((distance * self.timing[1])/1000)
        self.robot.stop()

    def clockwise(self, angle) -> None:
        left, right = tuple(self.wheel_values[2])
        self.robot.go_diff(left, right, 1, 0)
        sleep((angle * self.timing[2])/1000)
        self.robot.stop()

    def anticlockwise(self, angle) -> None:
        left, right = tuple(self.wheel_values[3])
        self.robot.go_diff(left, right, 0, 1)
        sleep((angle * self.timing[3])/1000)
        self.robot.stop()

    def estimate_distance(self, bounding_box):
        average_left = bounding_box[0][0]

    def find_aruco(self, image):
        aruco_dict = cv2.aruco.Dictionary_get(cv2.aruco.DICT_6X6_250)
        aruco_params = cv2.aruco.DetectorParameters_create()
        corners, ids, _ = cv2.aruco.detectMarkers(
            image,
            aruco_dict,
            parameters=aruco_params
        )

        return zip(corners, ids)

    def draw_arucos(self, image, bounding_boxes):
        for bounding, n in bounding_boxes:
            bounding_box = [
                (int(x), int(y)) for x, y in bounding.reshape((4,2))
            ]
            print(bounding_box)

            cv2.line(image, bounding_box[0], bounding_box[1], (0,255,0), 2)
            cv2.line(image, bounding_box[1], bounding_box[2], (0,255,0), 2)
            cv2.line(image, bounding_box[2], bounding_box[3], (0,255,0), 2)
            cv2.line(image, bounding_box[3], bounding_box[0], (0,255,0), 2)

            cv2.putText(
                image,
                str(n[0]),
                (bounding_box[0][0] - 10, bounding_box[0][1] - 10),
                cv2.FONT_HERSHEY_COMPLEX,
                0.8,
                (0,255,0),
                2
            )

        return image

    def _decode_calibration_code(self) -> None:
        wheel_values_hex, timing_hex = tuple(self.calibration_code.split("_"))
        self.wheel_values = [
            [int(x, 16) for x in values.split(".")]
            for values in wheel_values_hex.split("-")
        ]
        self.timing = [int(x, 16) for x in timing_hex.split("-")]

    def _calibrate(self) -> None:
        values = [START_VALUES.copy() for _ in range(4)]
        cps = [0 for _ in range(4)]
        calibrated = [False for _ in range(4)]
        tests = [
            test_forward,
            test_back,
            test_clockwise,
            test_anticlockwise
        ]

        while not all(calibrated):
            print(calibrated, values)
            for i, function in enumerate(tests):
                calibrated[i] = False
                wheels = function(self.robot, *values[i])
                fraction = max(wheels)/min(wheels)
                if fraction <= THRESHOLD:
                    calibrated[i] = True
                elif wheels[0] < wheels[1]:
                    values[i][0] += 1
                else:
                    values[i][1] += 1

                cps[i] = wheels[0]/SLEEP_TIME

        time = [0 for _ in range(4)]

        cpc = 144/(3.14159*15) # wheel counts per cm

        # milliseconds per 10cm forward
        time[0] = int((1000 * cpc)/cps[0])
        # milliseconds per 10cm backwards
        time[1] = int((1000 * cpc)/cps[1])

        cpr = 3.1415*38 # 1 rotation in cm
        # milliseconds per 1 degrees clockwise
        time[2] = int((1000 * cpr * cpc)/(cps[2] * 360))
        # milliseconds per 1 degrees anticlockwise
        time[3] = int((1000 * cpr * cpc)/(cps[3] * 360))

        self.wheel_values = values
        self.timing = time

        values_hex = "-".join(
            [".".join([format(i, "x") for i in v]) for v in values]
        )
        time_hex = "-".join([format(i, "x") for i in time])
        self.calibration_code = f"{values_hex}_{time_hex}"
        print("Calibration code:", self.calibration_code)