From 868f174208101352b5c5174cfa9325c666eb29c8 Mon Sep 17 00:00:00 2001 From: "Gareth.Y" <91742523+linjiaye929@users.noreply.github.com> Date: Thu, 5 Nov 2026 07:53:22 +0900 Subject: [PATCH] Fix BinaryCodec temporary resource cleanup on errors --- compressai/utils/bench/codecs.py | 63 +++++++++-------- tests/test_bench_binary_codec.py | 115 +++++++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 32 deletions(-) create mode 100644 tests/test_bench_binary_codec.py diff --git a/compressai/utils/bench/codecs.py b/compressai/utils/bench/codecs.py index 254fc6be..b6055061 100644 --- a/compressai/utils/bench/codecs.py +++ b/compressai/utils/bench/codecs.py @@ -35,7 +35,7 @@ import sys import time -from tempfile import mkstemp +from tempfile import TemporaryDirectory, mkstemp from typing import Dict, List, Optional, Union import numpy as np @@ -262,37 +262,36 @@ def setup_args(cls, _parser): pass def _run_impl(self, in_filepath, quality): - fd0, png_filepath = mkstemp(suffix=".png") - fd1, out_filepath = mkstemp(suffix=self.fmt) - - # Encode - start = time.time() - run_command(self._get_encode_cmd(in_filepath, quality, out_filepath)) - enc_time = time.time() - start - size = filesize(out_filepath) - - # Decode - start = time.time() - run_command(self._get_decode_cmd(out_filepath, png_filepath)) - dec_time = time.time() - start - - # Read image - rec = read_image(png_filepath) - os.close(fd0) - os.remove(png_filepath) - os.close(fd1) - os.remove(out_filepath) - - img = self._load_img(in_filepath) - bpp_val = float(size) * 8 / (img.size[0] * img.size[1]) - - out = { - "bpp": bpp_val, - "encoding_time": enc_time, - "decoding_time": dec_time, - } - - return out, rec + with TemporaryDirectory() as tmpdir: + fd, png_filepath = mkstemp(suffix=".png", dir=tmpdir) + os.close(fd) + fd, out_filepath = mkstemp(suffix=self.fmt, dir=tmpdir) + os.close(fd) + + # Encode + start = time.time() + run_command(self._get_encode_cmd(in_filepath, quality, out_filepath)) + enc_time = time.time() - start + size = filesize(out_filepath) + + # Decode + start = time.time() + run_command(self._get_decode_cmd(out_filepath, png_filepath)) + dec_time = time.time() - start + + # Read image + rec = read_image(png_filepath) + + img = self._load_img(in_filepath) + bpp_val = float(size) * 8 / (img.size[0] * img.size[1]) + + out = { + "bpp": bpp_val, + "encoding_time": enc_time, + "decoding_time": dec_time, + } + + return out, rec def _get_encode_cmd(self, in_filepath, quality, out_filepath): raise NotImplementedError() diff --git a/tests/test_bench_binary_codec.py b/tests/test_bench_binary_codec.py new file mode 100644 index 00000000..a75b8f76 --- /dev/null +++ b/tests/test_bench_binary_codec.py @@ -0,0 +1,115 @@ +# Copyright (c) 2021-2025, InterDigital Communications, Inc +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted (subject to the limitations in the disclaimer +# below) provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# * Neither the name of InterDigital Communications, Inc nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. + +# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY +# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT +# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A +# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR +# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF +# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +import math +import os +import shutil +import subprocess +import tempfile + +from pathlib import Path + +import numpy as np +import pytest + +from PIL import Image + +from compressai.utils.bench.codecs import JPEG2000 + + +@pytest.fixture +def jpeg2000_codec(): + ffmpeg = shutil.which("ffmpeg") + if ffmpeg is None: + pytest.skip("JPEG2000 integration tests require ffmpeg") + encoders = subprocess.check_output([ffmpeg, "-hide_banner", "-encoders"], text=True) + if "libopenjpeg" not in encoders: + pytest.skip("JPEG2000 integration tests require libopenjpeg") + return JPEG2000(argparse.Namespace()) + + +@pytest.fixture +def codec_tmpdir(tmp_path, monkeypatch): + directory = tmp_path / "codec-temp" + directory.mkdir() + # Redirect only the real temporary-file directory, not the codec subprocess. + monkeypatch.setattr(tempfile, "tempdir", str(directory)) + return directory + + +def _temporary_fds(directory): + fd_directory = Path("/proc/self/fd") + if not fd_directory.is_dir(): + return None + descriptors = set() + for path in fd_directory.iterdir(): + try: + target = os.readlink(path) + except FileNotFoundError: + continue + if target.startswith(str(directory) + os.sep): + descriptors.add(path.name) + return descriptors + + +def test_binary_codec_cleans_up_after_encoder_failure( + tmp_path, codec_tmpdir, jpeg2000_codec +): + filepath = tmp_path / "small.png" + Image.new("RGB", (1, 1), (60, 120, 180)).save(filepath) + + original_fds = _temporary_fds(codec_tmpdir) + + # OpenJPEG's default six resolution levels cannot encode this tiny image. + for _ in range(2): + with pytest.raises(SystemExit) as error: + jpeg2000_codec.run(str(filepath), quality=32) + assert error.value.code == 1 + assert _temporary_fds(codec_tmpdir) == original_fds + assert list(codec_tmpdir.iterdir()) == [] + + +def test_binary_codec_cleans_up_after_success(tmp_path, codec_tmpdir, jpeg2000_codec): + y, x = np.indices((64, 64)) + image = np.stack((x * 4, y * 4, (x + y) * 2), axis=-1).astype(np.uint8) + filepath = tmp_path / "image.png" + Image.fromarray(image).save(filepath) + + original_fds = _temporary_fds(codec_tmpdir) + info, rec = jpeg2000_codec.run(str(filepath), quality=32, return_rec=True) + + assert rec.mode == "RGB" + assert rec.size == (64, 64) + assert info["bpp"] > 0 + assert info["encoding_time"] >= 0 + assert info["decoding_time"] >= 0 + assert math.isfinite(info["psnr-rgb"]) + assert _temporary_fds(codec_tmpdir) == original_fds + assert list(codec_tmpdir.iterdir()) == []