diff --git a/js/CanvasIO.js b/js/CanvasIO.js index b33dcf0..7a49495 100644 --- a/js/CanvasIO.js +++ b/js/CanvasIO.js @@ -2,7 +2,7 @@ import { createCanvas } from "./utils/CommonUtils.js"; import { createModuleLogger } from "./log_system/log_funcs.js"; import { showErrorNotification } from "./utils/NotificationUtils.js"; import { webSocketManager } from "./utils/WebSocketManager.js"; -import { scaleImageToFit, createImageFromSource, tensorToImageData, createImageFromImageData } from "./utils/ImageUtils.js"; +import { scaleImageToFit, loadImage, tensorToImageData, createImageFromImageData } from "./utils/ImageUtils.js"; const log = createModuleLogger('CanvasIO'); export class CanvasIO { constructor(canvas) { @@ -544,7 +544,7 @@ export class CanvasIO { throw new Error("Could not create mask context"); ctx.putImageData(maskImageData, 0, 0); // Convert to HTMLImageElement - const maskImg = await createImageFromSource(maskCanvas.toDataURL()); + const maskImg = await loadImage(maskCanvas.toDataURL()); // Respect fit_on_add (scale to output area) const widgets = this.canvas.node.widgets; const fitOnAddWidget = widgets ? widgets.find((w) => w.name === "fit_on_add") : null; @@ -681,7 +681,7 @@ export class CanvasIO { log.info(`Processing batch of ${batch.length} images from backend`); for (let i = 0; i < batch.length; i++) { const imgData = batch[i]; - const img = await createImageFromSource(imgData.data); + const img = await loadImage(imgData.data); // Add image to canvas with unique name await this.canvas.canvasLayers.addLayerWithImage(img, { name: `Batch Image ${i + 1}` }, addMode, this.canvas.outputAreaBounds); log.debug(`Added batch image ${i + 1}/${batch.length} from backend`); @@ -692,7 +692,7 @@ export class CanvasIO { } else if (inputData.input_image) { // Handle single image (backward compatibility) - const img = await createImageFromSource(inputData.input_image); + const img = await loadImage(inputData.input_image); // Add image to canvas at output area position await this.canvas.canvasLayers.addLayerWithImage(img, {}, addMode, this.canvas.outputAreaBounds); log.info("Single input image added as new layer to canvas"); @@ -710,7 +710,7 @@ export class CanvasIO { if (allowMask && !maskLoaded && hasMaskInput && inputData.input_mask) { log.info("Processing input mask"); // Load mask image - const maskImg = await createImageFromSource(inputData.input_mask); + const maskImg = await loadImage(inputData.input_mask); // Determine if we should fit the mask or use it at original size const fitOnAddWidget2 = this.canvas.node.widgets.find((w) => w.name === "fit_on_add"); const shouldFit = fitOnAddWidget2 && fitOnAddWidget2.value; @@ -867,12 +867,7 @@ export class CanvasIO { const result = await response.json(); if (result.success && result.image_data) { log.info("Latest image received, adding to canvas."); - const img = new Image(); - await new Promise((resolve, reject) => { - img.onload = resolve; - img.onerror = reject; - img.src = result.image_data; - }); + const img = await loadImage(result.image_data); await this.canvas.canvasLayers.addLayerWithImage(img, {}, 'fit'); log.info("Latest image imported and placed on canvas successfully."); return true; @@ -896,7 +891,7 @@ export class CanvasIO { log.info(`Received ${result.images.length} new images, adding to canvas.`); const newLayers = []; for (const imageData of result.images) { - const img = await createImageFromSource(imageData); + const img = await loadImage(imageData); let processedImage = img; // If there's a custom shape, clip the image to that shape if (this.canvas.outputAreaShape && this.canvas.outputAreaShape.isClosed) { @@ -944,6 +939,6 @@ export class CanvasIO { ctx.closePath(); ctx.fill(); // Create a new image from the clipped canvas - return await createImageFromSource(canvas.toDataURL()); + return await loadImage(canvas.toDataURL()); } } diff --git a/js/utils/ImageUtils.js b/js/utils/ImageUtils.js index c925b3f..4c8adf9 100644 --- a/js/utils/ImageUtils.js +++ b/js/utils/ImageUtils.js @@ -146,12 +146,7 @@ export const tensorToImage = withErrorHandling(async function (tensor) { imageData.data[pixelIndex + 3] = 255; } ctx.putImageData(imageData, 0, 0); - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } throw new Error("Canvas context not available"); }, 'tensorToImage'); @@ -169,12 +164,7 @@ export const resizeImage = withErrorHandling(async function (image, maxWidth, ma ctx.imageSmoothingEnabled = true; ctx.imageSmoothingQuality = 'high'; ctx.drawImage(image, 0, 0, newWidth, newHeight); - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } throw new Error("Canvas context not available"); }, 'resizeImage'); @@ -199,11 +189,8 @@ export const base64ToImage = withErrorHandling(function (base64) { if (!base64) { throw createValidationError("Base64 string is required"); } - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = () => reject(new Error("Failed to load image from base64")); - img.src = base64; + return loadImage(base64).catch(() => { + throw new Error("Failed to load image from base64"); }); }, 'base64ToImage'); export function isValidImage(image) { @@ -225,9 +212,12 @@ export function getImageInfo(image) { area: width * height }; } -export function createImageFromSource(source) { +export function loadImage(source, options = {}) { return new Promise((resolve, reject) => { const img = new Image(); + if (options.crossOrigin !== undefined) { + img.crossOrigin = options.crossOrigin; + } img.onload = () => resolve(img); img.onerror = (err) => reject(err); img.src = source; @@ -240,12 +230,7 @@ export const createEmptyImage = withErrorHandling(function (width, height, color ctx.fillStyle = color; ctx.fillRect(0, 0, width, height); } - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } throw new Error("Canvas context not available"); }, 'createEmptyImage'); @@ -259,41 +244,7 @@ export async function convertToImage(source) { if (source instanceof HTMLImageElement) { return source; // Already an image } - const image = new Image(); - image.src = source.toDataURL(); - await new Promise((resolve, reject) => { - image.onload = () => resolve(); - image.onerror = reject; - }); - return image; -} -/** - * Creates a mask from image source for use in mask editor - * Consolidated from mask_utils.create_mask_from_image_src() - * @param imageSrc - Image source (URL or data URL) - * @returns Promise returning Image object - */ -export function createMaskFromImageSrc(imageSrc) { - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = imageSrc; - }); -} -/** - * Converts canvas to Image for use as mask - * Consolidated from mask_utils.canvas_to_mask_image() - * @param canvas - Canvas to convert - * @returns Promise returning Image object - */ -export function canvasToMaskImage(canvas) { - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(source.toDataURL()); } /** * Scales an image to fit within specified bounds while maintaining aspect ratio @@ -310,12 +261,7 @@ export async function scaleImageToFit(image, targetWidth, targetHeight) { if (!ctx) throw new Error("Could not create scaled image context"); ctx.drawImage(image, 0, 0, scaledWidth, scaledHeight); - return new Promise((resolve, reject) => { - const scaledImg = new Image(); - scaledImg.onload = () => resolve(scaledImg); - scaledImg.onerror = reject; - scaledImg.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } /** * Unified tensor to image data conversion @@ -392,5 +338,5 @@ export async function createImageFromImageData(imageData) { if (!ctx) throw new Error("Could not create canvas context"); ctx.putImageData(imageData, 0, 0); - return await createImageFromSource(canvas.toDataURL()); + return await loadImage(canvas.toDataURL()); } diff --git a/js/utils/mask_utils.js b/js/utils/mask_utils.js index 9d0c1b7..ce026cd 100644 --- a/js/utils/mask_utils.js +++ b/js/utils/mask_utils.js @@ -145,6 +145,3 @@ export const start_mask_editor_auto = withErrorHandling(function (canvasInstance } canvasInstance.startMaskEditor(null, true); }, 'start_mask_editor_auto'); -// Duplikowane funkcje zostały przeniesione do ImageUtils.ts: -// - create_mask_from_image_src -> createMaskFromImageSrc -// - canvas_to_mask_image -> canvasToMaskImage diff --git a/python/image_serialization.py b/python/image_serialization.py new file mode 100644 index 0000000..ce8c08a --- /dev/null +++ b/python/image_serialization.py @@ -0,0 +1,32 @@ +"""Shared image serialization helpers for LayerForge backend boundaries.""" + +import base64 +import io + +from PIL import Image + + +def _bytes_to_data_url(data: bytes, mime_type: str) -> str: + encoded = base64.b64encode(data).decode("ascii") + return f"data:{mime_type};base64,{encoded}" + + +def pil_to_data_url( + image: Image.Image, + *, + image_format: str = "PNG", + mime_type: str = "image/png", +) -> str: + """Serialize a PIL image using the existing PNG data-URL contract.""" + buffer = io.BytesIO() + image.save(buffer, format=image_format) + return _bytes_to_data_url(buffer.getvalue(), mime_type) + + +def file_to_data_url(path: str, *, mime_type: str = "image/png") -> str: + """Read an image file without transforming its bytes into a data URL.""" + with open(path, "rb") as image_file: + return _bytes_to_data_url(image_file.read(), mime_type) + + +__all__ = ["file_to_data_url", "pil_to_data_url"] diff --git a/python/image_utils.py b/python/image_utils.py index 5632360..ca818cd 100644 --- a/python/image_utils.py +++ b/python/image_utils.py @@ -5,6 +5,8 @@ import io import numpy as np from PIL import Image + +from .image_serialization import pil_to_data_url from torchvision import transforms from .node import log @@ -61,10 +63,7 @@ def convert_tensor_to_base64(tensor, alpha_mask=None, original_alpha=None): else: img = Image.fromarray(img_array, mode="RGB") - buffer = io.BytesIO() - img.save(buffer, format="PNG") - img_str = base64.b64encode(buffer.getvalue()).decode() - return f"data:image/png;base64,{img_str}" + return pil_to_data_url(img) except Exception as error: log.error(f"Error in convert_tensor_to_base64: {error}") log.debug(f"Tensor shape: {tensor.shape}, dtype: {tensor.dtype}") diff --git a/python/node.py b/python/node.py index 25401ea..b6b607d 100644 --- a/python/node.py +++ b/python/node.py @@ -15,6 +15,8 @@ import torch import torch.nn.functional as F from PIL import Image +from .image_serialization import pil_to_data_url + try: from .log_system import create_module_logger @@ -252,10 +254,7 @@ class LayerForgeNode: if batch_size == 1: img_np = (input_image.squeeze(0).cpu().numpy() * 255).astype(np.uint8) pil_img = Image.fromarray(img_np, "RGB") - buffered = io.BytesIO() - pil_img.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - input_data["input_image"] = f"data:image/png;base64,{img_str}" + input_data["input_image"] = pil_to_data_url(pil_img) input_data["input_image_width"] = pil_img.width input_data["input_image_height"] = pil_img.height log.debug(f"Stored single input image: {pil_img.width}x{pil_img.height}") @@ -264,12 +263,9 @@ class LayerForgeNode: for index in range(batch_size): img_np = (input_image[index].cpu().numpy() * 255).astype(np.uint8) pil_img = Image.fromarray(img_np, "RGB") - buffered = io.BytesIO() - pil_img.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() images_array.append( { - "data": f"data:image/png;base64,{img_str}", + "data": pil_to_data_url(pil_img), "width": pil_img.width, "height": pil_img.height, } @@ -290,10 +286,7 @@ class LayerForgeNode: mask_np = (input_mask.cpu().numpy() * 255).astype(np.uint8) pil_mask = Image.fromarray(mask_np, "L") - mask_buffered = io.BytesIO() - pil_mask.save(mask_buffered, format="PNG") - mask_str = base64.b64encode(mask_buffered.getvalue()).decode() - input_data["input_mask"] = f"data:image/png;base64,{mask_str}" + input_data["input_mask"] = pil_to_data_url(pil_mask) log.debug(f"Stored input mask: {pil_mask.width}x{pil_mask.height}") input_data["fit_on_add"] = fit_on_add @@ -437,10 +430,7 @@ class LayerForgeNode: def get_cached_image(self): if self.cached_image: - buffered = io.BytesIO() - self.cached_image.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - return f"data:image/png;base64,{img_str}" + return pil_to_data_url(self.cached_image) return None diff --git a/python/routes.py b/python/routes.py index 48b2eab..76f5755 100644 --- a/python/routes.py +++ b/python/routes.py @@ -1,7 +1,5 @@ """LayerForge HTTP and WebSocket route registration.""" -import base64 -import io import os import time @@ -10,6 +8,7 @@ from PIL import Image from server import PromptServer from .matting import register_matting_routes +from .image_serialization import file_to_data_url, pil_to_data_url from .node import log @@ -99,16 +98,10 @@ def register_routes(node_class): } if cache_data["image"] is not None: - buffered = io.BytesIO() - cache_data["image"].save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - response_data["data"]["image"] = f"data:image/png;base64,{img_str}" + response_data["data"]["image"] = pil_to_data_url(cache_data["image"]) if cache_data["mask"] is not None: - mask_buffer = io.BytesIO() - cache_data["mask"].save(mask_buffer, format="PNG") - mask_str = base64.b64encode(mask_buffer.getvalue()).decode() - response_data["data"]["mask"] = f"data:image/png;base64,{mask_str}" + response_data["data"]["mask"] = pil_to_data_url(cache_data["mask"]) return web.json_response(response_data) except Exception as error: @@ -122,9 +115,7 @@ def register_routes(node_class): latest_image_paths = node_class.get_latest_images(since_timestamp / 1000.0) images_data = [] for image_path in latest_image_paths: - with open(image_path, "rb") as image_file: - encoded_string = base64.b64encode(image_file.read()).decode("utf-8") - images_data.append(f"data:image/png;base64,{encoded_string}") + images_data.append(file_to_data_url(image_path)) return web.json_response({"success": True, "images": images_data}) except Exception as error: @@ -137,11 +128,9 @@ def register_routes(node_class): try: latest_image_path = node_class.get_latest_image() if latest_image_path: - with open(latest_image_path, "rb") as image_file: - encoded_string = base64.b64encode(image_file.read()).decode() return web.json_response({ "success": True, - "image_data": f"data:image/png;base64,{encoded_string}", + "image_data": file_to_data_url(latest_image_path), }) return web.json_response({ @@ -183,13 +172,10 @@ def register_routes(node_class): with Image.open(file_path) as image: if image.mode != "RGB": image = image.convert("RGB") - buffered = io.BytesIO() - image.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") log.info(f"Successfully loaded image from path: {file_path}") return web.json_response({ "success": True, - "image_data": f"data:image/png;base64,{img_str}", + "image_data": pil_to_data_url(image), "width": image.width, "height": image.height, }) diff --git a/src/CanvasIO.ts b/src/CanvasIO.ts index cb656e0..c6d87c1 100644 --- a/src/CanvasIO.ts +++ b/src/CanvasIO.ts @@ -2,7 +2,7 @@ import { createCanvas } from "./utils/CommonUtils.js"; import { createModuleLogger } from "./log_system/log_funcs.js"; import { showErrorNotification } from "./utils/NotificationUtils.js"; import { webSocketManager } from "./utils/WebSocketManager.js"; -import { scaleImageToFit, createImageFromSource, tensorToImageData, createImageFromImageData } from "./utils/ImageUtils.js"; +import { scaleImageToFit, loadImage, tensorToImageData, createImageFromImageData } from "./utils/ImageUtils.js"; import type { Canvas } from './Canvas'; import type { Layer, Shape } from './types'; @@ -620,7 +620,7 @@ export class CanvasIO { ctx.putImageData(maskImageData, 0, 0); // Convert to HTMLImageElement - const maskImg = await createImageFromSource(maskCanvas.toDataURL()); + const maskImg = await loadImage(maskCanvas.toDataURL()); // Respect fit_on_add (scale to output area) const widgets = this.canvas.node.widgets; @@ -768,7 +768,7 @@ export class CanvasIO { for (let i = 0; i < batch.length; i++) { const imgData = batch[i]; - const img = await createImageFromSource(imgData.data); + const img = await loadImage(imgData.data); // Add image to canvas with unique name await this.canvas.canvasLayers.addLayerWithImage( @@ -787,7 +787,7 @@ export class CanvasIO { } else if (inputData.input_image) { // Handle single image (backward compatibility) - const img = await createImageFromSource(inputData.input_image); + const img = await loadImage(inputData.input_image); // Add image to canvas at output area position await this.canvas.canvasLayers.addLayerWithImage( @@ -812,7 +812,7 @@ export class CanvasIO { log.info("Processing input mask"); // Load mask image - const maskImg = await createImageFromSource(inputData.input_mask); + const maskImg = await loadImage(inputData.input_mask); // Determine if we should fit the mask or use it at original size const fitOnAddWidget2 = this.canvas.node.widgets.find((w) => w.name === "fit_on_add"); @@ -994,12 +994,7 @@ export class CanvasIO { if (result.success && result.image_data) { log.info("Latest image received, adding to canvas."); - const img = new Image(); - await new Promise((resolve, reject) => { - img.onload = resolve; - img.onerror = reject; - img.src = result.image_data; - }); + const img = await loadImage(result.image_data); await this.canvas.canvasLayers.addLayerWithImage(img, {}, 'fit'); log.info("Latest image imported and placed on canvas successfully."); @@ -1025,7 +1020,7 @@ export class CanvasIO { const newLayers: (Layer | null)[] = []; for (const imageData of result.images) { - const img = await createImageFromSource(imageData); + const img = await loadImage(imageData); let processedImage = img; @@ -1079,6 +1074,6 @@ export class CanvasIO { ctx.fill(); // Create a new image from the clipped canvas - return await createImageFromSource(canvas.toDataURL()); + return await loadImage(canvas.toDataURL()); } } diff --git a/src/utils/ImageUtils.ts b/src/utils/ImageUtils.ts index 040c504..63be8eb 100644 --- a/src/utils/ImageUtils.ts +++ b/src/utils/ImageUtils.ts @@ -193,12 +193,7 @@ export const tensorToImage = withErrorHandling(async function (tensor: Tensor): ctx.putImageData(imageData, 0, 0); - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } throw new Error("Canvas context not available"); }, 'tensorToImage'); @@ -221,12 +216,7 @@ export const resizeImage = withErrorHandling(async function (image: HTMLImageEle ctx.imageSmoothingQuality = 'high'; ctx.drawImage(image, 0, 0, newWidth, newHeight); - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } throw new Error("Canvas context not available"); }, 'resizeImage'); @@ -257,11 +247,8 @@ export const base64ToImage = withErrorHandling(function (base64: string): Promis throw createValidationError("Base64 string is required"); } - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = () => reject(new Error("Failed to load image from base64")); - img.src = base64; + return loadImage(base64).catch(() => { + throw new Error("Failed to load image from base64"); }); }, 'base64ToImage'); @@ -288,9 +275,12 @@ export function getImageInfo(image: HTMLImageElement | HTMLCanvasElement): {widt }; } -export function createImageFromSource(source: string): Promise { +export function loadImage(source: string, options: { crossOrigin?: string } = {}): Promise { return new Promise((resolve, reject) => { const img = new Image(); + if (options.crossOrigin !== undefined) { + img.crossOrigin = options.crossOrigin; + } img.onload = () => resolve(img); img.onerror = (err) => reject(err); img.src = source; @@ -306,12 +296,7 @@ export const createEmptyImage = withErrorHandling(function (width: number, heigh ctx.fillRect(0, 0, width, height); } - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } throw new Error("Canvas context not available"); }, 'createEmptyImage'); @@ -327,45 +312,7 @@ export async function convertToImage(source: HTMLCanvasElement | HTMLImageElemen return source; // Already an image } - const image = new Image(); - image.src = source.toDataURL(); - - await new Promise((resolve, reject) => { - image.onload = () => resolve(); - image.onerror = reject; - }); - - return image; -} - -/** - * Creates a mask from image source for use in mask editor - * Consolidated from mask_utils.create_mask_from_image_src() - * @param imageSrc - Image source (URL or data URL) - * @returns Promise returning Image object - */ -export function createMaskFromImageSrc(imageSrc: string): Promise { - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = imageSrc; - }); -} - -/** - * Converts canvas to Image for use as mask - * Consolidated from mask_utils.canvas_to_mask_image() - * @param canvas - Canvas to convert - * @returns Promise returning Image object - */ -export function canvasToMaskImage(canvas: HTMLCanvasElement): Promise { - return new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => resolve(img); - img.onerror = (err) => reject(err); - img.src = canvas.toDataURL(); - }); + return loadImage(source.toDataURL()); } /** @@ -385,12 +332,7 @@ export async function scaleImageToFit(image: HTMLImageElement, targetWidth: numb ctx.drawImage(image, 0, 0, scaledWidth, scaledHeight); - return new Promise((resolve, reject) => { - const scaledImg = new Image(); - scaledImg.onload = () => resolve(scaledImg); - scaledImg.onerror = reject; - scaledImg.src = canvas.toDataURL(); - }); + return loadImage(canvas.toDataURL()); } /** @@ -473,5 +415,5 @@ export async function createImageFromImageData(imageData: ImageData): Promise createMaskFromImageSrc -// - canvas_to_mask_image -> canvasToMaskImage diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index 421a1ed..88c1ce3 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -197,6 +197,11 @@ def test_base64_image_conversion_preserves_rgb_and_alpha(layerforge_runtime): roundtrip = node_module.convert_tensor_to_base64(rgb_tensor) assert roundtrip.startswith("data:image/png;base64,") + roundtrip_image = Image.open(io.BytesIO(base64.b64decode(roundtrip.split(",", 1)[1]))) + assert roundtrip_image.mode == "RGB" + assert roundtrip_image.size == (2, 2) + assert roundtrip_image.getpixel((0, 0)) == (255, 0, 0) + def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, monkeypatch): import torch diff --git a/tests/test_image_serialization.py b/tests/test_image_serialization.py new file mode 100644 index 0000000..afd92da --- /dev/null +++ b/tests/test_image_serialization.py @@ -0,0 +1,35 @@ +import base64 +import io + +from PIL import Image + +from python.image_serialization import file_to_data_url, pil_to_data_url + + +def _decode_data_url(data_url): + header, encoded = data_url.split(",", 1) + return header, base64.b64decode(encoded) + + +def test_pil_to_data_url_preserves_png_mode_dimensions_and_pixels(): + image = Image.new("RGB", (2, 1), (255, 0, 128)) + + data_url = pil_to_data_url(image) + header, encoded = _decode_data_url(data_url) + decoded = Image.open(io.BytesIO(encoded)) + + assert header == "data:image/png;base64" + assert decoded.mode == "RGB" + assert decoded.size == (2, 1) + assert decoded.getpixel((0, 0)) == (255, 0, 128) + + +def test_file_to_data_url_preserves_file_bytes_and_mime_type(tmp_path): + path = tmp_path / "image.png" + path.write_bytes(b"not-transformed-image-bytes") + + data_url = file_to_data_url(str(path), mime_type="image/png") + header, encoded = _decode_data_url(data_url) + + assert header == "data:image/png;base64" + assert encoded == b"not-transformed-image-bytes" diff --git a/tests/test_image_utils.mjs b/tests/test_image_utils.mjs index dee8cf8..ae0b616 100644 --- a/tests/test_image_utils.mjs +++ b/tests/test_image_utils.mjs @@ -4,6 +4,7 @@ import test from 'node:test'; import { applyMaskToImageData, convertImageData, + loadImage, prepareImageForCanvas, tensorToImageData, validateImageData, @@ -92,3 +93,44 @@ test('tensorToImageData supports grayscale masks without a browser canvas', () = } } }); + +test('image loading resolves loaded images and preserves load errors', async () => { + const originalImage = globalThis.Image; + globalThis.Image = class { + constructor() { + this.onload = null; + this.onerror = null; + this.width = 4; + this.height = 3; + } + + set src(value) { + this.source = value; + queueMicrotask(() => { + if (value === 'bad-image') { + this.onerror?.(new Error('image decode failed')); + } else { + this.onload?.(); + } + }); + } + }; + + try { + const image = await loadImage('good-image'); + assert.equal(image.source, 'good-image'); + assert.equal(image.width, 4); + assert.equal(image.height, 3); + + await assert.rejects( + () => loadImage('bad-image'), + error => error.message === 'image decode failed' + ); + } finally { + if (originalImage === undefined) { + delete globalThis.Image; + } else { + globalThis.Image = originalImage; + } + } +});