megadetector 5.0.15__py3-none-any.whl → 5.0.17__py3-none-any.whl

megadetector/data_management/yolo_to_coco.py

@@ -29,9 +29,9 @@ from megadetector.data_management.yolo_output_to_md_output import read_classes_f
 
 def _filename_to_image_id(fn):
     """
-    Image IDs can't have spaces in them, replae spaces with underscores
+    Image IDs can't have spaces in them, replace spaces with underscores
     """
-    return fn.replace(' ','_')
+    return fn.replace(' ','_').replace('\\','/')
 
 
 def _process_image(fn_abs,input_folder,category_id_to_name):
@@ -40,7 +40,9 @@ def _process_image(fn_abs,input_folder,category_id_to_name):
     """
     
     # Create the image object for this image
-    fn_relative = os.path.relpath(fn_abs,input_folder)
+    #
+    # Always use forward slashes in image filenames and IDs
+    fn_relative = os.path.relpath(fn_abs,input_folder).replace('\\','/')
     image_id = _filename_to_image_id(fn_relative)
     
     # This is done in a separate loop now
@@ -51,7 +53,7 @@ def _process_image(fn_abs,input_folder,category_id_to_name):
     # image_ids.add(image_id)
     
     im = {}
-    im['file_name'] = fn_relative    
+    im['file_name'] = fn_relative
     im['id'] = image_id
     
     annotations_this_image = []
@@ -393,7 +395,8 @@ def yolo_to_coco(input_folder,
                  pool_type='thread',
                  recursive=True,
                  exclude_string=None,
-                 include_string=None):
+                 include_string=None,
+                 overwrite_handling='overwrite'):
     """
     Converts a YOLO-formatted dataset to a COCO-formatted dataset.
     
@@ -427,6 +430,8 @@ def yolo_to_coco(input_folder,
         recursive (bool, optional): whether to recurse into [input_folder]
         exclude_string (str, optional): exclude any images whose filename contains a string
         include_string (str, optional): include only images whose filename contains a string
+        overwrite_handling (bool, optional): behavior if output_file exists ('load', 'overwrite', or 
+            'error')
     
     Returns:
         dict: COCO-formatted data, the same as what's written to [output_file]
@@ -441,7 +446,21 @@ def yolo_to_coco(input_folder,
         ('no_annotations','empty_annotations','skip','error'), \
             'Unrecognized empty image handling spec: {}'.format(empty_image_handling)
      
-            
+    if (output_file is not None) and os.path.isfile(output_file):
+        
+            if overwrite_handling == 'overwrite':
+                print('Warning: output file {} exists, over-writing'.format(output_file))
+            elif overwrite_handling == 'load':
+                print('Output file {} exists, loading and returning'.format(output_file))
+                with open(output_file,'r') as f:
+                    d = json.load(f)
+                return d
+            elif overwrite_handling == 'error':
+                raise ValueError('Output file {} exists'.format(output_file))
+            else:
+                raise ValueError('Unrecognized overwrite_handling value: {}'.format(overwrite_handling))
+                
+                
     ## Read class names
     
     category_id_to_name = load_yolo_class_list(class_name_file)

megadetector/detection/process_video.py

@@ -32,11 +32,14 @@ from megadetector.visualization import visualize_detector_output
 from megadetector.utils.ct_utils import args_to_object
 from megadetector.utils.path_utils import insert_before_extension, clean_path
 from megadetector.detection.video_utils import video_to_frames
+from megadetector.detection.video_utils import run_callback_on_frames
+from megadetector.detection.video_utils import run_callback_on_frames_for_folder
 from megadetector.detection.video_utils import frames_to_video
 from megadetector.detection.video_utils import frame_results_to_video_results
 from megadetector.detection.video_utils import _add_frame_numbers_to_results
 from megadetector.detection.video_utils import video_folder_to_frames
 from megadetector.detection.video_utils import default_fourcc
+from megadetector.detection.run_detector import load_detector
 
 
 #%% Classes
@@ -76,7 +79,8 @@ class ProcessVideoOptions:
         
         #: Should we render a video with detection boxes?
         #:
-        #: Only supported when processing a single video, not a folder.
+        #: If processing a folder, this renders each input video to a separate
+        #: video with detection boxes.
         self.render_output_video = False
         
         #: If we are rendering boxes to a new video, should we keep the temporary
@@ -142,6 +146,10 @@ class ProcessVideoOptions:
         #: For debugging only, stop processing after a certain number of frames.
         self.debug_max_frames = -1
         
+        #: For debugging only, force on-disk frame extraction, even if it wouldn't otherwise be
+        #: necessary
+        self.force_on_disk_frame_extraction = False
+        
         #: File containing non-standard categories, typically only used if you're running a non-MD
         #: detector.
         self.class_mapping_filename = None
@@ -158,6 +166,11 @@ class ProcessVideoOptions:
         
         #: Enable image augmentation
         self.augment = False
+        
+        #: By default, a video with no frames (or no frames retrievable with the current parameters)
+        #: is an error, this makes it a warning.  This would apply if you request, e.g., the 100th
+        #: frame from each video, but a video only has 50 frames.
+        self.allow_empty_videos = False
     
 # ...class ProcessVideoOptions
 
@@ -206,6 +219,9 @@ def _clean_up_rendered_frames(options,rendering_output_folder,detected_frame_fil
     If necessary, delete rendered frames and/or the entire rendering output folder.
     """
     
+    if rendering_output_folder is None:
+        return
+    
     caller_provided_rendering_output_folder = (options.frame_rendering_folder is not None)
     
     # (Optionally) delete the temporary directory we used for rendered detection images
@@ -254,6 +270,9 @@ def _clean_up_extracted_frames(options,frame_output_folder,frame_filenames):
     If necessary, delete extracted frames and/or the entire temporary frame folder.
     """
     
+    if frame_output_folder is None:
+        return
+    
     caller_provided_frame_output_folder = (options.frame_folder is not None)
                                            
     if not options.keep_extracted_frames:
@@ -274,6 +293,9 @@ def _clean_up_extracted_frames(options,frame_output_folder,frame_filenames):
             # ...otherwise just delete the frames, but leave the folder in place
             else:
                 
+                if frame_filenames is None:
+                    return
+                
                 if options.force_extracted_frame_folder_deletion:
                     assert caller_provided_frame_output_folder
                     print('Warning: force_extracted_frame_folder_deletion supplied with a ' + \
@@ -322,66 +344,118 @@ def process_video(options):
     caller_provided_frame_output_folder = (options.frame_folder is not None)
     caller_provided_rendering_output_folder = (options.frame_rendering_folder is not None)
     
-    # This does not create any folders, just defines temporary folder names in 
-    # case we need them.
-    temporary_folder_info = _select_temporary_output_folders(options)
-        
-    if (caller_provided_frame_output_folder):
-        frame_output_folder = options.frame_folder
-    else:
-        frame_output_folder = temporary_folder_info['frame_output_folder']
-   
-    os.makedirs(frame_output_folder, exist_ok=True)
-
-
-    ## Extract frames
-    
-    frame_filenames, Fs = video_to_frames(
-                            options.input_video_file, 
-                            frame_output_folder, 
-                            every_n_frames=options.frame_sample, 
-                            overwrite=(not options.reuse_frames_if_available),
-                            quality=options.quality, 
-                            max_width=options.max_width, 
-                            verbose=options.verbose,
-                            frames_to_extract=options.frames_to_extract)
-
-    image_file_names = frame_filenames
-    if options.debug_max_frames > 0:
-        image_file_names = image_file_names[0:options.debug_max_frames]
-    
-    if options.model_file == 'no_detection':
-        assert options.keep_extracted_frames, \
-            'Internal error: keep_extracted_frames not set, but no model specified'
-        return
-    
+    frame_output_folder = None
+    frame_filenames = None
     
-    ## Run MegaDetector
-    
-    if options.reuse_results_if_available and \
-        os.path.isfile(options.output_json_file):
+    # If we should re-use existing results, and the output file exists, don't bother running MD
+    if (options.reuse_results_if_available and os.path.isfile(options.output_json_file)):
+            
             print('Loading results from {}'.format(options.output_json_file))
             with open(options.output_json_file,'r') as f:
                 results = json.load(f)
-    else:
-        results = run_detector_batch.load_and_run_detector_batch(
-            options.model_file, 
-            image_file_names,
-            confidence_threshold=options.json_confidence_threshold,
-            n_cores=options.n_cores,
-            class_mapping_filename=options.class_mapping_filename,
-            quiet=True,
-            augment=options.augment,
-            image_size=options.image_size)
-        
-        _add_frame_numbers_to_results(results)
+    
+    # Run MD in memory if we don't need to generate frames
+    #
+    # Currently if we're generating an output video, we need to generate frames on disk first.
+    elif (not options.keep_extracted_frames and \
+          not options.render_output_video and \
+          not options.force_on_disk_frame_extraction):
+        
+        # Run MegaDetector in memory
+        
+        if options.verbose:
+            print('Running MegaDetector in memory for {}'.format(options.input_video_file))
+            
+        if options.frame_folder is not None:
+            print('Warning: frame_folder specified, but keep_extracted_frames is ' + \
+                  'not; no raw frames will be written')
+        
+        detector = load_detector(options.model_file)
+        
+        def frame_callback(image_np,image_id):
+            return detector.generate_detections_one_image(image_np,
+                                                          image_id,
+                                                          detection_threshold=options.json_confidence_threshold,
+                                                          augment=options.augment)
+        
+        frame_results = run_callback_on_frames(options.input_video_file, 
+                                               frame_callback,
+                                               every_n_frames=options.frame_sample, 
+                                               verbose=options.verbose, 
+                                               frames_to_process=options.frames_to_extract)
+        
+        frame_results['results'] = _add_frame_numbers_to_results(frame_results['results'])
         
         run_detector_batch.write_results_to_file(
-            results, options.output_json_file,
-            relative_path_base=frame_output_folder,
+            frame_results['results'], 
+            options.output_json_file,
+            relative_path_base=None,
             detector_file=options.model_file,
-            custom_metadata={'video_frame_rate':Fs})
+            custom_metadata={'video_frame_rate':frame_results['frame_rate']})
 
+    # Extract frames and optionally run MegaDetector on those frames
+    else:
+                
+        if options.verbose:
+            print('Extracting frames for {}'.format(options.input_video_file))
+            
+        # This does not create any folders, just defines temporary folder names in 
+        # case we need them.
+        temporary_folder_info = _select_temporary_output_folders(options)
+            
+        if (caller_provided_frame_output_folder):
+            frame_output_folder = options.frame_folder
+        else:
+            frame_output_folder = temporary_folder_info['frame_output_folder']
+       
+        os.makedirs(frame_output_folder, exist_ok=True)
+    
+    
+        ## Extract frames
+        
+        frame_filenames, Fs = video_to_frames(
+                                options.input_video_file, 
+                                frame_output_folder, 
+                                every_n_frames=options.frame_sample, 
+                                overwrite=(not options.reuse_frames_if_available),
+                                quality=options.quality, 
+                                max_width=options.max_width, 
+                                verbose=options.verbose,
+                                frames_to_extract=options.frames_to_extract,
+                                allow_empty_videos=options.allow_empty_videos)
+    
+        image_file_names = frame_filenames
+        if options.debug_max_frames > 0:
+            image_file_names = image_file_names[0:options.debug_max_frames]
+                
+        ## Run MegaDetector on those frames
+        
+        if options.model_file != 'no_detection':
+        
+            if options.verbose:
+                print('Running MD for {}'.format(options.input_video_file))
+                
+            results = run_detector_batch.load_and_run_detector_batch(
+                options.model_file, 
+                image_file_names,
+                confidence_threshold=options.json_confidence_threshold,
+                n_cores=options.n_cores,
+                class_mapping_filename=options.class_mapping_filename,
+                quiet=True,
+                augment=options.augment,
+                image_size=options.image_size)
+        
+            results = _add_frame_numbers_to_results(results)
+        
+            run_detector_batch.write_results_to_file(
+                results, 
+                options.output_json_file,
+                relative_path_base=frame_output_folder,
+                detector_file=options.model_file,
+                custom_metadata={'video_frame_rate':Fs})
+            
+    # ...if we are/aren't keeping raw frames on disk
+        
     
     ## (Optionally) render output video
     
@@ -470,81 +544,141 @@ def process_video_folder(options):
     # case we need them.
     temporary_folder_info = _select_temporary_output_folders(options)
     
+    frame_output_folder = None
+    image_file_names = None
+    video_filename_to_fs = {}
     
-    ## Split every video into frames
+    # Run MD in memory if we don't need to generate frames
+    #
+    # Currently if we're generating an output video, we need to generate frames on disk first.
+    if (not options.keep_extracted_frames and \
+        not options.render_output_video and \
+        not options.force_on_disk_frame_extraction):
+                
+        if options.verbose:
+            print('Running MegaDetector in memory for folder {}'.format(options.input_video_file))
+            
+        if options.frame_folder is not None:
+            print('Warning: frame_folder specified, but keep_extracted_frames is ' + \
+                  'not; no raw frames will be written')
+        
+        detector = load_detector(options.model_file)
+        
+        def frame_callback(image_np,image_id):
+            return detector.generate_detections_one_image(image_np,
+                                                          image_id,
+                                                          detection_threshold=options.json_confidence_threshold,
+                                                          augment=options.augment)
+        
+        md_results = run_callback_on_frames_for_folder(input_video_folder=options.input_video_file, 
+                                                       frame_callback=frame_callback,
+                                                       every_n_frames=options.frame_sample, 
+                                                       verbose=options.verbose)
+        
+        video_results = md_results['results']
+        
+        for i_video,video_filename in enumerate(md_results['video_filenames']):
+            assert video_filename not in video_filename_to_fs
+            video_filename_to_fs[video_filename] = md_results['frame_rates'][i_video]
+        
+        all_frame_results = []
+        
+        # r = video_results[0]
+        for frame_results in video_results:
+            _add_frame_numbers_to_results(frame_results)
+            all_frame_results.extend(frame_results)
+        
+        run_detector_batch.write_results_to_file(
+            all_frame_results, 
+            frames_json,
+            relative_path_base=None,
+            detector_file=options.model_file)
     
-    if caller_provided_frame_output_folder:
-        frame_output_folder = options.frame_folder
     else:
-        frame_output_folder = temporary_folder_info['frame_output_folder']
-        
-    os.makedirs(frame_output_folder, exist_ok=True)
-
-    print('Extracting frames')
-    
-    frame_filenames, Fs, video_filenames = \
-        video_folder_to_frames(input_folder=options.input_video_file,
-                               output_folder_base=frame_output_folder, 
-                               recursive=options.recursive, 
-                               overwrite=(not options.reuse_frames_if_available),
-                               n_threads=options.n_cores,
-                               every_n_frames=options.frame_sample,
-                               verbose=options.verbose,
-                               quality=options.quality,
-                               max_width=options.max_width,
-                               frames_to_extract=options.frames_to_extract)
-    
-    print('Extracted frames for {} videos'.format(len(set(video_filenames))))
-    image_file_names = list(itertools.chain.from_iterable(frame_filenames))
-    
-    if len(image_file_names) == 0:
-        if len(video_filenames) == 0:
-            print('No videos found in folder {}'.format(options.input_video_file))
+        
+        ## Split every video into frames
+        
+        if options.verbose:
+            print('Extracting frames for folder {}'.format(options.input_video_file))
+            
+        if caller_provided_frame_output_folder:
+            frame_output_folder = options.frame_folder
         else:
-            print('No frames extracted from folder {}, this may be due to an '\
-                  'unsupported video codec'.format(options.input_video_file))
-        return
-
-    if options.debug_max_frames is not None and options.debug_max_frames > 0:
-        image_file_names = image_file_names[0:options.debug_max_frames]
+            frame_output_folder = temporary_folder_info['frame_output_folder']
+            
+        os.makedirs(frame_output_folder, exist_ok=True)
+    
+        frame_filenames, Fs, video_filenames = \
+            video_folder_to_frames(input_folder=options.input_video_file,
+                                   output_folder_base=frame_output_folder, 
+                                   recursive=options.recursive, 
+                                   overwrite=(not options.reuse_frames_if_available),
+                                   n_threads=options.n_cores,
+                                   every_n_frames=options.frame_sample,
+                                   verbose=options.verbose,
+                                   quality=options.quality,
+                                   max_width=options.max_width,
+                                   frames_to_extract=options.frames_to_extract,
+                                   allow_empty_videos=options.allow_empty_videos)
         
-    if options.model_file == 'no_detection':
-        assert options.keep_extracted_frames, \
-            'Internal error: keep_extracted_frames not set, but no model specified'
-        return
-    
-    
-    ## Run MegaDetector on the extracted frames
-    
-    if options.reuse_results_if_available and \
-        os.path.isfile(frames_json):
-            print('Bypassing inference, loading results from {}'.format(frames_json))
-            results = None
-    else:
-        print('Running MegaDetector')
-        results = run_detector_batch.load_and_run_detector_batch(
-            options.model_file, 
-            image_file_names,
-            confidence_threshold=options.json_confidence_threshold,
-            n_cores=options.n_cores,
-            class_mapping_filename=options.class_mapping_filename,
-            quiet=True,
-            augment=options.augment,
-            image_size=options.image_size)
-    
-        _add_frame_numbers_to_results(results)
+        for i_video,video_filename in enumerate(video_filenames):
+            assert video_filename not in video_filename_to_fs
+            video_filename_to_fs[video_filename] = Fs[i_video]
+                
+        print('Extracted frames for {} videos'.format(len(set(video_filenames))))
+        image_file_names = list(itertools.chain.from_iterable(frame_filenames))
         
-        run_detector_batch.write_results_to_file(
-            results, frames_json,
-            relative_path_base=frame_output_folder,
-            detector_file=options.model_file,
-            custom_metadata={'video_frame_rate':Fs})
+        if len(image_file_names) == 0:
+            if len(video_filenames) == 0:
+                print('No videos found in folder {}'.format(options.input_video_file))
+            else:
+                print('No frames extracted from folder {}, this may be due to an '\
+                      'unsupported video codec'.format(options.input_video_file))
+            return
     
+        if options.debug_max_frames is not None and options.debug_max_frames > 0:
+            image_file_names = image_file_names[0:options.debug_max_frames]
+            
+        if options.model_file == 'no_detection':
+            assert options.keep_extracted_frames, \
+                'Internal error: keep_extracted_frames not set, but no model specified'
+            return
+        
+        
+        ## Run MegaDetector on the extracted frames
+        
+        if options.reuse_results_if_available and \
+            os.path.isfile(frames_json):
+                print('Bypassing inference, loading results from {}'.format(frames_json))
+                with open(frames_json,'r') as f:
+                    results = json.load(f)
+        else:
+            print('Running MegaDetector')
+            results = run_detector_batch.load_and_run_detector_batch(
+                options.model_file, 
+                image_file_names,
+                confidence_threshold=options.json_confidence_threshold,
+                n_cores=options.n_cores,
+                class_mapping_filename=options.class_mapping_filename,
+                quiet=True,
+                augment=options.augment,
+                image_size=options.image_size)
+        
+            _add_frame_numbers_to_results(results)
+            
+            run_detector_batch.write_results_to_file(
+                results, 
+                frames_json,
+                relative_path_base=frame_output_folder,
+                detector_file=options.model_file)
+        
+    # ...if we're running MD on in-memory frames vs. extracting frames to disk
     
     ## Convert frame-level results to video-level results
 
     print('Converting frame-level results to video-level results')
-    frame_results_to_video_results(frames_json,video_json)
+    frame_results_to_video_results(frames_json,video_json,
+                                   video_filename_to_frame_rate=video_filename_to_fs)
 
 
     ## (Optionally) render output videos
@@ -646,13 +780,13 @@ def process_video_folder(options):
 
 def options_to_command(options):
     """
-    Convert a ProcessVideoOptions obejct to a corresponding command line.
+    Convert a ProcessVideoOptions object to a corresponding command line.
     
     Args:
         options (ProcessVideoOptions): the options set to render as a command line
         
     Returns:
-        str: the command line coresponding to [options]
+        str: the command line corresponding to [options]
         
     :meta private:
     """
@@ -725,8 +859,8 @@ if False:
     #%% Process a folder of videos
     
     model_file = 'MDV5A'
-    # input_dir = r'g:\temp\test-videos'
-    input_dir = r'G:\temp\md-test-package\md-test-images\video-samples'
+    input_dir = r'g:\temp\test-videos'
+    # input_dir = r'G:\temp\md-test-package\md-test-images\video-samples'
     output_base = r'g:\temp\video_test'
     frame_folder = os.path.join(output_base,'frames')
     rendering_folder = os.path.join(output_base,'rendered-frames')
@@ -744,25 +878,26 @@ if False:
     options.recursive = True
     options.reuse_frames_if_available = False
     options.reuse_results_if_available = False
-    options.quality = 90
+    options.quality = None # 90
     options.frame_sample = 10
-    options.max_width = 1280
+    options.max_width = None # 1280
     options.n_cores = 4
     options.verbose = True
-    options.render_output_video = True    
+    options.render_output_video = False  
     options.frame_folder = frame_folder
     options.frame_rendering_folder = rendering_folder    
-    options.keep_extracted_frames = True
-    options.keep_rendered_frames = True
+    options.keep_extracted_frames = False
+    options.keep_rendered_frames = False
     options.force_extracted_frame_folder_deletion = False
     options.force_rendered_frame_folder_deletion = False
     options.fourcc = 'mp4v'
+    options.force_on_disk_frame_extraction = True
     # options.rendering_confidence_threshold = 0.15
     
     cmd = options_to_command(options); print(cmd)
         
     # import clipboard; clipboard.copy(cmd)
-    # process_video_folder(options)
+    process_video_folder(options)
         
     
     #%% Process a single video
@@ -988,6 +1123,10 @@ def main():
     parser.add_argument('--augment',
                         action='store_true',
                         help='Enable image augmentation')
+    
+    parser.add_argument('--allow_empty_videos',
+                        action='store_true',
+                        help='By default, videos with no retrievable frames cause an error, this makes it a warning')
             
     if len(sys.argv[1:]) == 0:
         parser.print_help()
@@ -1000,6 +1139,10 @@ def main():
     if os.path.isdir(options.input_video_file):
         process_video_folder(options)
     else:
+        assert os.path.isfile(options.input_video_file), \
+            '{} is not a valid file or folder name'.format(options.input_video_file)
+        assert not options.recursive, \
+            '--recursive is only meaningful when processing a folder'
         process_video(options)
 
 if __name__ == '__main__':

megadetector/detection/pytorch_detector.py

@@ -2,7 +2,7 @@
 
 pytorch_detector.py
 
-Module to run MegaDetector v5, a PyTorch YOLOv5 animal detection model.
+Module to run MegaDetector v5.
 
 """
 
@@ -131,7 +131,7 @@ class PTDetector:
         try:
             self.model = PTDetector._load_model(model_path, self.device)
         except Exception as e:
-            # In a very estoeric scenario where an old version of YOLOv5 is used to run
+            # In a very esoteric scenario where an old version of YOLOv5 is used to run
             # newer models, we run into an issue because the "Model" class became
             # "DetectionModel".  New YOLOv5 code handles this case by just setting them
             # to be the same, so doing that via monkey-patch doesn't seem *that* rude.
@@ -180,7 +180,8 @@ class PTDetector:
             
         return model
 
-    def generate_detections_one_image(self, img_original, 
+    def generate_detections_one_image(self, 
+                                      img_original, 
                                       image_id='unknown', 
                                       detection_threshold=0.00001, 
                                       image_size=None,
@@ -190,7 +191,8 @@ class PTDetector:
         Applies the detector to an image.
 
         Args:
-            img_original (Image): the PIL Image object with EXIF rotation taken into account
+            img_original (Image): the PIL Image object (or numpy array) on which we should run the 
+                detector, with EXIF rotation already handled.
             image_id (str, optional): a path to identify the image; will be in the "file" field 
                 of the output object
             detection_threshold (float, optional): only detections above this confidence threshold 
@@ -209,20 +211,20 @@ class PTDetector:
                 - 'failure' (a failure string, or None if everything went fine)
         """
 
-        result = {
-            'file': image_id
-        }
+        result = {'file': image_id }
         detections = []
         max_conf = 0.0
 
         if detection_threshold is None:
+            
             detection_threshold = 0
             
         try:
             
-            img_original = np.asarray(img_original)
+            if not isinstance(img_original,np.ndarray):                
+                img_original = np.asarray(img_original)
 
-            # padded resize
+            # Padded resize
             target_size = PTDetector.IMAGE_SIZE
             
             # Image size can be an int (which translates to a square target size) or (h,w)
@@ -267,9 +269,9 @@ class PTDetector:
 
             # NMS
             if self.device == 'mps':
-                # As of v1.13.0.dev20220824, nms is not implemented for MPS.
+                # As of PyTorch 1.13.0.dev20220824, nms is not implemented for MPS.
                 #
-                # Send prediction back to the CPU to fix.
+                # Send predictions back to the CPU for NMS.
                 pred = non_max_suppression(prediction=pred.cpu(), conf_thres=detection_threshold)
             else: 
                 pred = non_max_suppression(prediction=pred, conf_thres=detection_threshold)