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

megadetector/detection/tf_detector.py

@@ -152,6 +152,9 @@ class TFDetector:
         assert image_size is None, 'Image sizing not supported for TF detectors'
         assert not skip_image_resizing, 'Image sizing not supported for TF detectors'
         
+        if detection_threshold is None:
+            detection_threshold = 0
+            
         result = { 'file': image_id }
         
         try:

megadetector/detection/video_utils.py

@@ -18,6 +18,7 @@ from multiprocessing.pool import ThreadPool
 from multiprocessing.pool import Pool
 from tqdm import tqdm
 from functools import partial
+from inspect import signature
 
 from megadetector.utils import path_utils    
 from megadetector.visualization import visualization_utils as vis_utils
@@ -92,6 +93,8 @@ def find_videos(dirname,
     if convert_slashes:
         files = [fn.replace('\\', '/') for fn in files]
     
+    files = [fn for fn in files if os.path.isfile(fn)]
+    
     return find_video_strings(files)
 
 
@@ -118,8 +121,11 @@ def frames_to_video(images, Fs, output_file_name, codec_spec=default_fourcc):
         codec_spec = 'h264'
         
     if len(images) == 0:
+        print('Warning: no frames to render')
         return
 
+    os.makedirs(os.path.dirname(output_file_name),exist_ok=True)
+    
     # Determine the width and height from the first image
     frame = cv2.imread(images[0])
     cv2.imshow('video',frame)
@@ -164,7 +170,8 @@ def _frame_number_to_filename(frame_number):
 
 
 def video_to_frames(input_video_file, output_folder, overwrite=True, 
-                    every_n_frames=None, verbose=False):
+                    every_n_frames=None, verbose=False, quality=None,
+                    max_width=None):
     """
     Renders frames from [input_video_file] to a .jpg in [output_folder].
     
@@ -179,6 +186,9 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
         every_n_frames (int, optional): sample every Nth frame starting from the first frame;
             if this is None or 1, every frame is extracted
         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]
     
     Returns:
         tuple: length-2 tuple containing (list of frame filenames,frame rate)
@@ -194,36 +204,58 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
     if overwrite == False:
         
         missing_frame_number = None
+        missing_frame_filename = None
         frame_filenames = []
+        found_existing_frame = False
         
         for frame_number in range(0,n_frames):
             
             if every_n_frames is not None:
-                if frame_number % every_n_frames != 0:
+                if (frame_number % every_n_frames) != 0:
                     continue
             
             frame_filename = _frame_number_to_filename(frame_number)
             frame_filename = os.path.join(output_folder,frame_filename)
             frame_filenames.append(frame_filename)
             if os.path.isfile(frame_filename):
+                found_existing_frame = True
                 continue
             else:
                 missing_frame_number = frame_number
+                missing_frame_filename = frame_filename
                 break
     
+        if verbose and missing_frame_number is not None:
+            print('Missing frame {} ({}) for video {}'.format(
+                missing_frame_number,
+                missing_frame_filename,
+                input_video_file))
+            
         # OpenCV seems to over-report the number of frames by 1 in some cases, or fails
         # to read the last frame; either way, I'm allowing one missing frame.
         allow_last_frame_missing = True
         
-        if missing_frame_number is None or \
-            (allow_last_frame_missing and (missing_frame_number == n_frames-1)):
+        # This doesn't have to mean literally the last frame number, it just means that if
+        # we find this frame or later, we consider the video done
+        last_expected_frame_number = n_frames-1
+        if every_n_frames is not None:
+            last_expected_frame_number -= (every_n_frames*2)
+            
+        # 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 \
+            (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:
-            pass
-            # print("Rendering video {}, couldn't find frame {}".format(
-            #    input_video_file,missing_frame_number))
+            # 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(
+                    input_video_file,
+                    missing_frame_number,
+                    missing_frame_filename,
+                    last_expected_frame_number))
     
     # ...if we need to check whether to skip this video entirely
         
@@ -232,6 +264,28 @@ 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.
+    imwrite_patched = False
+    n_imwrite_parameters = None
+    
+    try:
+        # calling signature() on the native cv2.imwrite function will
+        # fail, so an exception here is a good thing.  In fact I don't think
+        # there's a case where this *succeeds* and the number of parameters
+        # is wrong.
+        sig = signature(cv2.imwrite)
+        n_imwrite_parameters = len(sig.parameters)
+    except Exception:
+        pass
+    
+    if (n_imwrite_parameters is not None) and (n_imwrite_parameters < 3):
+        imwrite_patched = True
+        if verbose and (quality is not None):
+            print('Warning: quality value supplied, but YOLOv5 has mucked with cv2.imwrite, ignoring quality')
+            
     # for frame_number in tqdm(range(0,n_frames)):
     for frame_number in range(0,n_frames):
 
@@ -246,6 +300,25 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
             if frame_number % every_n_frames != 0:
                 continue
             
+        # Has resizing been requested?
+        if max_width is not None:
+        
+            # image.shape is h/w/dims            
+            input_shape = image.shape
+            assert input_shape[2] == 3
+            input_width = input_shape[1]
+            
+            # Is resizing necessary?
+            if input_width > max_width:
+                
+                scale = max_width / input_width
+                assert scale <= 1.0
+                
+                # INTER_AREA is recommended for size reduction
+                image = cv2.resize(image, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
+        
+        # ...if we need to deal with resizing
+        
         frame_filename = _frame_number_to_filename(frame_number)
         frame_filename = os.path.join(output_folder,frame_filename)
         frame_filenames.append(frame_filename)
@@ -256,9 +329,18 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
         else:
             try:
                 if frame_filename.isascii():
-                    cv2.imwrite(os.path.normpath(frame_filename),image)
+                    
+                    if quality is None or imwrite_patched:
+                        cv2.imwrite(os.path.normpath(frame_filename),image)
+                    else:                        
+                        cv2.imwrite(os.path.normpath(frame_filename),image,
+                                    [int(cv2.IMWRITE_JPEG_QUALITY), quality])
                 else:
-                    is_success, im_buf_arr = cv2.imencode('.jpg', image)
+                    if quality is None:
+                        is_success, im_buf_arr = cv2.imencode('.jpg', image)
+                    else:
+                        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
+                        is_success, im_buf_arr = cv2.imencode('.jpg', image, encode_param)
                     im_buf_arr.tofile(frame_filename)
                 assert os.path.isfile(frame_filename), \
                     'Output frame {} unavailable'.format(frame_filename)
@@ -269,7 +351,8 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
                 print('Error on frame {} of {}: {}'.format(frame_number,n_frames,str(e)))
 
     if verbose:
-        print('\nExtracted {} of {} frames'.format(len(frame_filenames),n_frames))
+        print('\nExtracted {} of {} frames for {}'.format(
+            len(frame_filenames),n_frames,input_video_file))
 
     vidcap.release()    
     return frame_filenames,Fs
@@ -277,7 +360,8 @@ def video_to_frames(input_video_file, output_folder, overwrite=True,
 # ...def video_to_frames(...)
 
 
-def _video_to_frames_for_folder(relative_fn,input_folder,output_folder_base,every_n_frames,overwrite,verbose):
+def _video_to_frames_for_folder(relative_fn,input_folder,output_folder_base,
+                                every_n_frames,overwrite,verbose,quality,max_width):
     """
     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.
@@ -295,7 +379,7 @@ def _video_to_frames_for_folder(relative_fn,input_folder,output_folder_base,ever
     # 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)
+                                         verbose=verbose,quality=quality,max_width=max_width)
     
     return frame_filenames,fs
 
@@ -303,7 +387,8 @@ def _video_to_frames_for_folder(relative_fn,input_folder,output_folder_base,ever
 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):
+                           verbose=False, parallelization_uses_threads=True,
+                           quality=None, max_width=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.
@@ -321,7 +406,10 @@ def video_folder_to_frames(input_folder, output_folder_base,
         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]
+        
     Returns:
         tuple: a length-3 tuple containing:
             - list of lists of frame filenames; the Nth list of frame filenames corresponds to 
@@ -352,7 +440,7 @@ 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)
+                                            every_n_frames,overwrite,verbose,quality,max_width)
             frame_filenames_by_video.append(frame_filenames)
             fs_by_video.append(fs)
     else:
@@ -367,7 +455,9 @@ def video_folder_to_frames(input_folder, output_folder_base,
                                              output_folder_base=output_folder_base,
                                              every_n_frames=every_n_frames,
                                              overwrite=overwrite,
-                                             verbose=verbose)
+                                             verbose=verbose,
+                                             quality=quality,
+                                             max_width=max_width)
         results = list(tqdm(pool.imap(
             partial(process_video_with_options),input_files_relative_paths), 
                             total=len(input_files_relative_paths)))
@@ -385,13 +475,15 @@ class FrameToVideoOptions:
     frame_results_to_video_results()    
     """
     
-    #: One-indexed indicator of which frame-level confidence value to use to determine detection confidence
-    #: for the whole video, i.e. "1" means "use the confidence value from the highest-confidence frame"
-    nth_highest_confidence = 1
-    
-    #: What to do if a file referred to in a .json results file appears not to be a 
-    #: video; can be 'error' or 'skip_with_warning'
-    non_video_behavior = 'error'
+    def __init__(self):
+            
+        #: One-indexed indicator of which frame-level confidence value to use to determine detection confidence
+        #: for the whole video, i.e. "1" means "use the confidence value from the highest-confidence frame"
+        self.nth_highest_confidence = 1
+        
+        #: What to do if a file referred to in a .json results file appears not to be a 
+        #: 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):
@@ -421,7 +513,7 @@ def frame_results_to_video_results(input_file,output_file,options=None):
     
     ## Break into videos
     
-    video_to_frames = defaultdict(list) 
+    video_to_frame_info = defaultdict(list) 
     
     # im = images[0]
     for im in tqdm(images):
@@ -437,25 +529,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))
-        video_to_frames[video_name].append(im)
+        video_to_frame_info[video_name].append(im)
     
     print('Found {} unique videos in {} frame-level results'.format(
-        len(video_to_frames),len(images)))
+        len(video_to_frame_info),len(images)))
     
     output_images = []
     
     ## For each video...
     
-    # video_name = list(video_to_frames.keys())[0]
-    for video_name in tqdm(video_to_frames):
+    # video_name = list(video_to_frame_info.keys())[0]
+    for video_name in tqdm(video_to_frame_info):
         
-        frames = video_to_frames[video_name]
+        frames = video_to_frame_info[video_name]
         
         all_detections_this_video = []
         
         # frame = frames[0]
         for frame in frames:
-            if frame['detections'] is not None:
+            if ('detections' in frame) and (frame['detections'] is not None):
                 all_detections_this_video.extend(frame['detections'])
             
         # At most one detection for each category for the whole video

megadetector/postprocessing/categorize_detections_by_size.py

@@ -22,20 +22,22 @@ class SizeCategorizationOptions:
     Options used to parameterize categorize_detections_by_size().
     """
     
-    #: Thresholds to use for separation, as a fraction of the image size.
-    #:
-    #: Should be sorted from smallest to largest.
-    size_thresholds = [0.95]
-    
-    #: List of category numbers to use in separation; uses all categories if None
-    categories_to_separate = None
-    
-    #: Dimension to use for thresholding; can be "size", "width", or "height"
-    measurement = 'size'
-    
-    #: Categories to assign to thresholded ranges; should have the same length as
-    #: "size_thresholds".
-    size_category_names = ['large_detection']
+    def __init__(self):
+        
+        #: Thresholds to use for separation, as a fraction of the image size.
+        #:
+        #: Should be sorted from smallest to largest.
+        self.size_thresholds = [0.95]
+        
+        #: List of category numbers to use in separation; uses all categories if None
+        self.categories_to_separate = None
+        
+        #: Dimension to use for thresholding; can be "size", "width", or "height"
+        self.measurement = 'size'
+        
+        #: Categories to assign to thresholded ranges; should have the same length as
+        #: "size_thresholds".
+        self.size_category_names = ['large_detection']
     
                                                  
 #%% Main functions