diff --git a/crab/actions/file_actions.py b/crab/actions/file_actions.py index f876631..4574d31 100644 --- a/crab/actions/file_actions.py +++ b/crab/actions/file_actions.py @@ -13,13 +13,70 @@ # =========== Copyright 2024 @ CAMEL-AI.org. All Rights Reserved. =========== import base64 from io import BytesIO +from typing import Tuple, Optional from PIL import Image from crab.core import action +def resize_image(image: Image.Image, ratio: float) -> Image.Image: + """ + Resize the given image by the specified ratio while maintaining aspect ratio. + + :param image: The original image. + :param ratio: The ratio to resize the image. Must be between 0 and 1 (exclusive). + :return: The resized image. + """ + if not (0 < ratio <= 1): + raise ValueError("Ratio must be between 0 and 1 (exclusive).") + + new_width = int(image.width * ratio) + new_height = int(image.height * ratio) + + return image.resize((new_width, new_height), Image.LANCZOS) # 使用 LANCZOS 替代 ANTIALIAS + +def compress_image(image: Image.Image, quality: int) -> BytesIO: + """ + Compress the given image to the specified quality. + + :param image: The original image. + :param quality: The quality level (1-100). + :return: A BytesIO object containing the compressed image. + """ + output = BytesIO() + image.save(output, format="JPEG", quality=quality) + output.seek(0) + return output + @action -def save_base64_image(image: str, path: str = "image.png") -> None: - image = Image.open(BytesIO(base64.b64decode(image))) - image.save(path) +def save_base64_image(image: str, path: str = "image.png", ratio: Optional[float] = None, quality: Optional[int] = None) -> None: + """ + Save a base64 encoded image to a file, optionally resizing and compressing it. + + :param image: The base64 encoded image. + :param path: The file path to save the image. + :param ratio: The ratio to resize the image. Must be between 0 and 1 (exclusive). If None, the original size is used. + :param quality: The quality level (1-100) for image compression. If None, no compression is applied. + """ + # Decode base64 string + image_data = base64.b64decode(image) + image = Image.open(BytesIO(image_data)) + + if ratio: + image = resize_image(image, ratio) + + # Determine the output format based on the file extension + format = path.split('.')[-1].upper() + if format == 'JPG': + format = 'JPEG' + + # Convert to RGB if saving as JPEG + if format == 'JPEG': + image = image.convert('RGB') + + # Save the image + if quality: + image.save(path, format=format, quality=quality) + else: + image.save(path, format=format) diff --git a/test/core/test_images/test_compressed_image.jpg b/test/core/test_images/test_compressed_image.jpg new file mode 100644 index 0000000..2fe7ace Binary files /dev/null and b/test/core/test_images/test_compressed_image.jpg differ diff --git a/test/core/test_images/test_image.png b/test/core/test_images/test_image.png new file mode 100644 index 0000000..72628c9 Binary files /dev/null and b/test/core/test_images/test_image.png differ diff --git a/test/core/test_images/test_image_original.png b/test/core/test_images/test_image_original.png new file mode 100644 index 0000000..91c51ab Binary files /dev/null and b/test/core/test_images/test_image_original.png differ diff --git a/test/core/test_images/test_resized_compressed_image.jpg b/test/core/test_images/test_resized_compressed_image.jpg new file mode 100644 index 0000000..71a075d Binary files /dev/null and b/test/core/test_images/test_resized_compressed_image.jpg differ diff --git a/test/core/test_images/test_resized_image.png b/test/core/test_images/test_resized_image.png new file mode 100644 index 0000000..8754626 Binary files /dev/null and b/test/core/test_images/test_resized_image.png differ diff --git a/test/core/test_resize_image.py b/test/core/test_resize_image.py new file mode 100644 index 0000000..7c66d20 --- /dev/null +++ b/test/core/test_resize_image.py @@ -0,0 +1,126 @@ +# =========== Copyright 2024 @ CAMEL-AI.org. All Rights Reserved. =========== +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# =========== Copyright 2024 @ CAMEL-AI.org. All Rights Reserved. =========== + +import unittest +from io import BytesIO +from PIL import Image +import base64 +import os +import traceback + +from crab.actions.file_actions import save_base64_image +from crab.core.models import Action + +class TestSaveBase64Image(unittest.TestCase): + + # 设置类变量 + RESIZE_RATIO = 0.2 + COMPRESSION_QUALITY = 20 + + def setUp(self): + # 创建 test_images 文件夹(如果不存在) + self.test_images_dir = "test_images" + os.makedirs(self.test_images_dir, exist_ok=True) + + # 使用一个具体的图片文件 + image_path = "test_images/test_image.png" # 请替换为实际的图片路径 + with open(image_path, "rb") as image_file: + self.image_base64 = base64.b64encode(image_file.read()).decode('utf-8') + + def tearDown(self): + # 不删除生成的文件,以便查看结果 + pass + + def run_test_with_error_logging(self, test_method): + try: + print(f"Current working directory: {os.getcwd()}") + test_method() + except Exception as e: + print(f"Error in {test_method.__name__}:") + print(traceback.format_exc()) + raise + + def _test_save_image_without_resize_or_compression(self): + output_path = os.path.join(self.test_images_dir, "test_image_original.png") + try: + print("Calling save_base64_image...") + action_obj = save_base64_image(self.image_base64, path=output_path) + print(f"save_base64_image returned: {action_obj}") + + if isinstance(action_obj, Action): + result = action_obj.run() + print(f"Action execution result: {result}") + else: + print("Unexpected return type from save_base64_image") + except Exception as e: + print(f"Exception: {e}") + raise + + print(f"Checking if file exists: {os.path.abspath(output_path)}") + if os.path.exists(output_path): + print("File exists") + file_size = os.path.getsize(output_path) + print(f"File size: {file_size} bytes") + else: + print("File does not exist") + + self.assertTrue(os.path.exists(output_path), f"File {output_path} was not created") + with Image.open(output_path) as saved_image: + print(f"Original image size: {saved_image.size}") + print(f"Original image format: {saved_image.format}") + + def _test_save_image_with_resize(self): + output_path = os.path.join(self.test_images_dir, "test_resized_image.png") + action_obj = save_base64_image(self.image_base64, path=output_path, ratio=self.RESIZE_RATIO) + if isinstance(action_obj, Action): + action_obj.run() + self.assertTrue(os.path.exists(output_path), f"File {output_path} was not created") + with Image.open(output_path) as saved_image: + print(f"Resized image size: {saved_image.size}") + print(f"Resized image format: {saved_image.format}") + + def _test_save_image_with_compression(self): + output_path = os.path.join(self.test_images_dir, "test_compressed_image.jpg") + action_obj = save_base64_image(self.image_base64, path=output_path, quality=self.COMPRESSION_QUALITY) + if isinstance(action_obj, Action): + action_obj.run() + self.assertTrue(os.path.exists(output_path), f"File {output_path} was not created") + with Image.open(output_path) as saved_image: + print(f"Compressed image size: {saved_image.size}") + print(f"Compressed image format: {saved_image.format}") + + def _test_save_image_with_resize_and_compression(self): + output_path = os.path.join(self.test_images_dir, "test_resized_compressed_image.jpg") + action_obj = save_base64_image(self.image_base64, path=output_path, ratio=self.RESIZE_RATIO, quality=self.COMPRESSION_QUALITY) + if isinstance(action_obj, Action): + action_obj.run() + self.assertTrue(os.path.exists(output_path), f"File {output_path} was not created") + with Image.open(output_path) as saved_image: + print(f"Resized and compressed image size: {saved_image.size}") + print(f"Resized and compressed image format: {saved_image.format}") + + def test_save_image_without_resize_or_compression(self): + self.run_test_with_error_logging(self._test_save_image_without_resize_or_compression) + + def test_save_image_with_resize(self): + self.run_test_with_error_logging(self._test_save_image_with_resize) + + def test_save_image_with_compression(self): + self.run_test_with_error_logging(self._test_save_image_with_compression) + + def test_save_image_with_resize_and_compression(self): + self.run_test_with_error_logging(self._test_save_image_with_resize_and_compression) + +if __name__ == '__main__': + unittest.main()