Spaces:
Runtime error
Runtime error
File size: 3,615 Bytes
999c5c9 |
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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
"""
DeepLabCut Toolbox (deeplabcut.org)
© A. & M. Mathis Labs
Licensed under GNU Lesser General Public License v3.0
"""
import serial
import struct
import time
import numpy as np
from dlclive.processor import Processor, KalmanFilterPredictor
class IzzyJump(Processor):
def __init__(self, com="", lik_thresh=0.5, baudrate=int(9600), **kwargs):
super().__init__()
self.ser = serial.Serial(com, baudrate, timeout=0)
self.lik_thresh = lik_thresh
self.led_times = []
self.last_light = 0
def close_serial(self):
self.ser.close()
def switch_led(self, val, frame_time):
### check status of led ###
self.ser.write(b"R")
led_byte = b""
led_status = None
while (len(led_byte) != 0) or (led_status is None):
led_byte = self.ser.read()
if len(led_byte) > 0:
led_status = ord(led_byte)
if led_status != val:
ctime = time.time()
if ctime - self.last_light > 0.25:
self.ser.write(b"L")
self.last_light = ctime
self.led_times.append((val, frame_time, ctime))
def process(self, pose, **kwargs):
### bodyparts
# 0. nose
# 1. L-eye
# 2. R-eye
# 3. L-ear
# 4. R-ear
# 5. Throat
# 6. Withers
# 7. Tailset
# 8. L-front-paw
# 9. R-front-paw
# 10. L-front-wrist
# 11. R-front-wrist
# 12. L-front-elbow
# 13. R-front-elbow
# ...
l_elbow = pose[12, 1] if pose[12, 2] > self.lik_thresh else None
r_elbow = pose[13, 1] if pose[13, 2] > self.lik_thresh else None
elbows = [l_elbow, r_elbow]
this_elbow = (
min([e for e in elbows if e is not None])
if any([e is not None for e in elbows])
else None
)
withers = pose[6, 1] if pose[6, 2] > self.lik_thresh else None
if kwargs["record"]:
if withers is not None and this_elbow is not None:
if this_elbow < withers:
self.switch_led(True, kwargs["frame_time"])
else:
self.switch_led(False, kwargs["frame_time"])
return pose
def save(self, filename):
### save stim on and stim off times
if filename[-4:] != ".npy":
filename += ".npy"
arr = np.array(self.led_times, dtype=float)
try:
np.save(filename, arr)
save_code = True
except Exception:
save_code = False
return save_code
class IzzyJumpKF(KalmanFilterPredictor, IzzyJump):
def __init__(
self,
com="",
lik_thresh=0.5,
baudrate=int(9600),
adapt=True,
forward=0.003,
fps=30,
nderiv=2,
priors=[1, 1],
initial_var=1,
process_var=1,
dlc_var=4,
):
super().__init__(
adapt=adapt,
forward=forward,
fps=fps,
nderiv=nderiv,
priors=priors,
initial_var=initial_var,
process_var=process_var,
dlc_var=dlc_var,
com=com,
lik_thresh=lik_thresh,
baudrate=baudrate,
)
def process(self, pose, **kwargs):
future_pose = KalmanFilterPredictor.process(self, pose, **kwargs)
final_pose = IzzyJump.process(self, future_pose, **kwargs)
return final_pose
def save(self, filename):
return IzzyJump.save(self, filename)
|