forgot to add frame variable to arguments

generate_yolo_regions.py

@@ -57,9 +57,9 @@ if __name__ == "__main__":
             sub_region = sub_region.split('-')
             start_loc = int(sub_region[0])
             end_loc = int(sub_region[1]) + 1
-            for frame in range(start_loc, end_loc - 5, 5):
-                vc_sub = vc[:, frame:frame+5, :, :]
-                if vc_sub.size(1) < 5:
+            for frame in range(start_loc, end_loc - args.num_frames, args.num_frames):
+                vc_sub = vc[:, frame:frame+args.num_frames, :, :]
+                if vc_sub.size(1) < args.num_frames:
                     continue
 
                 for clip in get_yolo_regions(yolo_model, vc_sub, is_right, image_width, image_height):
@@ -72,8 +72,8 @@ if __name__ == "__main__":
                 sub_region = sub_region.split('-')
                 start_loc = int(sub_region[0])
                 if len(sub_region) == 1:
-                    vc_sub = vc[:, start_loc:start_loc+5, :, :]
-                    if vc_sub.size(1) < 5:
+                    vc_sub = vc[:, start_loc:start_loc+args.num_frames, :, :]
+                    if vc_sub.size(1) < args.num_frames:
                         continue
 
                     for clip in get_yolo_regions(yolo_model, vc_sub, is_right, image_width, image_height):
@@ -86,26 +86,26 @@ if __name__ == "__main__":
                         end_loc = vc.size(1)
                     else:
                         end_loc = int(end_loc)
-                    for frame in range(start_loc, end_loc - 5, 5):
-                        vc_sub = vc[:, frame:frame+5, :, :]
+                    for frame in range(start_loc, end_loc - args.num_frames, args.num_frames):
+                        vc_sub = vc[:, frame:frame+args.num_frames, :, :]
 #                                         cv2.imwrite('test.png', vc_sub[0, 0, :, :].unsqueeze(2).numpy())
-                        if vc_sub.size(1) < 5:
+                        if vc_sub.size(1) < args.num_frames:
                             continue
                         for clip in get_yolo_regions(yolo_model, vc_sub, is_right, image_width, image_height):
                             clip = transforms(clip)
                             tv.io.write_video(os.path.join(args.output_dir, 'positive_yolo' + str(output_count) + '.mp4'), clip.swapaxes(0,1).swapaxes(1,2).swapaxes(2,3).numpy(), fps=20)
                             output_count += 1
     elif label == 'Positives':
-        vc_sub = vc[:, -5:, :, :]
-        if not vc_sub.size(1) < 5:
+        vc_sub = vc[:, -args.num_frames:, :, :]
+        if not vc_sub.size(1) < args.num_frames:
             for clip in get_yolo_regions(yolo_model, vc_sub, is_right, image_width, image_height):
                 clip = transforms(clip)
                 tv.io.write_video(os.path.join(args.output_dir, 'positive_yolo' + str(output_count) + '.mp4'), clip.swapaxes(0,1).swapaxes(1,2).swapaxes(2,3).numpy(), fps=20)
                 output_count += 1
     elif label == 'Negatives':
-        for j in range(0, vc.size(1) - 5, 5):
-            vc_sub = vc[:, j:j+5, :, :]
-            if not vc_sub.size(1) < 5:
+        for j in range(0, vc.size(1) - args.num_frames, args.num_frames):
+            vc_sub = vc[:, j:j+args.num_frames, :, :]
+            if not vc_sub.size(1) < args.num_frames:
                 for clip in get_yolo_regions(yolo_model, vc_sub, is_right, image_width, image_height):
                     clip = transforms(clip)
                     tv.io.write_video(os.path.join(args.output_dir, 'negative_yolo' + str(output_count) + '.mp4'), clip.swapaxes(0,1).swapaxes(1,2).swapaxes(2,3).numpy(), fps=20)