MagWire/collect_robot_data.py at main · dbiton/MagWire · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import json
from math import radians
import math
from time import sleep, time
from multiprocessing import Process
from hilbertcurve.hilbertcurve import HilbertCurve
from matplotlib import animation, pyplot as plt
from robot_controller import RobotController
import numpy as np
import os
try:
    import winsound
except ImportError:
    def playsound(frequency, duration):
        os.system(f'speaker-test -t sine -f {frequency}')
else:
    def playsound(frequency,duration):
        winsound.Beep(frequency,duration)

r = RobotController()

def hilbert_curve_3D(order):
    dimension = 3
    hilb = HilbertCurve(p=order, n=dimension)
    side_length = 2**order
    curve_points = []
    for dist in range(side_length**dimension):
        coords = hilb.point_from_distance(dist)
        x = coords[0] / (side_length - 1)
        y = coords[1] / (side_length - 1)
        z = coords[2] / (side_length - 1)
        curve_points.append((x, y, z))
    return curve_points


def trilinear_map_8pts(pt, corners):
    x, y, z = pt
    v000 = np.array(corners[0])
    v100 = np.array(corners[1])
    v010 = np.array(corners[2])
    v110 = np.array(corners[3])
    v001 = np.array(corners[4])
    v101 = np.array(corners[5])
    v011 = np.array(corners[6])
    v111 = np.array(corners[7])
    mapped = (
        (1 - x) * (1 - y) * (1 - z) * v000
        + x * (1 - y) * (1 - z) * v100
        + (1 - x) * y * (1 - z) * v010
        + x * y * (1 - z) * v110
        + (1 - x) * (1 - y) * z * v001
        + x * (1 - y) * z * v101
        + (1 - x) * y * z * v011
        + x * y * z * v111
    )
    return mapped

def space_filling_curve_3D(corners, order):
    hilbert_points = hilbert_curve_3D(order)
    mapped_points = []
    for pt in hilbert_points:
        mapped_pt = trilinear_map_8pts(pt, corners)
        mapped_points.append(tuple(mapped_pt))  # convert np.array -> tuple
    return mapped_points

def to_rad(p):
    return [radians(a) for a in p]

def log_data():
    with open("robot_pos.txt", "w") as f:
        while True:
            res = {"time": time(), "pos": r.get_config()}
            sleep(1 / 100)
            json.dump(res, f)

from footage_parser import FootageParser


if __name__ == "__main__":
    process_robot = Process(target=log_data)
    process_robot.start()

    # Bottom 4 corners
    v000 = (-700/1000,   0/1000,   40.0/1000)
    v100 = (-900/1000,   0/1000,   40.0/1000)
    v010 = (-900/1000,   200/1000,   40.0/1000)
    v110 = (-700/1000,   200/1000,   40.0/1000)

    # Top 4 corners
    v001 = (-700/1000,   0/1000,   40/1000)
    v101 = (-900/1000,   0/1000,   40/1000)
    v011 = (-900/1000,   200/1000,   40/1000)
    v111 = (-700/1000,   200/1000,   40/1000)

    corners = [v000, v100, v010, v110, v001, v101, v011, v111]

    acceleration = 0.2
    min_velocity = 0.1
    max_velocity = 0.8
    orientation = [0, math.pi, 0]


    steps_velocity = 8
    for order in range(1, 4):
      for velocity in np.linspace(min_velocity, max_velocity, steps_velocity):
        curve_points = space_filling_curve_3D(corners, order=order)
        curve_pts = np.array(curve_points)
        xs, ys, zs = curve_pts[:, 0], curve_pts[:, 1], curve_pts[:, 2]

        waypoints = []
        for curve_point in curve_points:
          waypoint = list(curve_point) + orientation + [velocity, acceleration]
          waypoints.append(waypoint)

        r.move_waypoints(waypoints)

        frequency = 432
        sleep(10)
    process.join()