-
Notifications
You must be signed in to change notification settings - Fork 10
/
DeepLabStream.py
executable file
·945 lines (840 loc) · 33.9 KB
/
DeepLabStream.py
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
#!/usr/bin/env python
"""
DeepLabStream
© J.Schweihoff, M. Loshakov
University Bonn Medical Faculty, Germany
https://github.com/SchwarzNeuroconLab/DeepLabStream
Licensed under GNU General Public License v3.0
"""
import multiprocessing as mp
import os
import sys
import time
from importlib.util import find_spec
import click
import cv2
import numpy as np
import pandas as pd
from utils.generic import VideoManager, GenericManager
from utils.configloader import (
RESOLUTION,
FRAMERATE,
OUT_DIR,
MODEL_NAME,
MULTI_CAM,
STACK_FRAMES,
ANIMALS_NUMBER,
FLATTEN_MA,
PASS_SEPARATE,
STREAMS,
STREAMING_SOURCE,
MODEL_ORIGIN,
CROP,
CROP_X,
CROP_Y,
USE_DLSTREAM_POSTURE_DETECTION,
)
from utils.plotter import plot_bodyparts, plot_metadata_frame
from utils.poser import (
load_deeplabcut,
load_dpk,
load_dlc_live,
load_sleap,
get_pose,
calculate_skeletons,
find_local_peaks_new,
get_ma_pose,
)
def create_video_files(directory, devices, resolution, framerate, codec):
"""
Function to create videofiles for every available camera at once
:param directory: directory to output files to
:param devices: list of camera ids/names
:param resolution: resolution in which to save the video
:param framerate: framerate in which to save the video
:param codec: codec in which to save the video
:return: dictionary with files
"""
files = {}
for device in devices:
file_name = (
"VideoOutput" + device + "-" + time.strftime("%d%m%Y-%H%M%S") + ".avi"
)
output_file = os.path.join(directory, file_name)
out = cv2.VideoWriter(output_file, codec, framerate, resolution)
files[device] = out
return files
def create_row(
index, animal_skeletons, experiment_status, experiment_trial, start_time=None
):
"""
Create a pd.Series for each frame from each camera with joints position
:param experiment_trial: current trial name
:param experiment_status: current experiment status
:param index: frame index
:param animal_skeletons: skeletons for that frame
:param start_time: (optional) starting time point for Time column
"""
row_dict = {}
# creating joints columns
for num, animal in enumerate(animal_skeletons):
for joint, value in animal.items():
(
row_dict[("Animal{}".format(num + 1), joint, "x")],
row_dict[("Animal{}".format(num + 1), joint, "y")],
) = value
# optional time column
if start_time is not None:
row_dict[("Time", "", "")] = round(time.time() - start_time, 3)
# experiment columns
row_dict[("Experiment", "Status", "")] = experiment_status
if experiment_trial is None and experiment_status:
row_dict[("Experiment", "Trial", "")] = "InterTrial"
else:
row_dict[("Experiment", "Trial", "")] = experiment_trial
row = pd.Series(row_dict, name=index)
return row
def show_stream(frames):
"""
Shows some number of stream frames, depending on cameras quantity
:param frames: dictionary of frames in format of {camera:frame}
"""
if STACK_FRAMES:
if len(frames.values()) >= 2:
stack = np.hstack(frames.values())
cv2.imshow("stacked stream", stack.copy())
else:
for camera in frames:
cv2.imshow(str(camera) + " stream", frames[camera])
def cls():
"""
Clear the screen
"""
os.system("cls" if os.name == "nt" else "clear")
class DeepLabStream:
"""
Class for managing everything stream-related
"""
def __init__(self):
"""
Initializing the DeepLabStream class with some predefined variables
(more information for each below)
"""
self._camera_manager = (
self.set_camera_manager()
) # camera manager, used to get frames from the camera
self._video_codec = cv2.VideoWriter_fourcc(
*"DIVX"
) # codec in which we output the videofiles
self._start_time = None # time of initialization
self._data_row = {
camera: {} for camera in self.cameras
} # dictionary for creating row of data for each frame
self._data_output = {} # dictionary for storing rows for dataframes
self._stored_frames = {} # dictionary for storing frames
self._dlc_running = False # has DeepLabCut started?
self._experiment_running = False # has experiment started?
self._recording_running = False # has recording started?
self._video_files = None
self._multiprocessing = None # variable for multiprocessing tools
self._experiment = self.set_up_experiment()
self.frame_index = 0
self._fps_counter = []
self._fps = 0
self.greetings()
@staticmethod
def set_camera_manager():
"""
Trying to load each present camera manager, if installed
Then checking for connected cameras and choosing the one with at least some cameras connected
! Camera managers cannot be mixed
:return: the chosen camera manager
"""
def select_camera_manager():
"""
Function to select from all available camera managers
"""
manager_list = []
# loading realsense manager, if installed
if find_spec("pyrealsense2") is not None:
from utils.realsense import RealSenseManager
realsense_manager = RealSenseManager()
manager_list.append(realsense_manager)
# loading basler manager, if installed
if find_spec("pypylon") is not None:
from utils.pylon import PylonManager
pylon_manager = PylonManager()
manager_list.append(pylon_manager)
def check_for_cameras(camera_manager):
"""
Helper method to get cameras, connected to that camera manager
"""
cameras = camera_manager.get_connected_devices()
if cameras:
print(
"Found {} {} camera(s)!".format(
len(cameras), camera_manager.get_name()
)
)
return True
else:
return False
# checking for connected cameras for all installed managers
for manager in manager_list:
if check_for_cameras(manager):
return manager
else:
# if no camera is found, try generic openCV manager
generic_manager = GenericManager()
return generic_manager
MANAGER_SOURCE = {"video": VideoManager, "camera": select_camera_manager}
# loading WebCam manager, if installed
if find_spec("pyzmq") is not None:
from utils.webcam import WebCamManager
MANAGER_SOURCE["ipwebcam"] = WebCamManager()
# initialize selected manager
camera_manager = MANAGER_SOURCE.get(STREAMING_SOURCE)()
if camera_manager is not None:
return camera_manager
else:
raise ValueError(
f"Streaming source {STREAMING_SOURCE} is not a valid option. \n"
f'Please choose from "video", "camera" or "ipwebcam". Make sure that if you are using "ipwebcam" you installed the additional dependencies.'
)
@property
def cameras(self):
"""
Used to dynamically get every connected camera from camera manager
"""
return self._camera_manager.get_connected_devices()
@property
def enabled_cameras(self):
"""
Used to dynamically get every enabled camera from camera manager
"""
return self._camera_manager.get_enabled_devices()
##########################
# start cameras and stream
##########################
def enable_streams(self, streams: list):
"""
Enabling streams if they are recognized
:param streams: list of stream names
"""
available_streams = ["color", "depth", "infrared"]
for stream in streams:
if stream in available_streams:
self._camera_manager.enable_stream(RESOLUTION, FRAMERATE, stream)
else:
print(
"Stream type {} not found!\n"
"Allowed stream types: {}".format(stream, available_streams)
)
def start_cameras(self, streams: list, multi_cam: bool = True):
"""
Enabling streams and starting cameras with enabled streams
If multicam then use all available cameras, else prompt for each
:param streams: list
:param multi_cam: bool
"""
# enabling streams
self.enable_streams(streams)
# starting cameras
# self._start_time = time.time()
if multi_cam:
print("Starting all available cameras")
self._camera_manager.enable_all_devices()
else:
if len(self.cameras) != 0:
print("Starting first available camera {}".format(self.cameras[0]))
self._camera_manager.enable_device(self.cameras[0])
else:
print("No cameras found.")
print("Exiting script...")
sys.exit()
# creating output for dataframe
# self.create_output()
def stop_cameras(self):
print("Stopping all cameras")
self._camera_manager.stop()
#######################
# video outputting part
#######################
def start_recording(self):
self.create_videofiles()
self._recording_running = True
def create_videofiles(self):
"""
Create video files dictionary by cameras names
"""
#check if video was cropped
if CROP:
SAVE_RESOLUTION = (CROP_X[1] - CROP_X[0], CROP_Y[1] - CROP_Y[0])
else:
SAVE_RESOLUTION = RESOLUTION
self._video_files = create_video_files(
OUT_DIR, self.enabled_cameras, SAVE_RESOLUTION, FRAMERATE, self._video_codec
)
def write_video(self, frames: dict, index: int):
"""
Make actual videos from frames
:param frames: dict with frames, sorted by camera
:param index: int frame number
"""
for camera in frames:
font = cv2.FONT_HERSHEY_SIMPLEX
# puts the frame index in the top-left corner
cv2.putText(frames[camera], str(index), (1, 15), font, 0.5, (0, 0, 255))
self._video_files[camera].write(frames[camera])
######################
# setting up DLC usage
######################
@staticmethod
def get_pose_mp(input_q, output_q):
"""
Process to be used for each camera/DLC stream of analysis
Designed to be run in an infinite loop
:param input_q: index and corresponding frame
:param output_q: index and corresponding analysis
"""
if MODEL_ORIGIN in ("DLC", "MADLC"):
config, sess, inputs, outputs = load_deeplabcut()
while True:
if input_q.full():
index, frame = input_q.get()
start_time = time.time()
if MODEL_ORIGIN == "DLC":
scmap, locref, pose = get_pose(
frame, config, sess, inputs, outputs
)
if USE_DLSTREAM_POSTURE_DETECTION:
""" This is a legacy function that was used in earlier versions"""
peaks = find_local_peaks_new(
scmap, locref, ANIMALS_NUMBER, config
)
# Use the line below to use raw DLC output rather then DLStream optimization
else:
peaks = pose
if MODEL_ORIGIN == "MADLC":
peaks = get_ma_pose(frame, config, sess, inputs, outputs)
analysis_time = time.time() - start_time
output_q.put((index, peaks, analysis_time))
elif MODEL_ORIGIN == "DLC-LIVE":
dlc_live = load_dlc_live()
while True:
if input_q.full():
index, frame = input_q.get()
start_time = time.time()
if not dlc_live.is_initialized:
peaks = dlc_live.init_inference(frame)
else:
peaks = dlc_live.get_pose(frame)
analysis_time = time.time() - start_time
output_q.put((index, peaks, analysis_time))
elif MODEL_ORIGIN == "DEEPPOSEKIT":
predict_model = load_dpk()
while True:
if input_q.full():
index, frame = input_q.get()
start_time = time.time()
frame = frame[..., 1][..., None]
st_frame = np.stack([frame])
prediction = predict_model.predict(
st_frame, batch_size=1, verbose=True
)
peaks = prediction[0, :, :2]
analysis_time = time.time() - start_time
output_q.put((index, peaks, analysis_time))
elif MODEL_ORIGIN == "SLEAP":
sleap_model = load_sleap()
while True:
if input_q.full():
index, frame = input_q.get()
start_time = time.time()
# Make sure image is (1, height, width, channels) and uint8
# (height, width) -> (height, width, 1)
frame = np.expand_dims(frame, axis=-1) if frame.ndim == 2 else frame
# (height, width, channels) -> (1, height, width, channels)
frame = np.expand_dims(frame, axis=0) if frame.ndim == 3 else frame
# predict_on_batch is MUCH faster as it does not retrace the model graph for same size inputs
pred = sleap_model.predict_on_batch(frame)
try:
peaks = pred["instance_peaks"][0] # (n_poses, n_nodes, 2)
except KeyError:
# necessary for old sleap versions where single_instance models have different key naming
peaks = pred["peaks"]
analysis_time = time.time() - start_time
output_q.put((index, peaks, analysis_time))
else:
raise ValueError(f"Model origin {MODEL_ORIGIN} not available.")
@staticmethod
def create_mp_tools(devices):
"""
Creating easy to use dictionaries for our multiprocessing needs
:param devices: list of cameras for each we should create a separate process
:return: dictionary with process and queues for each device
"""
device_mps = {}
for device in devices:
# creating queues
device_mps[device] = {"input": mp.Queue(1), "output": mp.Queue(1)}
# creating process
process = mp.Process(
target=DeepLabStream.get_pose_mp,
args=(device_mps[device]["input"], device_mps[device]["output"]),
name=device,
)
device_mps[device]["process"] = process
return device_mps
def set_up_multiprocessing(self):
"""
Creating multiprocessing tools
"""
self._multiprocessing = self.create_mp_tools(self.enabled_cameras)
def start_dlc(self):
"""
Starting DLC in stand-alone process for each available camera
"""
for camera in self.enabled_cameras:
self._multiprocessing[camera]["process"].start()
self._dlc_running = True
#####################
# working with frames
#####################
def get_frames(self) -> tuple:
"""
Get new frameset from each camera and make color frames and infrared frames useful
:return: color_frames, depth_maps, infrared_frames
"""
c_frames, d_maps, i_frames = self._camera_manager.get_frames()
for camera in c_frames:
c_frames[camera] = np.asanyarray(c_frames[camera])
if CROP:
c_frames[camera] = c_frames[camera][
CROP_Y[0] : CROP_Y[1], CROP_X[0] : CROP_X[1]
].copy()
for camera in i_frames:
i_frames[camera] = np.asanyarray(i_frames[camera])
return c_frames, d_maps, i_frames
def input_frames_for_analysis(self, frames: tuple, index: int):
"""
Passing the frame and its index to DLC process
:param frames: frameset from camera manager
:param index: index of a frameset
"""
if self._dlc_running:
c_frames, d_maps, i_frames = frames
for camera in self._multiprocessing:
if self._multiprocessing[camera]["input"].empty():
# passes color frame to analysis
frame = c_frames[camera]
frame_time = time.time()
self._multiprocessing[camera]["input"].put((index, frame))
if d_maps:
self.store_frames(
camera, frame, d_maps[camera], frame_time, index
)
else:
self.store_frames(camera, frame, None, frame_time, index)
def get_analysed_frames(self) -> tuple:
"""
The main magic is happening here
Getting the data from DLC processes
Plotting the data
Checking data in experiments
Gathering data to a series
"""
if self._dlc_running:
analysed_frames = {}
analysis_time = None
if CROP:
frame_height = CROP_Y[1] - CROP_Y[0]
frame_width = CROP_X[1] - CROP_X[0]
else:
frame_width, frame_height = RESOLUTION
for camera in self._multiprocessing:
if self._multiprocessing[camera]["output"].full():
if self._start_time is None:
self._start_time = time.time() # getting the first frame here
# Getting the analysed data
analysed_index, peaks, analysis_time = self._multiprocessing[
camera
]["output"].get()
skeletons = calculate_skeletons(peaks, ANIMALS_NUMBER)
print(
"", end="\r", flush=True
) # this is the line you should not remove
analysed_frame, depth_map, input_time = self.get_stored_frames(
camera, analysed_index
)
delay_time = time.time() - input_time
# Calculating FPS and plotting the data on frame
self.calculate_fps(analysis_time if analysis_time != 0 else 0.01)
frame_time = time.time() - self._start_time
analysed_image = plot_metadata_frame(
plot_bodyparts(analysed_frame, skeletons),
frame_width,
frame_height,
self._fps,
frame_time,
)
# Experiments
if (
self._experiment.experiment_finished
and self._experiment_running
):
self._experiment_running = False
if (
self._experiment_running
and not self._experiment.experiment_finished
):
if ANIMALS_NUMBER > 1 and not FLATTEN_MA and not PASS_SEPARATE:
self._experiment.check_skeleton(analysed_image, skeletons)
else:
for skeleton in skeletons:
self._experiment.check_skeleton(
analysed_image, skeleton
)
# Gathering data as pd.Series for output
if self._data_output:
self.append_row(
camera,
analysed_index,
skeletons,
self._experiment_running,
self._experiment.get_trial(),
self._start_time,
)
analysed_frames[camera] = analysed_image
return analysed_frames, analysis_time
def store_frames(self, camera: str, c_frame, d_map, frame_time: float, index: int):
"""
Store frames currently sent for analysis in index based dictionary
:param camera: camera name
:param c_frame: color frame
:param d_map: depth map
:param frame_time: inputting time of frameset
:param index: index of frame that is currently analysed
"""
if camera in self._stored_frames.keys():
self._stored_frames[camera][index] = c_frame, d_map, frame_time
else:
self._stored_frames[camera] = {}
self._stored_frames[camera][index] = c_frame, d_map, frame_time
def get_stored_frames(self, camera: str, index: int):
"""
Retrieve frames currently sent for analysis, retrieved frames will be removed (popped) from the dictionary
:param camera: camera name
:param index: index of analysed frame
:return:
"""
c_frame, d_map, frame_time = self._stored_frames[camera].pop(index, None)
return c_frame, d_map, frame_time
def convert_depth_map_to_image(self, d_map):
"""
Colorize depth map using build in camera manager colorizer
:param d_map: depth map
:return: colorized image of a depth map
"""
colorized_depth_frame = self._camera_manager.colorize_depth_frame(d_map)
return colorized_depth_frame
def calculate_fps(self, current_analysis_time):
"""
Calculates average FPS for 10 frames
:param current_analysis_time: instant analysis time for one frame
"""
current_fps = 1 / current_analysis_time
self._fps_counter.append(current_fps)
if len(self._fps_counter) == 10:
self._fps = np.average(self._fps_counter)
self._fps_counter.clear()
##################
# experiments part
##################
@staticmethod
def set_up_experiment():
from experiments.utils.exp_setup import setup_experiment
experiment = setup_experiment()
return experiment
def start_experiment(self):
if not self._experiment_running:
self._experiment.start_experiment()
self._experiment_running = True
#########################
# after streaming is done
#########################
def stop_recording(self):
# finishing with the videos
if self._recording_running:
print("Saving the video")
for file in self._video_files.values():
file.release()
self._recording_running = False
self._video_files = None
print("Video saved")
def stop_dlc(self):
# cleaning up the dlc processes
if self._dlc_running:
for camera in self._multiprocessing:
# finishing the process
self._multiprocessing[camera]["process"].terminate()
# closing all the Queues
self._multiprocessing[camera]["input"].close()
self._multiprocessing[camera]["output"].close()
self._dlc_running = False
self._multiprocessing = None
self._start_time = None
# writing database
if self._data_output:
self.create_dataframes()
def stop_experiment(self):
# stopping the experiment
if self._experiment_running:
self._experiment.stop_experiment()
self._experiment_running = False
self._experiment = self.set_up_experiment()
def finish_streaming(self):
"""
Clean up after ourselves
"""
self.stop_experiment()
self.stop_dlc()
self.stop_recording()
cv2.destroyAllWindows()
#####################
# working with pandas
#####################
def create_output(self):
"""
Create lists to contain serialized skeletons
"""
for camera in self.enabled_cameras:
self._data_output[camera] = []
def append_row(
self,
camera,
index,
animal_skeletons,
experiment_status,
experiment_trial,
start_time=None,
):
"""
Create a pd.Series for each frame from each camera with joints position and store it
:param experiment_trial: current trial name
:param experiment_status: current experiment status
:param camera: camera name
:param index: frame index
:param animal_skeletons: skeletons for that frame
:param start_time: (optional) starting time point for Time column
"""
row = create_row(
index, animal_skeletons, experiment_status, experiment_trial, start_time
)
self._data_output[camera].append(row)
def create_dataframes(self):
"""
Outputting dataframes to csv
"""
for num, camera in enumerate(self._data_output):
print("Saving database for device {}".format(camera))
df = pd.DataFrame(self._data_output[camera])
df.index.name = "Frame"
df.to_csv(
OUT_DIR
+ "/DataOutput{}".format(camera)
+ "-"
+ time.strftime("%d%m%Y-%H%M%S")
+ ".csv",
sep=";",
)
print("Database saved")
######
# meta
######
@staticmethod
def greetings():
from utils.configloader import EXP_NAME, EXP_ORIGIN
print("This is DeepLabStream")
print("Developed by: Jens Schweihoff and Matvey Loshakov")
print(f"Initializing {EXP_ORIGIN.lower()} experiment: {EXP_NAME}...")
def get_camera_manager(self):
return self._camera_manager
def get_video_files(self):
return self._video_files
def get_multiprocessing_tools(self):
return self._multiprocessing
def get_enabled_cameras(self):
return self.enabled_cameras
def get_connected_cameras(self):
return self._camera_manager.get_connected_devices()
def dlc_status(self):
return self._dlc_running
def experiment_status(self):
return self._experiment_running
def recording_status(self):
return self._recording_running
def get_fps(self):
return self._fps
def get_start_time(self):
return self._start_time
# testing part
@click.command()
@click.option("--dlc-enabled", "dlc_enabled", is_flag=True)
@click.option("--benchmark-enabled", "benchmark_enabled", is_flag=True)
@click.option("--recording-enabled", "recording_enabled", is_flag=True)
@click.option("--data-output-enabled", "data_output_enabled", is_flag=True)
def start_deeplabstream(
dlc_enabled, benchmark_enabled, recording_enabled, data_output_enabled
):
if not dlc_enabled and benchmark_enabled:
print("Cannot benchmark with DLC turned off")
print("Please enable DLC with --dlc-enabled flag")
if not dlc_enabled and data_output_enabled:
print("Cannot output data with DLC turned off")
print("Please enable DLC with --dlc-enabled flag")
# initializing DeepLabStream
print("Your current config:")
print("Resolution : {}".format(RESOLUTION))
print("Framerate : {}".format(FRAMERATE))
print("Output directory : {}".format(OUT_DIR))
print("Use multiple cameras : {}".format(MULTI_CAM))
print("DLC enabled: {}".format(dlc_enabled))
print("Benchmarking enabled: {}".format(benchmark_enabled))
print("Data output enabled: {}".format(data_output_enabled))
print("Start with current config? (y/n)")
if input().lower() != "y":
print("Please edit config and restart the script.")
print("Exiting script...")
sys.exit()
print("Starting DeepLabStream manager")
stream_manager = DeepLabStream()
# starting streams
print("Starting cameras with defined above streams")
stream_manager.start_cameras(STREAMS, MULTI_CAM)
if dlc_enabled:
# starting DeepLabCut
print("Starting DeepLabCut")
stream_manager.set_up_multiprocessing()
stream_manager.start_dlc()
if benchmark_enabled:
# benchmarking tools
fps_data = []
analysis_time_data = []
whole_loop_time_data = []
tracking_accuracy_counter = 0
def describe_dataset(dataset, name):
"""
Function to describe dataset and print out the results
"""
average = np.average(dataset)
print("Average {0} time {1:6.5f}".format(name, average))
maximum = np.max(dataset)
print("Maximum {0} time {1:6.5f}".format(name, maximum))
minimum = np.min(dataset)
print("Minimum {0} time {1:6.5f}".format(name, minimum))
standart_deviation = np.std(dataset)
print(
"Standard deviation {0} time {1:6.5f}".format(name, standart_deviation)
)
def show_benchmark_statistics():
"""
Outputting all stream benchmark statistics
"""
analysis_full_time = time.time() - start_time
print("Full analysis time {0:4.2f}".format(analysis_full_time))
##################################################################
avg_fps = np.average(fps_data)
print("Average FPS {0:4.2f}".format(avg_fps))
frame_count = len(fps_data)
print("Got total {0:.0f}".format(frame_count))
tracking_accuracy = tracking_accuracy_counter / frame_count * 100
print("Tracking accuracy {0:.3f}%".format(tracking_accuracy))
##################################################################
describe_dataset(analysis_time_data, "analysis")
describe_dataset(whole_loop_time_data, "whole_loop")
if recording_enabled:
stream_manager.start_recording()
if data_output_enabled:
stream_manager.create_output()
got_first_analysed_frame = False
while True:
loop_time = time.time() # start of the loop
all_frames = stream_manager.get_frames()
color_frames, depth_maps, infrared_frames = all_frames
if recording_enabled:
stream_manager.write_video(color_frames, stream_manager.frame_index)
if dlc_enabled:
###########################################################
# Analysis part
# outputting the frames
res_frames, res_time = stream_manager.get_analysed_frames()
# inputting the frames
stream_manager.input_frames_for_analysis(
all_frames, stream_manager.frame_index
)
###########################################################
# Benchmarking part
if res_time is not None:
if benchmark_enabled and got_first_analysed_frame:
analysis_time_data.append(res_time)
else:
if benchmark_enabled and got_first_analysed_frame:
tracking_accuracy_counter += 1
###########################################################
# streaming the stream
if res_frames:
if not got_first_analysed_frame and benchmark_enabled:
got_first_analysed_frame = True
show_stream(res_frames)
else:
show_stream(color_frames)
###########################################################
# finishing the loop
stream_manager.frame_index += 1
whole_loop = time.time() - loop_time
current_fps = stream_manager.get_fps()
if benchmark_enabled:
if got_first_analysed_frame:
fps_data.append(current_fps)
whole_loop_time_data.append(whole_loop)
# exit clauses
if cv2.waitKey(1) & 0xFF == ord("q"):
start_time = stream_manager.get_start_time()
stream_manager.finish_streaming()
if benchmark_enabled:
print("Benchmark statistics:")
show_benchmark_statistics()
if recording_enabled:
stream_manager.stop_recording()
break
if benchmark_enabled:
if len(analysis_time_data) > 3000:
start_time = stream_manager.get_start_time()
stream_manager.finish_streaming()
print("Benchmark statistics:")
show_benchmark_statistics()
if recording_enabled:
stream_manager.stop_recording()
break
elif got_first_analysed_frame:
print(
"[{0}/3000] Benchmarking in progress".format(
len(analysis_time_data)
)
)
if benchmark_enabled:
model_parts = MODEL_NAME.split("_")
if len(model_parts) == 3:
short_model = model_parts[0] + "_" + model_parts[2]
else:
short_model = MODEL_NAME
# the best way to save files
np.savetxt(
f"{OUT_DIR}/{short_model}_framerate_{FRAMERATE}_resolution_{RESOLUTION[0]}_{RESOLUTION[1]}.txt",
np.transpose([fps_data, whole_loop_time_data]),
)
if __name__ == "__main__":
mp.freeze_support()
cls()
start_deeplabstream()