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

megadetector/data_management/read_exif.py

@@ -16,7 +16,7 @@ path.  No attempt is made to be consistent in format across the two approaches.
 import os
 import subprocess
 import json
-from datetime import datetime
+from datetime import date, datetime
 
 from multiprocessing.pool import ThreadPool as ThreadPool
 from multiprocessing.pool import Pool as Pool
@@ -64,7 +64,7 @@ class ReadExifOptions:
         #:
         #: A useful set of tags one might want to limit queries for:
         #:
-        #: options.tags_to_include = ['DateTime','Model','Make','ExifImageWidth','ExifImageHeight','DateTime',
+        #: options.tags_to_include = ['DateTime','Model','Make','ExifImageWidth','ExifImageHeight',
         #: 'DateTimeOriginal','Orientation']    
         self.tags_to_include = None
         
@@ -103,7 +103,7 @@ class ExifResultsToCCTOptions:
     
         #: Function for extracting location information, should take a string
         #: and return a string.  Defaults to ct_utils.image_file_to_camera_folder.  If
-        #: this is None, uses folder names as locations.
+        #: this is None, location is written as "unknown".
         self.filename_to_location_function = image_file_to_camera_folder
     
 
@@ -689,7 +689,7 @@ def exif_results_to_cct(exif_results,cct_output_file=None,options=None):
         
         # By default we assume that each leaf-node folder is a location
         if options.filename_to_location_function is None:
-            im['location'] = os.path.dirname(exif_result['file_name'])
+            im['location'] = 'unknown'
         else:
             im['location'] = options.filename_to_location_function(exif_result['file_name'])            
             
@@ -738,9 +738,15 @@ def exif_results_to_cct(exif_results,cct_output_file=None,options=None):
     d['annotations'] = []
     d['categories'] = []
     
+    def json_serialize_datetime(obj):
+        if isinstance(obj, (datetime, date)):
+            return obj.isoformat()
+        raise TypeError('Object {} (type {}) not serializable'.format(
+            str(obj),type(obj)))
+        
     if cct_output_file is not None:
         with open(cct_output_file,'w') as f:
-            json.dump(d,indent=1)
+            json.dump(d,f,indent=1,default=json_serialize_datetime)
     
     return d
 

megadetector/detection/process_video.py

@@ -12,8 +12,6 @@ writing them to disk.  The upside, though, is that this approach allows you to r
 detection elimination after running MegaDetector, and it allows allows more efficient re-use
 of frames if you end up running MD more than once, or running multiple versions of MD.
 
-TODO: optionally skip writing frames to disk, and process frames in memory.
-
 """
 
 #%% Imports
@@ -36,6 +34,7 @@ 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 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
 
@@ -50,6 +49,11 @@ class ProcessVideoOptions:
     def __init__(self):
         
         #: Can be a model filename (.pt or .pb) or a model name (e.g. "MDV5A")
+        #: 
+        #: Use the string "no_detection" to indicate that you only want to extract frames,
+        #: not run a model.  If you do this, you almost definitely want to set 
+        #: keep_extracted_frames to "True", otherwise everything in this module is a no-op.
+        #: I.e., there's no reason to extract frames, do nothing with them, then delete them.
         self.model_file = 'MDV5A'
         
         #: Video (of folder of videos) to process
@@ -66,7 +70,7 @@ class ProcessVideoOptions:
         #: if this is None
         self.frame_folder = None
         
-        # Folder to use for rendered frames (if rendering output video); will use a folder 
+        #: Folder to use for rendered frames (if rendering output video); will use a folder 
         #: in system temp space if this is None
         self.frame_rendering_folder = None
         
@@ -111,6 +115,10 @@ class ProcessVideoOptions:
         #: fourcc code to use for writing videos; only relevant if render_output_video is True
         self.fourcc = None
     
+        #: force a specific frame rate for output videos; only relevant if render_output_video 
+        #: is True
+        self.rendering_fs = None
+        
         #: Confidence threshold to use for writing videos with boxes, only relevant if
         #: if render_output_video is True.  Defaults to choosing a reasonable threshold
         #: based on the model version.
@@ -121,9 +129,13 @@ class ProcessVideoOptions:
         
         #: Sample every Nth frame; set to None (default) or 1 to sample every frame.  Typically
         #: we sample down to around 3 fps, so for typical 30 fps videos, frame_sample=10 is a 
-        #: typical value.
+        #: typical value.  Mutually exclusive with [frames_to_extract].
         self.frame_sample = None
         
+        #: Extract a specific set of frames (list of ints, or a single int).  Mutually exclusive with
+        #: [frame_sample].
+        self.frames_to_extract = None
+        
         #: Number of workers to use for parallelization; set to <= 1 to disable parallelization
         self.n_cores = 1
     
@@ -138,7 +150,14 @@ class ProcessVideoOptions:
         self.quality = 90
         
         #: Resize frames so they're at most this wide
-        self.max_width = 1600
+        self.max_width = None
+        
+        #: Run the model at this image size (don't mess with this unless you know what you're
+        #: getting into)
+        self.image_size = None
+        
+        #: Enable image augmentation
+        self.augment = False
     
 # ...class ProcessVideoOptions
 
@@ -278,7 +297,8 @@ def _clean_up_extracted_frames(options,frame_output_folder,frame_filenames):
 
 def process_video(options):
     """
-    Process a single video through MD, optionally writing a new video with boxes
+    Process a single video through MD, optionally writing a new video with boxes.
+    Can also be used just to split a video into frames, without running a model.
     
     Args: 
         options (ProcessVideoOptions): all the parameters used to control this process,
@@ -294,6 +314,10 @@ def process_video(options):
     if options.render_output_video and (options.output_video_file is None):
         options.output_video_file = options.input_video_file + '.detections.mp4'
 
+    if options.model_file == 'no_detection' and not options.keep_extracted_frames:
+        print('Warning: you asked for no detection, but did not specify keep_extracted_frames, this is a no-op')
+        return
+    
     # Track whether frame and rendering folders were created by this script
     caller_provided_frame_output_folder = (options.frame_folder is not None)
     caller_provided_rendering_output_folder = (options.frame_rendering_folder is not None)
@@ -309,15 +333,31 @@ def process_video(options):
    
     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)
+                            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
+    
+    
+    ## Run MegaDetector
+    
     if options.reuse_results_if_available and \
         os.path.isfile(options.output_json_file):
             print('Loading results from {}'.format(options.output_json_file))
@@ -325,12 +365,17 @@ def process_video(options):
                 results = json.load(f)
     else:
         results = run_detector_batch.load_and_run_detector_batch(
-            options.model_file, image_file_names,
+            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)
-    
+            quiet=True,
+            augment=options.augment,
+            image_size=options.image_size)
+        
+        _add_frame_numbers_to_results(results)
+        
         run_detector_batch.write_results_to_file(
             results, options.output_json_file,
             relative_path_base=frame_output_folder,
@@ -357,14 +402,20 @@ def process_video(options):
             confidence_threshold=options.rendering_confidence_threshold)
 
         # Combine into a video
-        if options.frame_sample is None:
+        if options.rendering_fs is not None:
+            rendering_fs = options.rendering_fs
+        elif options.frame_sample is None:
             rendering_fs = Fs
         else:
+            # If the original video was 30fps and we sampled every 10th frame, 
+            # render at 3fps
             rendering_fs = Fs / options.frame_sample
             
         print('Rendering {} frames to {} at {} fps (original video {} fps)'.format(
             len(detected_frame_files), options.output_video_file,rendering_fs,Fs))
-        frames_to_video(detected_frame_files, rendering_fs, options.output_video_file, 
+        frames_to_video(detected_frame_files, 
+                        rendering_fs, 
+                        options.output_video_file, 
                         codec_spec=options.fourcc)
         
         # Possibly clean up rendered frames
@@ -381,7 +432,13 @@ def process_video(options):
 
 def process_video_folder(options):
     """
-    Process a folder of videos through MD
+    Process a folder of videos through MD. Can also be used just to split a folder of
+    videos into frames, without running a model.
+    
+    When this function is used to run MD, two .json files will get written, one with 
+    an entry for each *frame* (identical to what's created by process_video()), and
+    one with an entry for each *video* (which is more suitable for, e.g., reading into
+    Timelapse).
     
     Args: 
         options (ProcessVideoOptions): all the parameters used to control this process,
@@ -393,13 +450,17 @@ def process_video_folder(options):
     assert os.path.isdir(options.input_video_file), \
         '{} is not a folder'.format(options.input_video_file)
            
-    assert options.output_json_file is not None, \
-        'When processing a folder, you must specify an output .json file'
-                         
-    assert options.output_json_file.endswith('.json')
-    video_json = options.output_json_file
-    frames_json = options.output_json_file.replace('.json','.frames.json')
-    os.makedirs(os.path.dirname(video_json),exist_ok=True)
+    if options.model_file == 'no_detection' and not options.keep_extracted_frames:
+        print('Warning: you asked for no detection, but did not specify keep_extracted_frames, this is a no-op')
+        return
+    
+    if options.model_file != 'no_detection':
+        assert options.output_json_file is not None, \
+            'When processing a folder, you must specify an output .json file'      
+        assert options.output_json_file.endswith('.json')
+        video_json = options.output_json_file
+        frames_json = options.output_json_file.replace('.json','.frames.json')
+        os.makedirs(os.path.dirname(video_json),exist_ok=True)
     
     # Track whether frame and rendering folders were created by this script
     caller_provided_frame_output_folder = (options.frame_folder is not None)
@@ -420,6 +481,7 @@ def process_video_folder(options):
     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, 
@@ -429,8 +491,10 @@ def process_video_folder(options):
                                every_n_frames=options.frame_sample,
                                verbose=options.verbose,
                                quality=options.quality,
-                               max_width=options.max_width)
+                               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:
@@ -443,6 +507,11 @@ def process_video_folder(options):
 
     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
@@ -454,12 +523,17 @@ def process_video_folder(options):
     else:
         print('Running MegaDetector')
         results = run_detector_batch.load_and_run_detector_batch(
-            options.model_file, image_file_names,
+            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)
+            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,
@@ -518,9 +592,13 @@ def process_video_folder(options):
             
             video_fs = Fs[i_video]
             
-            if options.frame_sample is None:                
+            if options.rendering_fs is not None:
+                rendering_fs = options.rendering_fs
+            elif options.frame_sample is None:                
                 rendering_fs = video_fs
             else:
+                # If the original video was 30fps and we sampled every 10th frame, 
+                # render at 3fps                
                 rendering_fs = video_fs / options.frame_sample
             
             input_video_file_relative = os.path.relpath(input_video_file_abs,options.input_video_file)
@@ -547,7 +625,10 @@ def process_video_folder(options):
             # Create the output video            
             print('Rendering detections for video {} to {} at {} fps (original video {} fps)'.format(
                 input_video_file_relative,video_output_file,rendering_fs,video_fs))
-            frames_to_video(video_frame_files, rendering_fs, video_output_file, codec_spec=options.fourcc)
+            frames_to_video(video_frame_files, 
+                            rendering_fs, 
+                            video_output_file, 
+                            codec_spec=options.fourcc)
                 
         # ...for each video
         
@@ -607,6 +688,14 @@ def options_to_command(options):
         cmd += ' --n_cores ' + str(options.n_cores)
     if options.frame_sample is not None:
         cmd += ' --frame_sample ' + str(options.frame_sample)
+    if options.frames_to_extract is not None:
+        cmd += ' --frames_to_extract '
+        if isinstance(options.frames_to_extract,int):
+            frames_to_extract = [options.frames_to_extract]
+        else:
+            frames_to_extract = options.frames_to_extract
+        for frame_number in frames_to_extract:
+            cmd += ' {}'.format(frame_number)
     if options.debug_max_frames is not None:
         cmd += ' --debug_max_frames ' + str(options.debug_max_frames)
     if options.class_mapping_filename is not None:
@@ -631,16 +720,20 @@ def options_to_command(options):
 
 if False:    
         
+    pass
+
     #%% Process a folder of videos
     
     model_file = 'MDV5A'
-    input_dir = r'g:\temp\test-videos'
+    # 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')
     output_json_file = os.path.join(output_base,'video-test.json')
     output_video_folder = os.path.join(output_base,'output_videos')
     
+    
     print('Processing folder {}'.format(input_dir))
     
     options = ProcessVideoOptions()    
@@ -654,32 +747,28 @@ if False:
     options.quality = 90
     options.frame_sample = 10
     options.max_width = 1280
-    options.n_cores = 5
+    options.n_cores = 4
     options.verbose = True
-    options.render_output_video = True
-    
-    options.frame_folder = None # frame_folder
-    options.frame_rendering_folder = None # rendering_folder
-    
-    options.keep_extracted_frames = False
-    options.keep_rendered_frames = False
-    options.force_extracted_frame_folder_deletion = True
-    options.force_rendered_frame_folder_deletion = True
-    
-    # options.confidence_threshold = 0.15
+    options.render_output_video = True    
+    options.frame_folder = frame_folder
+    options.frame_rendering_folder = rendering_folder    
+    options.keep_extracted_frames = True
+    options.keep_rendered_frames = True
+    options.force_extracted_frame_folder_deletion = False
+    options.force_rendered_frame_folder_deletion = False
     options.fourcc = 'mp4v'
+    # options.rendering_confidence_threshold = 0.15
     
     cmd = options_to_command(options); print(cmd)
         
-    import clipboard; clipboard.copy(cmd)
-    
-    if False:
-        process_video_folder(options)
+    # import clipboard; clipboard.copy(cmd)
+    # process_video_folder(options)
         
     
     #%% Process a single video
 
     fn = r'g:\temp\test-videos\person_and_dog\DSCF0056.AVI'
+    assert os.path.isfile(fn)
     model_file = 'MDV5A'
     input_video_file = fn
     
@@ -687,39 +776,89 @@ if False:
     frame_folder = os.path.join(output_base,'frames')
     rendering_folder = os.path.join(output_base,'rendered-frames')
     output_json_file = os.path.join(output_base,'video-test.json')
-    output_video_file = os.path.join(output_base,'output_videos.mp4')
+    output_video_file = os.path.join(output_base,'output_video.mp4')
     
     options = ProcessVideoOptions()
     options.model_file = model_file
     options.input_video_file = input_video_file
     options.render_output_video = True
     options.output_video_file = output_video_file
-    
-    options.verbose = True
-    
+    options.output_json_file = output_json_file    
+    options.verbose = True    
     options.quality = 75
-    options.frame_sample = None # 10
-    options.max_width = 600
-    
-    options.frame_folder = None # frame_folder
-    options.frame_rendering_folder = None # rendering_folder
-    
+    options.frame_sample = 10
+    options.max_width = 1600    
+    options.frame_folder = frame_folder
+    options.frame_rendering_folder = rendering_folder    
     options.keep_extracted_frames = False
     options.keep_rendered_frames = False
     options.force_extracted_frame_folder_deletion = True
-    options.force_rendered_frame_folder_deletion = True
-    
-    # options.confidence_threshold = 0.15
+    options.force_rendered_frame_folder_deletion = True    
     options.fourcc = 'mp4v'
+    # options.rendering_confidence_threshold = 0.15
     
     cmd = options_to_command(options); print(cmd)
     
-    import clipboard; clipboard.copy(cmd)
-    
-    if False:        
-        process_video(options)    
+    # import clipboard; clipboard.copy(cmd)    
+    process_video(options)
             
     
+    #%% Extract specific frames from a single video, no detection
+
+    fn = r'g:\temp\test-videos\person_and_dog\DSCF0064.AVI'
+    assert os.path.isfile(fn)
+    model_file = 'no_detection'
+    input_video_file = fn
+    
+    output_base = r'g:\temp\video_test'
+    frame_folder = os.path.join(output_base,'frames')
+    output_video_file = os.path.join(output_base,'output_videos.mp4')
+    
+    options = ProcessVideoOptions()
+    options.model_file = model_file
+    options.input_video_file = input_video_file    
+    options.verbose = True    
+    options.quality = 90
+    options.frame_sample = None
+    options.frames_to_extract = [0,100]
+    options.max_width = None    
+    options.frame_folder = frame_folder
+    options.keep_extracted_frames = True
+    
+    cmd = options_to_command(options); print(cmd)
+    
+    # import clipboard; clipboard.copy(cmd)
+    process_video(options)    
+        
+        
+    #%% Extract specific frames from a folder, no detection
+
+    fn = r'g:\temp\test-videos\person_and_dog'
+    assert os.path.isdir(fn)
+    model_file = 'no_detection'
+    input_video_file = fn
+    
+    output_base = r'g:\temp\video_test'
+    frame_folder = os.path.join(output_base,'frames')
+    output_video_file = os.path.join(output_base,'output_videos.mp4')
+    
+    options = ProcessVideoOptions()
+    options.model_file = model_file
+    options.input_video_file = input_video_file    
+    options.verbose = True    
+    options.quality = 90
+    options.frame_sample = None
+    options.frames_to_extract = [0,100]
+    options.max_width = None    
+    options.frame_folder = frame_folder
+    options.keep_extracted_frames = True
+    
+    cmd = options_to_command(options); print(cmd)
+    
+    # import clipboard; clipboard.copy(cmd)
+    process_video(options)    
+
+        
 #%% Command-line driver
 
 def main():
@@ -731,7 +870,8 @@ def main():
         'producing a new video with detections annotated'))
 
     parser.add_argument('model_file', type=str,
-                        help='MegaDetector model file (.pt or .pb) or model name (e.g. "MDV5A")')
+                        help='MegaDetector model file (.pt or .pb) or model name (e.g. "MDV5A"), '\
+                             'or the string "no_detection" to run just frame extraction')
 
     parser.add_argument('input_video_file', type=str,
                         help='video file (or folder) to process')
@@ -788,22 +928,33 @@ def main():
                            'whether other files were present in the folder.')
         
     parser.add_argument('--rendering_confidence_threshold', type=float,
-                        default=None, help="don't render boxes with confidence below this threshold (defaults to choosing based on the MD version)")
+                        default=None, 
+                        help="don't render boxes with confidence below this threshold (defaults to choosing based on the MD version)")
+
+    parser.add_argument('--rendering_fs', type=float,
+                        default=None, 
+                        help='force a specific frame rate for output videos (only relevant when using '\
+                             '--render_output_video) (defaults to the original frame rate)')
 
     parser.add_argument('--json_confidence_threshold', type=float,
-                        default=0.0, help="don't include boxes in the .json file with confidence "\
+                        default=default_options.json_confidence_threshold, 
+                        help="don't include boxes in the .json file with confidence "\
                             'below this threshold (default {})'.format(
                                 default_options.json_confidence_threshold))
 
     parser.add_argument('--n_cores', type=int,
-                        default=1, help='Number of cores to use for frame separation and detection. '\
+                        default=default_options.n_cores,
+                        help='Number of cores to use for frame separation and detection. '\
                             'If using a GPU, this option will be respected for frame separation but '\
                             'ignored for detection.  Only relevant to frame separation when processing '\
-                            'a folder.')
+                            'a folder.  Default {}.'.format(default_options.n_cores))
 
     parser.add_argument('--frame_sample', type=int,
                         default=None, help='process every Nth frame (defaults to every frame)')
 
+    parser.add_argument('--frames_to_extract', nargs='+', type=int,
+                        default=None, help='extract specific frames (one or more ints)')
+
     parser.add_argument('--quality', type=int,
                         default=default_options.quality, 
                         help='JPEG quality for extracted frames (defaults to {})'.format(
@@ -828,7 +979,16 @@ def main():
     parser.add_argument('--verbose', action='store_true',
                         help='Enable additional debug output')
     
-        
+    parser.add_argument('--image_size',
+                        type=int,
+                        default=None,
+                        help=('Force image resizing to a specific integer size on the long '\
+                              'axis (not recommended to change this)'))    
+    
+    parser.add_argument('--augment',
+                        action='store_true',
+                        help='Enable image augmentation')
+            
     if len(sys.argv[1:]) == 0:
         parser.print_help()
         parser.exit()

megadetector/detection/pytorch_detector.py

@@ -180,9 +180,12 @@ class PTDetector:
             
         return model
 
-    def generate_detections_one_image(self, img_original, image_id='unknown', 
-                                      detection_threshold=0.00001, image_size=None,
-                                      skip_image_resizing=False):
+    def generate_detections_one_image(self, img_original, 
+                                      image_id='unknown', 
+                                      detection_threshold=0.00001, 
+                                      image_size=None,
+                                      skip_image_resizing=False,
+                                      augment=False):
         """
         Applies the detector to an image.
 
@@ -192,11 +195,11 @@ class PTDetector:
                 of the output object
             detection_threshold (float, optional): only detections above this confidence threshold 
                 will be included in the return value
-            image_size (tuple, optional): image size to use for inference, only mess with this
-                if (a) you're using a model other than MegaDetector or (b) you know what you're
-                doing
-            skip_image_resizing (bool, optional): whether to skip internal image resizing (and rely on external 
-                resizing)
+            image_size (tuple, optional): image size to use for inference, only mess with this if 
+                (a) you're using a model other than MegaDetector or (b) you know what you're getting into
+            skip_image_resizing (bool, optional): whether to skip internal image resizing (and rely on 
+                external resizing)
+            augment (bool, optional): enable (implementation-specific) image augmentation
 
         Returns:
             dict: a dictionary with the following fields:
@@ -242,8 +245,10 @@ class PTDetector:
             if skip_image_resizing:
                 img = img_original
             else:
-                letterbox_result = letterbox(img_original, new_shape=target_size,
-                                stride=PTDetector.STRIDE, auto=True)
+                letterbox_result = letterbox(img_original, 
+                                             new_shape=target_size,
+                                             stride=PTDetector.STRIDE, 
+                                             auto=True)
                 img = letterbox_result[0]                
             
             # HWC to CHW; PIL Image is RGB already
@@ -258,7 +263,7 @@ class PTDetector:
             if len(img.shape) == 3:  
                 img = torch.unsqueeze(img, 0)
 
-            pred: list = self.model(img)[0]
+            pred = self.model(img,augment=augment)[0]
 
             # NMS
             if self.device == 'mps':