Skip to content

Don't queue original images after preprocessing; read image metadata on the loaders - #260

Open
putdanil wants to merge 1 commit into
agentmorris:mainfrom
putdanil:drop-processed-image-copies
Open

putdanil wants to merge 1 commit into
agentmorris:mainfrom
putdanil:drop-processed-image-copies

Conversation

@putdanil

Copy link
Copy Markdown
Contributor

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() returns img_original_shape instead of img_original and img_original_pil. It is recorded after the optional resize, so it equals the img_original.shape that the old code read. _process_batch_group() uses it.
  • RFDETRDetector.preprocess_image() returns only img_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_base is a module constant, so the loaders use it directly.
  • Video frames are queued with empty metadata. _run_detector_on_videos_with_queue() never requests metadata, and frames get the smaller dicts through preprocess_image().
  • run_md_and_speciesnet.py and run_tiled_inference.py need no changes.

Errors while reading metadata

The consumer doesn't catch exceptions from get_image_datetime() or read_pil_exif(). Today, with batch size > 1, such an exception lands in _process_batch()'s try, and the whole batch is reported as FAILURE_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, as agent-notes/improved_exception_handling.md also points out: an item that fails to pickle is dropped by the queue's feeder thread, and q.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 of gps-tests images 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:

image main this PR
2048×1536 (3 MP, LILA) 22.6 MB 3.7 MB
4032×3024 (12 MP, gps-tests) 76.9 MB 3.7 MB
5472×3648 (20 MP, special-cases) 123 MB 3.4 MB

run_detector_batch CLI 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:

images configuration main this PR
500 × 5.2 MP (Snapshot Karoo, 2592×2000) no queue, --batch_size 8, redwood 71 s 73 s
--use_image_queue --preprocess_on_image_queue, redwood 113 s 38 s
same, --batch_size 8, redwood 76–77 s, 5.3 GB 26 s, 4.5 GB
same, --batch_size 8, spruce 73–77 s 15–19 s
200 × 8–21 MP (LILA, mostly 3840×2160 and 4000×3000) no queue, --batch_size 8, redwood 51 s 50 s
--use_image_queue --preprocess_on_image_queue, redwood 73 s 20 s
same, --batch_size 8, redwood 59–62 s, 6.5 GB 17 s, 4.7 GB
same, --batch_size 8, spruce 62 s 15 s

Excluding start-up (from runs on 100 and 500, or 40 and 200 images), redwood with --batch_size 8 goes 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-images and special-cases), gps-tests, 100 LILA images, and 10 edge cases I made from a gps-tests image: 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 all unless noted.

The images list is identical to main's, including the width, height, datetime and EXIF fields, for:

  • no queue, batch size 1 and 8;
  • queue at batch size 1: loading only; preprocessing on the queue with processes and with threads; preprocessing without the metadata flags;
  • queue at batch size 8 with --loader_workers 1, which forms batches in a fixed order: loading only; preprocessing with processes and with threads;
  • MDv5a (68 images): classic mode without the queue and with preprocessing on the queue at batch size 1; preprocessing at batch size 8 with one loader, in classic and default mode;
  • CPU only: --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_config to 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.
  • New test_image_metadata_with_image_queue() in run_detector_batch.py. It uses _FakeDetector, now with a preprocess_image() that returns only {'file': ...}, like PTDetector now 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 with KeyError: '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 for corrupt-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_mode is 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; typos reports nothing in the changed files.

I left TODO.md and the agent note as they are.


Prepared with help from Claude Code; I reviewed the diff and the measurements.

🤖 Generated with Claude Code

…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant