Repository navigation
Conversation
…on the loaders With --preprocess_on_image_queue, loaders queued the whole preprocess_image() dict, which for PTDetector held the original image twice (a numpy array and a PIL image), most of the pickled bytes for each image. Inference only needs the original's shape, and the consumer only read the PIL copy for image size, timestamp, and EXIF output. - PTDetector.preprocess_image() returns img_original_shape instead of the two originals, and _process_batch_group() uses it. - RFDETRDetector.preprocess_image() returns only the numpy original, which it runs inference on. - Loaders read the requested metadata fields (_read_image_metadata()) and queue them with each image as [file,image,producer_id,image_metadata]. Without the queue, the same function reads them where the image is loaded. Video frames are queued with no metadata. - An exception while reading metadata is passed along as a traceback string and raised once inference has succeeded, so it has the same effect as before: the batch fails with batch size > 1, the run stops with batch size 1. Adds test_image_metadata_with_image_queue(), which uses a fake detector. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This implements the TODO item "Avoid carrying original images around when they aren't necessary" (agent note), along the lines of the note, and supersedes #256. With
--use_image_queue --preprocess_on_image_queue, loaders now queue only what the consumer uses: the preprocessed image, the original image's shape, and the metadata fields that were requested.Changes
PTDetector.preprocess_image()returnsimg_original_shapeinstead ofimg_originalandimg_original_pil. It is recorded after the optional resize, so it equals theimg_original.shapethat the old code read._process_batch_group()uses it.RFDETRDetector.preprocess_image()returns onlyimg_original, the numpy array it runs inference on; the PIL copy is gone. The loaders don't need to know which detector they feed._read_image_metadata()reads the requested fields (width/height, timestamp, EXIF tags) with the code the consumer used before. It runs where the image is loaded: on the loaders when the queue is enabled, with or without preprocessing on the queue, so there is one path; otherwise in_process_batch()and_process_image(). Queue items are now[file, image, producer_id, image_metadata], and_run_detector_with_image_queue()passes the three flags to the loaders.exif_options_baseis a module constant, so the loaders use it directly._run_detector_on_videos_with_queue()never requests metadata, and frames get the smaller dicts throughpreprocess_image().run_md_and_speciesnet.pyandrun_tiled_inference.pyneed no changes.Errors while reading metadata
The consumer doesn't catch exceptions from
get_image_datetime()orread_pil_exif(). Today, with batch size > 1, such an exception lands in_process_batch()'stry, and the whole batch is reported asFAILURE_INFER. With batch size 1, it propagates out of_process_image()and stops the run. Either way, it only happens after inference succeeded.This PR keeps that behavior.
_read_image_metadata()returns the traceback as a string, and_add_image_metadata()raises it at the point where the old code read the metadata. It is a string because exception objects aren't always picklable, asagent-notes/improved_exception_handling.mdalso points out: an item that fails to pickle is dropped by the queue's feeder thread, andq.join()would then never return.I couldn't trigger such an exception with Pillow 12.3. None occurred on the MD test images, on
gps-tests, on the edge cases listed below, or on 2,300 copies ofgps-testsimages with random bytes written into their EXIF segment. If you'd rather handle this per image (e.g. leave the fields out and keep the detections), that's a small follow-up.Width and height come from the image that
load_image()returns, i.e. after EXIF rotation, as before. Test images with orientation 6 and 8 give swapped dimensions, the same as on main.Effect
Pickled size of one queued item, MDv1000-redwood:
gps-tests)special-cases)run_detector_batchCLI on an NVIDIA GPU + 8 vCPU (Xeon Gold 6248R VM), default 4 loaders. Whole run including model loading, and the peak resident memory of the consumer and loader processes together:--batch_size 8, redwood--use_image_queue --preprocess_on_image_queue, redwood--batch_size 8, redwood--batch_size 8, spruce--batch_size 8, redwood--use_image_queue --preprocess_on_image_queue, redwood--batch_size 8, redwood--batch_size 8, spruceExcluding start-up (from runs on 100 and 500, or 40 and 200 images), redwood with
--batch_size 8goes from 7.5 to 33 images/s on the 5 MP set and from 4.0 to 32 on the 8–21 MP set. On main, the preprocessing queue is slower than no queue on these images, as the note also saw on an RTX 3050. With this PR, a whole run takes about a third of the no-queue time (excluding start-up, 3.8× and 5.5× the no-queue rate). The memory figures add up resident memory across processes, so shared libraries count once per process; the difference between the two columns is what's meaningful.On 2–3 MP images (the 500-image LILA mix from #256), the copies no longer limit this machine: with
--batch_size 8, redwood does 31.3 images/s on main and 32.1 with this PR (one run each), spruce 35–36 vs. 38–46 (two runs each). In #256, before #255 was merged, main did 8.0 images/s on these images with redwood. The spread between runs on this machine is up to about 20%.Equivalence
168 images: the MD test images (including
corrupt-imagesandspecial-cases),gps-tests, 100 LILA images, and 10 edge cases I made from agps-testsimage: EXIF orientation 2, 3, 6 and 8, no EXIF, grayscale with EXIF, RGBA PNG, out-of-range EXIF and GPS IFD offsets, and a zeroed IFD. MDv1000-redwood, with--include_image_size --include_image_timestamp --include_exif_tags allunless noted.The
imageslist is identical to main's, including the width, height, datetime and EXIF fields, for:--loader_workers 1, which forms batches in a fixed order: loading only; preprocessing with processes and with threads;--ncores 4; queue with preprocessing.With the default 4 loaders and batch size 8, main isn't deterministic: this PR differs from main on 8–12 of the 168 images, and two runs of main differ on 2–12. Confidence values move by 0.001, which sometimes adds or drops a detection right at the output threshold.
RF-DETR. No RF-DETR MegaDetector model is public, so I used rfdetr 1.11.2's COCO RF-DETR Nano weights and added
model_configto the checkpoint. On the same 168 images, the output is identical to main's for all 7 configurations: no queue at batch size 1 and 8, queue at batch size 8, preprocessing on the queue at batch size 1 and 8 with processes, at 8 with threads, and at 8 without the metadata flags. RF-DETR still gets its original image. With TF32 off, two runs of main agree exactly here, even at batch size 8.Video is covered by md_tests, below.
Tests
pytest -vv: 24 passed, 4 skipped on main; 25 passed, 4 skipped on this branch.test_image_metadata_with_image_queue()inrun_detector_batch.py. It uses_FakeDetector, now with apreprocess_image()that returns only{'file': ...}, likePTDetectornow does. The test checks that size, timestamp and EXIF fields are the same without the queue, with the queue, and with preprocessing on the queue, at batch size 1 and 3, including a rotated image, an image without EXIF, and an unreadable file. On main it fails withKeyError: 'img_original_pil'.python -m megadetector.utils.md_tests --max_coord_error 0.01 --max_conf_error 0.01 --skip_download_tests --warning_mode(GPU, CPU and video): all 35 steps run on both, and the comparison lines against the expected results are the same. Of the 44 output files, 32 are identical to main's. Another 11 are video outputs that differ only in the path inside the failure message forcorrupt-video.mp4, which contains each run's scratch folder. The last one is the worker-side preprocessing with batched inference, the non-deterministic case above: 1 of 47 images, 0.152 vs. 0.153.--warning_modeis still needed for the video folder step, which fails the same way on main (see Don't queue full-resolution originals when preprocessing on the image queue #256).ruff check .is clean;typosreports nothing in the changed files.I left
TODO.mdand the agent note as they are.Prepared with help from Claude Code; I reviewed the diff and the measurements.
🤖 Generated with Claude Code