Extract video frames in load dataset

opendm/config.py

@@ -658,7 +658,7 @@ def config(argv=None, parser=None):
                         action=StoreValue,
                         default=500,
                         metavar='<positive integer>',
-                        help='Maximum number of frames to extract from video files for processing. Default: %(default)s')
+                        help='Maximum number of frames to extract from video files for processing. Set to 0 for no limit. Default: %(default)s')
 
     parser.add_argument('--video-resolution',
                         type=int,

opendm/context.py

@@ -41,6 +41,7 @@ settings_path = os.path.join(root_path, 'settings.yaml')
 
 # Define supported image extensions
 supported_extensions = {'.jpg','.jpeg','.png', '.tif', '.tiff', '.bmp'}
+supported_video_extensions = {'.mp4', '.mov'}
 
 # Define the number of cores
 num_cores = multiprocessing.cpu_count()

opendm/video/video2dataset.py

@@ -135,7 +135,7 @@ class Video2Dataset:
         if self.f is not None:
             self.f.close()
 
-        if self.parameters.limit is not None and self.global_idx >= self.parameters.limit:
+        if self.parameters.limit is not None and self.parameters.limit > 0 and self.global_idx >= self.parameters.limit:
             log.ODM_INFO("Limit of {} frames reached, trimming dataset".format(self.parameters.limit))
             output_file_paths = limit_files(output_file_paths, self.parameters.limit)
 

stages/dataset.py

@@ -15,6 +15,7 @@ from opendm import ai
 from opendm.skyremoval.skyfilter import SkyFilter
 from opendm.bgfilter import BgFilter
 from opendm.concurrency import parallel_map
+from opendm.video.video2dataset import Parameters, Video2Dataset
 
 def save_images_database(photos, database_file):
     with open(database_file, 'w') as f:
@@ -58,22 +59,25 @@ class ODMLoadDatasetStage(types.ODM_Stage):
             except Exception as e:
                 log.ODM_WARNING("Cannot write benchmark file: %s" % str(e))
 
-        # check if the image filename is supported
-        def valid_image_filename(filename):
+        def valid_filename(filename, supported_extensions):
             (pathfn, ext) = os.path.splitext(filename)
-            return ext.lower() in context.supported_extensions and pathfn[-5:] != "_mask"
+            return ext.lower() in supported_extensions and pathfn[-5:] != "_mask"
 
         # Get supported images from dir
         def get_images(in_dir):
             entries = os.listdir(in_dir)
             valid, rejects = [], []
             for f in entries:
-                if valid_image_filename(f):
+                if valid_filename(f, context.supported_extensions):
                     valid.append(f)
                 else:
                     rejects.append(f)
             return valid, rejects
 
+        def search_video_files(in_dir):
+            entries = os.listdir(in_dir)
+            return [os.path.join(in_dir, f) for f in entries if valid_filename(f, context.supported_video_extensions)]
+
         def find_mask(photo_path, masks):
             (pathfn, ext) = os.path.splitext(os.path.basename(photo_path))
             k = "{}_mask".format(pathfn)
@@ -85,6 +89,8 @@ class ODMLoadDatasetStage(types.ODM_Stage):
                     return mask
                 else:
                     log.ODM_WARNING("Image mask {} has a space. Spaces are currently not supported for image masks.".format(mask))
+        
+
 
         # get images directory
         images_dir = tree.dataset_raw
@@ -100,6 +106,51 @@ class ODMLoadDatasetStage(types.ODM_Stage):
             if not os.path.exists(images_dir):
                 raise system.ExitException("There are no images in %s! Make sure that your project path and dataset name is correct. The current is set to: %s" % (images_dir, args.project_path))
 
+            # Check if we need to extract video frames
+            frames_db_file = os.path.join(images_dir, 'frames.json')
+            if not os.path.exists(frames_db_file) or self.rerun():
+                video_files = search_video_files(images_dir)
+
+                # If we're re-running the pipeline, and frames have been extracted during a previous run
+                # we need to remove those before re-extracting them
+                if len(video_files) > 0 and os.path.exists(frames_db_file) and self.rerun():
+                    log.ODM_INFO("Re-run, removing previously extracted video frames")
+                    frames = []
+                    try:
+                        with open(frames_db_file, 'r') as f:
+                            frames = json.loads(f.read())
+                    except Exception as e:
+                        log.ODM_WARNING("Cannot check previous video extraction: %s" % str(e))
+
+                    for f in frames:
+                        fp = os.path.join(images_dir, f)
+                        if os.path.isfile(fp):
+                            os.remove(fp)
+                
+                if len(video_files) > 0:
+                    log.ODM_INFO("Found video files (%s), extracting frames" % len(video_files))
+
+                    try:
+                        params = Parameters({
+                            "input": video_files,
+                            "output": images_dir,
+                            
+                            "blur_threshold": 0.2,
+                            "distance_threshold": 10, 
+                            "black_ratio_threshold": 0.98,
+                            "pixel_black_threshold": 0.30,
+                            "use_srt": True,
+                            "max_dimension": args.video_resolution,
+                            "limit": args.video_limit,
+                        })
+                        v2d = Video2Dataset(params)
+                        frames = v2d.ProcessVideo()
+
+                        with open(frames_db_file, 'w') as f:
+                            f.write(json.dumps([os.path.basename(f) for f in frames]))
+                    except Exception as e:
+                        log.ODM_WARNING("Could not extract video frames: %s" % str(e))
+
             files, rejects = get_images(images_dir)
             if files:
                 # create ODMPhoto list