| … | … | @@ -0,0 +1,295 @@ |
| 1 | import bpy, socket, math, ast, threading, queue |
| 2 | from enum import Enum |
| 3 | |
| 4 | ARM_NAME_PREFIX = 'arm.' |
| 5 | LIGHT_NAME_PREFIX = 'light.' |
| 6 | #QUAD_NAME_PREFIX = 'quad.' |
| 7 | BGR_PREFIX = 'bgr-' |
| 8 | LIGHT_PWM_MAX = 255 |
| 9 | LIGHT_PWM_RED_MAX = 255 |
| 10 | LIGHT_LUM_MAX = 1.0 |
| 11 | |
| 12 | ConnState = Enum('ConnState', 'none connecting connected') |
| 13 | class Connection: |
| 14 | def __init__(self, socket, state=ConnState.none): |
| 15 | self.socket = socket |
| 16 | self.state = state |
| 17 | |
| 18 | # map from IP address to Connection |
| 19 | sockets = {} |
| 20 | quad_last_command_dict = {} |
| 21 | # ip addresses to be connected to and added to sockets dict |
| 22 | sockets_queue = queue.Queue() |
| 23 | |
| 24 | # connect to address and store socket in persistent global map |
| 25 | def connect_socket(ip, port = 1337): |
| 26 | server_address = (ip, port) |
| 27 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 28 | sock.settimeout(1.0) |
| 29 | |
| 30 | print('connecting to', server_address) |
| 31 | sockets[ip].state = ConnState.connecting |
| 32 | try: |
| 33 | sock.connect(server_address) |
| 34 | print('connected to', server_address) |
| 35 | sock.sendall(b'speed 200gZ78\n') |
| 36 | sockets[ip].socket = sock |
| 37 | sockets[ip].state = ConnState.connected |
| 38 | return sock |
| 39 | |
| 40 | except socket.timeout: |
| 41 | print('error: connection timed out on', ip) |
| 42 | |
| 43 | class SocketThread(threading.Thread): |
| 44 | def __init__(self, name): |
| 45 | self.shouldExit = False |
| 46 | threading.Thread.__init__(self, name=name, daemon=True) |
| 47 | self.start() |
| 48 | |
| 49 | def run(self): |
| 50 | while True: |
| 51 | if self.shouldExit: return |
| 52 | |
| 53 | ip = sockets_queue.get() |
| 54 | if sockets[ip].state != ConnState.connected: |
| 55 | connect_socket(ip) |
| 56 | else: |
| 57 | print('skipping', ip) |
| 58 | sockets_queue.task_done() |
| 59 | |
| 60 | socket_thread = SocketThread('st') |
| 61 | |
| 62 | # find the LAMP attached to this arm |
| 63 | def getLamp(arm): |
| 64 | for child in arm.children: |
| 65 | if child.type == 'LAMP': |
| 66 | return child |
| 67 | |
| 68 | # returns robot arm LED intensity in integer units |
| 69 | def pwmFromLuminosity(lum, pwm_max=LIGHT_PWM_MAX): |
| 70 | pwm = int(round(abs(lum) / LIGHT_LUM_MAX * pwm_max)) |
| 71 | return pwm |
| 72 | |
| 73 | def luminosityFromPWM(pwm, pwm_max=LIGHT_PWM_MAX): |
| 74 | lum = pwm * LIGHT_LUM_MAX / pwm_max |
| 75 | return lum |
| 76 | |
| 77 | # returns a list of arm angles, starting from the base |
| 78 | # units are dynamixel units [200, 824] |
| 79 | def getAngles(arm): |
| 80 | |
| 81 | def dynamixel_from_degrees(angle): |
| 82 | return int(round((angle - 90) * 600.0/180 + 512)) |
| 83 | |
| 84 | # convert to degrees because they are easier to debug |
| 85 | rad_to_deg = 180.0/math.pi |
| 86 | up = arm.pose.bones['base'] |
| 87 | fo = arm.pose.bones['forearm'] |
| 88 | |
| 89 | up_q = up.matrix.to_quaternion() |
| 90 | fo_q = fo.matrix.to_quaternion() |
| 91 | fo_q.rotate(up_q.inverted()) |
| 92 | |
| 93 | # TODO get local orientation somehow instead of adding 180 and 90 |
| 94 | up_a = up.matrix.to_euler().x * rad_to_deg + 180 |
| 95 | fo_a = fo_q.to_euler().z * rad_to_deg + 90 |
| 96 | #up_a = 0 if up_a < 0 else up_a |
| 97 | #print(fo_q.to_euler(), fo_a) |
| 98 | |
| 99 | # convert to dynamixel units |
| 100 | angles = [dynamixel_from_degrees(a) for a in [up_a, fo_a]] |
| 101 | #print(angles) |
| 102 | return angles |
| 103 | |
| 104 | '''upperarm_angles = upperarm.matrix.to_euler() |
| 105 | forearm_angles = (upperarm.matrix.inverted() * forearm.matrix).to_euler() |
| 106 | print('upperarm angles:', upperarm_angles) |
| 107 | print('forearm angles: ', forearm_angles) |
| 108 | s |
| 109 | #angles[0] = arm.pose.bones[0].rotation_quaternion.to_euler()[0] |
| 110 | #angles[1] = arm.pose.bones[1].rotation_quaternion.to_euler()[1] |
| 111 | |
| 112 | angles = [x * 600/math.pi + 512 for x in angles] |
| 113 | ''' |
| 114 | |
| 115 | def robot_anim_handler(scene): |
| 116 | # TODO alter serial protocol to use negative IDs |
| 117 | #baseID = 1 |
| 118 | |
| 119 | # dict mapping ip to (dict mapping 1-based index to pwm string)s |
| 120 | quad_pwm_dicts = {} |
| 121 | |
| 122 | # find light arm objects |
| 123 | for o in bpy.data.objects: |
| 124 | if o.name.startswith(ARM_NAME_PREFIX): |
| 125 | ip = o.name[len(ARM_NAME_PREFIX):] |
| 126 | pwm = pwmFromLuminosity(getLamp(o).data.energy) |
| 127 | angles = getAngles(o) |
| 128 | |
| 129 | # assemble command string |
| 130 | s = 's' |
| 131 | for i in range(len(angles)): |
| 132 | s += ' ' + str(i+baseID) + ':' + str(angles[i]) |
| 133 | |
| 134 | s += '\nc ' + str(pwm) + '\n' |
| 135 | sendString(ip, s) |
| 136 | |
| 137 | elif o.name.startswith(LIGHT_NAME_PREFIX): |
| 138 | ip = o.name[len(LIGHT_NAME_PREFIX):] |
| 139 | pwm = pwmFromLuminosity(o.data.energy) |
| 140 | |
| 141 | # assemble command string |
| 142 | s = 'c ' + str(pwm) + '\n' |
| 143 | sendString(ip, s) |
| 144 | |
| 145 | elif o.name.startswith(BGR_PREFIX): |
| 146 | tokens = o.name.split('-') |
| 147 | index = int(tokens[1]) |
| 148 | ip = tokens[2] |
| 149 | |
| 150 | # these will be appended together |
| 151 | ss = str(pwmFromLuminosity(o.data.color[0])) + ' ' |
| 152 | ss += str(pwmFromLuminosity(o.data.color[1])) + ' ' |
| 153 | ss += str(pwmFromLuminosity(o.data.color[2], LIGHT_PWM_RED_MAX)) + ' ' |
| 154 | |
| 155 | try: pwm_dict = quad_pwm_dicts[ip] |
| 156 | except: |
| 157 | pwm_dict = {} |
| 158 | quad_pwm_dicts[ip] = pwm_dict |
| 159 | pwm_dict[index] = ss |
| 160 | |
| 161 | for ip,pwm_dict in quad_pwm_dicts.items(): |
| 162 | # assemble command string in order, starting with index 1 |
| 163 | s = 'c ' |
| 164 | for i in range(1, 5): |
| 165 | try: s += pwm_dict[i] |
| 166 | except: s += '0 0 0 ' # not found--use zeros for this light |
| 167 | s += '\n' |
| 168 | |
| 169 | # save the last command sent to an IP address |
| 170 | # sending commands every 30 times a second overwhelms the ESP32 |
| 171 | #global quad_last_command_dict |
| 172 | try: |
| 173 | if s in quad_last_command_dict[ip]: continue |
| 174 | except: pass |
| 175 | quad_last_command_dict[ip] = s |
| 176 | print(s) |
| 177 | sendString(ip, s) |
| 178 | |
| 179 | def sendString(ip, s): |
| 180 | if not ip: return |
| 181 | print(s.replace('\n', '_')) |
| 182 | print(ip + ": " + s) |
| 183 | # if this is the first time trying to talk to this device, create an entry |
| 184 | # and connect |
| 185 | conn = None |
| 186 | try: |
| 187 | conn = sockets[ip] |
| 188 | except KeyError: |
| 189 | print('not connected; queueing', ip) |
| 190 | sockets[ip] = Connection(None) |
| 191 | sockets_queue.put(ip) |
| 192 | |
| 193 | if conn and conn.state == ConnState.connected: |
| 194 | try: |
| 195 | print(s) |
| 196 | conn.socket.sendall(bytes(s, 'UTF-8')) |
| 197 | except (socket.timeout, BrokenPipeError, ConnectionResetError) as e: |
| 198 | print('lost connection; queueing', ip) |
| 199 | conn.state = ConnState.none |
| 200 | sockets_queue.put(ip) |
| 201 | |
| 202 | def saveToFile(filepath): |
| 203 | with open(filepath, 'w') as f: |
| 204 | arms = {} |
| 205 | |
| 206 | for arm in bpy.data.objects: |
| 207 | if not arm.name.startswith(ARM_NAME_PREFIX): continue |
| 208 | |
| 209 | ip = arm.name[len(ARM_NAME_PREFIX):] |
| 210 | angles = getAngles(arm) |
| 211 | light = pwmFromLuminosity(getLamp(arm).data.energy) |
| 212 | |
| 213 | dict = {'angles': angles, 'light': light} |
| 214 | arms[ip] = dict |
| 215 | |
| 216 | text = "{'LightArms': \n " + str(arms) + "\n}\n" |
| 217 | f.write(text) |
| 218 | |
| 219 | def loadFromFile(filepath): |
| 220 | with open(filepath, 'r') as f: |
| 221 | text = f.read() |
| 222 | json = ast.literal_eval(text) |
| 223 | arms = json['LightArms'] |
| 224 | |
| 225 | for ip, data in arms.items(): |
| 226 | name = ARM_NAME_PREFIX + ip |
| 227 | arm = bpy.data.objects[name] |
| 228 | |
| 229 | pwm = data['light'] |
| 230 | getLamp(arm).data.energy = luminosityFromPWM(pwm) |
| 231 | |
| 232 | angles = data['angles'] |
| 233 | # TODO_set angles |
| 234 | |
| 235 | ################################################################################### |
| 236 | # Menu in toolprops region |
| 237 | class RobotPanel(bpy.types.Panel): |
| 238 | bl_label = "Robot Scene" |
| 239 | bl_space_type = "VIEW_3D" |
| 240 | bl_region_type = "TOOL_PROPS" |
| 241 | |
| 242 | def draw(self, context): |
| 243 | self.layout.operator("circus.saver", text='Save Robot Positions') |
| 244 | #self.layout.operator("circus.loader", text='Load Robot Positions') |
| 245 | |
| 246 | class CircusSaver(bpy.types.Operator): |
| 247 | """Export Robot Data""" |
| 248 | bl_idname = "circus.saver" |
| 249 | bl_label = "Save the robots" |
| 250 | filepath = bpy.props.StringProperty(subtype="FILE_PATH") |
| 251 | |
| 252 | @classmethod |
| 253 | def poll(cls, context): |
| 254 | return context.object is not None |
| 255 | |
| 256 | def execute(self, context): |
| 257 | saveToFile(self.filepath) |
| 258 | return {'FINISHED'} |
| 259 | |
| 260 | def invoke(self, context, event): |
| 261 | context.window_manager.fileselect_add(self) |
| 262 | return {'RUNNING_MODAL'} |
| 263 | |
| 264 | class CircusLoader(bpy.types.Operator): |
| 265 | """Export Robot Data""" |
| 266 | bl_idname = "circus.loader" |
| 267 | bl_label = "Load the robots" |
| 268 | filepath = bpy.props.StringProperty(subtype="FILE_PATH") |
| 269 | |
| 270 | @classmethod |
| 271 | def poll(cls, context): |
| 272 | return context.object is not None |
| 273 | |
| 274 | def execute(self, context): |
| 275 | loadFromFile(self.filepath) |
| 276 | return {'FINISHED'} |
| 277 | |
| 278 | def invoke(self, context, event): |
| 279 | context.window_manager.fileselect_add(self) |
| 280 | return {'RUNNING_MODAL'} |
| 281 | |
| 282 | def unregister(): |
| 283 | bpy.utils.unregister_class(RobotPanel) |
| 284 | |
| 285 | if __name__ == "__main__": |
| 286 | # re-register this handler |
| 287 | bpy.app.handlers.frame_change_post.clear() |
| 288 | bpy.app.handlers.frame_change_post.append(robot_anim_handler) |
| 289 | |
| 290 | # create menu |
| 291 | #bpy.utils.register_module(__name__) |
| 292 | bpy.utils.register_class(RobotPanel) |
| 293 | bpy.utils.register_class(CircusLoader) |
| 294 | bpy.utils.register_class(CircusSaver) |
| 295 | |