Can we improve the memory usage of the slicing operation?

[akshayrb22]

Hi guys,

Working with large JPEG images ends up using a lot more RAM than needed. Specifically, slice_image() uses about 3.4 times the memory of the array it is slicing. That means get_sliced_prediction() on a large image needs several times the memory the image itself takes up. For a 1144 MP (39266x29140) scan, which is a 3.2 GiB array, that works out to about 11 GiB, and it raises MemoryError on a machine with enough RAM to hold the image several times over.

Measurements

Measured on 356299b, which was the tip of main at the time. slicing.py and utils/cv.py are unchanged since then, so the numbers still apply to 9f5a79d.

Each cell is the median of three runs, each in a new process. Peak memory varied by 0-4 MB between runs. A Python interpreter with import sahi uses 69 MB before any image is touched, so that is the floor for every number below.

input	RGB array	slice_image() peak memory
121 MP (11000²)	346 MB	1234 MB (3.57x)
256 MP (16000²)	732 MB	2522 MB (3.44x)

Where the memory goes

Measured step by step on the 256 MP image, reading live memory from /proc/self/statm alongside ru_maxrss, and checked against Pillow's source:

step	cost	why
interpreter + import sahi	69 MB	measured floor
PIL decode buffer	976 MB	Pillow stores RGB at 4 bytes per pixel, not 3. It adds a padding byte. Measured at exactly 4.00 bytes/px.
tobytes() chunk list	732 MB	np.asarray(pil_image) goes through __array_interface__, which calls tobytes()
b"".join(chunks)	732 MB	joins the chunks into a second full copy while the chunk list is still in memory
total	2509 MB	against 2522 MB measured

Image.tobytes() on its own costs exactly 2.00x the size of the image data, measured at 1467 MB added for a 732 MB array. OpenCV decodes straight into a 3 bytes/px array and never calls tobytes(). On the same file it peaks at 1527 MB, against Pillow's 2479 MB.

Two things that are not the cause

I checked both of these before concluding the above.

It is not the slices. slicing.py does image_pil_arr[tly:bry, tlx:brx], which is a numpy view, not a copy. For 100 slices: all views? True | distinct parents: 1 | owned (copied) bytes 0.0 MiB. The whole slice list is about 78 Python objects. Freeing the slices earlier saves nothing.

It is not .convert("RGB"). Skipping it changes peak memory by zero bytes. Its copy is freed before the peak happens.

Why batch_size does not help

batch_size only limits the np.ascontiguousarray(img) copies at predict.py:377. The full decoded array stays in memory for the whole inference loop, because the slice views all point into it. Changing the batch size does not release it.

One related thing: perform_standard_pred=True is the default (predict.py:183). After slicing, predict.py:398 calls get_prediction(image=image, ...) on the full image, which decodes it a second time and runs the model at full resolution. That may be worth its own discussion. I have left it out of scope for now.

How to reproduce this

The image has to be created in a different process from the one that measures, and the measuring process must not be started by the generating one.

ru_maxrss is a high water mark. It records the largest amount of memory a process ever used and never goes down. Creating a 16000x16000 image uses 1782 MB on its own, so if you create it in the process that measures, that 1782 MB is counted against slice_image. The single-process version prints 2766 MB instead of 2522 MB.

Running the measurement as a child process is not enough either. A child inherits its parent's high water mark on Linux, so a child of the process that created the image starts at 1782 MB and reports at least that. Two separate commands, as below, are fine because neither is the parent of the other.

# step 1 - create the image. Its memory goes away when this process exits.
python -c "import numpy as np; from PIL import Image; Image.MAX_IMAGE_PIXELS=None; \
Image.fromarray(np.random.randint(0,255,(16000,16000,3),dtype=np.uint8)).save('big.jpg', quality=85)"

# step 2 - measure in a process that has done nothing else
python -c "import resource; from sahi.slicing import slice_image; slice_image(image='big.jpg'); \
print(f'peak {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.0f} MB for a 732 MB image')"

Output on my machine: peak 2522 MB for a 732 MB image

The full benchmark I used for the table is here:
https://gist.github.com/akshayrb22/653268e052174d2d17120e78ff32d048

It measures the "before" column against a real checkout of upstream main, passed with --baseline-root, prints which sahi.slicing produced each column, takes the median over --repeats runs, and prints slice counts so you can check both paths produce the same slices. It creates its images in a separate process for the reason above.

Proposed direction

I have this working on a branch and I am splitting it into small PRs, opened one at a time (see #1407). In order:

1. Decode to an array instead of going through PIL. No new dependency and no change to any function signature. Peak memory goes from 1234 to 764 MB at 121 MP, and from 2522 to 1536 MB at 256 MP.
2. Decode PredictionResult.image only when it is read. Many callers never read it.
3. Produce slices one batch at a time so get_sliced_prediction iterates instead of building the full list first. That goes from 764 to 456 MB at 121 MP, on top of step 1. The image is still decoded in full underneath.
4. Read the image in horizontal bands off disk using libvips. This is the step that removes the full decode, which is what a gigapixel scan needs. It requires an optional dependency.

Question: steps 1 to 3 need nothing new. Step 4 needs pyvips. I would add it as an optional extra, pip install sahi[bigimage], with the current whole-image decode used whenever it is not installed, so nothing changes for anyone who does not install it.

Is an optional dependency acceptable here at all? If it is not, I will stop after step 3 rather than try to work around it.

[onuralpszr]

@akshayrb22 small Prs would be nice to see and also I need to check pyvips also I am curious to see PIL vs pyvips cases if possible to avoid new deps* I also made a new release can also help couple of cases in advance please re-check and test as well.

[akshayrb22]

Sure. I'll work on a comparison between PIL and pyvips and add it as a comment here with a link to the benchmark gist (I might take a day or so to get this out). A quick check on the latest release does show the postprocessing upgrades helping on the numpy backend, greedy_nmm on 20k crowded boxes drops from 232 MB to 12 MB, and it no longer scales with crowding. But that's on the tail end of the operation and the changes I am suggesting are on the decoding side (for the first PR atleast).

Also, what is your policy on tests - do you prefer to handle them yourself to ensure that the contributor doesn't miss out on anything? Or would you prefer that I come up with smaller tests?

[akshayrb22]

Here's the PIL vs pyvips comparison.

There are four ways of getting a JPEG off disk and into an HWC RGB uint8 array (with the libraries in sahi and pyvips as an extra) with nothing else happening in the process. Peak RSS, each cell in a fresh process, median of three runs. growth is (peak - that path's own import floor) / the full decoded array, so each library's import cost isn't charged to its decode. Importing cv2 costs more than importing pyvips and that has nothing to do with decoding.

image	path	file	full array	largest band	import floor	peak	growth	wall
16 MP	pil	6 MB	46 MB	n/a	34 MB	189 MB	3.38x	0.30 s
16 MP	opencv	6 MB	46 MB	n/a	45 MB	139 MB	2.05x	0.32 s
16 MP	vips	6 MB	46 MB	n/a	47 MB	178 MB	2.86x	0.35 s
16 MP	vips-band	6 MB	46 MB	12 MB	59 MB	134 MB	1.64x	0.39 s
64 MP	pil	25 MB	183 MB	n/a	34 MB	647 MB	3.35x	0.89 s
64 MP	opencv	25 MB	183 MB	n/a	45 MB	414 MB	2.01x	0.78 s
64 MP	vips	25 MB	183 MB	n/a	47 MB	308 MB	1.42x	0.68 s
64 MP	vips-band	25 MB	183 MB	23 MB	59 MB	234 MB	0.96x	0.96 s
256 MP	pil	100 MB	732 MB	n/a	34 MB	2480 MB	3.34x	3.21 s
256 MP	opencv	100 MB	732 MB	n/a	45 MB	1513 MB	2.00x	2.64 s
256 MP	vips	100 MB	732 MB	n/a	47 MB	973 MB	1.26x	1.92 s
256 MP	vips-band	100 MB	732 MB	47 MB	59 MB	386 MB	0.45x	3.12 s

• pil is Image.open().convert("RGB") then np.asarray(), which is what sahi does today.
• opencv is cv2.imread() with an in-place BGR to RGB swap, which is what the first PR does.
• vips is a whole-image pyvips read.
• vips-band reads top to bottom in 1024-row bands and holds one at a time, which is what step 4 would do. It never allocates the full array, so the full array column is what it avoided rather than what it used, and its growth shares the same denominator as the other rows.

All four decode to byte-identical arrays. --verify reports max abs diff 0 against the PIL reference, and the banded read matches the whole-image vips read exactly.

Benchmark script: https://gist.github.com/akshayrb22/b7daf2286e986730d8f085f1a1b7aec5

Two things I inferred from the table:

A pyvips whole-image read isn't the interesting part. It beats OpenCV at large sizes (1.26x against 2.00x at 256 MP) but it ends up costing a multiple of the decoded array, so it only moves the constant. If that were the only use for pyvips I wouldn't argue for a dependency to get it.

The band read is the only path whose cost isn't a multiple of the decoded array. Its growth column falls as the image grows, 1.64x to 0.96x to 0.45x, because what it actually costs is the compressed file plus a couple of bands, and a band is width x band_rows x 3 no matter how tall the image is. Extrapolating both curves to the 1144 MP scan that started this (3.2 GiB decoded, and real scans compress better than my test images): the whole-decode paths use up between 6-11 GB while the band read wants a few hundred MB. That gap is the main argument for using pyvips.

One measurement note, because it messed with my numbers and might mess with others - libvips maps its source file, so the compressed file size lands in RSS. My first randomly generated test images were per-pixel Gaussian noise, which compresses about 2x, so the 256 MP generated image was a 354 MB file and every vips row was inflated by that much. It made a band read look proportional to the image when it's proportional to the file. The script now generates content that compresses about 7x, near the low end of a real scan, and reports file size as its own column so the term stays visible. PIL and OpenCV are unaffected either way. Their peaks came out identical on both sets of randomly generated images, so they don't hold the compressed file resident.

If you have a large image you want to test out against this benchmark, --image <path> benchmarks that instead of the generated sweep.