Skip to content

Commit cdb1fdb

Browse files
gkane26gkane26
committed
add create labeled vid, fix proc settings in gui
1 parent e46b24f commit cdb1fdb

4 files changed

Lines changed: 69 additions & 47 deletions

File tree

‎dlclivegui/dlclivegui.py‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -833,7 +833,7 @@ def edit_proc(self):
833833
def_args = inspect.getargspec(self.proc_object)
834834
self.proc_param_names = def_args[0]
835835
self.proc_param_default_values = def_args[3]
836-
self.proc_param_default_types = [type(v) for v in def_args[3]]
836+
self.proc_param_default_types = [type(v) if type(v) is not list else [type(v[0])] for v in def_args[3]]
837837
for i in range(len(def_args[0])-len(def_args[3])):
838838
self.proc_param_default_values = ('',) + self.proc_param_default_values
839839
self.proc_param_default_types = [str] + self.proc_param_default_types
@@ -867,18 +867,6 @@ def edit_proc(self):
867867
self.cfg['processor_args'][self.dlc_proc_dir.get()][self.dlc_proc_name.get()] = proc_args_gui.get_values()
868868
self.save_config()
869869

870-
# proc_param_gui = Tk()
871-
# self.proc_param_values = []
872-
# for i in range(1,len(self.proc_param_names)):
873-
# Label(proc_param_gui, text=self.proc_param_names[i]+": ").grid(sticky="w", row=i, column=0)
874-
875-
876-
877-
# self.proc_param_values.append(StringVar(proc_param_gui, value=str(self.proc_param_default_values[i])))
878-
# Entry(proc_param_gui, textvariable=self.proc_param_values[i-1]).grid(sticky="nsew", row=i, column=1)
879-
880-
# Button(proc_param_gui, text="Update", command=lambda: self.start_proc(proc_param_gui)).grid(sticky="nsew", row=i+1, column=1)
881-
882870

883871
def init_session(self):
884872

@@ -1047,7 +1035,7 @@ def createGUI(self):
10471035
### select camera ###
10481036

10491037
# camera entry
1050-
Label(self.window, text="Camera "+": ").grid(sticky="w", row=cur_row, column=0)
1038+
Label(self.window, text="Camera: ").grid(sticky="w", row=cur_row, column=0)
10511039
self.camera_name = StringVar(self.window)
10521040
self.camera_entry = Combobox(self.window, textvariable=self.camera_name)
10531041
cam_names = self.get_camera_names()

‎dlclivegui/process.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -144,10 +144,6 @@ def _pose_loop(self):
144144
else:
145145
self.q_to_process.write(cmd)
146146

147-
# print(f"POSE RATE = {int(1 / (time.time()-end_time))} / POSE LATENCY = {(pose_time-frame_time):0.03f}")
148-
149-
end_time = time.time()
150-
151147

152148
def start_record(self, timeout=5):
153149

‎dlclivegui/video.py‎

Lines changed: 66 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,23 @@
1212
import cv2
1313
import colorcet as cc
1414
from PIL import ImageColor
15-
16-
17-
def create_labeled_video(video_file, ts_file, dlc_file, out_file=None, save_images=False, cut=(0, np.Inf), cmap='bgy', radius=3, lik_thresh=0.5, display=False):
15+
from tqdm import tqdm
16+
17+
18+
def create_labeled_video(video_file,
19+
ts_file,
20+
dlc_file,
21+
out_file=None,
22+
save_images=False,
23+
cut=(0, np.Inf),
24+
cmap='bgy',
25+
radius=3,
26+
lik_thresh=0.5,
27+
write_ts=False,
28+
write_scale=2,
29+
display=False,
30+
progress=True,
31+
label=True):
1832
""" Create a labeled video from DeepLabCut-live-GUI recording
1933
2034
Parameters
@@ -27,6 +41,8 @@ def create_labeled_video(video_file, ts_file, dlc_file, out_file=None, save_imag
2741
path to DeepLabCut file
2842
out_file : str, optional
2943
path for output file. If None, output file will be "'video_file'_LABELED.avi". by default None. If NOn
44+
save_images : bool, optional
45+
boolean flag to save still images in a folder
3046
cut : tuple, optional
3147
time of video to use. Will only save labeled video for time after cut[0] and before cut[1], by default (0, np.Inf)
3248
cmap : str, optional
@@ -37,6 +53,8 @@ def create_labeled_video(video_file, ts_file, dlc_file, out_file=None, save_imag
3753
likelihood threshold to plot keypoints, by default 0.5
3854
display : bool, optional
3955
boolean flag to display images as video is written, by default False
56+
progress : bool, optional
57+
boolean flag to display progress bar
4058
4159
Raises
4260
------
@@ -48,23 +66,31 @@ def create_labeled_video(video_file, ts_file, dlc_file, out_file=None, save_imag
4866
cam_frame_times = np.load(ts_file)
4967
n_frames = cam_frame_times.size
5068

69+
70+
lab = "LABELED" if label else "UNLABELED"
5171
if out_file:
52-
out_file = f"{out_file}_VIDEO_LABELED.avi"
53-
out_times_file = f"{out_file}_TS_LABELED.npy"
72+
out_file = f"{out_file}_VIDEO_{lab}.avi"
73+
out_times_file = f"{out_file}_TS_{lab}.npy"
5474
else:
55-
out_file = f"{os.path.splitext(video_file)[0]}_LABELED.avi"
56-
out_times_file = f"{os.path.splitext(ts_file)[0]}_LABELED.npy"
75+
out_file = f"{os.path.splitext(video_file)[0]}_{lab}.avi"
76+
out_times_file = f"{os.path.splitext(ts_file)[0]}_L{lab}.npy"
5777

5878
if save_images:
5979
im_dir = os.path.splitext(out_file)[0]
60-
os.makedirs(im_dir)
80+
os.makedirs(im_dir, exist_ok=True)
6181

6282
fourcc = cv2.VideoWriter_fourcc(*'DIVX')
6383
fps = cap.get(cv2.CAP_PROP_FPS)
6484
im_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
6585
vwriter = cv2.VideoWriter(out_file, fourcc, fps, im_size)
6686
label_times = []
6787

88+
if write_ts:
89+
ts_font = cv2.FONT_HERSHEY_PLAIN
90+
ts_coord = (0, im_size[1])
91+
ts_color = (255, 255, 255)
92+
ts_size = 2
93+
6894
poses = pd.read_hdf(dlc_file)
6995
pose_times = poses['pose_time']
7096
poses = poses.melt(id_vars=['frame_time', 'pose_time'])
@@ -73,7 +99,23 @@ def create_labeled_video(video_file, ts_file, dlc_file, out_file=None, save_imag
7399
all_colors = getattr(cc, cmap)
74100
colors = [ImageColor.getcolor(c, "RGB")[::-1] for c in all_colors[::int(len(all_colors)/bodyparts.size)]]
75101

76-
for i in range(cam_frame_times.size):
102+
ind = 0
103+
vid_time = 0
104+
while vid_time < cut[0]:
105+
106+
cur_time = cam_frame_times[ind]
107+
vid_time = cur_time - cam_frame_times[0]
108+
ret, frame = cap.read()
109+
ind += 1
110+
111+
if not ret:
112+
raise Exception(f"Could not read frame = {ind+1} at time = {cur_time-cam_frame_times[0]}.")
113+
114+
115+
frame_times_sub = cam_frame_times[(cam_frame_times-cam_frame_times[0] > cut[0]) & (cam_frame_times-cam_frame_times[0] < cut[1])]
116+
iterator = tqdm(range(ind, ind+frame_times_sub.size)) if progress else range(ind, ind+frame_times_sub.size)
117+
118+
for i in iterator:
77119

78120
cur_time = cam_frame_times[i]
79121
vid_time = cur_time - cam_frame_times[0]
@@ -82,31 +124,29 @@ def create_labeled_video(video_file, ts_file, dlc_file, out_file=None, save_imag
82124
if not ret:
83125
raise Exception(f"Could not read frame = {i+1} at time = {cur_time-cam_frame_times[0]}.")
84126

85-
if vid_time > cut[1]:
86-
87-
break
88-
89-
elif vid_time > cut[0]:
90-
91-
cur_pose_time = pose_times[np.where(pose_times - cur_time > 0)[0][0]]
92-
this_pose = poses[poses['pose_time']==cur_pose_time]
127+
cur_pose_time = pose_times[np.where(pose_times - cur_time > 0)[0][0]]
128+
this_pose = poses[poses['pose_time']==cur_pose_time]
93129

130+
if label:
94131
for j in range(bodyparts.size):
95132
this_bp = this_pose[this_pose['bodyparts'] == bodyparts[j]]['value'].values
96133
if this_bp[2] > lik_thresh:
97134
x = int(this_bp[0])
98135
y = int(this_bp[1])
99136
frame = cv2.circle(frame, (x, y), radius, colors[j], thickness=-1)
100137

101-
if display:
102-
cv2.imshow('DLC Live Labeled Video', frame)
103-
cv2.waitKey(1)
104-
105-
vwriter.write(frame)
106-
label_times.append(cur_time)
107-
if save_images:
108-
new_file = f"{im_dir}/frame_{i}.png"
109-
cv2.imwrite(new_file, frame)
138+
if write_ts:
139+
frame = cv2.putText(frame, f"{vid_time:0.3f}", ts_coord, ts_font, write_scale, ts_color, ts_size)
140+
141+
if display:
142+
cv2.imshow('DLC Live Labeled Video', frame)
143+
cv2.waitKey(1)
144+
145+
vwriter.write(frame)
146+
label_times.append(cur_time)
147+
if save_images:
148+
new_file = f"{im_dir}/frame_{i}.png"
149+
cv2.imwrite(new_file, frame)
110150

111151
if display:
112152
cv2.destroyAllWindows()

‎setup.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
#!/usr/bin/env python3
2-
# -*- coding: utf-8 -*-
31
"""
42
DeepLabCut Toolbox (deeplabcut.org)
53
© A. & M. Mathis Labs
@@ -23,7 +21,7 @@
2321
long_description_content_type="text/markdown",
2422
url="https://github.com/DeepLabCut/DeepLabCut-live-GUI",
2523
python_requires = '>=3.5, <3.8',
26-
install_requires=['deeplabcut-live', 'pyserial', 'pandas', 'tables', 'multiprocess', 'imutils', 'pillow'],
24+
install_requires=['deeplabcut-live', 'pyserial', 'pandas', 'tables', 'multiprocess', 'imutils', 'pillow', 'tqdm'],
2725
packages=setuptools.find_packages(),
2826
include_package_data=True,
2927
classifiers=(

0 commit comments

Comments
 (0)