Skip to content

Commit f64aa7b

Browse files
committed
initial commit
1 parent 75b4555 commit f64aa7b

File tree

5 files changed

+157
-13
lines changed

5 files changed

+157
-13
lines changed

DeepLabStream.py

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,9 @@
2020
from utils.generic import VideoManager, WebCamManager, GenericManager
2121
from utils.configloader import RESOLUTION, FRAMERATE, OUT_DIR, MODEL_NAME, MULTI_CAM, STACK_FRAMES, \
2222
ANIMALS_NUMBER, STREAMS, STREAMING_SOURCE, MODEL_ORIGIN, CROP, CROP_X, CROP_Y
23-
from utils.plotter import plot_bodyparts, plot_metadata_frame
24-
from utils.poser import load_deeplabcut, load_dpk, load_dlc_live, get_pose, calculate_skeletons,\
25-
find_local_peaks_new, get_ma_pose
26-
23+
from utils.plotter import plot_bodyparts,plot_metadata_frame
24+
from utils.poser import load_deeplabcut,load_dpk,load_dlc_live,get_pose,calculate_skeletons, \
25+
find_local_peaks_new,get_ma_pose,load_sleap
2726

2827
def create_video_files(directory, devices, resolution, framerate, codec):
2928
"""
@@ -317,6 +316,20 @@ def get_pose_mp(input_q, output_q):
317316
peaks = prediction[0, :, :2]
318317
analysis_time = time.time() - start_time
319318
output_q.put((index,peaks,analysis_time))
319+
320+
elif MODEL_ORIGIN == 'SLEAP':
321+
sleap_model = load_sleap()
322+
while True:
323+
if input_q.full():
324+
index, frame = input_q.get()
325+
start_time = time.time()
326+
frame = frame[:, :, ::-1]
327+
#this is weird, but i without it, it does not seem to work...
328+
frames = np.array([frame])
329+
prediction = sleap_model.predict(frames[[0]], batch_size=1)
330+
peaks = prediction['instance_peaks'][0, :]
331+
analysis_time = time.time() - start_time
332+
output_q.put((index,peaks,analysis_time))
320333
else:
321334
raise ValueError(f'Model origin {MODEL_ORIGIN} not available.')
322335

@@ -426,8 +439,13 @@ def get_analysed_frames(self) -> tuple:
426439
self._experiment_running = False
427440

428441
if self._experiment_running and not self._experiment.experiment_finished:
429-
for skeleton in skeletons:
430-
self._experiment.check_skeleton(analysed_image, skeleton)
442+
#TODO: Update to work for multiple animal and single animal experiments
443+
#Shift responsibility to experiments
444+
if ANIMALS_NUMBER > 1:
445+
self._experiment.check_skeleton(analysed_image,skeletons)
446+
else:
447+
for skeleton in skeletons:
448+
self._experiment.check_skeleton(analysed_image, skeleton)
431449

432450
# Gathering data as pd.Series for output
433451
if self._data_output:

experiments/custom/experiments.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,106 @@
1717
from experiments.custom.stimulation import show_visual_stim_img,laser_switch
1818

1919

20+
class ExampleSocialExperiment:
21+
"""
22+
Simple class to contain all of the experiment properties
23+
Uses multiprocess to ensure the best possible performance and
24+
to showcase that it is possible to work with any type of equipment, even timer-dependent
25+
"""
26+
def __init__(self):
27+
self.experiment_finished = False
28+
self._process = ExampleProtocolProcess()
29+
self._green_point = (313, 552)
30+
self._radius = 80
31+
self._event = None
32+
self._current_trial = None
33+
self._trial_count = {trial: 0 for trial in self._trials}
34+
self._trial_timers = {trial: Timer(10) for trial in self._trials}
35+
self._exp_timer = Timer(600)
36+
37+
def check_skeleton(self, frame, skeletons):
38+
"""
39+
Checking each passed animal skeleton for a pre-defined set of conditions
40+
Outputting the visual representation, if exist
41+
Advancing trials according to inherent logic of an experiment
42+
:param frame: frame, on which animal skeleton was found
43+
:param skeletons: skeletons, consisting of multiple joints of an animal
44+
"""
45+
self.check_exp_timer() # checking if experiment is still on
46+
for trial in self._trial_count:
47+
# checking if any trial hit a predefined cap
48+
if self._trial_count[trial] >= 10:
49+
self.stop_experiment()
50+
51+
if not self.experiment_finished:
52+
result, response = False, None
53+
for trial in self._trials:
54+
# check for all trials if condition is met
55+
result_list = []
56+
for skeleton in skeletons:
57+
result, response = self._trials[trial]['trigger'](skeleton=skeleton)
58+
if result:
59+
break
60+
plot_triggers_response(frame, response)
61+
if result:
62+
if self._current_trial is None:
63+
if not self._trial_timers[trial].check_timer():
64+
self._current_trial = trial
65+
self._trial_timers[trial].reset()
66+
self._trial_count[trial] += 1
67+
print(trial, self._trial_count[trial])
68+
else:
69+
if self._current_trial == trial:
70+
self._current_trial = None
71+
self._trial_timers[trial].start()
72+
73+
self._process.set_trial(self._current_trial)
74+
return result, response
75+
76+
@property
77+
def _trials(self):
78+
"""
79+
Defining the trials
80+
"""
81+
green_roi = RegionTrigger('circle', self._green_point, self._radius * 2 + 7.5, 'bp1')
82+
trials = {'Greenbar_whiteback': dict(trigger=green_roi.check_skeleton,
83+
count=0)}
84+
return trials
85+
86+
def check_exp_timer(self):
87+
"""
88+
Checking the experiment timer
89+
"""
90+
if not self._exp_timer.check_timer():
91+
print("Experiment is finished")
92+
print("Time ran out.")
93+
self.stop_experiment()
94+
95+
def start_experiment(self):
96+
"""
97+
Start the experiment
98+
"""
99+
self._process.start()
100+
if not self.experiment_finished:
101+
self._exp_timer.start()
102+
103+
def stop_experiment(self):
104+
"""
105+
Stop the experiment and reset the timer
106+
"""
107+
self.experiment_finished = True
108+
print('Experiment completed!')
109+
self._exp_timer.reset()
110+
# don't forget to end the process!
111+
self._process.end()
112+
113+
def get_trial(self):
114+
"""
115+
Check which trial is going on right now
116+
"""
117+
return self._current_trial
118+
119+
20120
class ExampleExperiment:
21121
"""
22122
Simple class to contain all of the experiment properties

experiments/custom/stimulus_process.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,16 +68,16 @@ def get_start_time(self):
6868

6969
def example_protocol_run(condition_q: mp.Queue):
7070
current_trial = None
71-
dmod_device = DigitalModDevice('Dev1/PFI0')
71+
#dmod_device = DigitalModDevice('Dev1/PFI0')
7272
while True:
7373
if condition_q.full():
7474
current_trial = condition_q.get()
7575
if current_trial is not None:
76-
show_visual_stim_img(img_type=current_trial, name='inside')
77-
dmod_device.toggle()
76+
show_visual_stim_img(type=current_trial, name='inside')
77+
#dmod_device.toggle()
7878
else:
7979
show_visual_stim_img(name='inside')
80-
dmod_device.turn_off()
80+
#dmod_device.turn_off()
8181

8282
if cv2.waitKey(1) & 0xFF == ord('q'):
8383
break

settings.ini

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,10 @@ CAMERA_SOURCE = 0
99
STREAMING_SOURCE = video
1010

1111
[Pose Estimation]
12-
#possible origins are: DLC, DLC-LIVE, MADLC, DEEPPOSEKIT
12+
#possible origins are: SLEAP, DLC, DLC-LIVE,MADLC, DEEPPOSEKIT
1313
MODEL_ORIGIN = ORIGIN
14+
#takes path to model or models (in case of SLEAP topdown, bottom up) in style "string" or "string , string", without ""
15+
# E.g.: MODEL_PATH = D:\SLEAP\models\baseline_model.centroids , D:\SLEAP\models\baseline_model.topdown
1416
MODEL_PATH = PATH_TO_MODEL
1517
MODEL_NAME = MODEL_NAME
1618
; only used in DLC-LIVE and DeepPoseKit for now; if left empty or to short, auto-naming will be enabled in style bp1, bp2 ...

utils/poser.py

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,10 @@
5656
from dlclive import DLCLive
5757
from utils.configloader import MODEL_PATH
5858

59+
elif MODEL_ORIGIN == 'SLEAP':
60+
from sleap import load_model
61+
from utils.configloader import MODEL_PATH
62+
5963

6064
def load_deeplabcut():
6165
"""
@@ -230,13 +234,13 @@ def get_ma_pose(image, config, session, inputs, outputs):
230234

231235
return pose
232236

233-
234-
def calculate_ma_skeletons(pose: dict, animals_number: int, threshold = 0.1) -> list:
237+
def calculate_ma_skeletons(pose: dict, animals_number: int, threshold:float = 0.1) -> list:
235238
"""
236239
Creating skeletons from given pose in maDLC
237240
There could be no more skeletons than animals_number
238241
Only unique skeletons output
239242
"""
243+
240244
def filter_mapredictions(pose):
241245
detection = []
242246
conf = np.array(pose['confidence'])
@@ -288,6 +292,11 @@ def load_dpk():
288292
def load_dlc_live():
289293
return DLCLive(MODEL_PATH)
290294

295+
def load_sleap():
296+
model = load_model(MODEL_PATH)
297+
model.inference_model
298+
return model.inference_model
299+
291300
def flatten_maDLC_skeletons(skeletons):
292301
"""Flattens maDLC multi skeletons into one skeleton to simulate dlc output
293302
where animals are not identical e.g. for animals with different fur colors (SIMBA)"""
@@ -362,6 +371,16 @@ def calculate_skeletons_dlc_live(pose) -> list:
362371
skeletons = [transform_2skeleton(pose)]
363372
return skeletons
364373

374+
def calculate_sleap_skeletons(pose, animals_number)-> list:
375+
"""
376+
Creating skeleton from sleap output
377+
"""
378+
skeletons = []
379+
for animal in range(pose.shape[0]):
380+
skeleton = transform_2skeleton(pose[animal])
381+
skeletons.append(skeleton)
382+
return skeletons
383+
365384

366385
def calculate_skeletons(peaks: dict, animals_number: int) -> list:
367386
"""
@@ -383,9 +402,14 @@ def calculate_skeletons(peaks: dict, animals_number: int) -> list:
383402
raise ValueError('Multiple animals are currently not supported by DLC-LIVE.'
384403
' If you are using differently colored animals, please refer to the bodyparts directly.')
385404
animal_skeletons = calculate_skeletons_dlc_live(peaks)
405+
elif MODEL_ORIGIN == 'SLEAP':
406+
animal_skeletons = calculate_sleap_skeletons(peaks, animals_number)
407+
if FLATTEN_MA:
408+
animal_skeletons = flatten_maDLC_skeletons(animal_skeletons)
386409

387410
animal_skeletons = handle_missing_bp(animal_skeletons)
388411

389412
return animal_skeletons
390413

391414

415+

0 commit comments

Comments
 (0)