video: add crop, bug fix

Author: gkane26

dlclivegui/video.py

@@ -18,10 +18,11 @@
 def create_labeled_video(video_file,
                          ts_file,
                          dlc_file,
-                         out_file=None,
+                         out_dir=None,
                          save_images=False,
                          cut=(0, np.Inf),
-                         cmap='bgy',
+                         crop=None,
+                         cmap='bmy',
                          radius=3,
                          lik_thresh=0.5,
                          write_ts=False,
@@ -46,7 +47,7 @@ def create_labeled_video(video_file,
     cut : tuple, optional
         time of video to use. Will only save labeled video for time after cut[0] and before cut[1], by default (0, np.Inf)
     cmap : str, optional
-        a :package:`colorcet` colormap, by default 'bgy'
+        a :package:`colorcet` colormap, by default 'bmy'
     radius : int, optional
         radius for keypoints, by default 3
     lik_thresh : float, optional
@@ -68,26 +69,34 @@ def create_labeled_video(video_file,
 
 
     lab = "LABELED" if label else "UNLABELED"
-    if out_file:
-        out_file = f"{out_file}_VIDEO_{lab}.avi"
-        out_times_file = f"{out_file}_TS_{lab}.npy"
+    if out_dir:
+        out_file = f"{out_dir}/{os.path.splitext(os.path.basename(video_file))[0]}_{lab}.avi"
+        out_times_file = f"{out_dir}/{os.path.splitext(os.path.basename(ts_file))[0]}_{lab}.npy"
     else:
         out_file = f"{os.path.splitext(video_file)[0]}_{lab}.avi"
-        out_times_file = f"{os.path.splitext(ts_file)[0]}_L{lab}.npy"
+        out_times_file = f"{os.path.splitext(ts_file)[0]}_{lab}.npy"
+
+    os.makedirs(os.path.normpath(os.path.dirname(out_file)), exist_ok=True)
 
     if save_images:
         im_dir = os.path.splitext(out_file)[0]
         os.makedirs(im_dir, exist_ok=True)
 
+    im_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+    if crop is not None:
+        crop = np.max(np.vstack((crop, [0, im_size[1], 0, im_size[0]])), axis=0)
+        im_size = (crop[3]-crop[2], crop[1]-crop[0])
+
     fourcc = cv2.VideoWriter_fourcc(*'DIVX')
     fps = cap.get(cv2.CAP_PROP_FPS)
-    im_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
     vwriter = cv2.VideoWriter(out_file, fourcc, fps, im_size)
     label_times = []
-
+    
     if write_ts:
         ts_font = cv2.FONT_HERSHEY_PLAIN
-        ts_coord = (0, im_size[1])
+        ts_w = 0 if crop is None else crop[0]
+        ts_h = im_size[1] if crop is None else crop[1]
+        ts_coord = (ts_w, ts_h)
         ts_color = (255, 255, 255)
         ts_size = 2
 
@@ -114,7 +123,8 @@ def create_labeled_video(video_file,
 
     frame_times_sub = cam_frame_times[(cam_frame_times-cam_frame_times[0] > cut[0]) & (cam_frame_times-cam_frame_times[0] < cut[1])]
     iterator = tqdm(range(ind, ind+frame_times_sub.size)) if progress else range(ind, ind+frame_times_sub.size)
-    
+    this_pose = np.zeros((bodyparts.size, 3))
+
     for i in iterator:
 
         cur_time = cam_frame_times[i]
@@ -124,8 +134,10 @@ def create_labeled_video(video_file,
         if not ret:
             raise Exception(f"Could not read frame = {i+1} at time = {cur_time-cam_frame_times[0]}.")
 
-        cur_pose_time = pose_times[np.where(pose_times - cur_time > 0)[0][0]]
-        this_pose = poses[poses['pose_time']==cur_pose_time]
+        poses_before_index = np.where(pose_times < cur_time)[0]
+        if poses_before_index.size > 0:
+            cur_pose_time = pose_times[poses_before_index[-1]]
+            this_pose = poses[poses['pose_time']==cur_pose_time]
 
         if label:
             for j in range(bodyparts.size):
@@ -134,6 +146,9 @@ def create_labeled_video(video_file,
                     x = int(this_bp[0])
                     y = int(this_bp[1])
                     frame = cv2.circle(frame, (x, y), radius, colors[j], thickness=-1)
+        
+        if crop is not None:
+            frame = frame[crop[0]:crop[1], crop[2]:crop[3]]
 
         if write_ts:
             frame = cv2.putText(frame, f"{vid_time:0.3f}", ts_coord, ts_font, write_scale, ts_color, ts_size)
@@ -153,3 +168,45 @@ def create_labeled_video(video_file,
 
     vwriter.release()
     np.save(out_times_file, label_times)
+
+
+def main():
+
+    import argparse
+    import os
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('file', type=str)
+    parser.add_argument('-o', '--out-dir', type=str, default=None)
+    parser.add_argument('-s', '--save-images', action='store_true')
+    parser.add_argument('-u', '--cut', nargs='+', type=float, default=[0, np.Inf])
+    parser.add_argument('-c', '--crop', nargs='+', type=int, default=None)
+    parser.add_argument('-m', '--cmap', type=str, default='bmy')
+    parser.add_argument('-r', '--radius', type=int, default=3)
+    parser.add_argument('-l', '--lik-thresh', type=float, default=0.5)
+    parser.add_argument('-w', '--write-ts', action='store_true')
+    parser.add_argument('--write-scale', type=int, default=2)
+    parser.add_argument('-d', '--display', action='store_true')
+    parser.add_argument('--no-progress', action='store_false')
+    parser.add_argument('--no-label', action='store_false')
+    args = parser.parse_args()
+
+    vid_file = os.path.normpath(f"{args.file}_VIDEO.avi")
+    ts_file = os.path.normpath(f"{args.file}_TS.npy")
+    dlc_file = os.path.normpath(f"{args.file}_DLC.hdf5")
+
+    create_labeled_video(vid_file,
+                         ts_file,
+                         dlc_file,
+                         out_dir=args.out_dir,
+                         save_images=args.save_images,
+                         cut=tuple(args.cut),
+                         crop=args.crop,
+                         cmap=args.cmap,
+                         radius=args.radius,
+                         lik_thresh=args.lik_thresh,
+                         write_ts=args.write_ts,
+                         write_scale=args.write_scale,
+                         display=args.display,
+                         progress=args.no_progress,
+                         label=args.no_label)