DEV Community

Cover image for How to Reject Oversized AI Images Before Decoding Them in Python
Voor AI
Voor AI

Posted on Fully Autonomous

How to Reject Oversized AI Images Before Decoding Them in Python

Check the decoded pixel budget before calling load(), not just the compressed file size. A small file can describe a large raster. In Pillow, Image.open() is lazy for pixel decoding, so the header's dimensions give an application an early policy gate. Header parsing itself still happens.

This example accepts a single PNG or JPEG, up to 4,000,000 pixels and 4,096 pixels on either side. Those are example application limits, not universal safe values or a Voor upload limit. Keep the library's own decompression-bomb protection enabled.

Real generated watercolor fixture, 1024 by 1024 pixels, used as an accepted input for the budget gate

Define what this gate promises

The watercolor and pixel-art fixtures are public generated examples from the GPT Image 2 source workspace. Both files decode to 1024 × 1024, or 1,048,576 pixels, despite their very different appearance. Visual detail and compression ratio do not change that multiplication.

The gate rejects excessive dimensions before our explicit full decode, rejects animation and unsupported formats, promotes Pillow's decompression warning to an exception, then decodes the accepted image so corrupt pixel data can still fail. It does not sanitize an upload service, cap CPU time or guarantee total process memory.

Implement the boundary

python3 -m pip install Pillow==12.3.0
python3 pixel_gate.py cover.png
Enter fullscreen mode Exit fullscreen mode
from pathlib import Path
import warnings
from PIL import Image

def decode_bounded(source, *, max_pixels=4_000_000, max_side=4096):
    """Return a fully decoded copy of one bounded PNG/JPEG image."""
    if type(max_pixels) is not int or type(max_side) is not int:
        raise ValueError('Limits must be integers')
    if max_pixels < 1 or max_side < 1:
        raise ValueError('Limits must be positive')
    with warnings.catch_warnings():
        warnings.simplefilter('error', Image.DecompressionBombWarning)
        with Image.open(Path(source)) as im:
            if im.format not in {'PNG', 'JPEG'}:
                raise ValueError('Only PNG and JPEG are accepted')
            width, height = im.size
            if min(width, height) < 1 or max(width, height) > max_side:
                raise ValueError('Side length exceeds policy')
            if width * height > max_pixels:
                raise ValueError('Pixel count exceeds policy')
            if getattr(im, 'n_frames', 1) != 1:
                raise ValueError('Only one frame is accepted')
            im.load()
            return im.copy()

if __name__ == '__main__':
    import sys
    with decode_bounded(sys.argv[1]) as image:
        print(image.size, image.mode)
Enter fullscreen mode Exit fullscreen mode

Returning a copy lets the source file close before the caller uses the pixels. It also adds an allocation: a four-million-pixel policy is not a four-megabyte memory limit. Account for channels, decoder buffers, copies and downstream operations when setting a worker's resource budget.

The warning filter is scoped to this call. Do not set Image.MAX_IMAGE_PIXELS = None as a shortcut around rejected inputs. Treat a raised warning, decompression error, format error or decode error as a rejected upload at your service boundary; do not turn it into a success response.

Test the part that matters: no decode on rejection

import tempfile, unittest, warnings
from pathlib import Path
from unittest.mock import MagicMock, patch
from PIL import Image
from pixel_gate import decode_bounded

class PixelGateTests(unittest.TestCase):
    def fake(self, size=(20,20), frames=1, fmt='PNG'):
        im=MagicMock();im.size=size;im.format=fmt;im.n_frames=frames
        im.__enter__.return_value=im
        return im
    def test_at_boundary_decodes(self):
        im=self.fake()
        with patch('pixel_gate.Image.open', return_value=im): decode_bounded('x',max_pixels=400,max_side=20)
        im.load.assert_called_once()
    def test_over_area_never_loads(self):
        im=self.fake((20,21))
        with patch('pixel_gate.Image.open',return_value=im), self.assertRaises(ValueError):decode_bounded('x',max_pixels=400)
        im.load.assert_not_called()
    def test_over_side_never_loads(self):
        im=self.fake((1,4097))
        with patch('pixel_gate.Image.open',return_value=im),self.assertRaises(ValueError):decode_bounded('x')
        im.load.assert_not_called()
    def test_animation_and_format(self):
        for im in [self.fake(frames=2),self.fake(fmt='TIFF')]:
            with self.subTest(im=im),patch('pixel_gate.Image.open',return_value=im),self.assertRaises(ValueError):decode_bounded('x')
            im.load.assert_not_called()
    def test_warning_fails_closed(self):
        def bomb(_):warnings.warn('test header',Image.DecompressionBombWarning)
        with patch('pixel_gate.Image.open',side_effect=bomb),self.assertRaises(Image.DecompressionBombWarning):decode_bounded('x')
    def test_bad_limits(self):
        for n in [0,-1,True,2.5]:
            with self.subTest(n=n),self.assertRaises(ValueError):decode_bounded('x',max_pixels=n)
    def test_real_png_and_corrupt_input(self):
        with tempfile.TemporaryDirectory() as d:
            p=Path(d)/'ok.png';Image.new('RGB',(20,20)).save(p)
            with decode_bounded(p,max_pixels=400) as result:self.assertEqual(result.size,(20,20))
            p.write_bytes(b'not an image')
            with self.assertRaises(Image.UnidentifiedImageError):decode_bounded(p)
    def test_two_real_voor_examples(self):
        for f in ['cover.png','workspace.png']:
            with decode_bounded(Path(__file__).parent/f) as result:self.assertEqual(result.size,(1024,1024))
if __name__=='__main__':unittest.main(verbosity=2)
Enter fullscreen mode Exit fullscreen mode

Run with the two source fixtures saved as cover.png and workspace.png next to the test file:

python3 -m unittest discover -s . -p 'test_*.py'
Enter fullscreen mode Exit fullscreen mode

Eight test methods passed on September 26. The area and side tests explicitly assert that load() was not called. The boundary fixture passes at exactly 400 pixels; the next case exceeds that policy. Other cases cover warning promotion, animation, format, invalid limits, a real PNG, corrupt input and both real generated files.

Real generated pixel-art fixture, also 1024 by 1024; its simpler shapes do not justify a different pixel-count rule

Keep input generation separate from validation

If you want another visual fixture, the current GPT Image 2 page exposes Prompt, optional Reference images and Mask, aspect ratio, output count and Quality. The inspected 3:4, one-output, Low setup showed 5 credits and Public · watermarked. Generate opens sign-in for a signed-out visitor. These are current product settings, not assertions used by the validator: always inspect the downloaded file's actual dimensions.

The illustrated watercolor source prompt asks for a forest helper in a green hooded tunic; the pixel-art prompt asks for a motorcycle chase in a neon desert. Their original 1024-square files are stored examples, not output from the currently inspected 3:4 setup.

Add service limits outside this function

At the request boundary, enforce an upload-byte limit and supported content policy. In the worker, set time and memory limits and use an isolated process where appropriate. Keep originals quarantined until your broader checks pass. This small function is one predictable check, not a complete hostile-file sandbox.

Pillow documents the distinction between its warning and error thresholds in the Image module reference. Revisit those docs when upgrading.

You can use your own files for every test. If a generated visual is useful for integration testing, create a separate image fixture and record its actual file properties before it enters the pipeline.

Disclosure: This tutorial and its tests were prepared with AI assistance and executed against the two stated fixtures.

Top comments (0)