megadetector 5.0.13__py3-none-any.whl → 5.0.14__py3-none-any.whl

megadetector/detection/video_utils.py

@@ -9,6 +9,7 @@ Utilities for splitting, rendering, and assembling videos.
 #%% Constants, imports, environment
 
 import os
+import re
 import cv2
 import glob
 import json
@@ -98,7 +99,7 @@ def find_videos(dirname,
     return find_video_strings(files)
 
 
-#%% Function for rendering frames to video and vice-versa
+#%% Functions for rendering frames to video and vice-versa
 
 # http://tsaith.github.io/combine-images-into-a-video-with-python-3-and-opencv-3.html
 
@@ -169,9 +170,55 @@ def _frame_number_to_filename(frame_number):
     return 'frame{:06d}.jpg'.format(frame_number)
 
 
-def video_to_frames(input_video_file, output_folder, overwrite=True, 
-                    every_n_frames=None, verbose=False, quality=None,
-                    max_width=None):
+def _filename_to_frame_number(filename):
+    """
+    Extract the frame number from a filename that was created using 
+    _frame_number_to_filename.
+    
+    Args:
+        filename (str): a filename created with _frame_number_to_filename.
+    Returns:
+        int: the frame number extracted from [filename]
+    """
+    
+    filename = os.path.basename(filename)
+    match = re.search(r'frame(\d+)\.jpg', filename)
+    if match is None:
+        raise ValueError('{} does not appear to be a frame file'.format(filename))
+    frame_number = match.group(1)
+    try:
+        frame_number = int(frame_number)
+    except:
+        raise ValueError('Filename {} does contain a valid frame number'.format(filename))
+        
+    return frame_number
+
+
+def _add_frame_numbers_to_results(results):
+    """
+    Given the 'images' list from a set of MD results that was generated on video frames,
+    add a 'frame_number' field to each image.
+    
+    Args:
+        results (list): list of image dicts        
+    """
+    
+    # Add video-specific fields to the results
+    for im in results:
+        fn = im['file']
+        frame_number = _filename_to_frame_number(fn)
+        im['frame_number'] = frame_number
+    
+
+def video_to_frames(input_video_file, 
+                    output_folder, 
+                    overwrite=True, 
+                    every_n_frames=None, 
+                    verbose=False, 
+                    quality=None,
+                    max_width=None, 
+                    frames_to_extract=None,
+                    allow_empty_videos=False):
     """
     Renders frames from [input_video_file] to a .jpg in [output_folder].
     
@@ -184,11 +231,18 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
         output_folder (str): folder to put frame images in
         overwrite (bool, optional): whether to overwrite existing frame images
         every_n_frames (int, optional): sample every Nth frame starting from the first frame;
-            if this is None or 1, every frame is extracted
+            if this is None or 1, every frame is extracted. Mutually exclusive with
+            frames_to_extract.
         verbose (bool, optional): enable additional debug console output
         quality (int, optional): JPEG quality for frame output, from 0-100.  Defaults
             to the opencv default (typically 95).
         max_width (int, optional): resize frames to be no wider than [max_width]
+        frames_to_extract (list of int, optional): extract this specific set of frames;
+            mutually exclusive with every_n_frames.  If all values are beyond the length
+            of the video, no frames are extracted.  Can also be a single int, specifying
+            a single frame number.
+        allow_empty_videos (bool, optional): Just print a warning if a video appears to have no
+            frames (by default, this is an error).
     
     Returns:
         tuple: length-2 tuple containing (list of frame filenames,frame rate)
@@ -196,6 +250,14 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
     
     assert os.path.isfile(input_video_file), 'File {} not found'.format(input_video_file)
     
+    if isinstance(frames_to_extract,int):
+        frames_to_extract = [frames_to_extract]
+        
+    if (frames_to_extract is not None) and (every_n_frames is not None):
+        raise ValueError('frames_to_extract and every_n_frames are mutually exclusive')
+    
+    os.makedirs(output_folder,exist_ok=True)
+    
     vidcap = cv2.VideoCapture(input_video_file)
     n_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
     Fs = vidcap.get(cv2.CAP_PROP_FPS)
@@ -211,9 +273,17 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
         for frame_number in range(0,n_frames):
             
             if every_n_frames is not None:
+                assert frames_to_extract is None, \
+                    'Internal error: frames_to_extract and every_n_frames are exclusive'
                 if (frame_number % every_n_frames) != 0:
                     continue
             
+            if frames_to_extract is not None:
+                assert every_n_frames is None, \
+                    'Internal error: frames_to_extract and every_n_frames are exclusive'
+                if frame_number not in frames_to_extract:
+                    continue
+                
             frame_filename = _frame_number_to_filename(frame_number)
             frame_filename = os.path.join(output_folder,frame_filename)
             frame_filenames.append(frame_filename)
@@ -240,15 +310,23 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
         last_expected_frame_number = n_frames-1
         if every_n_frames is not None:
             last_expected_frame_number -= (every_n_frames*2)
+          
+        # When specific frames are requested, if anything is missing, reprocess the video
+        if (frames_to_extract is not None) and (missing_frame_number is not None):
             
+            pass
+        
         # If no frames are missing, or only frames very close to the end of the video are "missing",
         # skip this video
-        if (missing_frame_number is None) or \
+        elif (missing_frame_number is None) or \
             (allow_last_frame_missing and (missing_frame_number >= last_expected_frame_number)):
+                
             if verbose:
                 print('Skipping video {}, all output frames exist'.format(input_video_file))
             return frame_filenames,Fs
+        
         else:
+            
             # If we found some frames, but not all, print a message
             if verbose and found_existing_frame:
                 print("Rendering video {}, couldn't find frame {} ({}) of {}".format(
@@ -264,10 +342,14 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
 
     frame_filenames = []
 
-    # YOLOv5 does some totally bananas monkey-patching of opencv,
-    # which causes problems if we try to supply a third parameter to
-    # imwrite (to specify JPEG quality).  Detect this case, and ignore the quality 
-    # parameter if it looks like imwrite has been messed with.
+    # YOLOv5 does some totally bananas monkey-patching of opencv, which causes 
+    # problems  if we try to supply a third parameter to imwrite (to specify JPEG 
+    # quality).  Detect  this case, and ignore the quality  parameter if it looks 
+    # like imwrite has been messed with.
+    #
+    # See:
+    #
+    # https://github.com/ultralytics/yolov5/issues/7285
     imwrite_patched = False
     n_imwrite_parameters = None
     
@@ -299,6 +381,12 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
         if every_n_frames is not None:
             if frame_number % every_n_frames != 0:
                 continue
+
+        if frames_to_extract is not None:
+            if frame_number > max(frames_to_extract):
+                break
+            if frame_number not in frames_to_extract:
+                continue
             
         # Has resizing been requested?
         if max_width is not None:
@@ -350,6 +438,10 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
             except Exception as e:
                 print('Error on frame {} of {}: {}'.format(frame_number,n_frames,str(e)))
 
+    if len(frame_filenames) == 0:
+        raise Exception('Error: found no frames in file {}'.format(
+            input_video_file))
+        
     if verbose:
         print('\nExtracted {} of {} frames for {}'.format(
             len(frame_filenames),n_frames,input_video_file))
@@ -361,10 +453,12 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
 
 
 def _video_to_frames_for_folder(relative_fn,input_folder,output_folder_base,
-                                every_n_frames,overwrite,verbose,quality,max_width):
+                                every_n_frames,overwrite,verbose,quality,max_width,
+                                frames_to_extract):
     """
-    Internal function to call video_to_frames in the context of video_folder_to_frames; 
-    makes sure the right output folder exists, then calls video_to_frames.
+    Internal function to call video_to_frames for a single video in the context of 
+    video_folder_to_frames;  makes sure the right output folder exists, then calls 
+    video_to_frames.
     """    
     
     input_fn_absolute = os.path.join(input_folder,relative_fn)
@@ -379,7 +473,8 @@ def _video_to_frames_for_folder(relative_fn,input_folder,output_folder_base,
     # input_video_file = input_fn_absolute; output_folder = output_folder_video
     frame_filenames,fs = video_to_frames(input_fn_absolute,output_folder_video,
                                          overwrite=overwrite,every_n_frames=every_n_frames,
-                                         verbose=verbose,quality=quality,max_width=max_width)
+                                         verbose=verbose,quality=quality,max_width=max_width,
+                                         frames_to_extract=frames_to_extract)
     
     return frame_filenames,fs
 
@@ -388,7 +483,8 @@ def video_folder_to_frames(input_folder, output_folder_base,
                            recursive=True, overwrite=True,
                            n_threads=1, every_n_frames=None,
                            verbose=False, parallelization_uses_threads=True,
-                           quality=None, max_width=None):
+                           quality=None, max_width=None, 
+                           frames_to_extract=None):
     """
     For every video file in input_folder, creates a folder within output_folder_base, and 
     renders frame of that video to images in that folder.
@@ -402,13 +498,18 @@ def video_folder_to_frames(input_folder, output_folder_base,
         n_threads (int, optional): number of concurrent workers to use; set to <= 1 to disable
             parallelism
         every_n_frames (int, optional): sample every Nth frame starting from the first frame;
-            if this is None or 1, every frame is extracted
+            if this is None or 1, every frame is extracted. Mutually exclusive with 
+            frames_to_extract.
         verbose (bool, optional): enable additional debug console output
         parallelization_uses_threads (bool, optional): whether to use threads (True) or
             processes (False) for parallelization; ignored if n_threads <= 1
         quality (int, optional): JPEG quality for frame output, from 0-100.  Defaults
             to the opencv default (typically 95).
         max_width (int, optional): resize frames to be no wider than [max_width]
+        frames_to_extract (list of int, optional): extract this specific set of frames from
+            each video; mutually exclusive with every_n_frames.  If all values are beyond 
+            the length of a video, no frames are extracted. Can also be a single int, 
+            specifying a single frame number.
         
     Returns:
         tuple: a length-3 tuple containing:
@@ -440,7 +541,8 @@ def video_folder_to_frames(input_folder, output_folder_base,
         
             frame_filenames,fs = \
                 _video_to_frames_for_folder(input_fn_relative,input_folder,output_folder_base,
-                                            every_n_frames,overwrite,verbose,quality,max_width)
+                                            every_n_frames,overwrite,verbose,quality,max_width,
+                                            frames_to_extract)
             frame_filenames_by_video.append(frame_filenames)
             fs_by_video.append(fs)
     else:
@@ -457,7 +559,8 @@ def video_folder_to_frames(input_folder, output_folder_base,
                                              overwrite=overwrite,
                                              verbose=verbose,
                                              quality=quality,
-                                             max_width=max_width)
+                                             max_width=max_width,
+                                             frames_to_extract=frames_to_extract)
         results = list(tqdm(pool.imap(
             partial(process_video_with_options),input_files_relative_paths), 
                             total=len(input_files_relative_paths)))
@@ -485,7 +588,7 @@ class FrameToVideoOptions:
         #: video; can be 'error' or 'skip_with_warning'
         self.non_video_behavior = 'error'
     
-    
+
 def frame_results_to_video_results(input_file,output_file,options=None):
     """
     Given an MD results file produced at the *frame* level, corresponding to a directory 
@@ -511,6 +614,7 @@ def frame_results_to_video_results(input_file,output_file,options=None):
     images = input_data['images']
     detection_categories = input_data['detection_categories']
     
+    
     ## Break into videos
     
     video_to_frame_info = defaultdict(list) 
@@ -520,7 +624,9 @@ def frame_results_to_video_results(input_file,output_file,options=None):
         
         fn = im['file']
         video_name = os.path.dirname(fn)
+        
         if not is_video_file(video_name):
+            
             if options.non_video_behavior == 'error':
                 raise ValueError('{} is not a video file'.format(video_name))
             elif options.non_video_behavior == 'skip_with_warning':
@@ -529,13 +635,25 @@ def frame_results_to_video_results(input_file,output_file,options=None):
             else:
                 raise ValueError('Unrecognized non-video handling behavior: {}'.format(
                     options.non_video_behavior))
+        
+        # Attach video-specific fields to the output, specifically attach the frame
+        # number to both the video and each detection.  Only the frame number for the 
+        # canonical detection will end up in the video-level output file.
+        frame_number = _filename_to_frame_number(fn)
+        im['frame_number'] = frame_number
+        for detection in im['detections']: 
+            detection['frame_number'] = frame_number
+        
         video_to_frame_info[video_name].append(im)
     
+    # ...for each frame referred to in the results file
+    
     print('Found {} unique videos in {} frame-level results'.format(
         len(video_to_frame_info),len(images)))
     
     output_images = []
     
+    
     ## For each video...
     
     # video_name = list(video_to_frame_info.keys())[0]
@@ -594,37 +712,60 @@ def frame_results_to_video_results(input_file,output_file,options=None):
 # ...def frame_results_to_video_results(...)
 
 
-#%% Test driver
+#%% Test drivers
 
 if False:
 
+    pass
+
     #%% Constants
     
-    Fs = 30.01
-    confidence_threshold = 0.75
-    input_folder = 'z:\\'
-    frame_folder_base = r'e:\video_test\frames'
-    detected_frame_folder_base = r'e:\video_test\detected_frames'
-    rendered_videos_folder_base = r'e:\video_test\rendered_videos'
-    
-    results_file = r'results.json'
-    os.makedirs(detected_frame_folder_base,exist_ok=True)
-    os.makedirs(rendered_videos_folder_base,exist_ok=True)
-    
+    input_folder = r'G:\temp\usu-long\data'
+    frame_folder_base = r'g:\temp\usu-long-single-frames'
+    assert os.path.isdir(input_folder)
+        
     
     #%% Split videos into frames
         
     frame_filenames_by_video,fs_by_video,video_filenames = \
-        video_folder_to_frames(input_folder,frame_folder_base,recursive=True)
+        video_folder_to_frames(input_folder,
+                               frame_folder_base,
+                               recursive=True,
+                               overwrite=True,
+                               n_threads=10, 
+                               every_n_frames=None,
+                               verbose=True, 
+                               parallelization_uses_threads=True,
+                               quality=None, 
+                               max_width=None, 
+                               frames_to_extract=150)
+    
     
+    #%% Constants for detection tests
     
+    detected_frame_folder_base = r'e:\video_test\detected_frames'
+    rendered_videos_folder_base = r'e:\video_test\rendered_videos'
+    os.makedirs(detected_frame_folder_base,exist_ok=True)
+    os.makedirs(rendered_videos_folder_base,exist_ok=True)
+    results_file = r'results.json'
+    confidence_threshold = 0.75
+        
+    #%% Load detector output    
+    
+    with open(results_file,'r') as f:
+        detection_results = json.load(f)
+    detections = detection_results['images']
+    detector_label_map = detection_results['detection_categories']
+    for d in detections:
+        d['file'] = d['file'].replace('\\','/').replace('video_frames/','')
+
+
     #%% List image files, break into folders
     
     frame_files = path_utils.find_images(frame_folder_base,True)
     frame_files = [s.replace('\\','/') for s in frame_files]
     print('Enumerated {} total frames'.format(len(frame_files)))
     
-    Fs = 30.01
     # Find unique folders
     folders = set()
     # fn = frame_files[0]
@@ -634,16 +775,6 @@ if False:
     print('Found {} folders for {} files'.format(len(folders),len(frame_files)))
     
         
-    #%% Load detector output
-    
-    with open(results_file,'r') as f:
-        detection_results = json.load(f)
-    detections = detection_results['images']
-    detector_label_map = detection_results['detection_categories']
-    for d in detections:
-        d['file'] = d['file'].replace('\\','/').replace('video_frames/','')
-
-
     #%% Render detector frames
     
     # folder = list(folders)[0]

megadetector/postprocessing/convert_output_format.py

@@ -30,8 +30,11 @@ CONF_DIGITS = 3
 
 #%% Conversion functions
 
-def convert_json_to_csv(input_path,output_path=None,min_confidence=None,
-                        omit_bounding_boxes=False,output_encoding=None,
+def convert_json_to_csv(input_path,
+                        output_path=None,
+                        min_confidence=None,
+                        omit_bounding_boxes=False,
+                        output_encoding=None,
                         overwrite=True):
     """
     Converts a MD results .json file to a totally non-standard .csv format.
@@ -76,9 +79,9 @@ def convert_json_to_csv(input_path,output_path=None,min_confidence=None,
     # n_non_empty_detection_categories = len(annotation_constants.annotation_bbox_categories) - 1
     n_non_empty_detection_categories = annotation_constants.NUM_DETECTOR_CATEGORIES
     detection_category_column_names = []
-    assert annotation_constants.detector_bbox_categories[0] == 'empty'
+    assert annotation_constants.detector_bbox_category_id_to_name[0] == 'empty'
     for cat_id in range(1,n_non_empty_detection_categories+1):
-        cat_name = annotation_constants.detector_bbox_categories[cat_id]
+        cat_name = annotation_constants.detector_bbox_category_id_to_name[cat_id]
         detection_category_column_names.append('max_conf_' + cat_name)
     
     n_classification_categories = 0
@@ -370,6 +373,8 @@ def main():
     parser.add_argument('--output_path',type=str,default=None,
                         help='Output filename ending in .json or .csv (defaults to ' + \
                              'input file, with .json/.csv replaced by .csv/.json)')
+    parser.add_argument('--omit_bounding_boxes',action='store_true',
+                        help='Output bounding box text from .csv output (large and usually not useful)')        
     
     if len(sys.argv[1:]) == 0:
         parser.print_help()
@@ -386,9 +391,11 @@ def main():
             raise ValueError('Illegal input file extension')    
     
     if args.input_path.endswith('.csv') and args.output_path.endswith('.json'):
+        assert not args.omit_bounding_boxes, \
+            '--omit_bounding_boxes does not apply to csv --> json conversion'
         convert_csv_to_json(args.input_path,args.output_path)
     elif args.input_path.endswith('.json') and args.output_path.endswith('.csv'):
-        convert_json_to_csv(args.input_path,args.output_path)
+        convert_json_to_csv(args.input_path,args.output_path,omit_bounding_boxes=args.omit_bounding_boxes)
     else:
         raise ValueError('Illegal format combination')