| … | … | @@ -0,0 +1,399 @@ |
| 1 | import bpy, socket, math, threading, queue, json |
| 2 | from enum import Enum |
| 3 | from bpy.app.handlers import persistent |
| 4 | |
| 5 | # ----------------------- |
| 6 | # Constants |
| 7 | # ----------------------- |
| 8 | ARM_NAME_PREFIX = 'arm.' |
| 9 | LIGHT_NAME_PREFIX = 'light.' |
| 10 | BGR_PREFIX = 'bgr-' |
| 11 | LIGHT_PWM_MAX = 65534 |
| 12 | LIGHT_PWM_RED_MAX = 30000 |
| 13 | LIGHT_LUM_MAX = 1.0 |
| 14 | |
| 15 | # ----------------------- |
| 16 | # Networking |
| 17 | # ----------------------- |
| 18 | ConnState = Enum('ConnState', 'none connecting connected') |
| 19 | |
| 20 | import time as _time |
| 21 | |
| 22 | class Connection: |
| 23 | def __init__(self, sock=None, state=ConnState.none): |
| 24 | self.socket = sock |
| 25 | self.state = state |
| 26 | self.failures = 0 # consecutive failed connect attempts |
| 27 | self.next_retry_ts = 0.0 # wall-clock earliest retry; backoff target |
| 28 | |
| 29 | sockets = {} |
| 30 | quad_last_command_dict = {} |
| 31 | sockets_queue = queue.Queue() |
| 32 | |
| 33 | def connect_socket(ip, port=1337): |
| 34 | server_address = (ip, port) |
| 35 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 36 | sock.settimeout(3.0) # tolerate ESP32 hiccups during accept |
| 37 | print('Connecting to', server_address) |
| 38 | sockets[ip].state = ConnState.connecting |
| 39 | try: |
| 40 | sock.connect(server_address) |
| 41 | print('Connected to', server_address) |
| 42 | sock.sendall(b'speed 200gZ78\n') |
| 43 | sockets[ip].socket = sock |
| 44 | sockets[ip].state = ConnState.connected |
| 45 | sockets[ip].failures = 0 |
| 46 | sockets[ip].next_retry_ts = 0.0 |
| 47 | return sock |
| 48 | except (socket.timeout, OSError) as e: |
| 49 | # Reset state so future sendString calls re-queue. Apply exponential |
| 50 | # backoff so a dead unit doesn't generate a connect storm. |
| 51 | sockets[ip].failures += 1 |
| 52 | backoff = min(30.0, 2.0 ** min(sockets[ip].failures, 5)) |
| 53 | sockets[ip].next_retry_ts = _time.monotonic() + backoff |
| 54 | sockets[ip].state = ConnState.none |
| 55 | print(f'[ERROR] connection failed on {ip}: {e} (next retry in {backoff:.0f}s)') |
| 56 | try: sock.close() |
| 57 | except OSError: pass |
| 58 | |
| 59 | class SocketThread(threading.Thread): |
| 60 | def __init__(self): |
| 61 | self.shouldExit = False |
| 62 | threading.Thread.__init__(self, daemon=True) |
| 63 | self.start() |
| 64 | def run(self): |
| 65 | while True: |
| 66 | if self.shouldExit: return |
| 67 | ip = sockets_queue.get() |
| 68 | if sockets[ip].state != ConnState.connected: |
| 69 | connect_socket(ip) |
| 70 | sockets_queue.task_done() |
| 71 | |
| 72 | # Guard against module reload spawning orphan worker threads. |
| 73 | if 'socket_thread' not in globals(): |
| 74 | socket_thread = SocketThread() |
| 75 | |
| 76 | def sendString(ip, s): |
| 77 | if not ip: return |
| 78 | conn = sockets.get(ip) |
| 79 | if not conn: |
| 80 | sockets[ip] = Connection(None) |
| 81 | sockets_queue.put(ip) |
| 82 | return |
| 83 | if conn.state == ConnState.connected: |
| 84 | try: |
| 85 | conn.socket.sendall(bytes(s, 'UTF-8')) |
| 86 | print(f"[SEND] {ip} -> {s.strip()}") |
| 87 | except (socket.timeout, BrokenPipeError, ConnectionResetError, OSError) as e: |
| 88 | print(f"[ERROR] Lost connection to {ip}: {e}; re-queueing") |
| 89 | conn.state = ConnState.none |
| 90 | conn.failures = max(conn.failures, 1) |
| 91 | conn.next_retry_ts = _time.monotonic() # one immediate retry |
| 92 | sockets_queue.put(ip) |
| 93 | elif conn.state == ConnState.none: |
| 94 | # Honor backoff timer so dead IPs don't generate a connect storm. |
| 95 | if _time.monotonic() >= conn.next_retry_ts: |
| 96 | sockets_queue.put(ip) |
| 97 | |
| 98 | # ----------------------- |
| 99 | # Helpers |
| 100 | # ----------------------- |
| 101 | def getLamp(arm): |
| 102 | for child in arm.children: |
| 103 | if child.type == 'LIGHT': |
| 104 | return child |
| 105 | return None |
| 106 | |
| 107 | def pwmFromLuminosity(lum, pwm_max=LIGHT_PWM_MAX): |
| 108 | return int(round(abs(lum) / LIGHT_LUM_MAX * pwm_max)) |
| 109 | |
| 110 | def luminosityFromPWM(pwm, pwm_max=LIGHT_PWM_MAX): |
| 111 | return pwm * LIGHT_LUM_MAX / pwm_max |
| 112 | |
| 113 | def getAngles(arm): |
| 114 | def dynamixel_from_degrees(angle): |
| 115 | return int(round((angle - 90) * 600.0 / 180 + 512)) |
| 116 | rad_to_deg = 180.0 / math.pi |
| 117 | up = arm.pose.bones['base'] |
| 118 | fo = arm.pose.bones['forearm'] |
| 119 | up_q = up.matrix.to_quaternion() |
| 120 | fo_q = fo.matrix.to_quaternion() |
| 121 | fo_q.rotate(up_q.inverted()) |
| 122 | up_a = up.matrix.to_euler().x * rad_to_deg + 180 |
| 123 | fo_a = fo_q.to_euler().z * rad_to_deg + 90 |
| 124 | return [dynamixel_from_degrees(a) for a in [up_a, fo_a]] |
| 125 | |
| 126 | # ----------------------- |
| 127 | # Lazy Recompute Cache for Arms Only |
| 128 | # ----------------------- |
| 129 | last_arm_state = {} |
| 130 | quad_last_command_dict = {} |
| 131 | |
| 132 | # ----------------------- |
| 133 | # Playback / keyframe helpers |
| 134 | # ----------------------- |
| 135 | last_pwm2_target = {} # ip -> last 12-tuple of ints we sent in a pwm2 |
| 136 | |
| 137 | def is_playing(): |
| 138 | """True if any window is actively playing the timeline (not scrubbing).""" |
| 139 | wm = bpy.context.window_manager |
| 140 | if not wm: return False |
| 141 | for win in wm.windows: |
| 142 | if win.screen and win.screen.is_animation_playing: |
| 143 | return True |
| 144 | return False |
| 145 | |
| 146 | def frames_to_ms(scene, frames): |
| 147 | fps = scene.render.fps / scene.render.fps_base |
| 148 | return max(1, int(round(frames * 1000.0 / fps))) |
| 149 | |
| 150 | def get_animated_fcurves(ad): |
| 151 | """Return the FCurves driving a given AnimData block. |
| 152 | |
| 153 | Handles both the legacy (≤ 4.3) Action.fcurves shape and the Blender 4.4+ |
| 154 | Slotted Actions shape, where FCurves live on a per-slot channelbag and |
| 155 | must be fetched via bpy_extras.anim_utils.action_get_channelbag_for_slot. |
| 156 | Only the slot bound to *this* AnimData is returned (other slots in the |
| 157 | same Action belong to other datablocks).""" |
| 158 | if not ad or not ad.action: |
| 159 | return [] |
| 160 | action = ad.action |
| 161 | # Legacy compat path — pre-4.4 single-fcurves-on-Action. |
| 162 | if hasattr(action, 'fcurves'): |
| 163 | return list(action.fcurves) |
| 164 | # 4.4+: action.slots[*] + channelbag lookup. |
| 165 | slot = getattr(ad, 'action_slot', None) |
| 166 | if slot is None: |
| 167 | return [] |
| 168 | try: |
| 169 | from bpy_extras import anim_utils |
| 170 | cb = anim_utils.action_get_channelbag_for_slot(action, slot) |
| 171 | except Exception: |
| 172 | return [] |
| 173 | return list(cb.fcurves) if cb else [] |
| 174 | |
| 175 | def next_keyframe_time(fcurves, current_frame): |
| 176 | """Smallest keyframe.x strictly greater than current_frame across fcurves, |
| 177 | or None if no keyframe is ahead.""" |
| 178 | nxt = None |
| 179 | for fc in fcurves: |
| 180 | for kp in fc.keyframe_points: |
| 181 | x = kp.co.x |
| 182 | if x > current_frame and (nxt is None or x < nxt): |
| 183 | nxt = x |
| 184 | return nxt |
| 185 | |
| 186 | def pick_pwm2_args(total_ms): |
| 187 | """Decompose total fade ms into (numPeriods, periodMs), both 1..1000. |
| 188 | Total fade cap = 1,000,000 ms (~1000 s).""" |
| 189 | total_ms = max(1, min(1_000_000, int(total_ms))) |
| 190 | if total_ms <= 1000: |
| 191 | return (1, total_ms) |
| 192 | period_ms = (total_ms + 999) // 1000 # ceil so numPeriods stays <= 1000 |
| 193 | num_periods = max(1, total_ms // period_ms) |
| 194 | return (num_periods, period_ms) |
| 195 | |
| 196 | # ----------------------- |
| 197 | # Core Handler |
| 198 | # ----------------------- |
| 199 | @persistent |
| 200 | def robot_anim_handler(scene): |
| 201 | global last_arm_state, quad_last_command_dict |
| 202 | quad_pwm_dicts = {} |
| 203 | |
| 204 | for obj in bpy.data.objects: |
| 205 | if obj.name.startswith(ARM_NAME_PREFIX): |
| 206 | ip = obj.name[len(ARM_NAME_PREFIX):] |
| 207 | angles = getAngles(obj) |
| 208 | lamp = getLamp(obj) |
| 209 | energy = lamp.data.energy if lamp else 0.0 |
| 210 | pwm = pwmFromLuminosity(energy) |
| 211 | |
| 212 | cur_state = (tuple(angles), pwm) |
| 213 | if last_arm_state.get(obj.name) == cur_state: |
| 214 | print(f"[SKIP] {obj.name} unchanged") |
| 215 | continue |
| 216 | last_arm_state[obj.name] = cur_state |
| 217 | |
| 218 | s = "s " + " ".join(f"{i+1}:{a}" for i,a in enumerate(angles)) |
| 219 | s += f"\npwm {pwm}\n" |
| 220 | print(f"[SEND] {obj.name} -> {s.strip()}") |
| 221 | sendString(ip, s) |
| 222 | |
| 223 | elif obj.name.startswith(LIGHT_NAME_PREFIX): |
| 224 | ip = obj.name[len(LIGHT_NAME_PREFIX):] |
| 225 | pwm = pwmFromLuminosity(obj.data.energy) |
| 226 | s = f"pwm {pwm}\n" |
| 227 | print(f"[SEND] {obj.name} -> {s.strip()}") |
| 228 | sendString(ip, s) |
| 229 | |
| 230 | elif obj.name.startswith(BGR_PREFIX): |
| 231 | # bgr-INDEX-IP — maxsplit=2 so hostnames containing '-' survive. |
| 232 | tokens = obj.name.split('-', 2) |
| 233 | if len(tokens) < 3: continue |
| 234 | index = int(tokens[1]) |
| 235 | ip = tokens[2] |
| 236 | pwm_r = pwmFromLuminosity(obj.data.color[0], LIGHT_PWM_RED_MAX) |
| 237 | pwm_g = pwmFromLuminosity(obj.data.color[1]) |
| 238 | pwm_b = pwmFromLuminosity(obj.data.color[2]) |
| 239 | cur_str = f"{pwm_b} {pwm_g} {pwm_r} " |
| 240 | # Keep the obj reference so playback can re-evaluate fcurves at the |
| 241 | # next keyframe (target ≠ current interpolated value). |
| 242 | quad_pwm_dicts.setdefault(ip, {})[index] = (obj, cur_str) |
| 243 | |
| 244 | scene = bpy.context.scene |
| 245 | cur_frame = scene.frame_current_final |
| 246 | playing = is_playing() |
| 247 | |
| 248 | for ip, slot_dict in quad_pwm_dicts.items(): |
| 249 | # Current-frame plain-pwm payload, used for scrub/pause fallback. |
| 250 | cur_vals = [] |
| 251 | for i in range(1, 5): |
| 252 | if i in slot_dict: |
| 253 | cur_vals += slot_dict[i][1].split() |
| 254 | else: |
| 255 | cur_vals += ["0", "0", "0"] |
| 256 | |
| 257 | if playing: |
| 258 | # Index each slot's color fcurves by array_index (0=R 1=G 2=B). |
| 259 | slot_fcurves = {} |
| 260 | all_fcurves = [] |
| 261 | for idx, (obj, _) in slot_dict.items(): |
| 262 | ad = obj.data.animation_data if obj.data else None |
| 263 | curves = [fc for fc in get_animated_fcurves(ad) |
| 264 | if fc.data_path == 'color'] |
| 265 | all_fcurves += curves |
| 266 | slot_fcurves[idx] = {fc.array_index: fc for fc in curves} |
| 267 | nxt = next_keyframe_time(all_fcurves, cur_frame) |
| 268 | if nxt is not None: |
| 269 | # Build target by evaluating each fcurve AT THE NEXT KEYFRAME. |
| 270 | # This is the value the firmware should fade toward; it's |
| 271 | # constant within a segment so dedup catches the next frame |
| 272 | # and we send exactly one pwm2 per segment. |
| 273 | target_vals = [] |
| 274 | for i in range(1, 5): |
| 275 | if i in slot_dict: |
| 276 | obj = slot_dict[i][0] |
| 277 | fcs = slot_fcurves[i] |
| 278 | def channel(ai, max_pwm): |
| 279 | if ai in fcs: |
| 280 | v = fcs[ai].evaluate(nxt) |
| 281 | else: |
| 282 | v = obj.data.color[ai] |
| 283 | return pwmFromLuminosity(v, max_pwm) |
| 284 | # BGR send order, matches firmware expectation. |
| 285 | target_vals += [ |
| 286 | channel(2, LIGHT_PWM_MAX), |
| 287 | channel(1, LIGHT_PWM_MAX), |
| 288 | channel(0, LIGHT_PWM_RED_MAX), |
| 289 | ] |
| 290 | else: |
| 291 | target_vals += [0, 0, 0] |
| 292 | target_tuple = tuple(target_vals) |
| 293 | if last_pwm2_target.get(ip) == target_tuple: |
| 294 | continue # same segment, already commanded |
| 295 | last_pwm2_target[ip] = target_tuple |
| 296 | fade_ms = frames_to_ms(scene, nxt - cur_frame) |
| 297 | num, per = pick_pwm2_args(fade_ms) |
| 298 | s = f"pwm2 {num} {per} " + " ".join(str(v) for v in target_vals) + "\n" |
| 299 | sendString(ip, s) |
| 300 | quad_last_command_dict.pop(ip, None) |
| 301 | continue |
| 302 | # no keyframes ahead — fall through to plain pwm |
| 303 | |
| 304 | # Scrubbing, paused, or no more keyframes: plain pwm (cancels any fade). |
| 305 | s = "pwm " + " ".join(cur_vals) + "\n" |
| 306 | if quad_last_command_dict.get(ip) == s: |
| 307 | continue |
| 308 | quad_last_command_dict[ip] = s |
| 309 | last_pwm2_target.pop(ip, None) |
| 310 | sendString(ip, s) |
| 311 | |
| 312 | # ----------------------- |
| 313 | # File Save/Load |
| 314 | # ----------------------- |
| 315 | def saveToFile(filepath): |
| 316 | arms = {} |
| 317 | for arm in bpy.data.objects: |
| 318 | if not arm.name.startswith(ARM_NAME_PREFIX): continue |
| 319 | ip = arm.name[len(ARM_NAME_PREFIX):] |
| 320 | angles = getAngles(arm) |
| 321 | lamp = getLamp(arm) |
| 322 | energy = pwmFromLuminosity(lamp.data.energy) if lamp else 0 |
| 323 | arms[ip] = {"angles": angles, "light": energy} |
| 324 | text = {"LightArms": arms} |
| 325 | with open(filepath, 'w') as f: |
| 326 | json.dump(text, f, indent=2) |
| 327 | print(f"[SAVE] Robot positions saved to {filepath}") |
| 328 | |
| 329 | def loadFromFile(filepath): |
| 330 | with open(filepath, 'r') as f: |
| 331 | text = json.load(f) |
| 332 | arms = text['LightArms'] |
| 333 | for ip,data in arms.items(): |
| 334 | name = ARM_NAME_PREFIX + ip |
| 335 | arm = bpy.data.objects[name] |
| 336 | lamp = getLamp(arm) |
| 337 | if lamp: |
| 338 | lamp.data.energy = luminosityFromPWM(data['light']) |
| 339 | print(f"[LOAD] Robot positions loaded from {filepath}") |
| 340 | # TODO: set angles if needed |
| 341 | |
| 342 | # ----------------------- |
| 343 | # UI Panel |
| 344 | # ----------------------- |
| 345 | class ROBOT_PT_scene(bpy.types.Panel): |
| 346 | bl_label = "Robot Scene" |
| 347 | bl_space_type = "VIEW_3D" |
| 348 | bl_region_type = "UI" |
| 349 | bl_category = "Robot" |
| 350 | def draw(self, context): |
| 351 | self.layout.operator("circus.saver", text='Save Robot Positions') |
| 352 | self.layout.operator("circus.loader", text='Load Robot Positions') |
| 353 | |
| 354 | class CircusSaver(bpy.types.Operator): |
| 355 | bl_idname = "circus.saver" |
| 356 | bl_label = "Save the robots" |
| 357 | filepath: bpy.props.StringProperty(subtype="FILE_PATH") |
| 358 | def execute(self, context): |
| 359 | saveToFile(self.filepath) |
| 360 | return {'FINISHED'} |
| 361 | def invoke(self, context, event): |
| 362 | context.window_manager.fileselect_add(self) |
| 363 | return {'RUNNING_MODAL'} |
| 364 | |
| 365 | class CircusLoader(bpy.types.Operator): |
| 366 | bl_idname = "circus.loader" |
| 367 | bl_label = "Load the robots" |
| 368 | filepath: bpy.props.StringProperty(subtype="FILE_PATH") |
| 369 | def execute(self, context): |
| 370 | loadFromFile(self.filepath) |
| 371 | return {'FINISHED'} |
| 372 | def invoke(self, context, event): |
| 373 | context.window_manager.fileselect_add(self) |
| 374 | return {'RUNNING_MODAL'} |
| 375 | |
| 376 | # ----------------------- |
| 377 | # Register Everything |
| 378 | # ----------------------- |
| 379 | _CLASSES = (ROBOT_PT_scene, CircusSaver, CircusLoader) |
| 380 | |
| 381 | def unregister(): |
| 382 | for cls in _CLASSES: |
| 383 | try: bpy.utils.unregister_class(cls) |
| 384 | except Exception: pass |
| 385 | # Strip *all* handlers named robot_anim_handler — catches stale instances |
| 386 | # left by previous script runs (different function-object identity, same |
| 387 | # name). identity-only remove() doesn't see those. |
| 388 | bpy.app.handlers.frame_change_post[:] = [ |
| 389 | h for h in bpy.app.handlers.frame_change_post |
| 390 | if getattr(h, '__name__', None) != 'robot_anim_handler' |
| 391 | ] |
| 392 | |
| 393 | if __name__ == "__main__": |
| 394 | unregister() |
| 395 | for cls in _CLASSES: |
| 396 | bpy.utils.register_class(cls) |
| 397 | if robot_anim_handler not in bpy.app.handlers.frame_change_post: |
| 398 | bpy.app.handlers.frame_change_post.append(robot_anim_handler) |
| 399 | print("[INFO] Robot script initialized") |