33 lines
903 B
Python
33 lines
903 B
Python
# test_save_image.py
|
|
# Test the save_image_node.py script
|
|
|
|
import torch
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from save_image_node import SaveImageHigherBitDepth
|
|
|
|
def main():
|
|
# Instantiate the class
|
|
test = SaveImageHigherBitDepth()
|
|
|
|
# Create a synthetic image tensor (RGB, 256x256) with values in [0,1]
|
|
img_tensor = torch.rand(3, 256, 256)
|
|
|
|
# Create extra_pnginfo with metadata
|
|
png_info = {}
|
|
png_info['Author'] = 'TestUser'
|
|
png_info['Description'] = 'Synthetic Image for Testing'
|
|
|
|
# Call the save_images method
|
|
test.save_images(
|
|
images=[img_tensor],
|
|
bit_depth='16',
|
|
format='PNG',
|
|
filename_prefix='TestImage',
|
|
extra_pnginfo=png_info,
|
|
new_icc_profile='XYZ',
|
|
# existing_icc_profile_path="/Users/aimeri/Downloads/Adobe ICC Profiles (end-user)/RGB/AdobeRGB1998.icc"
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|