Skip to content

Commit b81b5c5

Browse files
JensBlackJensBlack
committed
Merge branch 'master' into dev_classification
# Conflicts: # DeepLabStream.py # experiments/custom/triggers.py # requirements.txt # settings.ini # utils/configloader.py # utils/plotter.py # utils/poser.py
2 parents 7d54efa + ccb38d8 commit b81b5c5

File tree

9 files changed

+355
-87
lines changed

9 files changed

+355
-87
lines changed

‎DeepLabStream.py‎

Lines changed: 73 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,23 @@
66
https://github.com/SchwarzNeuroconLab/DeepLabStream
77
Licensed under GNU General Public License v3.0
88
"""
9-
import time
9+
import multiprocessing as mp
1010
import os
1111
import sys
12-
import multiprocessing as mp
12+
import time
1313
from importlib.util import find_spec
1414

15+
import click
1516
import cv2
1617
import numpy as np
1718
import pandas as pd
18-
import click
1919

20-
from utils.configloader import RESOLUTION, FRAMERATE, OUT_DIR, MODEL_NAME, MULTI_CAM, STACK_FRAMES, \
21-
ANIMALS_NUMBER, STREAMS, STREAMING_SOURCE
22-
from utils.poser import load_deeplabcut, get_pose, find_local_peaks_new, calculate_skeletons
20+
from utils.generic import VideoManager, WebCamManager, GenericManager
21+
from utils.configloader import RESOLUTION, FRAMERATE, OUT_DIR, MODEL_NAME, MULTI_CAM, STACK_FRAMES, \
22+
ANIMALS_NUMBER, STREAMS, STREAMING_SOURCE, MODEL_ORIGIN
2323
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
2426

2527

2628
def create_video_files(directory, devices, resolution, framerate, codec):
@@ -127,28 +129,19 @@ def set_camera_manager():
127129
:return: the chosen camera manager
128130
"""
129131

130-
if STREAMING_SOURCE.lower() == 'video':
131-
from utils.generic import VideoManager
132-
manager = VideoManager()
133-
return manager
134-
135-
elif STREAMING_SOURCE.lower() == 'ipwebcam':
136-
from utils.generic import WebCamManager
137-
manager = WebCamManager()
138-
return manager
139-
140-
elif STREAMING_SOURCE.lower() == 'camera':
132+
def select_camera_manager():
133+
"""
134+
Function to select from all available camera managers
135+
"""
141136
manager_list = []
142137
# loading realsense manager, if installed
143-
realsense = find_spec("pyrealsense2") is not None
144-
if realsense:
138+
if find_spec("pyrealsense2") is not None:
145139
from utils.realsense import RealSenseManager
146140
realsense_manager = RealSenseManager()
147141
manager_list.append(realsense_manager)
148142

149143
# loading basler manager, if installed
150-
pylon = find_spec("pypylon") is not None
151-
if pylon:
144+
if find_spec("pypylon") is not None:
152145
from utils.pylon import PylonManager
153146
pylon_manager = PylonManager()
154147
manager_list.append(pylon_manager)
@@ -170,9 +163,19 @@ def check_for_cameras(camera_manager):
170163
return manager
171164
else:
172165
# if no camera is found, try generic openCV manager
173-
from utils.generic import GenericManager
174166
generic_manager = GenericManager()
175167
return generic_manager
168+
169+
MANAGER_SOURCE = {
170+
'video': VideoManager,
171+
'ipwebcam': WebCamManager,
172+
'camera': select_camera_manager
173+
}
174+
175+
# initialize selected manager
176+
camera_manager = MANAGER_SOURCE.get(STREAMING_SOURCE)()
177+
if camera_manager is not None:
178+
return camera_manager
176179
else:
177180
raise ValueError(f'Streaming source {STREAMING_SOURCE} is not a valid option. \n'
178181
f'Please choose from "video", "camera" or "ipwebcam".')
@@ -272,13 +275,46 @@ def get_pose_mp(input_q, output_q):
272275
:param input_q: index and corresponding frame
273276
:param output_q: index and corresponding analysis
274277
"""
275-
config, sess, inputs, outputs = load_deeplabcut()
276-
while True:
277-
if input_q.full():
278-
index, frame = input_q.get()
279-
scmap, locref, pose = get_pose(frame, config, sess, inputs, outputs)
280-
peaks = find_local_peaks_new(scmap, locref, ANIMALS_NUMBER, config)
281-
output_q.put((index, peaks))
278+
279+
if MODEL_ORIGIN in ('DLC', 'MADLC'):
280+
config, sess, inputs, outputs = load_deeplabcut()
281+
while True:
282+
if input_q.full():
283+
index, frame = input_q.get()
284+
if MODEL_ORIGIN == 'DLC':
285+
scmap, locref, pose = get_pose(frame, config, sess, inputs, outputs)
286+
# TODO: Remove alterations to original
287+
peaks = find_local_peaks_new(scmap, locref, ANIMALS_NUMBER, config)
288+
# peaks = pose
289+
if MODEL_ORIGIN == 'MADLC':
290+
peaks = get_ma_pose(frame, config, sess, inputs, outputs)
291+
292+
output_q.put((index, peaks))
293+
294+
elif MODEL_ORIGIN == 'DLC-LIVE':
295+
dlc_live = load_dlc_live()
296+
while True:
297+
if input_q.full():
298+
index, frame = input_q.get()
299+
if not dlc_live.is_initialized:
300+
peaks = dlc_live.init_inference(frame)
301+
else:
302+
peaks = dlc_live.get_pose(frame)
303+
304+
output_q.put((index, peaks))
305+
306+
elif MODEL_ORIGIN == 'DEEPPOSEKIT':
307+
predict_model = load_dpk()
308+
while True:
309+
if input_q.full():
310+
index, frame = input_q.get()
311+
frame = frame[..., 1][..., None]
312+
st_frame = np.stack([frame])
313+
prediction = predict_model.predict(st_frame, batch_size=1, verbose=True)
314+
peaks = prediction[0, :, :2]
315+
output_q.put((index, peaks))
316+
else:
317+
raise ValueError(f'Model origin {MODEL_ORIGIN} not available.')
282318

283319
@staticmethod
284320
def create_mp_tools(devices):
@@ -721,10 +757,14 @@ def show_benchmark_statistics():
721757
print("[{0}/3000] Benchmarking in progress".format(len(analysis_time_data)))
722758

723759
if benchmark_enabled:
724-
import re
725-
short_model = re.split('[-_]', MODEL_NAME)
726-
short_model = short_model[0] + '_' + short_model[2]
727-
np.savetxt(f'{OUT_DIR}/{short_model}_framerate_{FRAMERATE}_resolution_{RESOLUTION[0]}_{RESOLUTION[1]}.txt', np.transpose([fps_data, whole_loop_time_data]))
760+
model_parts = MODEL_NAME.split('_')
761+
if len(model_parts) == 3:
762+
short_model = model_parts[0] + '_' + model_parts[2]
763+
else:
764+
short_model = MODEL_NAME
765+
# the best way to save files
766+
np.savetxt(f'{OUT_DIR}/{short_model}_framerate_{FRAMERATE}_resolution_{RESOLUTION[0]}_{RESOLUTION[1]}.txt',
767+
np.transpose([fps_data, whole_loop_time_data]))
728768

729769

730770
if __name__ == '__main__':

‎Readme.md‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,10 @@
99
[![License: GPL v3](https://img.shields.io/badge/License-GPL%20v3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0)
1010
[![Twitter Follow](https://img.shields.io/twitter/follow/SNeuroconnect.svg?label=SNeuroconnect&style=social)](https://twitter.com/SNeuroconnect)
1111

12-
DeepLabStream is a python based multi-purpose tool that enables the realtime tracking of animals and manipulation of experiments.
13-
Our toolbox is adapted from the previously published [DeepLabCut](https://github.com/AlexEMG/DeepLabCut) ([Mathis et al., 2018](https://www.nature.com/articles/s41593-018-0209-y)) and expands on its core capabilities.
14-
DeepLabStreams core feature is the real-time analysis using any type of camera-based video stream (incl. multiple streams). Building onto that, we designed a full experimental closed-loop toolkit. It enables running experimental protocols that are dependent on a constant stream of bodypart positions and feedback activation of several input/output devices. It's capabilities range from simple region of interest (ROI) based triggers to headdirection or behavior dependent stimulation.
12+
DeepLabStream is a python based multi-purpose tool that enables the realtime tracking and manipulation of animals during ongoing experiments.
13+
Our toolbox was orginally adapted from the previously published [DeepLabCut](https://github.com/AlexEMG/DeepLabCut) ([Mathis et al., 2018](https://www.nature.com/articles/s41593-018-0209-y)) and expanded on its core capabilities, but is now able to utilize a variety of different network architectures for online pose estimation
14+
([DLC + maDLC](https://github.com/AlexEMG/DeepLabCut), [DLC-Live](https://github.com/DeepLabCut/DeepLabCut-live), [DeepPosekit's](https://github.com/jgraving/DeepPoseKit) StackedDenseNet, StackedHourGlass and [LEAP](https://github.com/murthylab/sleap)).
15+
DeepLabStreams core feature is the utilization of real-time tracking to orchestrate closed-loop experiments. This can be achieved using any type of camera-based video stream (incl. multiple streams). It enables running experimental protocols that are dependent on a constant stream of bodypart positions and feedback activation of several input/output devices. It's capabilities range from simple region of interest (ROI) based triggers to headdirection or behavior dependent stimulation.
1516

1617
![DLS_Stim](docs/DLSSTim_example.gif)
1718

@@ -21,7 +22,15 @@ DeepLabStreams core feature is the real-time analysis using any type of camera-b
2122

2223
### Read the preprint: [Schweihoff et al, 2019](https://doi.org/10.1101/2019.12.20.884478).
2324

24-
### 1. [Installation & Testing](https://github.com/SchwarzNeuroconLab/DeepLabStream/wiki/Installation-&-Testing)
25+
### Contributing
26+
27+
If you have feature requests or questions regarding the design of experiments join our [slack group](https://join.slack.com/t/dlstream/shared_invite/zt-jpy2olk1-CuJu0ZylGg_SLbO7zBkcrg)!
28+
29+
We are constantly working to update and increase the capabilities of DLStream.
30+
We welcome all feedback and input from your side.
31+
Also, do not hesitate to contact us for collaborations.
32+
33+
### 1. [Updated Installation & Testing](https://github.com/SchwarzNeuroconLab/DeepLabStream/wiki/Installation-&-Testing)
2534

2635
### 2. [How to use DLStream GUI](https://github.com/SchwarzNeuroconLab/DeepLabStream/wiki/How-to-use-DLStream)
2736

@@ -108,8 +117,9 @@ This project is licensed under the GNU General Public License v3.0. Note that th
108117

109118
## Authors
110119

111-
Lead Researcher: Jens Schweihoff, [email protected]
120+
Developed by:
121+
- Jens Schweihoff, [email protected]
112122

113-
Lead Developer: Matvey Loshakov, [email protected]
123+
- Matvey Loshakov, [email protected]
114124

115125
Corresponding Author: Martin Schwarz, [email protected]

‎requirements.txt‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
numba
2-
gpiozero
3-
pigpio
4-
pyserial
5-
nidaqmx>=0.5.7
6-
click>=7.0
7-
opencv-python>=3.4.5.20
2+
gpiozero==1.5.1
3+
pigpio==1.78
4+
pyserial==3.5
5+
nidaqmx==0.5.7
6+
click==7.1.2
7+
opencv-python==3.4.5.20
8+
opencv-contrib-python==4.4.0.46
89
numpy>=1.14.5
9-
pandas>=0.21.0
10-
matplotlib>=3.0.3
11-
scikit-image>=0.14.2
12-
scipy>=1.1.0
10+
pandas==1.1.4
11+
matplotlib==3.0.3
12+
scikit-image==0.17.2
13+
scipy==1.5.4
14+
pyzmq==20.0.0

‎settings.ini‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,26 @@
11
[Streaming]
2-
RESOLUTION = 320, 256
2+
RESOLUTION = 848, 480
33
FRAMERATE = 30
44
OUTPUT_DIRECTORY = /Output
55
#if you have connected multiple cameras (USB), you will need to select the number OpenCV has given them.
66
#Default is "0", which takes the first available camera.
77
CAMERA_SOURCE = 0
8-
#you can use "camera", "ipwebcam" or "video" to select your input source
9-
STREAMING_SOURCE = video
8+
#you can use camera, ipwebcam or video to select your input source
9+
STREAMING_SOURCE = camera
1010

1111
[Pose Estimation]
12-
#possible origins are: DLC
12+
#possible origins are: DLC, DLC-LIVE, MADLC, DEEPPOSEKIT
1313
MODEL_ORIGIN = ORIGIN
1414
MODEL_PATH = PATH_TO_MODEL
1515
MODEL_NAME = MODEL_NAME
16-
; Currently used to translate pose estimation into postures with bodypart names ... not necessary for classic DLC networks
16+
; only used in DLC-LIVE and DeepPoseKit for now; if left empty or to short, auto-naming will be enabled in style bp1, bp2 ...
1717
ALL_BODYPARTS = bp1, bp2, bp3, bp4
1818

1919
[Experiment]
2020
#Available parameters are "CUSTOM" and "BASE"
21-
EXP_ORIGIN = BASE/CUSTOM
21+
EXP_ORIGIN = CUSTOM
2222
#Name of the experiment config in /experiments/configs or name of the custom experiment in /experiments/custom/experiments.py
23-
EXP_NAME = CONFIG_NAME
23+
EXP_NAME = ExampleExperiment
2424
#if you want the experiment to be recorded as a raw video set this to "True".
2525
RECORD_EXP = True
2626

‎utils/VideoAnalyzer.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
from utils.configloader import VIDEO_SOURCE, OUT_DIR, ANIMALS_NUMBER
1919
from experiments.custom.experiments import ExampleExperiment
2020

21+
from deepposekit.models import load_model
22+
2123

2224
def create_dataframes(data_output):
2325
"""
@@ -43,7 +45,7 @@ def create_dataframes(data_output):
4345

4446
def start_videoanalyser():
4547
print("Starting DeepLabCut")
46-
config, sess, inputs, outputs = load_deeplabcut()
48+
model = load_model(r"D:\DeepPoseKit-Data-master\datasets\fly\best_model_densenet.h5")
4749

4850
experiment_enabled = False
4951
video_output = True

‎utils/generic.py‎

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,15 @@
55
https://github.com/SchwarzNeuroconLab/DeepLabStream
66
Licensed under GNU General Public License v3.0
77
"""
8+
import time
9+
import base64
810

911
import cv2
10-
from utils.configloader import CAMERA_SOURCE, VIDEO_SOURCE, RESOLUTION, FRAMERATE, PORT
11-
import time
1212
import numpy as np
13+
import zmq
14+
15+
from utils.configloader import CAMERA_SOURCE, VIDEO_SOURCE, RESOLUTION, FRAMERATE, PORT
16+
1317

1418
class GenericManager:
1519
"""
@@ -86,7 +90,6 @@ def get_name(self) -> str:
8690
return self._manager_name
8791

8892

89-
9093
class VideoManager(GenericManager):
9194

9295
"""
@@ -97,13 +100,12 @@ def __init__(self):
97100
Generic video manager from video files
98101
Uses pure opencv
99102
"""
100-
self._manager_name = "generic"
103+
super().__init__()
101104
self._camera = cv2.VideoCapture(VIDEO_SOURCE)
102105
self._camera_name = "Video"
103106
self.initial_wait = False
104107
self.last_frame_time = time.time()
105108

106-
107109
def get_frames(self) -> tuple:
108110
"""
109111
Collect frames for camera and outputs it in 'color' dictionary
@@ -116,7 +118,6 @@ def get_frames(self) -> tuple:
116118
infra_frames = {}
117119
ret, image = self._camera.read()
118120
self.last_frame_time = time.time()
119-
#print(ret)
120121
if ret:
121122
if not self.initial_wait:
122123
cv2.waitKey(1000)
@@ -127,6 +128,10 @@ def get_frames(self) -> tuple:
127128
if running_time <= 1 / FRAMERATE:
128129
sleepy_time = int(np.ceil(1000/FRAMERATE - running_time / 1000))
129130
cv2.waitKey(sleepy_time)
131+
else:
132+
# cycle the video for testing purposes
133+
self._camera.set(cv2.CAP_PROP_POS_FRAMES, 0)
134+
return self.get_frames()
130135

131136
return color_frames, depth_maps, infra_frames
132137

@@ -138,24 +143,23 @@ def __init__(self):
138143
Binds the computer to a ip address and starts listening for incoming streams.
139144
Adapted from StreamViewer.py https://github.com/CT83/SmoothStream
140145
"""
141-
import zmq
146+
super().__init__()
142147
self._context = zmq.Context()
143148
self._footage_socket = self._context.socket(zmq.SUB)
144149
self._footage_socket.bind('tcp://*:' + PORT)
145150
self._footage_socket.setsockopt_string(zmq.SUBSCRIBE, np.unicode(''))
146151

147-
self._manager_name = "generic"
148152
self._camera = None
149153
self._camera_name = "webcam"
150154
self.initial_wait = False
151155
self.last_frame_time = time.time()
152156

153-
def string_to_image(self, string):
154-
""" Taken from https://github.com/CT83/SmoothStream"""
157+
@ staticmethod
158+
def string_to_image(string):
159+
"""
160+
Taken from https://github.com/CT83/SmoothStream
161+
"""
155162

156-
import numpy as np
157-
import cv2
158-
import base64
159163
img = base64.b64decode(string)
160164
npimg = np.fromstring(img, dtype=np.uint8)
161165
return cv2.imdecode(npimg, 1)

0 commit comments

Comments
 (0)