DeepLabCut-live/example_processors/DogJumpLED/izzy_jump.py at main · DeepLabCut/DeepLabCut-live · 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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
"""
DeepLabCut Toolbox (deeplabcut.org)
© A. & M. Mathis Labs

Licensed under GNU Lesser General Public License v3.0
"""

import time

import numpy as np
import serial

from dlclive.processor import KalmanFilterPredictor, Processor


class IzzyJump(Processor):
    def __init__(self, com="", lik_thresh=0.5, baudrate=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=9600,
        adapt=True,
        forward=0.003,
        fps=30,
        nderiv=2,
        priors=None,
        initial_var=1,
        process_var=1,
        dlc_var=4,
    ):

        if priors is None:
            priors = [1, 1]
        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)